diff --git a/README.md b/README.md index 32f6ce56..c774fa66 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,12 @@ Run the package benchmark suite with: make test_bench ``` +The PostgreSQL driver has a bounded, driver-wide Cypher compilation cache for both text Cypher and programmatic graph +queries. Its default capacity is 256 entries; callers that need a different capacity can use +`pg.NewDriverWithOptions`. The process-wide `pg.SetOptimizedTranslation` switch selects baseline translation for +newly started compilations across every PostgreSQL driver in the process. See +[PostgreSQL translation](docs/postgresql_translation.md#translation-cache) for the cache contract and rollback controls. + Use `cmd/benchdiff` to compare benchmarks between two committed refs without changing the active worktree: ```bash diff --git a/docs/postgresql_translation.md b/docs/postgresql_translation.md index c05d2c94..34e613d0 100644 --- a/docs/postgresql_translation.md +++ b/docs/postgresql_translation.md @@ -59,6 +59,51 @@ Substring and suffix predicates are not promoted to blanket schema indexes. Post parameter/property forms that lower to helper functions remain outside the hard index-match contract until their lowering changes. +## Translation Cache + +The PostgreSQL driver keeps one bounded, driver-wide compilation cache with a default capacity of 256 rendered SQL +statements. It serves both `Transaction.Query` and the legacy programmatic query builder used by relationship and node +queries. A warm builder hit avoids optimization, lowering-plan construction, PostgreSQL AST construction, and SQL +rendering; it still builds the request AST, deterministically names its runtime parameters, renders the canonical +Cypher cache identity, and assembles its debug comment. + +Entries are partitioned by the SHA-256 digest of trimmed Cypher text, target graph ID, sorted parameter names and PostgreSQL type shapes, a +translation-key format/policy identity, and the schema generation. Parameter values are never retained or used as key +material, and the cache key does not retain the source text. The retained data is limited to the SQL string and generated-parameter-to-Cypher-parameter source names; +the cache does not retain caller maps or values, ASTs, contexts, transactions, connections, rows, or connection +strings. Structural literals remain part of the cache identity. Generated parameters that cannot be reconstructed from +request-local named sources, translation errors, disabled/closed cache state, and source text over 64 KiB bypass +retention. + +Configure the bounded cache through the PostgreSQL driver constructor: + +```go +database := pg.NewDriverWithOptions(0, pool, pg.DriverOptions{ + TranslationCacheEntries: 0, // disables cache lookup and retention +}) +``` + +For a process-wide rollback, disable optimized translation directly and restore the prior state when appropriate: + +```go +previous := pg.SetOptimizedTranslation(false) +defer pg.SetOptimizedTranslation(previous) +``` + +When disabled, newly started PostgreSQL compilations bypass cache lookup and retention and translate the original AST +with no PostgreSQL rewrite rules or lowering decisions. The setting is atomic and applies to every PostgreSQL driver in +the process; each compilation snapshots it at entry, so already-running compilations continue with their selected path. +Cached optimized entries remain dormant and are reusable after re-enabling. A zero cache capacity still runs optimized +translation without retaining entries. The default remains optimized translation with a 256-entry cache. + +Concurrent cacheable misses for the same key share one translation. Waiters behind a failed, canceled, or non-cacheable +build continue independently rather than serializing behind repeated failed work. Statistics are available through +`(*pg.Driver).TranslationCacheStats()` and expose aggregate hits, build-leader misses, coalesced requests, bypasses, +unoptimized compilations, insertions, evictions, binding/build failures, live size, capacity, and generation; they never +expose query or value data. Successful schema assertions and kind refreshes advance the schema generation and discard completed entries. +External schema or type changes cannot be detected automatically; recreate the driver/pool after those changes. Driver +close retires the cache before the PostgreSQL pool closes. + ## Validation Workflow Optimizer changes should include focused optimizer/lowering tests, SQL-shape translation tests, and backend-equivalent diff --git a/drivers/pg/compiler.go b/drivers/pg/compiler.go index a971684d..90157b0d 100644 --- a/drivers/pg/compiler.go +++ b/drivers/pg/compiler.go @@ -6,22 +6,84 @@ import ( "github.com/specterops/dawgs/cypher/frontend" "github.com/specterops/dawgs/cypher/models/cypher" + cypherFormat "github.com/specterops/dawgs/cypher/models/cypher/format" "github.com/specterops/dawgs/cypher/models/pgsql/translate" + "github.com/specterops/dawgs/cypher/models/walk" + "github.com/specterops/dawgs/query" ) +const builderParameterPrefix = "__dawgs_builder_p" + +var builderCommentNewlines = strings.NewReplacer("\r\n", "\n-- ", "\r", "\n-- ", "\n", "\n-- ") + +// preparedRegularQuery holds a canonical view of a builder query. The source +// AST remains untouched on cache hits; cold compilation creates its own copy. +// Parameter values remain request-local and are never retained by the cache. +type preparedRegularQuery struct { + query *cypher.RegularQuery + source string + commentSource string + parameters map[string]any +} + +func prepareRegularQuery(regularQuery *cypher.RegularQuery) (preparedRegularQuery, error) { + namer := query.NewParameterNamerWithPrefix(builderParameterPrefix) + if err := walk.Cypher(regularQuery, namer); err != nil { + return preparedRegularQuery{}, err + } else if source, err := cypherFormat.RegularQueryWithParameterSequence(regularQuery, false, namer.Symbols); err != nil { + return preparedRegularQuery{}, err + } else if commentSource, err := cypherFormat.RegularQueryWithParameterSequence(regularQuery, true, namer.Symbols); err != nil { + return preparedRegularQuery{}, err + } else { + return preparedRegularQuery{ + query: regularQuery, + source: strings.TrimSpace(source), + commentSource: strings.TrimSpace(commentSource), + parameters: namer.Parameters, + }, nil + } +} + +func (s preparedRegularQuery) translationQuery() (*cypher.RegularQuery, error) { + owned := cypher.Copy(s.query) + rewriter := query.NewParameterRewriterWithPrefix(builderParameterPrefix) + if err := walk.Cypher(owned, rewriter); err != nil { + return nil, err + } + + return owned, nil +} + +// Both PostgreSQL Cypher entry points use these methods so a builder query and +// its text equivalent share the same cache and cacheability rules. func (s *SchemaManager) compileText(ctx context.Context, source string, parameters map[string]any, graphID int32) (string, map[string]any, error) { return s.compile(ctx, strings.TrimSpace(source), parameters, graphID, func() (*cypher.RegularQuery, error) { return frontend.ParseCypher(frontend.NewContext(), source) }) } +func (s *SchemaManager) compileRegularQuery(ctx context.Context, prepared preparedRegularQuery, graphID int32) (string, map[string]any, error) { + return s.compile(ctx, prepared.source, prepared.parameters, graphID, func() (*cypher.RegularQuery, error) { + return prepared.translationQuery() + }) +} + func (s *SchemaManager) compile(ctx context.Context, source string, parameters map[string]any, graphID int32, parse func() (*cypher.RegularQuery, error)) (string, map[string]any, error) { - translationCache := s.translationCacheProvider.TranslationCache() + var ( + translationCache = s.translationCacheProvider.TranslationCache() + optimized = OptimizedTranslationEnabled() + translationOptions = translate.Options{ + OptimizerMode: translate.OptimizerDisabled, + } + ) + if optimized { + translationOptions.OptimizerMode = translate.OptimizerEnabled + } build := func() (string, translationCacheBuildResult, error) { if regularQuery, err := parse(); err != nil { return "", translationCacheBuildResult{}, err - } else if translated, parameterSources, err := translate.TranslateWithOptionsAndParameterSources(ctx, regularQuery, s, parameters, graphID, translate.DefaultOptions()); err != nil { + } else if translated, parameterSources, err := translate.TranslateWithOptionsAndParameterSources(ctx, regularQuery, s, parameters, graphID, translationOptions); err != nil { return "", translationCacheBuildResult{}, err } else if sqlQuery, err := translate.Translated(translated); err != nil { return "", translationCacheBuildResult{}, err @@ -33,6 +95,14 @@ func (s *SchemaManager) compile(ctx context.Context, source string, parameters m } } + if !optimized { + return translationCache.BuildUnoptimized(build) + } + key := translationCache.Key(source, graphID, parameters) return translationCache.GetOrBuildContext(ctx, key, parameters, build) } + +func commentRegularQuery(source, sql string) string { + return "-- " + builderCommentNewlines.Replace(source) + "\n" + sql +} diff --git a/drivers/pg/compiler_test.go b/drivers/pg/compiler_test.go new file mode 100644 index 00000000..abf7ea4b --- /dev/null +++ b/drivers/pg/compiler_test.go @@ -0,0 +1,446 @@ +package pg + +import ( + "context" + "testing" + "time" + + "github.com/specterops/dawgs/cypher/models/cypher" + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/query" + "github.com/stretchr/testify/require" +) + +func TestPrepareRegularQueryUsesStableNamesAndDoesNotMutateBuilder(t *testing.T) { + firstBuilder := query.NewBuilder(nil) + firstBuilder.Apply( + query.Where(query.Equals(query.NodeProperty("objectid"), graph.ID(1))), + query.Returning(query.Node()), + ) + + first, err := firstBuilder.Build(false) + require.NoError(t, err) + + firstPrepared, err := prepareRegularQuery(first) + require.NoError(t, err) + require.Same(t, first, firstPrepared.query) + + secondBuilder := query.NewBuilder(nil) + secondBuilder.Apply( + query.Where(query.Equals(query.NodeProperty("objectid"), graph.ID(2))), + query.Returning(query.Node()), + ) + + second, err := secondBuilder.Build(false) + require.NoError(t, err) + + secondPrepared, err := prepareRegularQuery(second) + require.NoError(t, err) + + require.Equal(t, firstPrepared.source, secondPrepared.source) + require.Equal(t, map[string]any{ + builderParameterPrefix + "0": graph.ID(1), + }, firstPrepared.parameters) + require.Equal(t, map[string]any{ + builderParameterPrefix + "0": graph.ID(2), + }, secondPrepared.parameters) + require.NotContains(t, firstPrepared.source, "1") + require.NotContains(t, secondPrepared.source, "2") + + // Re-preparing the original builder query proves that preparation did not + // leak request-local generated names into it. + preparedAgain, err := prepareRegularQuery(first) + require.NoError(t, err) + require.Equal(t, firstPrepared.source, preparedAgain.source) +} + +func TestPrepareRegularQueryRedactsLiteralsInSQLComments(t *testing.T) { + builder := query.NewBuilder(nil) + builder.Apply(query.Returning(query.Literal("sensitive-value"))) + + regularQuery, err := builder.Build(false) + require.NoError(t, err) + prepared, err := prepareRegularQuery(regularQuery) + require.NoError(t, err) + require.Contains(t, prepared.source, "sensitive-value") + require.NotContains(t, prepared.commentSource, "sensitive-value") + require.Contains(t, prepared.commentSource, "$STRIPPED") + + commented := commentRegularQuery(prepared.commentSource, "select 1") + require.NotContains(t, commented, "sensitive-value") + require.Contains(t, commented, "$STRIPPED") +} + +func TestBuilderPreparedQueriesShareTranslationCacheKey(t *testing.T) { + cache := newTranslationCache(1) + first := preparedRegularQuery{ + source: "MATCH (n) WHERE n.id = $__dawgs_builder_p0 RETURN n", + parameters: map[string]any{ + builderParameterPrefix + "0": graph.ID(1), + }, + } + second := preparedRegularQuery{ + source: first.source, + parameters: map[string]any{ + builderParameterPrefix + "0": graph.ID(2), + }, + } + key := cache.Key(first.source, 1, first.parameters) + + _, bindings, err := cache.GetOrBuild(key, first.parameters, cacheableBuild( + "select @p0", + map[string]any{ + "p0": uint64(1), + }, + map[string]string{ + "p0": builderParameterPrefix + "0", + }, + )) + require.NoError(t, err) + require.Equal(t, map[string]any{ + "p0": uint64(1), + }, bindings) + + builds := 0 + _, bindings, err = cache.GetOrBuild(cache.Key(second.source, 1, second.parameters), second.parameters, func() (string, translationCacheBuildResult, error) { + builds++ + return "", translationCacheBuildResult{}, nil + }) + require.NoError(t, err) + require.Zero(t, builds) + require.Equal(t, map[string]any{ + "p0": uint64(2), + }, bindings) +} + +func buildPreparedNodeLookup(t *testing.T, id graph.ID) preparedRegularQuery { + t.Helper() + builder := query.NewBuilder(nil) + builder.Apply( + query.Where(query.Equals(query.NodeProperty("objectid"), id)), + query.Returning(query.Node()), + ) + regularQuery, err := builder.Build(false) + require.NoError(t, err) + prepared, err := prepareRegularQuery(regularQuery) + require.NoError(t, err) + return prepared +} + +func TestCompileRegularQueryCachesBuilderShapeAndRebindsValues(t *testing.T) { + setOptimizedTranslationForTest(t, true) + manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 2, + }) + first := buildPreparedNodeLookup(t, graph.ID(1)) + second := buildPreparedNodeLookup(t, graph.ID(2)) + + firstSQL, firstBindings, err := manager.compileRegularQuery(context.Background(), first, 7) + require.NoError(t, err) + secondSQL, secondBindings, err := manager.compileRegularQuery(context.Background(), second, 7) + require.NoError(t, err) + + require.Equal(t, firstSQL, secondSQL) + require.NotEqual(t, firstBindings, secondBindings) + require.Len(t, firstBindings, 1) + require.Len(t, secondBindings, 1) + for name, firstValue := range firstBindings { + require.Equal(t, uint64(1), firstValue) + require.Equal(t, uint64(2), secondBindings[name]) + } + stats := manager.translationCache.Stats() + require.Equal(t, int64(1), stats.Misses) + require.Equal(t, int64(1), stats.Hits) + require.Equal(t, int64(1), stats.Insertions) +} + +func compilePreparedBuilderQuery(t *testing.T, builder *query.Builder) preparedRegularQuery { + t.Helper() + + regularQuery, err := builder.Build(false) + require.NoError(t, err) + prepared, err := prepareRegularQuery(regularQuery) + require.NoError(t, err) + return prepared +} + +// requireWarmBindingsMatchCold compiles one builder request to populate a cache, +// then verifies that a same-shape request receives the exact bindings it would +// have received from a fresh translation. This protects the complete builder +// preparation -> cache -> translation path rather than only cache internals. +func requireWarmBindingsMatchCold(t *testing.T, first, second preparedRegularQuery) { + t.Helper() + setOptimizedTranslationForTest(t, true) + + warmManager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 8, + }) + coldManager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 0, + }) + + firstSQL, firstBindings, err := warmManager.compileRegularQuery(context.Background(), first, 7) + require.NoError(t, err) + warmSQL, warmBindings, err := warmManager.compileRegularQuery(context.Background(), second, 7) + require.NoError(t, err) + coldSQL, coldBindings, err := coldManager.compileRegularQuery(context.Background(), second, 7) + require.NoError(t, err) + + require.Equal(t, first.source, second.source) + require.Equal(t, firstSQL, warmSQL) + require.Equal(t, coldSQL, warmSQL) + require.Equal(t, coldBindings, warmBindings) + require.NotEqual(t, firstBindings, warmBindings, "a cache hit must not retain the first request's values") + require.Equal(t, int64(1), warmManager.translationCache.Stats().Misses) + require.Equal(t, int64(1), warmManager.translationCache.Stats().Hits) +} + +func TestCompileRegularQueryWarmBindingsMatchColdForMultipleParameters(t *testing.T) { + timestamp := time.Date(2026, 8, 26, 12, 0, 0, 0, time.UTC) + testCases := []struct { + name string + first func() *query.Builder + second func() *query.Builder + }{ + { + name: "nested scalar predicates preserve order", + first: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply(query.Where(query.And( + query.Equals(query.NodeProperty("id"), graph.ID(101)), + query.Or( + query.StringContains(query.NodeProperty("name"), "first-name"), + query.GreaterThan(query.NodeProperty("rank"), int64(303)), + ), + query.Equals(query.NodeProperty("enabled"), true), + )), query.Returning(query.Node())) + return builder + }, + second: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply(query.Where(query.And( + query.Equals(query.NodeProperty("id"), graph.ID(202)), + query.Or( + query.StringContains(query.NodeProperty("name"), "second-name"), + query.GreaterThan(query.NodeProperty("rank"), int64(404)), + ), + query.Equals(query.NodeProperty("enabled"), false), + )), query.Returning(query.Node())) + return builder + }, + }, + { + name: "ID collections and temporal values", + first: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply(query.Where(query.And( + query.InIDs(query.Node(), graph.ID(1), graph.ID(3)), + query.Before(query.NodeProperty("seen"), timestamp), + )), query.Returning(query.Node())) + return builder + }, + second: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply(query.Where(query.And( + query.InIDs(query.Node(), graph.ID(2), graph.ID(4), graph.ID(6)), + query.Before(query.NodeProperty("seen"), timestamp.Add(time.Hour)), + )), query.Returning(query.Node())) + return builder + }, + }, + { + name: "floating point values", + first: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply(query.Where(query.Equals(query.NodeProperty("score"), 1.25)), query.Returning(query.Node())) + return builder + }, + second: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply(query.Where(query.Equals(query.NodeProperty("score"), 9.75)), query.Returning(query.Node())) + return builder + }, + }, + { + name: "map properties", + first: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply(query.Create(query.NodePattern(nil, query.Parameter(map[string]any{"name": "first", "rank": int64(1)})))) + return builder + }, + second: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply(query.Create(query.NodePattern(nil, query.Parameter(map[string]any{"name": "second", "rank": int64(2)})))) + return builder + }, + }, + { + name: "set values and filter values", + first: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply( + query.Where(query.Equals(query.NodeProperty("id"), graph.ID(11))), + query.Updatef(func() graph.Criteria { return query.SetProperties(query.Node(), map[string]any{"value": "first"}) }), + ) + return builder + }, + second: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply( + query.Where(query.Equals(query.NodeProperty("id"), graph.ID(22))), + query.Updatef(func() graph.Criteria { return query.SetProperties(query.Node(), map[string]any{"value": "second"}) }), + ) + return builder + }, + }, + { + name: "relationship set values and filter values", + first: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply( + query.Where(query.InIDs(query.Relationship(), graph.ID(31))), + query.Updatef(func() graph.Criteria { + return query.SetProperties(query.Relationship(), map[string]any{"value": "first"}) + }), + ) + return builder + }, + second: func() *query.Builder { + builder := query.NewBuilder(nil) + builder.Apply( + query.Where(query.InIDs(query.Relationship(), graph.ID(32))), + query.Updatef(func() graph.Criteria { + return query.SetProperties(query.Relationship(), map[string]any{"value": "second"}) + }), + ) + return builder + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + first := compilePreparedBuilderQuery(t, testCase.first()) + second := compilePreparedBuilderQuery(t, testCase.second()) + requireWarmBindingsMatchCold(t, first, second) + }) + } +} + +func TestCompileRegularQueryPartitionsParameterTypesAndStructuralLiterals(t *testing.T) { + manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 8, + }) + + stringBuilder := query.NewBuilder(nil) + stringBuilder.Apply(query.Where(query.Equals(query.NodeProperty("value"), "one")), query.Returning(query.Node())) + stringQuery := compilePreparedBuilderQuery(t, stringBuilder) + intBuilder := query.NewBuilder(nil) + intBuilder.Apply(query.Where(query.Equals(query.NodeProperty("value"), int64(1))), query.Returning(query.Node())) + intQuery := compilePreparedBuilderQuery(t, intBuilder) + require.Equal(t, stringQuery.source, intQuery.source) + require.NotEqual(t, manager.translationCache.Key(stringQuery.source, 7, stringQuery.parameters), manager.translationCache.Key(intQuery.source, 7, intQuery.parameters)) + + limitOne := query.NewBuilder(nil) + limitOne.Apply(query.Where(query.Equals(query.NodeProperty("value"), "one")), query.Limit(1), query.Returning(query.Node())) + limitTwo := query.NewBuilder(nil) + limitTwo.Apply(query.Where(query.Equals(query.NodeProperty("value"), "one")), query.Limit(2), query.Returning(query.Node())) + limitOneQuery := compilePreparedBuilderQuery(t, limitOne) + limitTwoQuery := compilePreparedBuilderQuery(t, limitTwo) + require.NotEqual(t, limitOneQuery.source, limitTwoQuery.source) + require.NotEqual(t, manager.translationCache.Key(limitOneQuery.source, 7, limitOneQuery.parameters), manager.translationCache.Key(limitTwoQuery.source, 7, limitTwoQuery.parameters)) + + emptyIDs := map[string]any{"ids": []graph.ID{}} + populatedIDs := map[string]any{"ids": []graph.ID{1, 2}} + strings := map[string]any{"ids": []string{"1", "2"}} + require.Equal(t, manager.translationCache.Key("RETURN $ids", 7, emptyIDs), manager.translationCache.Key("RETURN $ids", 7, populatedIDs)) + require.NotEqual(t, manager.translationCache.Key("RETURN $ids", 7, populatedIDs), manager.translationCache.Key("RETURN $ids", 7, strings)) +} + +func TestPrepareRegularQueryRenamesExplicitAndRepeatedParameterSymbols(t *testing.T) { + build := func(value string) preparedRegularQuery { + parameter := cypher.NewParameter(builderParameterPrefix+"99", value) + builder := query.NewBuilder(nil) + builder.Apply( + query.Where(query.And( + query.Equals(query.NodeProperty("first"), parameter), + query.Equals(query.NodeProperty("second"), parameter), + )), + query.Returning(query.Node()), + ) + return compilePreparedBuilderQuery(t, builder) + } + + first := build("first") + second := build("second") + require.NotContains(t, first.source, builderParameterPrefix+"99") + require.Equal(t, first.source, second.source) + require.Equal(t, map[string]any{ + builderParameterPrefix + "0": "first", + builderParameterPrefix + "1": "first", + }, first.parameters) + requireWarmBindingsMatchCold(t, first, second) +} + +func TestCompileRegularQueryDisabledCacheDoesNotRetainTranslation(t *testing.T) { + setOptimizedTranslationForTest(t, true) + manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 0, + }) + + _, _, err := manager.compileRegularQuery(context.Background(), buildPreparedNodeLookup(t, graph.ID(1)), 7) + require.NoError(t, err) + _, _, err = manager.compileRegularQuery(context.Background(), buildPreparedNodeLookup(t, graph.ID(2)), 7) + require.NoError(t, err) + stats := manager.translationCache.Stats() + require.Zero(t, stats.Misses) + require.Zero(t, stats.Size) + require.Equal(t, int64(2), stats.Bypasses) +} + +func TestCompileRegularQueryUnoptimizedBypassesAndPreservesWarmEntries(t *testing.T) { + setOptimizedTranslationForTest(t, true) + manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 2, + }) + + first := buildPreparedNodeLookup(t, graph.ID(1)) + _, _, err := manager.compileRegularQuery(context.Background(), first, 7) + require.NoError(t, err) + warmStats := manager.translationCache.Stats() + + setOptimizedTranslationForTest(t, false) + _, _, err = manager.compileRegularQuery(context.Background(), buildPreparedNodeLookup(t, graph.ID(2)), 7) + require.NoError(t, err) + disabledStats := manager.translationCache.Stats() + require.Equal(t, warmStats.Hits, disabledStats.Hits) + require.Equal(t, warmStats.Misses, disabledStats.Misses) + require.Equal(t, warmStats.Insertions, disabledStats.Insertions) + require.Equal(t, warmStats.Bypasses+1, disabledStats.Bypasses) + require.Equal(t, warmStats.UnoptimizedCompilations+1, disabledStats.UnoptimizedCompilations) + + setOptimizedTranslationForTest(t, true) + _, _, err = manager.compileRegularQuery(context.Background(), buildPreparedNodeLookup(t, graph.ID(3)), 7) + require.NoError(t, err) + require.Equal(t, warmStats.Hits+1, manager.translationCache.Stats().Hits) +} + +func TestCompileTextUnoptimizedBypassesCache(t *testing.T) { + setOptimizedTranslationForTest(t, true) + manager := NewSchemaManagerWithOptions(nil, 0, DriverOptions{ + TranslationCacheEntries: 2, + }) + + _, _, err := manager.compileText(context.Background(), "MATCH (n) RETURN n", nil, 7) + require.NoError(t, err) + warmStats := manager.translationCache.Stats() + + setOptimizedTranslationForTest(t, false) + _, _, err = manager.compileText(context.Background(), "MATCH (n) RETURN n", nil, 7) + require.NoError(t, err) + disabledStats := manager.translationCache.Stats() + require.Equal(t, warmStats.Hits, disabledStats.Hits) + require.Equal(t, warmStats.Misses, disabledStats.Misses) + require.Equal(t, warmStats.Bypasses+1, disabledStats.Bypasses) + require.Equal(t, warmStats.UnoptimizedCompilations+1, disabledStats.UnoptimizedCompilations) +} diff --git a/drivers/pg/optimized_translation.go b/drivers/pg/optimized_translation.go new file mode 100644 index 00000000..79ecdfef --- /dev/null +++ b/drivers/pg/optimized_translation.go @@ -0,0 +1,25 @@ +package pg + +import "sync/atomic" + +var optimizedTranslationEnabled atomic.Bool + +func init() { + optimizedTranslationEnabled.Store(true) +} + +// SetOptimizedTranslation enables or disables optimized PostgreSQL Cypher +// translation for the current process. It returns the previous setting. +// +// Disabling optimization bypasses the translation cache and uses the baseline +// translation path for calls that begin after the setting changes. Existing +// cached entries remain available if optimization is enabled again. +func SetOptimizedTranslation(enabled bool) bool { + return optimizedTranslationEnabled.Swap(enabled) +} + +// OptimizedTranslationEnabled reports whether PostgreSQL Cypher translation +// uses optimization and the translation cache for newly started compilations. +func OptimizedTranslationEnabled() bool { + return optimizedTranslationEnabled.Load() +} diff --git a/drivers/pg/optimized_translation_test.go b/drivers/pg/optimized_translation_test.go new file mode 100644 index 00000000..576b19a2 --- /dev/null +++ b/drivers/pg/optimized_translation_test.go @@ -0,0 +1,48 @@ +package pg + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +func setOptimizedTranslationForTest(t testing.TB, enabled bool) { + t.Helper() + previous := SetOptimizedTranslation(enabled) + t.Cleanup(func() { + SetOptimizedTranslation(previous) + }) +} + +func TestOptimizedTranslationSwitch(t *testing.T) { + setOptimizedTranslationForTest(t, true) + require.True(t, OptimizedTranslationEnabled()) + + previous := SetOptimizedTranslation(false) + require.True(t, previous) + require.False(t, OptimizedTranslationEnabled()) + + previous = SetOptimizedTranslation(true) + require.False(t, previous) + require.True(t, OptimizedTranslationEnabled()) +} + +func TestOptimizedTranslationSwitchConcurrentAccess(t *testing.T) { + setOptimizedTranslationForTest(t, true) + var group sync.WaitGroup + for index := range 16 { + group.Add(1) + go func(enabled bool) { + defer group.Done() + for range 100 { + SetOptimizedTranslation(enabled) + _ = OptimizedTranslationEnabled() + } + }(index%2 == 0) + } + group.Wait() + + SetOptimizedTranslation(true) + require.True(t, OptimizedTranslationEnabled()) +} diff --git a/drivers/pg/query.go b/drivers/pg/query.go index c06a8eed..21367f2e 100644 --- a/drivers/pg/query.go +++ b/drivers/pg/query.go @@ -3,24 +3,21 @@ package pg import ( "context" - "github.com/specterops/dawgs/cypher/models/pgsql/translate" "github.com/specterops/dawgs/graph" "github.com/specterops/dawgs/query" ) type liveQuery struct { ctx context.Context - tx graph.Transaction - kindMapper KindMapper + tx *transaction graphIDResolver func() (int32, error) queryBuilder *query.Builder } -func newLiveQuery(ctx context.Context, tx graph.Transaction, kindMapper KindMapper, graphIDResolver func() (int32, error)) liveQuery { +func newLiveQuery(ctx context.Context, tx *transaction, graphIDResolver func() (int32, error)) liveQuery { return liveQuery{ ctx: ctx, tx: tx, - kindMapper: kindMapper, graphIDResolver: graphIDResolver, queryBuilder: query.NewBuilder(nil), } @@ -29,12 +26,14 @@ func newLiveQuery(ctx context.Context, tx graph.Transaction, kindMapper KindMapp func (s *liveQuery) runRegularQuery(allShortestPaths bool) graph.Result { if regularQuery, err := s.queryBuilder.Build(allShortestPaths); err != nil { return graph.NewErrorResult(err) + } else if prepared, err := prepareRegularQuery(regularQuery); err != nil { + return graph.NewErrorResult(err) } else if graphID, err := s.graphIDResolver(); err != nil { return graph.NewErrorResult(err) - } else if translation, err := translate.FromCypher(s.ctx, regularQuery, s.kindMapper, false, graphID); err != nil { + } else if sqlQuery, bindings, err := s.tx.schemaManager.compileRegularQuery(s.ctx, prepared, graphID); err != nil { return graph.NewErrorResult(err) } else { - return s.tx.Raw(translation.Statement, translation.Parameters) + return s.tx.Raw(commentRegularQuery(prepared.commentSource, sqlQuery), bindings) } } diff --git a/drivers/pg/relationship.go b/drivers/pg/relationship.go index 72cd3515..562d9cf0 100644 --- a/drivers/pg/relationship.go +++ b/drivers/pg/relationship.go @@ -52,11 +52,11 @@ func (s *relationshipQuery) Update(properties *graph.Properties) error { var updateStatements []graph.Criteria if modifiedProperties := properties.ModifiedProperties(); len(modifiedProperties) > 0 { - updateStatements = append(updateStatements, query.SetProperties(query.Node(), modifiedProperties)) + updateStatements = append(updateStatements, query.SetProperties(query.Relationship(), modifiedProperties)) } if deletedProperties := properties.DeletedProperties(); len(deletedProperties) > 0 { - updateStatements = append(updateStatements, query.DeleteProperties(query.Node(), deletedProperties...)) + updateStatements = append(updateStatements, query.DeleteProperties(query.Relationship(), deletedProperties...)) } return updateStatements diff --git a/drivers/pg/transaction.go b/drivers/pg/transaction.go index c99c06a3..ccfb10e2 100644 --- a/drivers/pg/transaction.go +++ b/drivers/pg/transaction.go @@ -180,7 +180,7 @@ func (s *transaction) UpdateNode(node *graph.Node) error { func (s *transaction) Nodes() graph.NodeQuery { return &nodeQuery{ - liveQuery: newLiveQuery(s.ctx, s, s.schemaManager, s.targetGraphID), + liveQuery: newLiveQuery(s.ctx, s, s.targetGraphID), } } @@ -258,7 +258,7 @@ func (s *transaction) UpdateRelationship(relationship *graph.Relationship) error func (s *transaction) Relationships() graph.RelationshipQuery { return &relationshipQuery{ - liveQuery: newLiveQuery(s.ctx, s, s.schemaManager, s.targetGraphID), + liveQuery: newLiveQuery(s.ctx, s, s.targetGraphID), } } diff --git a/integration/pgsql_translation_cache_test.go b/integration/pgsql_translation_cache_test.go new file mode 100644 index 00000000..55413de6 --- /dev/null +++ b/integration/pgsql_translation_cache_test.go @@ -0,0 +1,389 @@ +//go:build manual_integration + +package integration + +import ( + "fmt" + "testing" + + "github.com/specterops/dawgs/drivers/pg" + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ops" + "github.com/specterops/dawgs/query" + "github.com/stretchr/testify/require" +) + +func TestPostgreSQLFetchStartNodesUsesBuilderCompilationCache(t *testing.T) { + previous := pg.SetOptimizedTranslation(true) + t.Cleanup(func() { + pg.SetOptimizedTranslation(previous) + }) + + userKind := graph.StringKind("CacheUser") + groupKind := graph.StringKind("CacheGroup") + memberKind := graph.StringKind("CacheMember") + session := Open(t, Options{ + RequireDriver: pg.DriverName, + SkipIfNoConnection: true, + SkipIfDriverMismatch: true, + CleanupMode: CleanupGraph, + ExtraNodeKinds: graph.Kinds{userKind, groupKind}, + ExtraEdgeKinds: graph.Kinds{memberKind}, + }) + driver, ok := session.DB.(*pg.Driver) + require.True(t, ok) + + var ( + first *graph.Node + second *graph.Node + ) + require.NoError(t, session.DB.WriteTransaction(session.Ctx, func(tx graph.Transaction) error { + var err error + if first, err = tx.CreateNode(graph.AsProperties(map[string]any{"name": "first"}), userKind); err != nil { + return err + } + if second, err = tx.CreateNode(graph.AsProperties(map[string]any{"name": "second"}), userKind); err != nil { + return err + } + target, err := tx.CreateNode(graph.AsProperties(map[string]any{"name": "target"}), groupKind) + if err != nil { + return err + } + if _, err = tx.CreateRelationshipByIDs(first.ID, target.ID, memberKind, graph.NewProperties()); err != nil { + return err + } + _, err = tx.CreateRelationshipByIDs(second.ID, target.ID, memberKind, graph.NewProperties()) + return err + })) + + before := driver.TranslationCacheStats() + require.NoError(t, session.DB.ReadTransaction(session.Ctx, func(tx graph.Transaction) error { + firstNodes, err := ops.FetchStartNodes(tx.Relationships().Filter(query.InIDs(query.Start(), first.ID))) + if err != nil { + return err + } + secondNodes, err := ops.FetchStartNodes(tx.Relationships().Filter(query.InIDs(query.Start(), second.ID))) + if err != nil { + return err + } + if _, found := firstNodes[first.ID]; !found { + t.Fatalf("first result omitted node %d", first.ID) + } + if _, found := secondNodes[second.ID]; !found { + t.Fatalf("second result omitted node %d", second.ID) + } + return nil + })) + + after := driver.TranslationCacheStats() + require.Equal(t, before.Misses+1, after.Misses) + require.GreaterOrEqual(t, after.Hits, before.Hits+1) +} + +func TestPostgreSQLFetchStartNodesUnoptimizedBypassesCache(t *testing.T) { + previous := pg.SetOptimizedTranslation(false) + t.Cleanup(func() { + pg.SetOptimizedTranslation(previous) + }) + + userKind := graph.StringKind("BaselineCacheUser") + groupKind := graph.StringKind("BaselineCacheGroup") + memberKind := graph.StringKind("BaselineCacheMember") + session := Open(t, Options{ + RequireDriver: pg.DriverName, + SkipIfNoConnection: true, + SkipIfDriverMismatch: true, + CleanupMode: CleanupGraph, + ExtraNodeKinds: graph.Kinds{userKind, groupKind}, + ExtraEdgeKinds: graph.Kinds{memberKind}, + }) + driver, ok := session.DB.(*pg.Driver) + require.True(t, ok) + + var ( + first *graph.Node + second *graph.Node + ) + require.NoError(t, session.DB.WriteTransaction(session.Ctx, func(tx graph.Transaction) error { + var err error + if first, err = tx.CreateNode(graph.NewProperties(), userKind); err != nil { + return err + } + if second, err = tx.CreateNode(graph.NewProperties(), userKind); err != nil { + return err + } + target, err := tx.CreateNode(graph.NewProperties(), groupKind) + if err != nil { + return err + } + if _, err = tx.CreateRelationshipByIDs(first.ID, target.ID, memberKind, graph.NewProperties()); err != nil { + return err + } + _, err = tx.CreateRelationshipByIDs(second.ID, target.ID, memberKind, graph.NewProperties()) + return err + })) + + before := driver.TranslationCacheStats() + require.NoError(t, session.DB.ReadTransaction(session.Ctx, func(tx graph.Transaction) error { + if _, err := ops.FetchStartNodes(tx.Relationships().Filter(query.InIDs(query.Start(), first.ID))); err != nil { + return err + } + _, err := ops.FetchStartNodes(tx.Relationships().Filter(query.InIDs(query.Start(), second.ID))) + return err + })) + + after := driver.TranslationCacheStats() + require.Equal(t, before.Hits, after.Hits) + require.Equal(t, before.Misses, after.Misses) + require.Equal(t, before.Insertions, after.Insertions) + require.Equal(t, before.Bypasses+2, after.Bypasses) + require.Equal(t, before.UnoptimizedCompilations+2, after.UnoptimizedCompilations) +} + +func TestPostgreSQLRawCypherQueryRebindsCachedParameters(t *testing.T) { + previous := pg.SetOptimizedTranslation(true) + t.Cleanup(func() { + pg.SetOptimizedTranslation(previous) + }) + + nodeKind := graph.StringKind("RawCypherCacheNode") + session := Open(t, Options{ + RequireDriver: pg.DriverName, + SkipIfNoConnection: true, + SkipIfDriverMismatch: true, + CleanupMode: CleanupGraph, + ExtraNodeKinds: graph.Kinds{nodeKind}, + }) + driver, ok := session.DB.(*pg.Driver) + require.True(t, ok) + + var ( + first *graph.Node + second *graph.Node + ) + require.NoError(t, session.DB.WriteTransaction(session.Ctx, func(tx graph.Transaction) error { + var err error + if first, err = tx.CreateNode(graph.AsProperties(map[string]any{"name": "first"}), nodeKind); err != nil { + return err + } + second, err = tx.CreateNode(graph.AsProperties(map[string]any{"name": "second"}), nodeKind) + return err + })) + + before := driver.TranslationCacheStats() + require.NoError(t, session.DB.ReadTransaction(session.Ctx, func(tx graph.Transaction) error { + firstID, err := rawCypherNodeID(tx, "first") + if err != nil { + return err + } + secondID, err := rawCypherNodeID(tx, "second") + if err != nil { + return err + } + if firstID != first.ID || secondID != second.ID { + return fmt.Errorf("cached raw Cypher query returned incorrect node IDs: %d, %d", firstID, secondID) + } + return nil + })) + + after := driver.TranslationCacheStats() + require.Equal(t, before.Misses+1, after.Misses) + require.GreaterOrEqual(t, after.Hits, before.Hits+1) +} + +func TestPostgreSQLRawCypherQueryUnoptimizedBypassesCache(t *testing.T) { + previous := pg.SetOptimizedTranslation(false) + t.Cleanup(func() { + pg.SetOptimizedTranslation(previous) + }) + + nodeKind := graph.StringKind("RawCypherCacheBaselineNode") + session := Open(t, Options{ + RequireDriver: pg.DriverName, + SkipIfNoConnection: true, + SkipIfDriverMismatch: true, + CleanupMode: CleanupGraph, + ExtraNodeKinds: graph.Kinds{nodeKind}, + }) + driver, ok := session.DB.(*pg.Driver) + require.True(t, ok) + + before := driver.TranslationCacheStats() + require.NoError(t, session.DB.ReadTransaction(session.Ctx, func(tx graph.Transaction) error { + for _, name := range []string{"first", "second"} { + result := tx.Query("RETURN $name", map[string]any{"name": name}) + if !result.Next() { + result.Close() + return result.Error() + } + result.Close() + } + return nil + })) + + after := driver.TranslationCacheStats() + require.Equal(t, before.Hits, after.Hits) + require.Equal(t, before.Misses, after.Misses) + require.Equal(t, before.Bypasses+2, after.Bypasses) + require.Equal(t, before.UnoptimizedCompilations+2, after.UnoptimizedCompilations) +} + +func rawCypherNodeID(tx graph.Transaction, name string) (graph.ID, error) { + result := tx.Query("MATCH (n:RawCypherCacheNode) WHERE n.name = $name RETURN n", map[string]any{"name": name}) + defer result.Close() + if !result.Next() { + return 0, result.Error() + } + + var node graph.Node + if err := result.Scan(&node); err != nil { + return 0, err + } + return node.ID, result.Error() +} + +func TestPostgreSQLNodeUpdateRebindsCachedBuilderParameters(t *testing.T) { + previous := pg.SetOptimizedTranslation(true) + t.Cleanup(func() { + pg.SetOptimizedTranslation(previous) + }) + + nodeKind := graph.StringKind("CacheUpdateNode") + session := Open(t, Options{ + RequireDriver: pg.DriverName, + SkipIfNoConnection: true, + SkipIfDriverMismatch: true, + CleanupMode: CleanupGraph, + ExtraNodeKinds: graph.Kinds{nodeKind}, + }) + driver, ok := session.DB.(*pg.Driver) + require.True(t, ok) + + var ( + first *graph.Node + second *graph.Node + ) + require.NoError(t, session.DB.WriteTransaction(session.Ctx, func(tx graph.Transaction) error { + var err error + if first, err = tx.CreateNode(graph.AsProperties(map[string]any{"name": "before-first"}), nodeKind); err != nil { + return err + } + second, err = tx.CreateNode(graph.AsProperties(map[string]any{"name": "before-second"}), nodeKind) + return err + })) + + before := driver.TranslationCacheStats() + require.NoError(t, session.DB.WriteTransaction(session.Ctx, func(tx graph.Transaction) error { + firstProperties := graph.NewProperties().Set("name", "after-first") + if err := tx.Nodes().Filter(query.InIDs(query.Node(), first.ID)).Update(firstProperties); err != nil { + return err + } + + secondProperties := graph.NewProperties().Set("name", "after-second") + return tx.Nodes().Filter(query.InIDs(query.Node(), second.ID)).Update(secondProperties) + })) + after := driver.TranslationCacheStats() + require.Equal(t, before.Misses+1, after.Misses) + require.GreaterOrEqual(t, after.Hits, before.Hits+1) + + require.NoError(t, session.DB.ReadTransaction(session.Ctx, func(tx graph.Transaction) error { + updatedFirst, err := tx.Nodes().Filter(query.InIDs(query.Node(), first.ID)).First() + if err != nil { + return err + } + updatedSecond, err := tx.Nodes().Filter(query.InIDs(query.Node(), second.ID)).First() + if err != nil { + return err + } + firstName, err := updatedFirst.Properties.Get("name").String() + if err != nil { + return err + } + secondName, err := updatedSecond.Properties.Get("name").String() + if err != nil { + return err + } + if firstName != "after-first" || secondName != "after-second" { + return fmt.Errorf("cached update bound wrong values: got %q and %q", firstName, secondName) + } + return nil + })) +} + +func TestPostgreSQLRelationshipUpdateRebindsCachedBuilderParameters(t *testing.T) { + previous := pg.SetOptimizedTranslation(true) + t.Cleanup(func() { + pg.SetOptimizedTranslation(previous) + }) + + nodeKind := graph.StringKind("CacheUpdateRelationshipNode") + relationshipKind := graph.StringKind("CacheUpdateRelationship") + session := Open(t, Options{ + RequireDriver: pg.DriverName, + SkipIfNoConnection: true, + SkipIfDriverMismatch: true, + CleanupMode: CleanupGraph, + ExtraNodeKinds: graph.Kinds{nodeKind}, + ExtraEdgeKinds: graph.Kinds{relationshipKind}, + }) + driver, ok := session.DB.(*pg.Driver) + require.True(t, ok) + + var ( + first *graph.Relationship + second *graph.Relationship + ) + require.NoError(t, session.DB.WriteTransaction(session.Ctx, func(tx graph.Transaction) error { + start, err := tx.CreateNode(graph.NewProperties(), nodeKind) + if err != nil { + return err + } + end, err := tx.CreateNode(graph.NewProperties(), nodeKind) + if err != nil { + return err + } + secondEnd, err := tx.CreateNode(graph.NewProperties(), nodeKind) + if err != nil { + return err + } + if first, err = tx.CreateRelationshipByIDs(start.ID, end.ID, relationshipKind, graph.AsProperties(map[string]any{"name": "before-first"})); err != nil { + return err + } + second, err = tx.CreateRelationshipByIDs(start.ID, secondEnd.ID, relationshipKind, graph.AsProperties(map[string]any{"name": "before-second"})) + return err + })) + + before := driver.TranslationCacheStats() + require.NoError(t, session.DB.WriteTransaction(session.Ctx, func(tx graph.Transaction) error { + if err := tx.Relationships().Filter(query.InIDs(query.Relationship(), first.ID)).Update(graph.NewProperties().Set("name", "after-first")); err != nil { + return err + } + return tx.Relationships().Filter(query.InIDs(query.Relationship(), second.ID)).Update(graph.NewProperties().Set("name", "after-second")) + })) + after := driver.TranslationCacheStats() + require.Equal(t, before.Misses+1, after.Misses) + require.GreaterOrEqual(t, after.Hits, before.Hits+1) + + require.NoError(t, session.DB.ReadTransaction(session.Ctx, func(tx graph.Transaction) error { + updatedFirst, err := tx.Relationships().Filter(query.InIDs(query.Relationship(), first.ID)).First() + if err != nil { + return err + } + updatedSecond, err := tx.Relationships().Filter(query.InIDs(query.Relationship(), second.ID)).First() + if err != nil { + return err + } + firstName, err := updatedFirst.Properties.Get("name").String() + if err != nil { + return err + } + secondName, err := updatedSecond.Properties.Get("name").String() + if err != nil { + return err + } + if firstName != "after-first" || secondName != "after-second" { + return fmt.Errorf("cached relationship update bound wrong values: got %q and %q", firstName, secondName) + } + return nil + })) +}