feat(pg): add cfg memory control to driver BED-7082

drivers/pg/driver.go

@@ -7,6 +7,7 @@ import (
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/specterops/dawgs/graph"
+	"github.com/specterops/dawgs/util/size"
 )
 
 var (
@@ -40,10 +41,10 @@ type Driver struct {
 	*SchemaManager
 }
 
-func NewDriver(pool *pgxpool.Pool) *Driver {
+func NewDriver(graphQueryMemoryLimit size.Size, pool *pgxpool.Pool) *Driver {
 	return &Driver{
 		pool:          pool,
-		SchemaManager: NewSchemaManager(pool),
+		SchemaManager: NewSchemaManager(pool, graphQueryMemoryLimit),
 	}
 }
 

drivers/pg/manager.go

@@ -12,6 +12,7 @@ import (
 	"github.com/specterops/dawgs/drivers/pg/model"
 	"github.com/specterops/dawgs/drivers/pg/query"
 	"github.com/specterops/dawgs/graph"
+	"github.com/specterops/dawgs/util/size"
 )
 
 type KindMapper interface {
@@ -32,23 +33,25 @@ 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
+	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
 }
 
-func NewSchemaManager(pool *pgxpool.Pool) *SchemaManager {
+func NewSchemaManager(pool *pgxpool.Pool, graphQueryMemoryLimit size.Size) *SchemaManager {
 	return &SchemaManager{
-		pool:            pool,
-		hasDefaultGraph: false,
-		graphs:          map[string]model.Graph{},
-		kindsByID:       map[graph.Kind]int16{},
-		kindIDsByKind:   map[int16]graph.Kind{},
-		lock:            &sync.RWMutex{},
+		pool:                  pool,
+		hasDefaultGraph:       false,
+		graphs:                map[string]model.Graph{},
+		kindsByID:             map[graph.Kind]int16{},
+		kindIDsByKind:         map[int16]graph.Kind{},
+		lock:                  &sync.RWMutex{},
+		graphQueryMemoryLimit: graphQueryMemoryLimit,
 	}
 }
 

drivers/pg/pg.go

@@ -83,6 +83,6 @@ func NewPool(connectionString string) (*pgxpool.Pool, error) {
 
 func init() {
 	dawgs.Register(DriverName, func(ctx context.Context, cfg dawgs.Config) (graph.Database, error) {
-		return NewDriver(cfg.Pool), nil
+		return NewDriver(cfg.GraphQueryMemoryLimit, cfg.Pool), nil
 	})
 }

drivers/pg/transaction.go

@@ -87,7 +87,7 @@ func (s *transaction) driver() driver {
 }
 
 func (s *transaction) GraphQueryMemoryLimit() size.Size {
-	return size.Gibibyte
+	return s.schemaManager.graphQueryMemoryLimit
 }
 
 func (s *transaction) WithGraph(schema graph.Graph) graph.Transaction {