update db pooling to handle iam rds auth [DRAFT]#1804
update db pooling to handle iam rds auth [DRAFT]#1804MitchellAmundsen wants to merge 3 commits intomainfrom
Conversation
WalkthroughPostgreSQL pool creation and database opening were changed to accept configuration objects instead of raw connection strings; GORM now opens over a pgx v5 stdlib-backed pool. Graph service initialization and test harness functions were updated to propagate the configuration. PostgreSQLConnectionString now always builds an AWS RDS IAM token-based DSN. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Caller as App / CLI
participant Cfg as config (cfg)
participant PG as pg.NewPool
participant DSN as cfg.Database.PostgreSQLConnectionString()
participant PoolDB as stdlib.OpenDBFromPool
participant GORM as gorm.Open
participant Dawgs as dawgs.Open
Caller->>Cfg: NewDefaultConfiguration()
Cfg-->>Caller: cfg
Caller->>PG: NewPool(cfg.Database)
PG-->>Caller: pool
Caller->>PoolDB: OpenDBFromPool(pool)
PoolDB-->>Caller: *sql.DB
Caller->>GORM: Open(postgres.New({Conn: *sql.DB}))
GORM-->>Caller: *gorm.DB
Caller->>DSN: Build AWS RDS IAM DSN
DSN-->>Caller: postgresql://... (token)
Caller->>Dawgs: Open({Pool: pool, ConnectionString: dsn, ...})
Dawgs-->>Caller: graph DB handle
sequenceDiagram
autonumber
actor App as Graph App
participant Cfg as NewDefaultConfiguration
participant Init as initializeGraphDatabase
participant Svc as GraphService
App->>Cfg: NewDefaultConfiguration()
Cfg-->>App: cfg
App->>Init: initializeGraphDatabase(ctx, cfg, connectionStr)
Init-->>App: graphDB
App->>Svc: InitializeService(ctx, cfg, graphDB)
Svc-->>App: ok / error
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🧰 Additional context used🪛 golangci-lint (2.5.0)cmd/api/src/config/config.go[major] 37-37: could not import github.com/aws/aws-sdk-go-v2/config (cmd/api/src/config/config.go:37:2: no required module provides package github.com/aws/aws-sdk-go-v2/config; to add it: (typecheck) [major] 38-38: could not import github.com/aws/aws-sdk-go-v2/feature/rds/auth (cmd/api/src/config/config.go:38:2: no required module provides package github.com/aws/aws-sdk-go-v2/feature/rds/auth; to add it: (typecheck) 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}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cmd/api/src/services/entrypoint.go (1)
47-51: Signature update LGTM; pending migration of legacy string-based callsThe signature change to
OpenDatabase(cfg config.Configuration)incmd/api/src/services/entrypoint.gois correct, but several remaining call sites still invokedatabase.OpenDatabasewith a raw DSN string. These must be updated to pass aconfig.Configurationinstead:
- cmd/api/src/test/lab/fixtures/postgres.go:
Line 36:database.OpenDatabase(config.Database.PostgreSQLConnectionString())- cmd/api/src/test/integration/database.go:
Line 62:return database.OpenDatabase(connConf.URL())- cmd/api/src/daemons/datapipe/datapipe_integration_test.go:
Line 74:database.OpenDatabase(connConf.URL())- cmd/api/src/services/graphify/graphify_integration_test.go:
Line 68:database.OpenDatabase(connConf.URL())- cmd/api/src/database/database_integration_test.go:
Line 53:database.OpenDatabase(connConf.URL())Please migrate each of these calls to use your new configuration-based wiring (e.g. load or construct a
config.Configurationfrom the URL/DSN, then pass that intoOpenDatabase(cfg)). Once updated, tests and fixtures will align with the new pool-based setup.packages/go/graphify/graph/graph.go (1)
409-421: Prevent pool creation failure when cfg.Database.Connection is empty.Default the
cfg.Database.Connectionfrom the providedpostgresConnectionflag/env value when it’s not set in the config.-func initializeGraphDatabase(ctx context.Context, postgresConnection string, cfg config.Configuration) (graph.Database, error) { - if pool, err := pg.NewPool(cfg.Database); err != nil { +func initializeGraphDatabase(ctx context.Context, postgresConnection string, cfg config.Configuration) (graph.Database, error) { + // Fall back to provided connection string if the config doesn't specify one + if cfg.Database.Connection == "" { + cfg.Database.Connection = postgresConnection + } + + if pool, err := pg.NewPool(cfg.Database); err != nil { return nil, fmt.Errorf("error creating postgres connection: %w", err) } else if database, err := dawgs.Open(ctx, pg.DriverName, dawgs.Config{ GraphQueryMemoryLimit: size.Gibibyte, ConnectionString: postgresConnection, Pool: pool, }); err != nil { return nil, fmt.Errorf("error connecting to database: %w", err) } else { return database, nil }
🧹 Nitpick comments (4)
cmd/api/src/config/config.go (1)
36-38: Alias AWS imports to avoid name shadowing with localconfigand commonauthpackagesImporting
github.com/aws/aws-sdk-go-v2/configasconfiginside packageconfigis confusing and error-prone. Same forfeature/rds/authvs internalauthpackages used elsewhere. Prefer explicit aliases.- "github.com/aws/aws-sdk-go-v2/config" - "github.com/aws/aws-sdk-go-v2/feature/rds/auth" + awsconfig "github.com/aws/aws-sdk-go-v2/config" + rdsauth "github.com/aws/aws-sdk-go-v2/feature/rds/auth"cmd/api/src/cmd/dawgs-harness/main.go (2)
48-56: Use the provided context; avoid context.TODO in RunTestSuite.Propagate
ctxtodawgs.Openfor proper cancellation and timeouts.- if connection, err := dawgs.Open(context.TODO(), driverName, dawgs.Config{ + if connection, err := dawgs.Open(ctx, driverName, dawgs.Config{
54-56: Make pool creation resilient to missing DSN in cfg.Optional: Align the pool config with the
connectionStrflag ifcfg.Connectionis empty, preventing surprising failures when the harness is used without a fully populated config.- if driverName == pg.DriverName { - pool, err = pg.NewPool(cfg) + if driverName == pg.DriverName { + if cfg.Connection == "" { + cfg.Connection = connectionStr + } + pool, err = pg.NewPool(cfg)packages/go/graphify/graph/graph.go (1)
203-208: Log the underlying error when default config creation fails.Include the
errin the log entry to aid troubleshooting before exiting.- if err != nil { - slog.Error("Error creating new default configuration") - os.Exit(1) - } + if err != nil { + slog.Error("Error creating new default configuration", slog.Any("error", err)) + os.Exit(1) + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (7)
cmd/api/src/api/tools/pg.go(1 hunks)cmd/api/src/bootstrap/util.go(1 hunks)cmd/api/src/cmd/dawgs-harness/main.go(3 hunks)cmd/api/src/config/config.go(3 hunks)cmd/api/src/database/db.go(2 hunks)cmd/api/src/services/entrypoint.go(1 hunks)packages/go/graphify/graph/graph.go(4 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (4)
cmd/api/src/database/db.go (3)
cmd/api/src/test/integration/database.go (1)
OpenDatabase(43-53)cmd/api/src/config/config.go (1)
Configuration(152-183)cmd/api/src/database/log.go (1)
GormLogAdapter(30-33)
packages/go/graphify/graph/graph.go (4)
cmd/api/src/config/config.go (1)
Configuration(152-183)cmd/api/src/database/db.go (2)
Database(68-184)OpenDatabase(221-241)cmd/api/src/config/default.go (1)
NewDefaultConfiguration(28-87)packages/go/stbernard/environment/environment.go (1)
PostgresConnectionVarName(30-30)
cmd/api/src/services/entrypoint.go (2)
cmd/api/src/database/db.go (1)
OpenDatabase(221-241)cmd/api/src/test/integration/database.go (1)
OpenDatabase(43-53)
cmd/api/src/cmd/dawgs-harness/main.go (2)
cmd/api/src/config/config.go (1)
DatabaseConfiguration(57-64)cmd/api/src/config/default.go (1)
NewDefaultConfiguration(28-87)
🪛 GitHub Check: run-analysis
cmd/api/src/config/config.go
[failure] 37-37:
no required module provides package github.com/aws/aws-sdk-go-v2/feature/rds/auth; to add it:
[failure] 36-36:
no required module provides package github.com/aws/aws-sdk-go-v2/config; to add it:
🪛 GitHub Check: run-tests
cmd/api/src/config/config.go
[failure] 37-37:
no required module provides package github.com/aws/aws-sdk-go-v2/feature/rds/auth; to add it:
[failure] 36-36:
no required module provides package github.com/aws/aws-sdk-go-v2/config; to add it:
🪛 GitHub Actions: Run UI Tests
cmd/api/src/config/config.go
[error] 36-36: Step 'go tool stbernard test -y' failed: no required module provides package github.com/aws/aws-sdk-go-v2/config; add the module to go.mod (e.g., 'go get github.com/aws/aws-sdk-go-v2/config' and run 'go mod tidy').
🪛 GitHub Actions: Static Code Analysis
cmd/api/src/config/config.go
[error] 36-36: no required module provides package github.com/aws/aws-sdk-go-v2/config; to add it:
🪛 GitHub Actions: Run Go Tests
cmd/api/src/config/config.go
[error] 36-36: Step 'go tool stbernard test -g -i' failed: no required module provides package github.com/aws/aws-sdk-go-v2/config; to add it:
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Build BloodHound Container Image / Build and Package Container
- GitHub Check: build-ui
🔇 Additional comments (7)
cmd/api/src/bootstrap/util.go (1)
95-114: No changes needed: PostgreSQL driver ignoresConnectionStringand uses only the poolWe verified that:
dawgs.Opensimply dispatches to the registered driver constructor with the fullConfigstruct (including bothConnectionStringandPool).- The PG driver’s
init()registration indrivers/pg/pg.gounconditionally doesthereby ignoringdawgs.Register(DriverName, func(ctx context.Context, cfg dawgs.Config) (graph.Database, error) { return NewDriver(cfg.Pool), nil })cfg.ConnectionStringentirely.- No downstream code in the PG driver ever reads
Config.ConnectionString—all connections go through the preconfiguredpgxpool.Pool.Since
ConnectionStringis never used for actual database operations whenPoolis provided, there is no risk of IAM-token drift or redundant token generation. The code is correct as is. You may optionally drop passing the DSN intoConfigfor clarity, but it is not required for correctness.cmd/api/src/api/tools/pg.go (1)
451-458: Use pool-only configuration for PostgreSQL driverTo avoid embedding one-time tokens in the DSN (and the attendant risk of token drift or redundant IAM work), drop the
ConnectionStringwhen you supply aPool:
- File:
cmd/api/src/api/tools/pg.go(around lines 451–458)- Replace the config block in your
dawgs.Opencall as follows:return dawgs.Open(s.ServerCtx, pg.DriverName, dawgs.Config{ GraphQueryMemoryLimit: size.Gibibyte, - ConnectionString: s.Cfg.Database.PostgreSQLConnectionString(), Pool: pool, })Please verify that the DAWGS PostgreSQL driver’s
Openfunction tolerates an empty or nilConnectionStringwhenPoolis provided.cmd/api/src/database/db.go (2)
29-39: Imports LGTM and appropriate for pgx v5 stdlib adapterThe stdlib adapter and driver imports look correct for wiring GORM on top of a pgx pool.
221-241: Pool-backed GORM wiring LGTM; please manually verify pool implementationI wasn’t able to locate the
pg.NewPooldefinition in the codebase, so please confirm that:
pg.NewPool(cfg.Database)configures per-connection IAM token generation (e.g. using aBeforeConnecthook) so that connections renewed after ~15 minutes will automatically acquire fresh tokens.- Sensible connection recycling parameters (MaxConnLifetime, MaxConnIdleTime, etc.) are set (ideally < 15 minutes) to ensure periodic refresh of credentials.
packages/go/graphify/graph/graph.go (3)
158-164: LGTM: InitializeService now opens DB from config.The switch to
database.OpenDatabase(cfg)aligns with the new pool creation path and centralizes DB init in config.
209-212: Ensure the pool builder has a DSN.
config.NewDefaultConfiguration()leavescfg.Database.Connectionempty. BecauseinitializeGraphDatabasecallspg.NewPool(cfg.Database), pool creation can fail without a DSN. See the fix proposed in the function below that defaults the DSN frompostgresConnection.
126-129: ConfirmGraphService.InitializeServicesignature update across all implementers
- I wasn’t able to find any concrete implementations of
in this repository (only the interface declaration at packages/go/graphify/graph/graph.go:124).func (… ) InitializeService(ctx context.Context, cfg config.Configuration, db graph.Database) error- Only one call site remains in this module (
packages/go/graphify/graph/graph.go:211).- If your
GraphServiceimplementations reside in other modules or vendored code, please verify that each implementation’s method signature now takes aconfig.Configuration(not astring) and that all call sites are updated accordingly to prevent compilation breakage.
| cfg := config.NewDefaultConfiguration(); err != nil { | ||
| return configuration, fmt.Errorf("failed to create default configuration: %w", err) | ||
| } |
There was a problem hiding this comment.
Fix config initialization: invalid syntax and illegal return from main.
config.NewDefaultConfiguration()returns(config.Configuration, error), but the current code assigns to a single identifier and then appends; err != nil, which will not compile.main()cannotreturn configuration, error.
Also, without wiring the CLI -pg DSN into the config, pg.NewPool(cfg) will likely fail because cfg.Database.Connection is empty by default.
Apply this diff:
- 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)
+ }
+ // Ensure the pool builder has a DSN when running from the harness
+ cfg.Database.Connection = pgConnectionStr📝 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.
| 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) | |
| } | |
| // Ensure the pool builder has a DSN when running from the harness | |
| cfg.Database.Connection = pgConnectionStr |
| case "both": | ||
| n4jTestSuite := execSuite(neo4j.DriverName, func() tests.TestSuite { | ||
| return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName) | ||
| return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg) |
There was a problem hiding this comment.
Pass DatabaseConfiguration, not Configuration, to RunTestSuite.
RunTestSuite expects config.DatabaseConfiguration, but calls are passing config.Configuration. This is a compile-time type mismatch.
Apply this diff:
- return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg)
+ return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg.Database)- return RunTestSuite(ctx, pgConnectionStr, pg.DriverName, cfg)
+ return RunTestSuite(ctx, pgConnectionStr, pg.DriverName, cfg.Database)- return RunTestSuite(ctx, pgConnectionStr, pg.DriverName, cfg)
+ return RunTestSuite(ctx, pgConnectionStr, pg.DriverName, cfg.Database)- return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg)
+ return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg.Database)Also applies to: 161-161, 169-169, 175-175
🤖 Prompt for AI Agents
In cmd/api/src/cmd/dawgs-harness/main.go around lines 152, 161, 169 and 175, the
calls to RunTestSuite are passing cfg (config.Configuration) but RunTestSuite
expects a config.DatabaseConfiguration; change each call to pass
cfg.DatabaseConfiguration instead of cfg so the correct type is supplied (e.g.,
replace the cfg argument with cfg.DatabaseConfiguration at those call sites).
| "github.com/aws/aws-sdk-go-v2/config" | ||
| "github.com/aws/aws-sdk-go-v2/feature/rds/auth" | ||
| ) |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Fix build by adding missing AWS SDK v2 modules and tidying go.mod
CI failures show unresolved imports for AWS SDK v2. Add modules and tidy.
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Add required AWS SDK v2 modules
go get github.com/aws/aws-sdk-go-v2/config@latest
go get github.com/aws/aws-sdk-go-v2/feature/rds/auth@latest
# Tidy modules
go mod tidy
# Sanity check: list all imports referencing aws-sdk-go-v2
rg -n 'aws-sdk-go-v2' -C1
# Build the packages that reference the new imports
go build ./cmd/api/...Length of output: 24109
Action Required: Fix remaining build errors after adding AWS SDK v2 modules
Adding the two AWS SDK v2 imports and running go mod tidy correctly updates your go.mod, but the build still fails due to unrelated code issues. You’ll need to:
- Address syntax errors in the “dawgs-harness” command:
• File: cmd/api/src/cmd/dawgs-harness/main.go
• Errors include unexpectedreturn, non-declaration statements outside functions, and methods missing receivers. Likely mismatched braces or misplaced code—please review recent changes around lines 140–170. - Correct the type passed to
pg.NewPool:
• File: cmd/api/src/database/db.go
• Line 229:cfg.Databaseis aDatabaseConfigurationstruct, butpg.NewPoolexpects astring. Pass the appropriate connection string field (e.g.cfg.Database.ConnectionString) instead of the entire struct.
Once those fixes are applied, re-run:
go build ./cmd/api/...to confirm the build succeeds.
🧰 Tools
🪛 GitHub Check: run-analysis
[failure] 37-37:
no required module provides package github.com/aws/aws-sdk-go-v2/feature/rds/auth; to add it:
[failure] 36-36:
no required module provides package github.com/aws/aws-sdk-go-v2/config; to add it:
🪛 GitHub Check: run-tests
[failure] 37-37:
no required module provides package github.com/aws/aws-sdk-go-v2/feature/rds/auth; to add it:
[failure] 36-36:
no required module provides package github.com/aws/aws-sdk-go-v2/config; to add it:
🪛 GitHub Actions: Run UI Tests
[error] 36-36: Step 'go tool stbernard test -y' failed: no required module provides package github.com/aws/aws-sdk-go-v2/config; add the module to go.mod (e.g., 'go get github.com/aws/aws-sdk-go-v2/config' and run 'go mod tidy').
🪛 GitHub Actions: Static Code Analysis
[error] 36-36: no required module provides package github.com/aws/aws-sdk-go-v2/config; to add it:
🪛 GitHub Actions: Run Go Tests
[error] 36-36: Step 'go tool stbernard test -g -i' failed: no required module provides package github.com/aws/aws-sdk-go-v2/config; to add it:
🤖 Prompt for AI Agents
In cmd/api/src/cmd/dawgs-harness/main.go around lines 140–170 there are syntax
errors (unexpected return, non-declaration statements outside functions, and
methods missing receivers) caused by mismatched or misplaced braces—inspect and
correct block boundaries so all returns and method definitions are inside
functions or proper types, restore missing receivers on methods, and ensure
braces align with recent edits; in cmd/api/src/database/db.go at line 229
replace the struct passed to pg.NewPool with the actual connection string field
from cfg.Database (e.g., cfg.Database.ConnectionString) so the function receives
a string as expected, then run go build ./cmd/api/... to verify.
| slog.Info("loading default config for rds auth") | ||
| cfg, err := config.LoadDefaultConfig(context.TODO()) | ||
| if err != nil { | ||
| panic("configuration error: " + err.Error()) | ||
| } | ||
|
|
||
| return s.Connection | ||
| dbinput := s.Address + ":5432" | ||
| slog.Info("requesting auth token") | ||
| authenticationToken, err := auth.BuildAuthToken(context.TODO(), dbinput, "us-east-1", s.Username, cfg.Credentials) | ||
| if err != nil { | ||
| panic("failed to create authentication token: " + err.Error()) | ||
| } | ||
| slog.Info("auth token successfully created") | ||
| encodedToken := url.QueryEscape(authenticationToken) | ||
|
|
||
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, encodedToken, s.Address, s.Database) | ||
| } |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
IAM RDS auth should be opt-in, parameterized, and error-returning; avoid panics and hardcoded region/port; enforce SSL
Current implementation:
- Unconditionally uses IAM, ignoring
s.Connectionands.Secret. - Hardcodes
us-east-1and port 5432 for the auth token signature. - Panics on errors within a configuration helper.
- Omits
sslmode, while IAM requires TLS. - Builds token once at process start; if the pool also builds tokens, you risk mismatch or redundant work.
Recommend:
- Fallback to
s.Connectionif provided (preserves existing behavior). - Parameterize region via env (
AWS_REGION/AWS_DEFAULT_REGION) and derive port froms.Addressif present. - Return a safe DSN with
sslmode=requirewhen IAM auth succeeds; otherwise, gracefully fall back to password DSN if available. - Keep token out of logs (you already do—good).
- Consider gating IAM behind a config flag in a follow-up; for now,
s.Connectionoverride is a minimal safety valve.
Apply this diff within this function and adjust AWS import aliases as suggested above:
func (s DatabaseConfiguration) PostgreSQLConnectionString() string {
- slog.Info("loading default config for rds auth")
- cfg, err := config.LoadDefaultConfig(context.TODO())
- if err != nil {
- panic("configuration error: " + err.Error())
- }
-
- dbinput := s.Address + ":5432"
- slog.Info("requesting auth token")
- authenticationToken, err := auth.BuildAuthToken(context.TODO(), dbinput, "us-east-1", s.Username, cfg.Credentials)
- if err != nil {
- panic("failed to create authentication token: " + err.Error())
- }
- slog.Info("auth token successfully created")
- encodedToken := url.QueryEscape(authenticationToken)
-
- return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, encodedToken, s.Address, s.Database)
+ // Preserve explicit connection string if provided
+ if s.Connection != "" {
+ return s.Connection
+ }
+
+ // Determine region from env with a sensible default
+ region := os.Getenv("AWS_REGION")
+ if region == "" {
+ region = os.Getenv("AWS_DEFAULT_REGION")
+ }
+ if region == "" {
+ region = "us-east-1"
+ }
+
+ // Determine host and port (default 5432)
+ host := s.Address
+ port := "5432"
+ if h, p, err := net.SplitHostPort(s.Address); err == nil {
+ host, port = h, p
+ }
+ hostPort := net.JoinHostPort(host, port)
+
+ slog.Info("Using AWS RDS IAM authentication for PostgreSQL", "host", host, "port", port, "region", region)
+
+ awsCfg, err := awsconfig.LoadDefaultConfig(context.Background())
+ if err != nil {
+ slog.Error("AWS configuration error", "err", err)
+ // Fallback preserves local dev/test
+ if s.Secret != "" {
+ return fmt.Sprintf("postgresql://%s:%s@%s/%s?sslmode=require",
+ url.QueryEscape(s.Username), url.QueryEscape(s.Secret), hostPort, s.Database)
+ }
+ return ""
+ }
+
+ token, err := rdsauth.BuildAuthToken(context.Background(), hostPort, region, s.Username, awsCfg.Credentials)
+ if err != nil {
+ slog.Error("Failed to create AWS RDS IAM token", "err", err)
+ if s.Secret != "" {
+ return fmt.Sprintf("postgresql://%s:%s@%s/%s?sslmode=require",
+ url.QueryEscape(s.Username), url.QueryEscape(s.Secret), hostPort, s.Database)
+ }
+ return ""
+ }
+
+ return fmt.Sprintf("postgresql://%s:%s@%s/%s?sslmode=require",
+ url.QueryEscape(s.Username), url.QueryEscape(token), hostPort, s.Database)
}Verification checklist:
- Confirm DAWGS/pg pool does per-connection token generation or refresh, not only once at startup.
- Confirm server cert validation is handled (sslmode=require here; consider verify-full + CA wiring in a later change).
- Confirm local dev/tests still work by setting
database.connectionin config to bypass IAM.
🏁 Script executed:
#!/bin/bash
set -e
# Locate the DatabaseConfiguration struct
file=$(rg -lP 'type\s+DatabaseConfiguration\b' -g '*.go')
echo "DatabaseConfiguration defined in: $file"
# Print the entire struct definition
sed -n '/type DatabaseConfiguration\s*struct\s*{/,/}/p' "$file"Length of output: 656
🏁 Script executed:
#!/bin/bash
set -e
echo "=== Import block and beginning of file ==="
sed -n '1,120p' cmd/api/src/config/config.go
echo ""
echo "=== PostgreSQLConnectionString function ==="
rg -nP -C5 'func \(s DatabaseConfiguration\) PostgreSQLConnectionString' cmd/api/src/config/config.goLength of output: 4129
🏁 Script executed:
#!/bin/bash
set -e
echo "=== Usage of PostgreSQLConnectionString in codebase ==="
rg -n "PostgreSQLConnectionString" -g '*.go' || echo "No occurrences found"
echo ""
echo "=== pgxpool usage ==="
rg -n "pgxpool" -g '*.go' || echo "No pgxpool usage found"
echo ""
echo "=== database/sql usage ==="
rg -n '"database/sql"' -g '*.go' || echo "No database/sql usage found"
echo ""
echo "=== sql.Open usage ==="
rg -n "sql.Open" -g '*.go' || echo "No sql.Open usage found"
echo ""
echo "=== 'postgres://' literal usage ==="
rg -n "postgres://" -g '*.go' || echo "No 'postgres://' literal found"Length of output: 2635
🏁 Script executed:
#!/bin/bash
set -e
echo "=== bootstrap/util.go around pgxpool and pool usage ==="
sed -n '1,200p' cmd/api/src/bootstrap/util.goLength of output: 4035
🏁 Script executed:
#!/bin/bash
set -e
echo "=== Display first 200 lines of config_test.go ==="
sed -n '1,200p' cmd/api/src/config/config_test.goLength of output: 6621
DatabaseConfiguration: Parameterize IAM auth, avoid panics, and honour overrides
The current PostgreSQLConnectionString implementation in cmd/api/src/config/config.go has several critical issues:
- It always tries RDS IAM auth (ignoring s.Connection and s.Secret), panics on load/config errors, and hardcodes “us-east-1” and port 5432.
- It emits a plaintext token in the URL and omits sslmode, which RDS IAM requires for TLS.
- It generates the token once at startup, but the pgx pool may open new connections later (resulting in expired tokens).
- Existing tests and the Neo4j path will break if overrides and error‐handling aren’t adjusted.
Locations to update:
- cmd/api/src/config/config.go:
func (s DatabaseConfiguration) PostgreSQLConnectionString() string - Imports in that file: alias
github.com/aws/aws-sdk-go-v2/configto avoid shadowing the local package, and alias the RDS auth package. - cmd/api/src/config/config_test.go: update tests to cover the new override (
s.Connection) and appended?sslmode=require.
Suggested in‐place refactor (adjust imports as noted):
--- a/cmd/api/src/config/config.go
+++ b/cmd/api/src/config/config.go
@@
-import (
- "github.com/aws/aws-sdk-go-v2/config"
- "github.com/aws/aws-sdk-go-v2/feature/rds/auth"
+import (
+ awsconfig "github.com/aws/aws-sdk-go-v2/config"
+ rdsauth "github.com/aws/aws-sdk-go-v2/feature/rds/auth"
)
func (s DatabaseConfiguration) PostgreSQLConnectionString() string {
- slog.Info("loading default config for rds auth")
- cfg, err := config.LoadDefaultConfig(context.TODO())
- if err != nil {
- panic("configuration error: " + err.Error())
- }
+ // Allow explicit override
+ if s.Connection != "" {
+ return s.Connection
+ }
+ // Pick AWS region: AWS_REGION → AWS_DEFAULT_REGION → fallback
+ region := os.Getenv("AWS_REGION")
+ if region == "" {
+ region = os.Getenv("AWS_DEFAULT_REGION")
+ }
+ if region == "" {
+ region = "us-east-1"
+ }
+ // Split host:port (default 5432)
+ host, port := s.Address, "5432"
+ if h, p, err := net.SplitHostPort(s.Address); err == nil {
+ host, port = h, p
+ }
+ hostPort := net.JoinHostPort(host, port)
+ slog.Info("Using AWS RDS IAM auth", "host", host, "port", port, "region", region)
+ awsCfg, err := awsconfig.LoadDefaultConfig(context.Background())
+ if err != nil {
+ slog.Error("AWS config load failed", "err", err)
+ goto passwordFallback
+ }
+
+ token, err := rdsauth.BuildAuthToken(context.Background(), hostPort, region, s.Username, awsCfg.Credentials)
+ if err != nil {
+ slog.Error("RDS IAM token generation failed", "err", err)
+ goto passwordFallback
+ }
+
+ // IAM succeeded—require TLS
+ return fmt.Sprintf(
+ "postgresql://%s:%s@%s/%s?sslmode=require",
+ url.QueryEscape(s.Username), url.QueryEscape(token),
+ hostPort, s.Database,
+ )
+passwordFallback:
+ if s.Secret != "" {
+ return fmt.Sprintf(
+ "postgresql://%s:%s@%s/%s?sslmode=require",
+ url.QueryEscape(s.Username), url.QueryEscape(s.Secret),
+ hostPort, s.Database,
+ )
+ }
+ // No valid auth available
+ return ""
}Checklist before merging:
- Update the import block as shown.
- Adjust config_test.go to assert on s.Connection override and the appended
?sslmode=require. - Verify pgx pool’s handling of static tokens; if it doesn’t refresh per connection, plan for a custom connector or token provider.
- In a follow-up, consider changing the signature to return (string, error) and adding sslmode=verify-full with CA.
📝 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.
| slog.Info("loading default config for rds auth") | |
| cfg, err := config.LoadDefaultConfig(context.TODO()) | |
| if err != nil { | |
| panic("configuration error: " + err.Error()) | |
| } | |
| return s.Connection | |
| dbinput := s.Address + ":5432" | |
| slog.Info("requesting auth token") | |
| authenticationToken, err := auth.BuildAuthToken(context.TODO(), dbinput, "us-east-1", s.Username, cfg.Credentials) | |
| if err != nil { | |
| panic("failed to create authentication token: " + err.Error()) | |
| } | |
| slog.Info("auth token successfully created") | |
| encodedToken := url.QueryEscape(authenticationToken) | |
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, encodedToken, s.Address, s.Database) | |
| } | |
| // --- in cmd/api/src/config/config.go --- | |
| import ( | |
| // ... other existing imports ... | |
| awsconfig "github.com/aws/aws-sdk-go-v2/config" | |
| rdsauth "github.com/aws/aws-sdk-go-v2/feature/rds/auth" | |
| "context" | |
| "fmt" | |
| "net" | |
| "os" | |
| "net/url" | |
| "github.com/go-logr/logr" // or wherever slog is from | |
| ) | |
| func (s DatabaseConfiguration) PostgreSQLConnectionString() string { | |
| // Honor explicit override | |
| if s.Connection != "" { | |
| return s.Connection | |
| } | |
| // Pick AWS region: AWS_REGION → AWS_DEFAULT_REGION → fallback | |
| region := os.Getenv("AWS_REGION") | |
| if region == "" { | |
| region = os.Getenv("AWS_DEFAULT_REGION") | |
| } | |
| if region == "" { | |
| region = "us-east-1" | |
| } | |
| // Split host:port (default port 5432) | |
| host, port := s.Address, "5432" | |
| if h, p, err := net.SplitHostPort(s.Address); err == nil { | |
| host, port = h, p | |
| } | |
| hostPort := net.JoinHostPort(host, port) | |
| slog.Info("Using AWS RDS IAM auth", "host", host, "port", port, "region", region) | |
| awsCfg, err := awsconfig.LoadDefaultConfig(context.Background()) | |
| if err != nil { | |
| slog.Error(err, "AWS config load failed") | |
| goto passwordFallback | |
| } | |
| token, err := rdsauth.BuildAuthToken(context.Background(), hostPort, region, s.Username, awsCfg.Credentials) | |
| if err != nil { | |
| slog.Error(err, "RDS IAM token generation failed") | |
| goto passwordFallback | |
| } | |
| // IAM succeeded—require TLS | |
| return fmt.Sprintf( | |
| "postgresql://%s:%s@%s/%s?sslmode=require", | |
| url.QueryEscape(s.Username), | |
| url.QueryEscape(token), | |
| hostPort, | |
| s.Database, | |
| ) | |
| passwordFallback: | |
| if s.Secret != "" { | |
| return fmt.Sprintf( | |
| "postgresql://%s:%s@%s/%s?sslmode=require", | |
| url.QueryEscape(s.Username), | |
| url.QueryEscape(s.Secret), | |
| hostPort, | |
| s.Database, | |
| ) | |
| } | |
| // No valid auth available | |
| return "" | |
| } |
🤖 Prompt for AI Agents
In cmd/api/src/config/config.go around lines 80-96, update
PostgreSQLConnectionString to honor s.Connection and s.Secret (only attempt RDS
IAM when s.Connection indicates IAM), avoid panics on config/auth errors (return
an error or an empty string and log instead — plan a follow-up to change the
signature to (string,error)), and stop hardcoding region and port by using
s.Region (default "us-east-1") and s.Port (default 5432); change imports to
alias AWS config (e.g. awsconfig) and the RDS auth package to avoid shadowing,
generate and url-escape the auth token only when using IAM, append
"?sslmode=require" to the returned connection string, and add a comment that
static tokens generated at startup may expire and require a connection-level
token provider (update config_test.go to assert s.Connection override and that
the returned string includes ?sslmode=require).
|
CLA Assistant Lite bot: I have read the CLA Document and I hereby sign the CLA You can retrigger this bot by commenting recheck in this Pull Request |
2 participants
Proof of concept for introducing RDS iam auth capabilities for PG. This is intended to be used to test IAM rds auth functionality.
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.
Summary by CodeRabbit