feat: update db pooling to handle iam rds auth, move DatabaseConfigur…#2582
Open
mamundsen-specter wants to merge 9 commits intomainfrom
Open
feat: update db pooling to handle iam rds auth, move DatabaseConfigur…#2582mamundsen-specter wants to merge 9 commits intomainfrom
mamundsen-specter wants to merge 9 commits intomainfrom
Conversation
…ation to DAWGS - BI-1306
Contributor
|
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:
📝 WalkthroughWalkthroughReplaces direct PostgreSQL connection-string usage with a shared Changes
Sequence Diagram(s)(omitted) Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 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]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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/test/integration/dawgs.go (1)
59-63:⚠️ Potential issue | 🔴 CriticalPool may connect to a different database than DAWGS in this test path.
pgtestdb.Custom(...)returns a per-test DB URL, but the pool is created fromcfg.Databasebeforecfg.Database.Connectionis updated. This can causePoolandConnectionStringto point at different DBs.Proposed fix
case pg.DriverName: connConf := pgtestdb.Custom(t, GetPostgresConfig(cfg), pgtestdb.NoopMigrator{}) + cfg.Database.Connection = connConf.URL() pool, err := pg.NewPool(cfg.Database) test.RequireNilErrf(t, err, "Failed to create new pgx pool: %v", err) graphDatabase, err = dawgs.Open(context.Background(), cfg.GraphDriver, dawgs.Config{ - ConnectionString: connConf.URL(), + ConnectionString: cfg.Database.Connection, Pool: pool, })🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmd/api/src/test/integration/dawgs.go` around lines 59 - 63, The test creates the pg pool from cfg.Database before cfg.Database.Connection is updated to the per-test URL, causing pg.NewPool and dawgs.Open to point at different DBs; fix by ensuring the pool is created using the same per-test connection string (i.e., update cfg.Database.Connection or obtain the test DB URL first) before calling pg.NewPool so that the Pool passed into dawgs.Open and the ConnectionString in dawgs.Config are identical (update the sequence around pg.NewPool, cfg.Database.Connection, and dawgs.Open to use the same URL).
🧹 Nitpick comments (2)
cmd/api/src/config/config.go (1)
34-34: Consider using lowercase alias for import.The import alias
Dawgsuses PascalCase, but Go convention prefers lowercase for import aliases (e.g.,dawgsordrivers). This also makes it visually distinct from type names.♻️ Suggested fix
- Dawgs "github.com/specterops/dawgs/drivers" + "github.com/specterops/dawgs/drivers"Then update usages to
drivers.DatabaseConfiguration.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmd/api/src/config/config.go` at line 34, Change the import alias from `Dawgs` to a lowercase alias (e.g., `drivers` or `dawgs`) and update all usages accordingly; for example, replace references like `Dawgs.DatabaseConfiguration` with `drivers.DatabaseConfiguration` (or `dawgs.DatabaseConfiguration`) and adjust any other occurrences of the `Dawgs` symbol in this package to the new lowercase alias.cmd/api/src/services/graphify/test/test.go (1)
104-116: LGTM - Test setup updated for configuration-driven pattern.The test correctly creates a minimal configuration and populates
cfg.Database.Connectionbefore passing topg.NewPool().Minor note: Line 114 could use
cfg.Database.Connectioninstead ofconnConf.URL()for consistency, though functionally identical.♻️ Optional consistency improvement
graphDB, err := dawgs.Open(ctx, pg.DriverName, dawgs.Config{ GraphQueryMemoryLimit: 1024 * 1024 * 1024 * 2, - ConnectionString: connConf.URL(), + ConnectionString: cfg.Database.Connection, Pool: pool, })🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmd/api/src/services/graphify/test/test.go` around lines 104 - 116, Replace the direct use of connConf.URL() in the dawgs.Open call with the already-populated cfg.Database.Connection for consistency: update the dawgs.Open invocation (the ConnectionString field in dawgs.Config) to use cfg.Database.Connection so the test consistently uses the same configuration source as pg.NewPool and avoids mixing sources.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cmd/api/src/daemons/changelog/changelog_e2e/main.go`:
- Around line 82-85: The error path in the startup uses
config.NewDefaultConfiguration() and currently calls slog.Error("Error creating
new default configuration") without the actual err; update the failure handling
for config.NewDefaultConfiguration() to include the err value in the log (e.g.,
pass err as a field or append its message) so the slog.Error call records the
underlying error before calling os.Exit(1).
---
Outside diff comments:
In `@cmd/api/src/test/integration/dawgs.go`:
- Around line 59-63: The test creates the pg pool from cfg.Database before
cfg.Database.Connection is updated to the per-test URL, causing pg.NewPool and
dawgs.Open to point at different DBs; fix by ensuring the pool is created using
the same per-test connection string (i.e., update cfg.Database.Connection or
obtain the test DB URL first) before calling pg.NewPool so that the Pool passed
into dawgs.Open and the ConnectionString in dawgs.Config are identical (update
the sequence around pg.NewPool, cfg.Database.Connection, and dawgs.Open to use
the same URL).
---
Nitpick comments:
In `@cmd/api/src/config/config.go`:
- Line 34: Change the import alias from `Dawgs` to a lowercase alias (e.g.,
`drivers` or `dawgs`) and update all usages accordingly; for example, replace
references like `Dawgs.DatabaseConfiguration` with
`drivers.DatabaseConfiguration` (or `dawgs.DatabaseConfiguration`) and adjust
any other occurrences of the `Dawgs` symbol in this package to the new lowercase
alias.
In `@cmd/api/src/services/graphify/test/test.go`:
- Around line 104-116: Replace the direct use of connConf.URL() in the
dawgs.Open call with the already-populated cfg.Database.Connection for
consistency: update the dawgs.Open invocation (the ConnectionString field in
dawgs.Config) to use cfg.Database.Connection so the test consistently uses the
same configuration source as pg.NewPool and avoids mixing sources.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository YAML (base), Organization UI (inherited)
Review profile: CHILL
Plan: Pro
Run ID: 80bd70f6-7c3a-407d-a490-57a3f7b5b320
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (21)
cmd/api/src/api/tools/pg.gocmd/api/src/bootstrap/util.gocmd/api/src/cmd/dawgs-harness/main.gocmd/api/src/config/config.gocmd/api/src/config/default.gocmd/api/src/daemons/changelog/changelog_e2e/main.gocmd/api/src/daemons/changelog/ingestion_integration_test.gocmd/api/src/daemons/datapipe/datapipe_integration_test.gocmd/api/src/database/database_integration_test.gocmd/api/src/database/db.gocmd/api/src/database/migration/extensions_integration_test.gocmd/api/src/queries/queries_integration_test.gocmd/api/src/services/entrypoint.gocmd/api/src/services/graphify/graphify_integration_test.gocmd/api/src/services/graphify/test/test.gocmd/api/src/test/integration/database.gocmd/api/src/test/integration/dawgs.gocmd/api/src/test/lab/fixtures/postgres.gogo.modpackages/go/analysis/analysis_integration_test.gopackages/go/graphify/graph/graph.go
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
cmd/api/src/database/database_integration_test.go (1)
53-56: Consider hoisting new initializations into the function’svarblock.The new declarations on Line 58 and Line 65 can be grouped with existing top-of-function declarations for consistency with repository style.
Suggested refactor
var ( - ctx = context.Background() - connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{}) + ctx = context.Background() + connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{}) + cfg config.Configuration + gormDB *gorm.DB + err error ) -cfg, err := config.NewDefaultConfiguration() +cfg, err = config.NewDefaultConfiguration() if err != nil { - t.Errorf("Failed to create default configuration") + t.Errorf("Failed to create default configuration") } cfg.Database.Connection = connConf.URL() -gormDB, err := database.OpenDatabase(cfg.Database) +gormDB, err = database.OpenDatabase(cfg.Database) require.NoError(t, err)As per coding guidelines: "Group variable initializations in a
var ( ... )block and hoist them to the top of the function when possible".Also applies to: 58-65
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmd/api/src/database/database_integration_test.go` around lines 53 - 56, Hoist the two new local declarations currently defined later in the function (the ones on Lines 58 and 65) into the existing top-of-function var block alongside ctx and connConf so all test locals are grouped; update the var ( ... ) block that contains ctx and connConf to include those new variable names and their initializations, and remove the later standalone declarations so only the single grouped var block declares these locals.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cmd/api/src/cmd/dawgs-harness/main.go`:
- Around line 146-149: The error path after calling
config.NewDefaultConfiguration() currently prints a generic message; update the
failure branch in main (where cfg, err := config.NewDefaultConfiguration() is
checked) to include the underlying err details in the output (e.g., format the
error into the log message or write to stderr) so the real cause is visible
before calling os.Exit(1).
In `@cmd/api/src/daemons/datapipe/datapipe_integration_test.go`:
- Around line 71-74: The test helper is using t.Errorf which allows execution to
continue with an invalid cfg; change the failure to stop immediately by
replacing the t.Errorf call after cfg, err := config.NewDefaultConfiguration()
with a fatal assertion (e.g., t.Fatalf or require.NoError from testify) so
setupIntegrationTestSuite exits on configuration errors; ensure you reference
cfg and config.NewDefaultConfiguration in the updated check and remove any
subsequent code that assumes cfg is valid when err != nil.
In `@cmd/api/src/database/database_integration_test.go`:
- Around line 58-61: The test currently logs an error but continues when
config.NewDefaultConfiguration() fails, leaving cfg invalid; update the test to
fail fast by replacing the t.Errorf call with t.Fatalf (or use a test helper
like require.NoError) so the test stops immediately on error when cfg is
nil/invalid after calling config.NewDefaultConfiguration().
---
Nitpick comments:
In `@cmd/api/src/database/database_integration_test.go`:
- Around line 53-56: Hoist the two new local declarations currently defined
later in the function (the ones on Lines 58 and 65) into the existing
top-of-function var block alongside ctx and connConf so all test locals are
grouped; update the var ( ... ) block that contains ctx and connConf to include
those new variable names and their initializations, and remove the later
standalone declarations so only the single grouped var block declares these
locals.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository YAML (base), Organization UI (inherited)
Review profile: CHILL
Plan: Pro
Run ID: d4b157e2-c37b-4f7d-9dad-e227bdbeb98c
📒 Files selected for processing (11)
cmd/api/src/cmd/dawgs-harness/main.gocmd/api/src/daemons/changelog/changelog_e2e/main.gocmd/api/src/daemons/changelog/ingestion_integration_test.gocmd/api/src/daemons/datapipe/datapipe_integration_test.gocmd/api/src/database/database_integration_test.gocmd/api/src/database/migration/extensions_integration_test.gocmd/api/src/queries/queries_integration_test.gocmd/api/src/services/graphify/graphify_integration_test.gocmd/api/src/services/graphify/test/test.gopackages/go/analysis/analysis_integration_test.gopackages/go/graphify/graph/graph.go
✅ Files skipped from review due to trivial changes (1)
- cmd/api/src/database/migration/extensions_integration_test.go
🚧 Files skipped from review as they are similar to previous changes (7)
- cmd/api/src/services/graphify/test/test.go
- cmd/api/src/daemons/changelog/changelog_e2e/main.go
- cmd/api/src/services/graphify/graphify_integration_test.go
- packages/go/analysis/analysis_integration_test.go
- cmd/api/src/daemons/changelog/ingestion_integration_test.go
- packages/go/graphify/graph/graph.go
- cmd/api/src/queries/queries_integration_test.go
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cmd/api/src/database/migration/extensions_integration_test.go (1)
56-60: Hoist configuration variable declaration into the existingvarblock.Line [56] introduces
cfgmid-function; this can be grouped with the existing declarations at the top for consistency with repo conventions.♻️ Suggested refactor
func setupIntegrationTestSuite(t *testing.T) IntegrationTestSuite { t.Helper() var ( - ctx = context.Background() - connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{}) - gormDB *gorm.DB - db *database.BloodhoundDB - err error + ctx = context.Background() + connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{}) + testConfiguration config.Configuration + gormDB *gorm.DB + db *database.BloodhoundDB + err error ) - cfg, err := config.NewDefaultConfiguration() + testConfiguration, err = config.NewDefaultConfiguration() require.NoError(t, err) - cfg.Database.Connection = connConf.URL() + testConfiguration.Database.Connection = connConf.URL() @@ - gormDB, err = database.OpenDatabase(cfg.Database) + gormDB, err = database.OpenDatabase(testConfiguration.Database)As per coding guidelines "Group variable initializations in a
var ( ... )block and hoist them to the top of the function when possible" and "Prefer rich variable names, for example:databaseInterfaceinstead ofdiordbi".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmd/api/src/database/migration/extensions_integration_test.go` around lines 56 - 60, The cfg variable is declared mid-function; hoist its initialization into the existing var block at the top of the test to follow the repo convention. Replace the inline cfg, err := config.NewDefaultConfiguration() and subsequent cfg.Database.Connection = connConf.URL() with cfg and err variables declared inside the existing var ( ... ) block, then assign cfg using config.NewDefaultConfiguration() and set cfg.Database.Connection = connConf.URL() before use (keep the function names config.NewDefaultConfiguration and connConf.URL to locate the change).packages/go/graphify/graph/graph.go (1)
207-212: Use a richer variable name and hoist declarations inRun.Line [207] introduces
dbcfgvia short declaration; prefer a descriptive name and top-level grouped declarations for readability and consistency.♻️ Suggested refactor
func (s *Command) Run() error { - ctx, cancel := context.WithCancel(context.Background()) + var ( + databaseConfiguration config.Configuration + err error + ) + + ctx, cancel := context.WithCancel(context.Background()) @@ - dbcfg, err := config.NewDefaultConfiguration() + databaseConfiguration, err = config.NewDefaultConfiguration() if err != nil { return fmt.Errorf("failed to create default configuration: %w", err) } - dbcfg.Database.Connection = s.env[environment.PostgresConnectionVarName] + databaseConfiguration.Database.Connection = s.env[environment.PostgresConnectionVarName] - if graphDB, err := initializeGraphDatabase(ctx, dbcfg); err != nil { + if graphDB, err := initializeGraphDatabase(ctx, databaseConfiguration); err != nil { return fmt.Errorf("error connecting to graphDB: %w", err) - } else if err := s.service.InitializeService(ctx, dbcfg, graphDB); err != nil { + } else if err := s.service.InitializeService(ctx, databaseConfiguration, graphDB); err != nil { return fmt.Errorf("error connecting to database: %w", err)As per coding guidelines "Group variable initializations in a
var ( ... )block and hoist them to the top of the function when possible" and "Prefer rich variable names, for example:databaseInterfaceinstead ofdiordbi".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/go/graphify/graph/graph.go` around lines 207 - 212, In Run, replace the short declaration dbcfg with a clearer name (e.g., databaseConfig or defaultDatabaseConfig) and hoist its declaration into a grouped var block at the top of the function along with err and any other locals; then call config.NewDefaultConfiguration() with an assignment to the hoisted variable (databaseConfig, err = config.NewDefaultConfiguration()), check err as before, and continue using databaseConfig.Database.Connection = s.env[environment.PostgresConnectionVarName] (update all references from dbcfg to the new name). Ensure you do not re-declare variables with := inside Run so the grouped declarations are used consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@cmd/api/src/database/migration/extensions_integration_test.go`:
- Around line 56-60: The cfg variable is declared mid-function; hoist its
initialization into the existing var block at the top of the test to follow the
repo convention. Replace the inline cfg, err := config.NewDefaultConfiguration()
and subsequent cfg.Database.Connection = connConf.URL() with cfg and err
variables declared inside the existing var ( ... ) block, then assign cfg using
config.NewDefaultConfiguration() and set cfg.Database.Connection =
connConf.URL() before use (keep the function names
config.NewDefaultConfiguration and connConf.URL to locate the change).
In `@packages/go/graphify/graph/graph.go`:
- Around line 207-212: In Run, replace the short declaration dbcfg with a
clearer name (e.g., databaseConfig or defaultDatabaseConfig) and hoist its
declaration into a grouped var block at the top of the function along with err
and any other locals; then call config.NewDefaultConfiguration() with an
assignment to the hoisted variable (databaseConfig, err =
config.NewDefaultConfiguration()), check err as before, and continue using
databaseConfig.Database.Connection =
s.env[environment.PostgresConnectionVarName] (update all references from dbcfg
to the new name). Ensure you do not re-declare variables with := inside Run so
the grouped declarations are used consistently.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository YAML (base), Organization UI (inherited)
Review profile: CHILL
Plan: Pro
Run ID: a9d19d98-bf8c-40b0-9d14-d342714f7b11
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (12)
cmd/api/src/cmd/dawgs-harness/main.gocmd/api/src/daemons/changelog/changelog_e2e/main.gocmd/api/src/daemons/changelog/ingestion_integration_test.gocmd/api/src/daemons/datapipe/datapipe_integration_test.gocmd/api/src/database/database_integration_test.gocmd/api/src/database/migration/extensions_integration_test.gocmd/api/src/queries/queries_integration_test.gocmd/api/src/services/graphify/graphify_integration_test.gocmd/api/src/services/graphify/test/test.gogo.modpackages/go/analysis/analysis_integration_test.gopackages/go/graphify/graph/graph.go
✅ Files skipped from review due to trivial changes (1)
- cmd/api/src/daemons/datapipe/datapipe_integration_test.go
🚧 Files skipped from review as they are similar to previous changes (8)
- cmd/api/src/services/graphify/graphify_integration_test.go
- cmd/api/src/queries/queries_integration_test.go
- cmd/api/src/database/database_integration_test.go
- cmd/api/src/daemons/changelog/ingestion_integration_test.go
- cmd/api/src/cmd/dawgs-harness/main.go
- cmd/api/src/daemons/changelog/changelog_e2e/main.go
- packages/go/analysis/analysis_integration_test.go
- cmd/api/src/services/graphify/test/test.go
mamundsen-specter
force-pushed
the
bi-1306
branch
2 times, most recently
from
2a42978 to
ab31174
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
https://specterops.atlassian.net/browse/BI-1348
https://specterops.atlassian.net/browse/BI-1306
Bloodhound changes to enable RDS IAM auth. Coincides with SpecterOps/DAWGS#32
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.
PoC tenants have been created. These are in the process of being updated and recreated.
Chore (a change that does not modify the application functionality)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to change)
Database Migrations
Screenshots (optional):
Types of changes
Checklist:
Summary by CodeRabbit