feat: update db pooling to handle iam rds auth, move DatabaseConfiguration to dawgs#65
feat: update db pooling to handle iam rds auth, move DatabaseConfiguration to dawgs#65mamundsen-specter merged 4 commits intomainfrom
Conversation
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 10 minutes and 31 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. WalkthroughAdds a new drivers.DatabaseConfiguration type with connection-string builders (Postgres and Neo4j) including optional RDS IAM auth; updates pg.NewPool to accept that config; and changes CLI and integration code to construct and pass DatabaseConfiguration. go.mod adds AWS SDK v2 deps. Changes
Sequence DiagramsequenceDiagram
participant App as Application
participant Cfg as DatabaseConfiguration
participant AWS as AWS SDK
participant DB as PostgreSQL Database
App->>Cfg: Create DatabaseConfiguration (EnableRDSIAMAuth?)
Cfg->>Cfg: PostgreSQLConnectionString()
alt EnableRDSIAMAuth == true
Cfg->>Cfg: LookupEndpoint() (CNAME resolution)
Cfg->>AWS: LoadDefaultConfig()
AWS-->>Cfg: AWS Config
Cfg->>AWS: BuildAuthToken(endpoint, user)
AWS-->>Cfg: IAM Auth Token
Cfg->>Cfg: Return URI with token as password
else
Cfg->>Cfg: Return formatted PostgreSQL URI
end
App->>DB: Call pg.NewPool(cfg) -> pool init uses cfg.PostgreSQLConnectionString()
DB-->>App: Pool ready
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@drivers/config.go`:
- Around line 68-70: The Neo4jConnectionString function currently interpolates
s.Username and s.Secret directly into the URI userinfo, which can break when
those contain reserved chars; update Neo4jConnectionString to escape credentials
like PostgreSQLConnectionString does (use url.QueryEscape or equivalent) for
s.Username and s.Secret before formatting the connection string so the returned
string is a valid/parseable Neo4j URI.
- Around line 76-85: The LookupEndpoint logic must handle addresses that already
include a port and use the context-aware CNAME API: change
DatabaseConfiguration.LookupEndpoint to accept a context.Context, split
s.Address with net.SplitHostPort to separate host and port (use ":5432" only
when no explicit port present), call net.DefaultResolver.LookupCNAME(ctx, host)
instead of net.LookupCNAME, trim any trailing '.' from the resolved host and
reattach the explicit or default port, and update callers to pass context (use
context.TODO() where no context is available and pass the existing poolCtx from
the PostgreSQL pool creation site).
- Around line 35-57: The RDS IAM token function currently swallows
AWS/config/auth errors and falls back to static credentials; change
DatabaseConfiguration.RDSIAMAuthConnectionString to return (string, error) and
return explicit errors on AWS config load, endpoint lookup, or
auth.BuildAuthToken failures instead of calling
defaultPostgreSQLConnectionString; update PostgreSQLConnectionString (and any
wrapper that calls RDSIAMAuthConnectionString) to accept and propagate the
error, and modify pg.NewPool callers (the two places calling
PostgreSQLConnectionString) to pass poolCtx and handle/return the error so IAM
failures bubble up rather than silently falling back.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 794c115f-1511-428f-85fa-ee235edf2ed5
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (6)
cmd/benchmark/main.gocmd/export/main.godrivers/config.godrivers/pg/pg.gogo.modintegration/harness.go
There was a problem hiding this comment.
♻️ Duplicate comments (3)
drivers/config.go (3)
36-37:⚠️ Potential issue | 🟠 MajorFail closed when IAM token generation fails.
Line 56 falls back to the default/static PostgreSQL connection string after AWS config or token-generation errors. With
EnableRDSIAMAuth, that can silently bypass IAM or hide misconfiguration until later. Return an error from the IAM/PostgreSQL connection-string path and propagate it throughdrivers/pg/pg.goinstead. Based on learnings,DatabaseConfiguration.PostgreSQLConnectionString()intentionally ignores theConnectionfield whenEnableRDSIAMAuthis true because IAM auth always takes precedence.🛡️ Proposed direction
-func (s DatabaseConfiguration) RDSIAMAuthConnectionString() string { - if cfg, err := awsConfig.LoadDefaultConfig(context.TODO()); err != nil { - slog.Error("AWS Config Loading Error", slog.String("err", err.Error())) +func (s DatabaseConfiguration) RDSIAMAuthConnectionString(ctx context.Context) (string, error) { + if cfg, err := awsConfig.LoadDefaultConfig(ctx); err != nil { + return "", fmt.Errorf("load AWS config for RDS IAM auth: %w", err) } else { // Must use instance endpoint with IAM auth var endpoint string if s.Endpoint != "" { endpoint = s.Endpoint } else { - endpoint = s.LookupEndpoint() + endpoint = s.LookupEndpoint(ctx) } slog.Info("Requesting RDS IAM Auth Token") - if authenticationToken, err := auth.BuildAuthToken(context.TODO(), endpoint, cfg.Region, s.Username, cfg.Credentials); err != nil { - slog.Error("RDS IAM Auth Token Request Error", slog.String("err", err.Error())) + if authenticationToken, err := auth.BuildAuthToken(ctx, endpoint, cfg.Region, s.Username, cfg.Credentials); err != nil { + return "", fmt.Errorf("build RDS IAM auth token: %w", err) } else { slog.Info("RDS IAM Auth Token Created") - return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(authenticationToken), endpoint, s.Database) + return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(authenticationToken), endpoint, s.Database), nil } } - - slog.Warn("Failed to create IAM auth token. Falling back to default Postgres connection string") - return s.defaultPostgreSQLConnectionString() }Also applies to: 48-49, 56-57
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@drivers/config.go` around lines 36 - 37, The current AWS IAM flow in DatabaseConfiguration.PostgreSQLConnectionString falls back to the static Connection string on AWS config or token-generation errors; instead fail closed by returning an error from the IAM path when awsConfig.LoadDefaultConfig or the RDS IAM token generation fails (do not silently return the default Connection), and update callers to propagate that error (notably the drivers/pg package entrypoint in drivers/pg/pg.go) so EnableRDSIAMAuth truly takes precedence; locate and modify the blocks around awsConfig.LoadDefaultConfig and the token-generation calls to return descriptive errors and ensure drivers/pg/pg.go surfaces those errors to its caller.
31-31:⚠️ Potential issue | 🟡 MinorBuild connection URLs with
url.URLandurl.UserPassword.
url.QueryEscapeis for query strings, not URI userinfo (Go net/url docs). These lines still leaveUsernameunescaped and can mis-encode credentials containing spaces; useurl.UserPasswordso PostgreSQL, IAM, and Neo4j credentials are encoded consistently.🔧 Proposed refactor
- return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database) + postgresURL := url.URL{ + Scheme: "postgresql", + User: url.UserPassword(s.Username, s.Secret), + Host: s.Address, + Path: "/" + strings.TrimPrefix(s.Database, "/"), + } + return postgresURL.String()- return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(authenticationToken), endpoint, s.Database) + postgresURL := url.URL{ + Scheme: "postgresql", + User: url.UserPassword(s.Username, authenticationToken), + Host: endpoint, + Path: "/" + strings.TrimPrefix(s.Database, "/"), + } + return postgresURL.String()- return fmt.Sprintf("neo4j://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database) + neo4jURL := url.URL{ + Scheme: "neo4j", + User: url.UserPassword(s.Username, s.Secret), + Host: s.Address, + Path: "/" + strings.TrimPrefix(s.Database, "/"), + } + return neo4jURL.String()Also applies to: 52-52, 70-70
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@drivers/config.go` at line 31, The connection URL assembly is using fmt.Sprintf with url.QueryEscape (which is for query strings) and leaves s.Username unescaped; update the code that returns the DSN (the line using fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database)) to build the URL via net/url: create a url.URL with Scheme "postgresql", Host set to s.Address, Path set to "/"+s.Database, and set User to url.UserPassword(s.Username, s.Secret) so both username and secret are encoded correctly; apply the same pattern for the other similar return lines that reference s.Username, s.Secret, s.Address, s.Database.
76-86:⚠️ Potential issue | 🟠 MajorDefault the endpoint port and use the context-aware DNS resolver.
Line 77 ignores
SplitHostPorterrors, soAddress: "db.example.com"leaves bothhostandportempty and can return an IAM endpoint without:5432. AWS documents RDS auth token endpoints as<host>:<port>(AWS SDK for Go v2 RDS utilities); Line 78 also triggers thenoctxlint becausenet.LookupCNAMEusescontext.Backgroundinternally (Go net docs).🔧 Proposed fix
-func (s DatabaseConfiguration) LookupEndpoint() string { - host, port, _ := net.SplitHostPort(s.Address) - if hostCName, err := net.LookupCNAME(host); err != nil { +func (s DatabaseConfiguration) LookupEndpoint(ctx context.Context) string { + host := s.Address + port := "5432" + if splitHost, splitPort, err := net.SplitHostPort(s.Address); err == nil { + host = splitHost + port = splitPort + } + + if hostCName, err := net.DefaultResolver.LookupCNAME(ctx, host); err != nil { slog.Warn("Error looking up CNAME for DB host. Using original address.", slog.String("err", err.Error())) - return s.Address } else { host = hostCName } // Instance endpoint always returns with a trailing '.' - return strings.TrimSuffix(host, ".") + ":" + port + return net.JoinHostPort(strings.TrimSuffix(host, "."), port) }Update callers to pass the available context, especially
drivers/pg/pg.go:75and theBeforeConnectpath.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@drivers/config.go`:
- Around line 36-37: The current AWS IAM flow in
DatabaseConfiguration.PostgreSQLConnectionString falls back to the static
Connection string on AWS config or token-generation errors; instead fail closed
by returning an error from the IAM path when awsConfig.LoadDefaultConfig or the
RDS IAM token generation fails (do not silently return the default Connection),
and update callers to propagate that error (notably the drivers/pg package
entrypoint in drivers/pg/pg.go) so EnableRDSIAMAuth truly takes precedence;
locate and modify the blocks around awsConfig.LoadDefaultConfig and the
token-generation calls to return descriptive errors and ensure drivers/pg/pg.go
surfaces those errors to its caller.
- Line 31: The connection URL assembly is using fmt.Sprintf with url.QueryEscape
(which is for query strings) and leaves s.Username unescaped; update the code
that returns the DSN (the line using fmt.Sprintf("postgresql://%s:%s@%s/%s",
s.Username, url.QueryEscape(s.Secret), s.Address, s.Database)) to build the URL
via net/url: create a url.URL with Scheme "postgresql", Host set to s.Address,
Path set to "/"+s.Database, and set User to url.UserPassword(s.Username,
s.Secret) so both username and secret are encoded correctly; apply the same
pattern for the other similar return lines that reference s.Username, s.Secret,
s.Address, s.Database.
mamundsen-specter
changed the title
|
@CodeRabbit review |
✅ Actions performedReview triggered.
|
2 participants
…ation to dawgs
Description
Recreated from previous PR: #32
Didn't have commit signing enabled for this repo and figured it would be easier to create a new branch/PR instead of attempting to squash merge commits.
Updates pg pooling to enable IAM auth. Moves DatabaseConfiguation to DAWGs repo from the BHCE repo.
Resolves: <TICKET_OR_ISSUE_NUMBER>
Type of Change
Testing
go test -tags manual_integration ./integration/...)Test tenant created with IAM enabled.
Screenshots (if appropriate):
Driver Impact
drivers/pg)drivers/neo4j)Checklist
go.mod/go.sumare up to date if dependencies changedSummary by CodeRabbit
New Features
Refactor