Skip to content

Commit 09aa07c

Browse files
mamundsen-spectermamundsen-specter
committed
update OpenDatabase and InitializeService to use DatabaseConfiguration, fixup tests
1 parent 5b9f5e5 commit 09aa07c

File tree

12 files changed

+73
-27
lines changed

12 files changed

+73
-27
lines changed

‎cmd/api/src/daemons/changelog/changelog_e2e/main.go‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import (
2727
"time"
2828

2929
"github.com/specterops/bloodhound/cmd/api/src/auth"
30+
"github.com/specterops/bloodhound/cmd/api/src/config"
3031
"github.com/specterops/bloodhound/cmd/api/src/daemons/changelog"
3132
"github.com/specterops/bloodhound/cmd/api/src/database"
3233
"github.com/specterops/bloodhound/packages/go/bhlog/attr"
@@ -75,7 +76,14 @@ func newHarness() *Harness {
7576
os.Exit(1)
7677
}
7778

78-
pool, err := pg.NewPool(connStr)
79+
cfg, err := config.NewDefaultConfiguration()
80+
if err != nil {
81+
slog.Error("Error creating new default configuration")
82+
os.Exit(1)
83+
}
84+
cfg.Database.Connection = connStr
85+
86+
pool, err := pg.NewPool(cfg.Database)
7987
if err != nil {
8088
slog.Error("Failed to connect", attr.Error(err))
8189
os.Exit(1)
@@ -87,7 +95,7 @@ func newHarness() *Harness {
8795
os.Exit(1)
8896
}
8997

90-
gormDB, err := database.OpenDatabase(connStr)
98+
gormDB, err := database.OpenDatabase(cfg.Database)
9199
if err != nil {
92100
slog.Error("Failed to open", attr.Error(err))
93101
os.Exit(1)

‎cmd/api/src/daemons/changelog/ingestion_integration_test.go‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929
_ "github.com/jackc/pgx/v5/stdlib"
3030
"github.com/peterldowns/pgtestdb"
3131
"github.com/specterops/bloodhound/cmd/api/src/auth"
32+
"github.com/specterops/bloodhound/cmd/api/src/config"
3233
"github.com/specterops/bloodhound/cmd/api/src/database"
3334
"github.com/specterops/bloodhound/cmd/api/src/model/appcfg"
3435
"github.com/specterops/bloodhound/cmd/api/src/test/integration/utils"
@@ -95,8 +96,14 @@ func setupIntegrationTest(t *testing.T) IntegrationTestSuite {
9596
connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{})
9697
)
9798

99+
cfg, err := config.NewDefaultConfiguration()
100+
if err != nil {
101+
t.Logf("Error creating new default configuration: %v", err)
102+
}
103+
cfg.Database.Connection = connConf.URL()
104+
98105
// Create connection pool
99-
pool, err := pg.NewPool(connConf.URL())
106+
pool, err := pg.NewPool(cfg.Database)
100107
require.NoError(t, err)
101108

102109
// Open graph database
@@ -111,7 +118,7 @@ func setupIntegrationTest(t *testing.T) IntegrationTestSuite {
111118
err = graphDB.AssertSchema(ctx, schema())
112119
require.NoError(t, err)
113120

114-
gormDB, err := database.OpenDatabase(connConf.URL())
121+
gormDB, err := database.OpenDatabase(cfg.Database)
115122
require.NoError(t, err)
116123

117124
db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver())

‎cmd/api/src/daemons/datapipe/datapipe_integration_test.go‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -68,11 +68,17 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes
6868
workDir = t.TempDir()
6969
)
7070

71+
cfg, err := config.NewDefaultConfiguration()
72+
if err != nil {
73+
t.Logf("Error creating new default configuration: %v", err)
74+
}
75+
cfg.Database.Connection = connConf.URL()
76+
7177
//#region Setup for dbs
72-
pool, err := pg.NewPool(connConf.URL())
78+
pool, err := pg.NewPool(cfg.Database)
7379
require.NoError(t, err)
7480

75-
gormDB, err := database.OpenDatabase(connConf.URL())
81+
gormDB, err := database.OpenDatabase(cfg.Database)
7682
require.NoError(t, err)
7783

7884
db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver())
@@ -107,9 +113,7 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes
107113
err = os.Mkdir(path.Join(workDir, "tmp"), 0755)
108114
require.NoError(t, err)
109115

110-
cfg := config.Configuration{
111-
WorkDir: workDir,
112-
}
116+
cfg.WorkDir = workDir
113117

114118
cl := changelog.NewChangelog(graphDB, db, changelog.DefaultOptions())
115119

‎cmd/api/src/database/database_integration_test.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import (
2828

2929
"github.com/peterldowns/pgtestdb"
3030
"github.com/specterops/bloodhound/cmd/api/src/auth"
31+
"github.com/specterops/bloodhound/cmd/api/src/config"
3132
"github.com/specterops/bloodhound/cmd/api/src/database"
3233
"github.com/specterops/bloodhound/cmd/api/src/test/integration/utils"
3334
"github.com/stretchr/testify/require"
@@ -54,9 +55,15 @@ func setupIntegrationTestSuite(t *testing.T) IntegrationTestSuite {
5455
connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{})
5556
)
5657

58+
cfg, err := config.NewDefaultConfiguration()
59+
if err != nil {
60+
t.Logf("Error creating new default configuration: %v", err)
61+
}
62+
cfg.Database.Connection = connConf.URL()
63+
5764
// #region Setup for dbs
5865

59-
gormDB, err := database.OpenDatabase(connConf.URL())
66+
gormDB, err := database.OpenDatabase(cfg.Database)
6067
require.NoError(t, err)
6168

6269
db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver())

‎cmd/api/src/database/db.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,14 @@ import (
2929
"github.com/gofrs/uuid"
3030
"github.com/jackc/pgx/v5/stdlib"
3131
"github.com/specterops/bloodhound/cmd/api/src/auth"
32-
"github.com/specterops/bloodhound/cmd/api/src/config"
3332
"github.com/specterops/bloodhound/cmd/api/src/database/migration"
3433
"github.com/specterops/bloodhound/cmd/api/src/model"
3534
"github.com/specterops/bloodhound/cmd/api/src/model/appcfg"
3635
"github.com/specterops/bloodhound/cmd/api/src/services/agi"
3736
"github.com/specterops/bloodhound/cmd/api/src/services/dataquality"
3837
"github.com/specterops/bloodhound/cmd/api/src/services/upload"
3938
"github.com/specterops/bloodhound/packages/go/bhlog/attr"
39+
"github.com/specterops/dawgs/drivers"
4040
"github.com/specterops/dawgs/drivers/pg"
4141
"gorm.io/driver/postgres"
4242
"gorm.io/gorm"
@@ -241,15 +241,15 @@ func (s *BloodhoundDB) Transaction(ctx context.Context, fn func(tx *BloodhoundDB
241241
}, opts...)
242242
}
243243

244-
func OpenDatabase(cfg config.Configuration) (*gorm.DB, error) {
244+
func OpenDatabase(cfg drivers.DatabaseConfiguration) (*gorm.DB, error) {
245245
gormConfig := &gorm.Config{
246246
Logger: &GormLogAdapter{
247247
SlowQueryErrorThreshold: time.Second * 30,
248248
SlowQueryWarnThreshold: time.Second * 10,
249249
},
250250
}
251251

252-
pool, err := pg.NewPool(cfg.Database)
252+
pool, err := pg.NewPool(cfg)
253253
if err != nil {
254254
return nil, err
255255
}

‎cmd/api/src/queries/queries_integration_test.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import (
2727
"testing"
2828

2929
"github.com/peterldowns/pgtestdb"
30+
"github.com/specterops/bloodhound/cmd/api/src/config"
3031
"github.com/specterops/bloodhound/cmd/api/src/migrations"
3132
"github.com/specterops/bloodhound/cmd/api/src/test/integration/utils"
3233
schema "github.com/specterops/bloodhound/packages/go/graphschema"
@@ -58,6 +59,12 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes
5859
defaultGraph = schema.DefaultGraph()
5960
)
6061

62+
cfg, err := config.NewDefaultConfiguration()
63+
if err != nil {
64+
t.Logf("Error creating new default configuration: %v", err)
65+
}
66+
cfg.Database.Connection = connConf.URL()
67+
6168
defaultGraph.Nodes.Add(graph.StringKind("Person"))
6269

6370
openGraphSchema := graph.Schema{
@@ -68,7 +75,7 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes
6875
}
6976

7077
//#region Setup for dbs
71-
pool, err := pg.NewPool(connConf.URL())
78+
pool, err := pg.NewPool(cfg.Database)
7279
require.NoError(t, err)
7380

7481
graphDB, err := dawgs.Open(ctx, pg.DriverName, dawgs.Config{

‎cmd/api/src/services/entrypoint.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ import (
4848

4949
// ConnectPostgres initializes a connection to PG, and returns errors if any
5050
func ConnectPostgres(cfg config.Configuration) (*database.BloodhoundDB, error) {
51-
if db, err := database.OpenDatabase(cfg); err != nil {
51+
if db, err := database.OpenDatabase(cfg.Database); err != nil {
5252
return nil, fmt.Errorf("error while attempting to create database connection: %w", err)
5353
} else {
5454
return database.NewBloodhoundDB(db, auth.NewIdentityResolver()), nil

‎cmd/api/src/services/graphify/graphify_integration_test.go‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -63,11 +63,17 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes
6363
workDir = t.TempDir()
6464
)
6565

66+
cfg, err := config.NewDefaultConfiguration()
67+
if err != nil {
68+
t.Logf("Error creating new default configuration: %v", err)
69+
}
70+
cfg.Database.Connection = connConf.URL()
71+
6672
//#region Setup for dbs
67-
pool, err := pg.NewPool(connConf.URL())
73+
pool, err := pg.NewPool(cfg.Database)
6874
require.NoError(t, err)
6975

70-
gormDB, err := database.OpenDatabase(connConf.URL())
76+
gormDB, err := database.OpenDatabase(cfg.Database)
7177
require.NoError(t, err)
7278

7379
db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver())
@@ -102,9 +108,7 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes
102108
err = os.Mkdir(path.Join(workDir, "tmp"), 0755)
103109
require.NoError(t, err)
104110

105-
cfg := config.Configuration{
106-
WorkDir: workDir,
107-
}
111+
cfg.WorkDir = workDir
108112

109113
return IntegrationTestSuite{
110114
Context: ctx,

‎cmd/api/src/test/integration/database.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,13 @@ func setupPGTestDB(t *testing.T, cfg config.Configuration) (*gorm.DB, error) {
6363
connConf = pgtestdb.Custom(t, GetPostgresConfig(cfg), pgtestdb.NoopMigrator{})
6464
)
6565

66-
return database.OpenDatabase(connConf.URL())
66+
cfg, err := config.NewDefaultConfiguration()
67+
if err != nil {
68+
t.Fatalf("Error creating new default configuration: %v", err)
69+
}
70+
cfg.Database.Connection = connConf.URL()
71+
72+
return database.OpenDatabase(cfg.Database)
6773
}
6874

6975
// GetPostgresConfig reads key/value pairs from the default integration

‎cmd/api/src/test/integration/dawgs.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,8 @@ func OpenGraphDB(t *testing.T, schema graph.Schema) graph.Database {
5656
switch cfg.GraphDriver {
5757
case pg.DriverName:
5858
connConf := pgtestdb.Custom(t, GetPostgresConfig(cfg), pgtestdb.NoopMigrator{})
59-
pool, err := pg.NewPool(connConf.URL())
59+
cfg.Database.Connection = connConf.URL()
60+
pool, err := pg.NewPool(cfg.Database)
6061
test.RequireNilErrf(t, err, "Failed to create new pgx pool: %v", err)
6162
graphDatabase, err = dawgs.Open(context.Background(), cfg.GraphDriver, dawgs.Config{
6263
ConnectionString: connConf.URL(),

0 commit comments

Comments
 (0)