update db pooling to handle iam rds auth - BI-1306#2356
update db pooling to handle iam rds auth - BI-1306#2356mamundsen-specter wants to merge 13 commits intomainfrom
Conversation
|
Howdy! Thank you for opening this pull request 🙇 It looks like your pull request title needs some adjustment. Details: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughDatabase and pool initialization now accept and propagate full configuration objects (Dawgs.DatabaseConfiguration / config.Configuration) instead of raw connection strings; function signatures, imports, and call sites were updated to use the new types. Changes
Sequence Diagram(s)sequenceDiagram
participant App as Application
participant Config as config.Configuration
participant PoolFactory as dawgs/pg PoolFactory
participant Stdlib as pgx stdlib
participant GORM as GORM (postgres driver)
App->>Config: obtain cfg (includes Dawgs.DatabaseConfiguration)
App->>PoolFactory: pg.NewPool(cfg.Database)
PoolFactory-->>Stdlib: provide *sql.DB (OpenDBFromPool)
Stdlib-->>GORM: provide Conn to postgres.New
GORM-->>App: return initialized *gorm.DB
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
coderabbitai
bot
added
the
api
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
cmd/api/src/cmd/dawgs-harness/main.go (1)
150-153:⚠️ Potential issue | 🔴 CriticalCritical: Missing
cfgargument in neo4jRunTestSuitecall.
RunTestSuitenow requires 4 arguments, but the neo4j call on line 152 only passes 3. This is a compile error.Proposed fix
n4jTestSuite := execSuite(neo4j.DriverName, func() tests.TestSuite { - return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName) + return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg.Database) })Note: The exact argument depends on resolving the type issue in
RunTestSuite's signature (line 48).packages/go/graphify/graph/graph.go (1)
412-423:⚠️ Potential issue | 🟠 MajorUse
cfg.Database.PostgreSQLConnectionString()instead of the rawpostgresConnectionvariable for consistency with the rest of the codebase.The
initializeGraphDatabasefunction creates a pool fromcfg.Database(which only hasMaxConcurrentSessions: 10set) while bypassing the connection string that should be retrieved fromcfg.Database.PostgreSQLConnectionString(). This deviates from the pattern incmd/api/src/api/tools/pg.goandcmd/api/src/bootstrap/util.go, wherepg.NewPool(cfg.Database)is followed bycfg.Database.PostgreSQLConnectionString()for the connection string.Either populate
cfg.Databasewith connection details before callinginitializeGraphDatabase, or refactor this function to use the standard pattern of callingpg.NewPool()with a fully-configuredcfg.Database.cmd/api/src/database/db.go (1)
243-263:⚠️ Potential issue | 🔴 CriticalCritical resource leak:
pgxpool.Poolis never closed. The pool created bypg.NewPool()at line 251 is wrapped withstdlib.OpenDBFromPool()but the underlying pool reference is discarded. Per pgx documentation,stdlib.OpenDBFromPooldoes not take ownership of pool closure—callingsql.DB.Close()only stops the database/sql wrapper, leaving the pool's connections permanently open. WhenBloodhoundDB.Close()callssqlDBRef.Close(), the pgxpool remains active and unreleased, preventing graceful shutdown and exhausting connection limits.Store the
*pgxpool.Poolreference in theBloodhoundDBstruct and callpool.Close()during shutdown in addition tosqlDBRef.Close().
🤖 Fix all issues with AI agents
In `@cmd/api/src/cmd/dawgs-harness/main.go`:
- Around line 145-148: The snippet uses invalid Go syntax and undefined
identifiers; replace the malformed line with a proper assignment and error
check: call config.NewDefaultConfiguration() as cfg, err :=
config.NewDefaultConfiguration(); then test if err != nil and handle it inside
main (e.g., log.Fatalf or fmt.Fprintf(os.Stderr, ...) followed by os.Exit(1))
instead of returning a non-existent configuration variable; ensure only cfg is
used thereafter.
- Line 48: The function RunTestSuite currently references the removed type
config.DatabaseConfiguration; change its parameter to use the new type from the
dawgs drivers package (e.g. import github.com/specterops/dawgs/drivers as dawgs
and update the signature to accept dawgs.DatabaseConfiguration), and update any
other unqualified usages (e.g. DatabaseConfiguration{} in
cmd/api/src/config/default.go) to the fully-qualified
dawgs.DatabaseConfiguration or switch those callers to pass config.Configuration
and read cfg.Database internally; ensure you add the dawgs import alias so
symbols resolve.
In `@packages/go/graphify/graph/graph.go`:
- Around line 206-214: The code calls config.NewDefaultConfiguration() and on
failure logs a message without the error and calls os.Exit(1), which prevents
deferred TeardownService from running; update the error handling in Run() so
that the slog/Error log includes the actual error value (err) and replace
os.Exit(1) with returning the error (e.g. return fmt.Errorf("creating default
configuration: %w", err)) so deferred s.service.TeardownService (and other
defers) can execute; modify the block around config.NewDefaultConfiguration()
accordingly.
🧹 Nitpick comments (1)
cmd/api/src/config/config.go (1)
34-34: Unconventional uppercase import aliasDawgs.Go convention is to use lowercase import aliases (e.g.,
dawgsdrvordrivers). The uppercaseDawgsreads like an exported identifier rather than a package alias.Suggested change
- Dawgs "github.com/specterops/dawgs/drivers" + dawgsdrv "github.com/specterops/dawgs/drivers"Then update usages accordingly (e.g.,
dawgsdrv.DatabaseConfiguration).
| cfg := config.NewDefaultConfiguration(); err != nil { | ||
| return configuration, fmt.Errorf("failed to create default configuration: %w", err) | ||
| } | ||
|
|
There was a problem hiding this comment.
Critical: This code will not compile — broken syntax and undefined identifiers.
Lines 145–147 contain invalid Go syntax:
config.NewDefaultConfiguration()returns(Configuration, error)but only one LHS variable is assigned.- The
;pseudo-if pattern is malformed — it's not within anifblock. return configuration, ...is invalid inmain()(no return values, andconfigurationis undefined).
Proposed fix
- cfg := config.NewDefaultConfiguration(); err != nil {
- return configuration, fmt.Errorf("failed to create default configuration: %w", err)
- }
+ cfg, err := config.NewDefaultConfiguration()
+ if err != nil {
+ fatalf("failed to create default configuration: %v", err)
+ }🤖 Prompt for AI Agents
In `@cmd/api/src/cmd/dawgs-harness/main.go` around lines 145 - 148, The snippet
uses invalid Go syntax and undefined identifiers; replace the malformed line
with a proper assignment and error check: call config.NewDefaultConfiguration()
as cfg, err := config.NewDefaultConfiguration(); then test if err != nil and
handle it inside main (e.g., log.Fatalf or fmt.Fprintf(os.Stderr, ...) followed
by os.Exit(1)) instead of returning a non-existent configuration variable;
ensure only cfg is used thereafter.
| dbcfg, err := config.NewDefaultConfiguration() | ||
| if err != nil { | ||
| slog.Error("Error creating new default configuration") | ||
| os.Exit(1) | ||
| } | ||
|
|
||
| if graphDB, err := initializeGraphDatabase(ctx, s.env[environment.PostgresConnectionVarName], dbcfg); err != nil { | ||
| return fmt.Errorf("error connecting to graphDB: %w", err) | ||
| } else if err := s.service.InitializeService(ctx, s.env[environment.PostgresConnectionVarName], graphDB); err != nil { | ||
| } else if err := s.service.InitializeService(ctx, dbcfg, graphDB); err != nil { |
There was a problem hiding this comment.
Missing error value in log message and abrupt os.Exit.
Line 208 logs "Error creating new default configuration" but does not include the actual error. Also, os.Exit(1) prevents the deferred TeardownService (line 202) from running. Consider returning the error from Run() instead.
Proposed fix
dbcfg, err := config.NewDefaultConfiguration()
if err != nil {
- slog.Error("Error creating new default configuration")
- os.Exit(1)
+ return fmt.Errorf("error creating default configuration: %w", err)
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| dbcfg, err := config.NewDefaultConfiguration() | |
| if err != nil { | |
| slog.Error("Error creating new default configuration") | |
| os.Exit(1) | |
| } | |
| if graphDB, err := initializeGraphDatabase(ctx, s.env[environment.PostgresConnectionVarName], dbcfg); err != nil { | |
| return fmt.Errorf("error connecting to graphDB: %w", err) | |
| } else if err := s.service.InitializeService(ctx, s.env[environment.PostgresConnectionVarName], graphDB); err != nil { | |
| } else if err := s.service.InitializeService(ctx, dbcfg, graphDB); err != nil { | |
| dbcfg, err := config.NewDefaultConfiguration() | |
| if err != nil { | |
| return fmt.Errorf("error creating default configuration: %w", err) | |
| } | |
| if graphDB, err := initializeGraphDatabase(ctx, s.env[environment.PostgresConnectionVarName], dbcfg); err != nil { | |
| return fmt.Errorf("error connecting to graphDB: %w", err) | |
| } else if err := s.service.InitializeService(ctx, dbcfg, graphDB); err != nil { |
🤖 Prompt for AI Agents
In `@packages/go/graphify/graph/graph.go` around lines 206 - 214, The code calls
config.NewDefaultConfiguration() and on failure logs a message without the error
and calls os.Exit(1), which prevents deferred TeardownService from running;
update the error handling in Run() so that the slog/Error log includes the
actual error value (err) and replace os.Exit(1) with returning the error (e.g.
return fmt.Errorf("creating default configuration: %w", err)) so deferred
s.service.TeardownService (and other defers) can execute; modify the block
around config.NewDefaultConfiguration() accordingly.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cmd/api/src/cmd/dawgs-harness/main.go (1)
155-167:⚠️ Potential issue | 🟠 MajorBug: Neo4j test suite receives
cfg.Databaseinstead ofcfg.Neo4J.Line 150 sets
cfg.Neo4J.Connection = neo4jConnectionStr, but on line 157 the neo4jRunTestSuitecall passescfg.Database— meaning the neo4j connection string configuration is unused and the wrongDatabaseConfigurationis passed. The same issue occurs on line 179.Proposed fix
n4jTestSuite := execSuite(neo4j.DriverName, func() tests.TestSuite { - return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg.Database) + return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg.Neo4J) })And for the
"neo4j"case (line 179):execSuite(neo4j.DriverName, func() tests.TestSuite { - return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg.Database) + return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg.Neo4J) })
🤖 Fix all issues with AI agents
In `@cmd/api/src/cmd/dawgs-harness/main.go`:
- Around line 146-149: The call to config.NewDefaultConfiguration() can fail but
the code currently creates an error value with fmt.Errorf and then continues
with a nil/zero cfg; replace that silent discard by checking err and terminating
with a fatal log (e.g., log.Fatalf or fmt.Fatalf or panic) so the program
doesn't continue with an invalid cfg; update the error handling right after cfg,
err := config.NewDefaultConfiguration() to log a descriptive message including
the wrapped err and exit (reference symbols: NewDefaultConfiguration, cfg, err).
🧹 Nitpick comments (3)
go.mod (1)
50-50: Pseudo-version dependency on dawgs — pin to a release tag before merging to main.The dawgs dependency uses a pseudo-version (
v0.4.8-0.20260211001809-7d3b4dd80186), which ties this repo to a specific unreleased commit. This is fine for development but should be updated to a proper tagged release before merging, to ensure reproducible builds and clear version tracking.cmd/api/src/config/default.go (1)
22-22: Uppercase import aliasDawgsis unconventional.Go convention uses lowercase import aliases (e.g.,
dawgsDriversor justdrivers). The uppercaseDawgsreads like an exported identifier rather than a package alias. Consider renaming for consistency with Go idioms.Suggested rename
- Dawgs "github.com/specterops/dawgs/drivers" + dawgsDrivers "github.com/specterops/dawgs/drivers"Then update usages accordingly (e.g.,
dawgsDrivers.DatabaseConfiguration).cmd/api/src/cmd/dawgs-harness/main.go (1)
49-60:connectionStrparameter is unused insideRunTestSuitefor pg path — consider removing redundancy.The
connectionStris passed separately butcfgalready contains the connection string (set at line 151). For the pg path,pg.NewPool(cfg)usescfgdirectly, makingconnectionStrredundant for that code path. It's still used fordawgs.Openon line 64, so this is more of a design observation — the connection string is duplicated in two places which could drift.
| cfg, err := config.NewDefaultConfiguration() | ||
| if err != nil { | ||
| fmt.Errorf("failed to create default configuration: %w", err) | ||
| } |
There was a problem hiding this comment.
Critical: Error from NewDefaultConfiguration() is silently discarded.
fmt.Errorf creates an error value but doesn't do anything with it — the program continues with a zero-value cfg on failure. This should use fatalf (or similar) to halt execution.
Proposed fix
cfg, err := config.NewDefaultConfiguration()
if err != nil {
- fmt.Errorf("failed to create default configuration: %w", err)
+ fatalf("failed to create default configuration: %v", err)
}🤖 Prompt for AI Agents
In `@cmd/api/src/cmd/dawgs-harness/main.go` around lines 146 - 149, The call to
config.NewDefaultConfiguration() can fail but the code currently creates an
error value with fmt.Errorf and then continues with a nil/zero cfg; replace that
silent discard by checking err and terminating with a fatal log (e.g.,
log.Fatalf or fmt.Fatalf or panic) so the program doesn't continue with an
invalid cfg; update the error handling right after cfg, err :=
config.NewDefaultConfiguration() to log a descriptive message including the
wrapped err and exit (reference symbols: NewDefaultConfiguration, cfg, err).
mamundsen-specter
force-pushed
the
bi-1348-main
branch
from
b15c0de to
5b9f5e5
Compare
1 participant
https://specterops.atlassian.net/browse/BI-1348
https://specterops.atlassian.net/browse/BI-1306
Bloodhound changes to enable RDS IAM auth. Coincides with DAWGs changes
Updates PostgresConnectionString() to fetch an IAM auth token and encode it before adding it to the password of the return postgres connection string.
Updates NewPool() calls coinciding with DAWGs newpool change
Updates OpenDatabase() to pass pgxpool to Gorm.
Motivation and Context
Enabling IAM auth allows us to get away from basic password authentication to the database. This frees us up from a infrastructure resource perspective and can move us towards better density in the cluster.
How Has This Been Tested?
PoC tenants have been created. These are in the process of being updated and recreated.
Checklist:
Summary by CodeRabbit
New Features
Refactor
Chores