BED-5467 Allow DB injection into DAWGS (#1314)

* chore: combine driver and schemamanager into a hyperstruct

* feat: move pgxpool creation outside of the driver with injection

* fix: infinite loop

* chore: add another lock to prevent deadlocks

* chore: partial revert

* feat: refactor schemamanager/driver to remove circular dependency

This embeds schemamanager into the driver, so that the transaction methods can live on schemamanager but be accessed from the driver layer.

This did necessitate moving batchWriteSize to a package level var instead of living on the driver.

This ended up being the simplest change to allow injecting a pgxpool into dawgs, and additional work should be done to better unwind the driver and schemamanager responsibilities

* chore: move of NewPool was too aggressive, moving back to pg driver package

* chore: oops

* fix: overapplied rename

* chore: fix infertypeargs flags

Author: superlinkx

Diff for: cmd/api/src/api/auth_internal_test.go

@@ -333,15 +333,15 @@ func TestValidateRequestSignature(t *testing.T) {
 
 		tmpFiles, err := os.ReadDir(authenticator.cfg.TempDirectory())
 		assert.NoError(t, err)
-		assert.Len(t, slicesext.Filter[fs.DirEntry](tmpFiles, func(file fs.DirEntry) bool {
+		assert.Len(t, slicesext.Filter(tmpFiles, func(file fs.DirEntry) bool {
 			return strings.HasPrefix(file.Name(), "bh-request-")
 		}), 1)
 
 		// Closing the body should remove the tmp file
 		req.Body.Close()
 		tmpFiles, err = os.ReadDir(os.TempDir())
 		assert.NoError(t, err)
-		assert.Len(t, slicesext.Filter[fs.DirEntry](tmpFiles, func(file fs.DirEntry) bool {
+		assert.Len(t, slicesext.Filter(tmpFiles, func(file fs.DirEntry) bool {
 			return strings.HasPrefix(file.Name(), "bh-request-")
 		}), 0)
 	})
@@ -376,7 +376,7 @@ func TestValidateRequestSignature(t *testing.T) {
 		// "small" payloads should not create a tmp file
 		tmpFiles, err := os.ReadDir(os.TempDir())
 		assert.NoError(t, err)
-		assert.Len(t, slicesext.Filter[fs.DirEntry](tmpFiles, func(file fs.DirEntry) bool {
+		assert.Len(t, slicesext.Filter(tmpFiles, func(file fs.DirEntry) bool {
 			return strings.HasPrefix(file.Name(), "bh-request-")
 		}), 0)
 	})

Diff for: cmd/api/src/api/tools/pg.go

@@ -83,7 +83,12 @@ func migrateTypes(ctx context.Context, neoDB, pgDB graph.Database) error {
 		return err
 	}
 
-	_, err := pgDB.(*pg.Driver).KindMapper().AssertKinds(ctx, append(neoNodeKinds, neoEdgeKinds...))
+	driver, ok := pgDB.(*pg.Driver)
+	if !ok {
+		return fmt.Errorf("current graph database is not a pg driver")
+	}
+
+	_, err := driver.KindMapper().AssertKinds(ctx, append(neoNodeKinds, neoEdgeKinds...))
 	return err
 }
 
@@ -443,15 +448,20 @@ func (s *PGMigrator) MigrationStatus(response http.ResponseWriter, request *http
 }
 
 func (s *PGMigrator) OpenPostgresGraphConnection() (graph.Database, error) {
-	return dawgs.Open(s.ServerCtx, pg.DriverName, dawgs.Config{
-		GraphQueryMemoryLimit: size.Gibibyte,
-		DriverCfg:             s.Cfg.Database.PostgreSQLConnectionString(),
-	})
+	if pool, err := pg.NewPool(s.Cfg.Database.PostgreSQLConnectionString()); err != nil {
+		return nil, err
+	} else {
+		return dawgs.Open(s.ServerCtx, pg.DriverName, dawgs.Config{
+			GraphQueryMemoryLimit: size.Gibibyte,
+			ConnectionString:      s.Cfg.Database.PostgreSQLConnectionString(),
+			Pool:                  pool,
+		})
+	}
 }
 
 func (s *PGMigrator) OpenNeo4jGraphConnection() (graph.Database, error) {
 	return dawgs.Open(s.ServerCtx, neo4j.DriverName, dawgs.Config{
 		GraphQueryMemoryLimit: size.Gibibyte,
-		DriverCfg:             s.Cfg.Neo4J.Neo4jConnectionString(),
+		ConnectionString:      s.Cfg.Neo4J.Neo4jConnectionString(),
 	})
 }

Diff for: cmd/api/src/bootstrap/util.go

@@ -22,6 +22,7 @@ import (
 	"log/slog"
 	"os"
 
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/specterops/bloodhound/bhlog"
 	"github.com/specterops/bloodhound/bhlog/level"
 	"github.com/specterops/bloodhound/dawgs"
@@ -75,34 +76,45 @@ func DefaultConfigFilePath() string {
 }
 
 func ConnectGraph(ctx context.Context, cfg config.Configuration) (*graph.DatabaseSwitch, error) {
-	var connectionString string
-
-	if driverName, err := tools.LookupGraphDriver(ctx, cfg); err != nil {
+	var (
+		connectionString string
+		pool             *pgxpool.Pool
+		err              error
+	)
+
+	driverName, err := tools.LookupGraphDriver(ctx, cfg)
+	if err != nil {
 		return nil, err
-	} else {
-		switch driverName {
-		case neo4j.DriverName:
-			slog.InfoContext(ctx, "Connecting to graph using Neo4j")
-			connectionString = cfg.Neo4J.Neo4jConnectionString()
+	}
 
-		case pg.DriverName:
-			slog.InfoContext(ctx, "Connecting to graph using PostgreSQL")
-			connectionString = cfg.Database.PostgreSQLConnectionString()
+	switch driverName {
+	case neo4j.DriverName:
+		slog.InfoContext(ctx, "Connecting to graph using Neo4j")
+		connectionString = cfg.Neo4J.Neo4jConnectionString()
 
-		default:
-			return nil, fmt.Errorf("unknown graphdb driver name: %s", driverName)
-		}
+	case pg.DriverName:
+		slog.InfoContext(ctx, "Connecting to graph using PostgreSQL")
+		connectionString = cfg.Database.PostgreSQLConnectionString()
 
-		if connectionString == "" {
-			return nil, fmt.Errorf("graph connection requires a connection url to be set")
-		} else if graphDatabase, err := dawgs.Open(ctx, driverName, dawgs.Config{
-			GraphQueryMemoryLimit: size.Size(cfg.GraphQueryMemoryLimit) * size.Gibibyte,
-			DriverCfg:             connectionString,
-		}); err != nil {
+		pool, err = pg.NewPool(connectionString)
+		if err != nil {
 			return nil, err
-		} else {
-			return graph.NewDatabaseSwitch(ctx, graphDatabase), nil
 		}
+
+	default:
+		return nil, fmt.Errorf("unknown graphdb driver name: %s", driverName)
+	}
+
+	if connectionString == "" {
+		return nil, fmt.Errorf("graph connection requires a connection url to be set")
+	} else if graphDatabase, err := dawgs.Open(ctx, driverName, dawgs.Config{
+		GraphQueryMemoryLimit: size.Size(cfg.GraphQueryMemoryLimit) * size.Gibibyte,
+		ConnectionString:      connectionString,
+		Pool:                  pool,
+	}); err != nil {
+		return nil, err
+	} else {
+		return graph.NewDatabaseSwitch(ctx, graphDatabase), nil
 	}
 }
 

Diff for: cmd/api/src/cmd/dawgs-harness/main.go

@@ -27,6 +27,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/jedib0t/go-pretty/v6/table"
 	"github.com/specterops/bloodhound/bhlog"
 	"github.com/specterops/bloodhound/dawgs"
@@ -44,9 +45,22 @@ func fatalf(format string, args ...any) {
 }
 
 func RunTestSuite(ctx context.Context, connectionStr, driverName string) tests.TestSuite {
+	var (
+		pool *pgxpool.Pool
+		err  error
+	)
+
+	if driverName == pg.DriverName {
+		pool, err = pg.NewPool(connectionStr)
+		if err != nil {
+			fatalf("Failed creating a new pgxpool: %s", err)
+		}
+	}
+
 	if connection, err := dawgs.Open(context.TODO(), driverName, dawgs.Config{
 		GraphQueryMemoryLimit: size.Gibibyte,
-		DriverCfg:             connectionStr,
+		ConnectionString:      connectionStr,
+		Pool:                  pool,
 	}); err != nil {
 		fatalf("Failed opening %s database: %v", driverName, err)
 	} else {
@@ -65,7 +79,7 @@ func RunTestSuite(ctx context.Context, connectionStr, driverName string) tests.T
 		}
 	}
 
-	panic(nil)
+	panic("unexpected error")
 }
 
 func newContext() context.Context {

Diff for: cmd/api/src/database/log.go

@@ -53,7 +53,7 @@ func (s *GormLogAdapter) Trace(ctx context.Context, begin time.Time, fc func() (
 	if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
 		sql, _ := fc()
 
-		slog.ErrorContext(ctx, "Database error", "query", sql, "error", err)
+		slog.ErrorContext(ctx, "Database error", "query", sql, "err", err)
 	} else {
 		elapsed := time.Since(begin)
 

Diff for: cmd/api/src/test/integration/dawgs.go

@@ -47,20 +47,24 @@ func OpenGraphDB(testCtrl test.Controller, schema graph.Schema) graph.Database {
 
 	switch cfg.GraphDriver {
 	case pg.DriverName:
+		pool, err := pg.NewPool(cfg.Database.PostgreSQLConnectionString())
+		test.RequireNilErrf(testCtrl, err, "Failed to create new pgx pool: %v", err)
 		graphDatabase, err = dawgs.Open(context.TODO(), cfg.GraphDriver, dawgs.Config{
-			DriverCfg: cfg.Database.PostgreSQLConnectionString(),
+			ConnectionString: cfg.Database.PostgreSQLConnectionString(),
+			Pool:             pool,
 		})
+		test.RequireNilErrf(testCtrl, err, "Failed connecting to graph database: %v", err)
 
 	case neo4j.DriverName:
 		graphDatabase, err = dawgs.Open(context.TODO(), cfg.GraphDriver, dawgs.Config{
-			DriverCfg: cfg.Neo4J.Neo4jConnectionString(),
+			ConnectionString: cfg.Neo4J.Neo4jConnectionString(),
 		})
+		test.RequireNilErrf(testCtrl, err, "Failed connecting to graph database: %v", err)
 
 	default:
 		testCtrl.Fatalf("unsupported graph driver name %s", cfg.GraphDriver)
 	}
 
-	test.RequireNilErrf(testCtrl, err, "Failed connecting to graph database: %v", err)
 	test.RequireNilErr(testCtrl, graphDatabase.AssertSchema(context.Background(), schema))
 
 	return graphDatabase

Diff for: packages/go/dawgs/dawgs.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"errors"
 
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/specterops/bloodhound/dawgs/graph"
 	"github.com/specterops/bloodhound/dawgs/util/size"
 )
@@ -38,7 +39,8 @@ func Register(driverName string, constructor DriverConstructor) {
 
 type Config struct {
 	GraphQueryMemoryLimit size.Size
-	DriverCfg             any
+	ConnectionString      string
+	Pool                  *pgxpool.Pool
 }
 
 func Open(ctx context.Context, driverName string, config Config) (graph.Database, error) {

Diff for: packages/go/dawgs/drivers/neo4j/neo4j.go

@@ -35,9 +35,7 @@ const (
 )
 
 func newNeo4jDB(_ context.Context, cfg dawgs.Config) (graph.Database, error) {
-	if connectionURLStr, typeOK := cfg.DriverCfg.(string); !typeOK {
-		return nil, fmt.Errorf("expected string for configuration type but got %T", cfg.DriverCfg)
-	} else if connectionURL, err := url.Parse(connectionURLStr); err != nil {
+	if connectionURL, err := url.Parse(cfg.ConnectionString); err != nil {
 		return nil, err
 	} else if connectionURL.Scheme != DriverName {
 		return nil, fmt.Errorf("expected connection URL scheme %s for Neo4J but got %s", DriverName, connectionURL.Scheme)

Diff for: packages/go/dawgs/drivers/neo4j/result_internal_test.go

@@ -50,28 +50,28 @@ func Test_mapValue(t *testing.T) {
 	mapTestCase[uint32, uint32](t, 0, 0)
 	mapTestCase[uint64, uint64](t, 0, 0)
 
-	mapTestCase[int, int](t, 0, 0)
+	mapTestCase(t, 0, 0) // Inferred int
 	mapTestCase[int8, int8](t, 0, 0)
 	mapTestCase[int16, int16](t, 0, 0)
 	mapTestCase[int32, int32](t, 0, 0)
 	mapTestCase[int64, int64](t, 0, 0)
 	mapTestCase[int64, graph.ID](t, 0, 0)
 
 	mapTestCase[float32, float32](t, 1.5, 1.5)
-	mapTestCase[float64, float64](t, 1.5, 1.5)
+	mapTestCase(t, 1.5, 1.5) // Inferred float64
 
-	mapTestCase[bool, bool](t, true, true)
-	mapTestCase[string, string](t, "test", "test")
+	mapTestCase(t, true, true)
+	mapTestCase(t, "test", "test")
 
-	mapTestCase[time.Time, time.Time](t, utcNow, utcNow)
-	mapTestCase[string, time.Time](t, utcNow.Format(time.RFC3339Nano), utcNow)
-	mapTestCase[int64, time.Time](t, utcNow.Unix(), time.Unix(utcNow.Unix(), 0))
-	mapTestCase[dbtype.Time, time.Time](t, dbtype.Time(utcNow), utcNow)
-	mapTestCase[dbtype.LocalTime, time.Time](t, dbtype.LocalTime(utcNow), utcNow)
-	mapTestCase[dbtype.Date, time.Time](t, dbtype.Date(utcNow), utcNow)
-	mapTestCase[dbtype.LocalDateTime, time.Time](t, dbtype.LocalDateTime(utcNow), utcNow)
+	mapTestCase(t, utcNow, utcNow)
+	mapTestCase(t, utcNow.Format(time.RFC3339Nano), utcNow)
+	mapTestCase(t, utcNow.Unix(), time.Unix(utcNow.Unix(), 0))
+	mapTestCase(t, dbtype.Time(utcNow), utcNow)
+	mapTestCase(t, dbtype.LocalTime(utcNow), utcNow)
+	mapTestCase(t, dbtype.Date(utcNow), utcNow)
+	mapTestCase(t, dbtype.LocalDateTime(utcNow), utcNow)
 
-	mapTestCase[[]any, []string](t, anyStringSlice, stringSlice)
-	mapTestCase[[]any, []graph.Kind](t, anyStringSlice, kindSlice)
-	mapTestCase[[]any, graph.Kinds](t, anyStringSlice, kinds)
+	mapTestCase(t, anyStringSlice, stringSlice)
+	mapTestCase(t, anyStringSlice, kindSlice)
+	mapTestCase(t, anyStringSlice, kinds)
 }