feat: cache reusable raw Cypher translations - BED-9469

bdd/utils_test.go

@@ -21,7 +21,7 @@ import (
 	"testing"
 
 	"github.com/cucumber/godog"
-	"github.com/cucumber/messages/go/v34"
+	messages "github.com/cucumber/messages/go/v34"
 	"github.com/specterops/dawgs/graph"
 	"github.com/stretchr/testify/require"
 )

drivers/pg/compiler.go

@@ -0,0 +1,38 @@
+package pg
+
+import (
+	"context"
+	"strings"
+
+	"github.com/specterops/dawgs/cypher/frontend"
+	"github.com/specterops/dawgs/cypher/models/cypher"
+	"github.com/specterops/dawgs/cypher/models/pgsql/translate"
+)
+
+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) 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()
+
+	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 {
+			return "", translationCacheBuildResult{}, err
+		} else if sqlQuery, err := translate.Translated(translated); err != nil {
+			return "", translationCacheBuildResult{}, err
+		} else {
+			return sqlQuery, translationCacheBuildResult{
+				parameters:       translated.Parameters,
+				parameterSources: parameterSources,
+			}, nil
+		}
+	}
+
+	key := translationCache.Key(source, graphID, parameters)
+	return translationCache.GetOrBuildContext(ctx, key, parameters, build)
+}

drivers/pg/driver.go

@@ -42,11 +42,38 @@ type Driver struct {
 	*SchemaManager
 }
 
+// DriverOptions configures driver-wide behavior. TranslationCacheEntries is a
+// count, not a per-connection setting; zero disables translation caching.
+type DriverOptions struct {
+	TranslationCacheEntries int
+}
+
+// DefaultDriverOptions returns the driver-wide PostgreSQL settings used by
+// NewDriver.
+func DefaultDriverOptions() DriverOptions {
+	return DriverOptions{
+		TranslationCacheEntries: translationCacheCapacity,
+	}
+}
+
 func NewDriver(graphQueryMemoryLimit size.Size, pool *pgxpool.Pool) *Driver {
+	return NewDriverWithOptions(graphQueryMemoryLimit, pool, DefaultDriverOptions())
+}
+
+// NewDriverWithOptions constructs a PostgreSQL driver with driver-wide options.
+func NewDriverWithOptions(graphQueryMemoryLimit size.Size, pool *pgxpool.Pool, options DriverOptions) *Driver {
+	options = normalizeDriverOptions(options)
 	return &Driver{
 		pool:          pool,
-		SchemaManager: NewSchemaManager(pool, graphQueryMemoryLimit),
+		SchemaManager: NewSchemaManagerWithOptions(pool, graphQueryMemoryLimit, options),
+	}
+}
+
+func normalizeDriverOptions(options DriverOptions) DriverOptions {
+	if options.TranslationCacheEntries < 0 {
+		options.TranslationCacheEntries = 0
 	}
+	return options
 }
 
 func (s *Driver) SetDefaultGraph(ctx context.Context, graphSchema graph.Graph) error {
@@ -96,10 +123,17 @@ func (s *Driver) BatchOperation(ctx context.Context, batchDelegate graph.BatchDe
 }
 
 func (s *Driver) Close(ctx context.Context) error {
+	s.translationCache.Close()
 	s.pool.Close()
 	return nil
 }
 
+// TranslationCacheStats returns aggregate PostgreSQL translation-cache
+// counters. It never exposes cached query text, SQL, or parameter data.
+func (s *Driver) TranslationCacheStats() TranslationCacheStats {
+	return s.translationCache.Stats()
+}
+
 func renderConfig(batchWriteSize int, pgxOptions pgx.TxOptions, userOptions []graph.TransactionOption) (*Config, error) {
 	graphCfg := graph.TransactionConfig{
 		DriverConfig: &Config{
@@ -175,7 +209,12 @@ func (s *Driver) RefreshKinds(ctx context.Context) error {
 
 	// Wipe this map to be rebuilt in the fetch call below
 	s.SchemaManager.kindIDsByKind = map[int16]graph.Kind{}
-	return s.SchemaManager.Fetch(ctx)
+	if err := s.SchemaManager.Fetch(ctx); err != nil {
+		return err
+	}
+
+	s.translationCache.Invalidate()
+	return nil
 }
 
 func (s *Driver) OptimizeStorage(ctx context.Context) error {

drivers/pg/manager.go

@@ -33,17 +33,37 @@ func KindMapperFromGraphDatabase(graphDB graph.Database) (KindMapper, error) {
 }
 
 type SchemaManager struct {
-	defaultGraph          model.Graph
-	pool                  *pgxpool.Pool
-	hasDefaultGraph       bool
-	graphs                map[string]model.Graph
-	kindsByID             map[graph.Kind]int16
-	kindIDsByKind         map[int16]graph.Kind
-	lock                  *sync.RWMutex
-	graphQueryMemoryLimit size.Size
+	defaultGraph             model.Graph
+	pool                     *pgxpool.Pool
+	hasDefaultGraph          bool
+	graphs                   map[string]model.Graph
+	kindsByID                map[graph.Kind]int16
+	kindIDsByKind            map[int16]graph.Kind
+	lock                     *sync.RWMutex
+	graphQueryMemoryLimit    size.Size
+	translationCache         *translationCache
+	translationCacheProvider translationCacheProvider
 }
 
 func NewSchemaManager(pool *pgxpool.Pool, graphQueryMemoryLimit size.Size) *SchemaManager {
+	return NewSchemaManagerWithOptions(pool, graphQueryMemoryLimit, DefaultDriverOptions())
+}
+
+// NewSchemaManagerWithTranslationCache permits an application to disable the
+// compilation cache (zero), or to choose a bounded driver-wide entry capacity.
+// Negative values disable the cache.
+func NewSchemaManagerWithTranslationCache(pool *pgxpool.Pool, graphQueryMemoryLimit size.Size, translationCacheEntries int) *SchemaManager {
+	return NewSchemaManagerWithOptions(pool, graphQueryMemoryLimit, DriverOptions{
+		TranslationCacheEntries: translationCacheEntries,
+	})
+}
+
+// NewSchemaManagerWithOptions constructs the shared compilation service with
+// a bounded translation cache.
+func NewSchemaManagerWithOptions(pool *pgxpool.Pool, graphQueryMemoryLimit size.Size, options DriverOptions) *SchemaManager {
+	options = normalizeDriverOptions(options)
+	translationCache := newTranslationCache(options.TranslationCacheEntries)
+
 	return &SchemaManager{
 		pool:                  pool,
 		hasDefaultGraph:       false,
@@ -52,6 +72,10 @@ func NewSchemaManager(pool *pgxpool.Pool, graphQueryMemoryLimit size.Size) *Sche
 		kindIDsByKind:         map[int16]graph.Kind{},
 		lock:                  &sync.RWMutex{},
 		graphQueryMemoryLimit: graphQueryMemoryLimit,
+		translationCache:      translationCache,
+		translationCacheProvider: sharedTranslationCacheProvider{
+			cache: translationCache,
+		},
 	}
 }
 
@@ -406,7 +430,12 @@ func (s *SchemaManager) AssertSchema(ctx context.Context, schema graph.Schema) e
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
-	return s.WriteTransaction(ctx, func(tx graph.Transaction) error {
+	if err := s.WriteTransaction(ctx, func(tx graph.Transaction) error {
 		return s.assertSchema(tx, schema)
-	}, OptionSetQueryExecMode(pgx.QueryExecModeSimpleProtocol))
+	}, OptionSetQueryExecMode(pgx.QueryExecModeSimpleProtocol)); err != nil {
+		return err
+	}
+
+	s.translationCache.Invalidate()
+	return nil
 }

drivers/pg/transaction.go

@@ -5,12 +5,10 @@ import (
 	"fmt"
 
 	"github.com/specterops/dawgs/cypher/models/pgsql"
-	"github.com/specterops/dawgs/cypher/models/pgsql/translate"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgxpool"
-	"github.com/specterops/dawgs/cypher/frontend"
 	"github.com/specterops/dawgs/drivers/pg/model"
 	"github.com/specterops/dawgs/graph"
 	"github.com/specterops/dawgs/query"
@@ -275,16 +273,15 @@ func (s *transaction) query(query string, parameters map[string]any) (pgx.Rows,
 }
 
 func (s *transaction) Query(query string, parameters map[string]any) graph.Result {
-	if parsedQuery, err := frontend.ParseCypher(frontend.NewContext(), query); err != nil {
-		return graph.NewErrorResult(err)
-	} else if graphTarget, err := s.getTargetGraph(); err != nil {
-		return graph.NewErrorResult(err)
-	} else if translated, err := translate.Translate(s.ctx, parsedQuery, s.schemaManager, parameters, graphTarget.ID); err != nil {
-		return graph.NewErrorResult(err)
-	} else if sqlQuery, err := translate.Translated(translated); err != nil {
+	if graphTarget, err := s.getTargetGraph(); err != nil {
 		return graph.NewErrorResult(err)
 	} else {
-		return s.Raw(sqlQuery, translated.Parameters)
+		sqlQuery, bindings, err := s.schemaManager.compileText(s.ctx, query, parameters, graphTarget.ID)
+		if err != nil {
+			return graph.NewErrorResult(err)
+		}
+
+		return s.Raw(sqlQuery, bindings)
 	}
 }