feat: move DatabaseConfiguration from Bloodhound repo to dawgs allowing for IAM auth option#57
feat: move DatabaseConfiguration from Bloodhound repo to dawgs allowing for IAM auth option#57mamundsen-specter wants to merge 1 commit intomainfrom
Conversation
WalkthroughA new Changes
Sequence Diagram(s)sequenceDiagram
participant CLI as CLI Tool
participant Config as DatabaseConfiguration
participant AWS as AWS SDK
participant RDS as RDS Database
participant Pool as pgxpool
rect rgba(100, 150, 200, 0.5)
Note over CLI,Pool: RDS IAM Auth Flow (EnableRDSIAMAuth = true)
CLI->>Config: Construct cfg (EnableRDSIAMAuth=true)
CLI->>Pool: NewPool(cfg)
Pool->>Config: PostgreSQLConnectionString()
Config->>AWS: Load SDK config / resolve host / request RDS IAM token
AWS-->>Config: Token (or error)
Config-->>Pool: Connection URI with token
Pool->>RDS: Connect using IAM token
RDS-->>Pool: Connected
end
rect rgba(150, 180, 100, 0.5)
Note over CLI,Pool: Standard PostgreSQL Flow (EnableRDSIAMAuth = false)
CLI->>Config: Construct cfg (EnableRDSIAMAuth=false)
CLI->>Pool: NewPool(cfg)
Pool->>Config: PostgreSQLConnectionString()
Config-->>Pool: Standard PostgreSQL URI
Pool->>RDS: Connect with username/password
RDS-->>Pool: Connected
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 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 |
mamundsen-specter
force-pushed
the
bi-1306
branch
from
5b67067 to
591a278
Compare
![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
🧹 Nitpick comments (3)
drivers/pg/pg.go (1)
64-66:MaxConcurrentSessionsfield is declared but not used.The
DatabaseConfigurationstruct includes aMaxConcurrentSessionsfield (seedrivers/config.go:21), but it's not used here. Consider wiring this field to the pool'sMaxConnssetting to make the pool configurable as intended by the TODO comment.♻️ Suggested implementation
// TODO: Min and Max connections for the pool should be configurable - poolCfg.MinConns = 5 - poolCfg.MaxConns = 50 + poolCfg.MinConns = 5 + if cfg.MaxConcurrentSessions > 0 { + poolCfg.MaxConns = int32(cfg.MaxConcurrentSessions) + } else { + poolCfg.MaxConns = 50 + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@drivers/pg/pg.go` around lines 64 - 66, The pool configuration hardcodes poolCfg.MinConns and poolCfg.MaxConns; wire DatabaseConfiguration.MaxConcurrentSessions into the pool by setting poolCfg.MaxConns = dbCfg.MaxConcurrentSessions (and keep or derive MinConns appropriately, e.g., min 1 or a configured value) inside the function that builds the pg pool (where poolCfg is assigned in drivers/pg/pg.go); ensure you reference the DatabaseConfiguration instance (e.g., dbCfg or cfg) and validate/clip MaxConcurrentSessions to sensible bounds before assigning to poolCfg.MaxConns.cmd/benchmark/main.go (1)
47-47: Consider declaringdbcfgcloser to its usage.Declaring
dbcfgin the flag variable block is slightly misleading since it's not a flag. Consider moving the declaration to where it's used (around line 67) to improve readability.♻️ Suggested refactor
var ( driver = flag.String("driver", "pg", "database driver (pg, neo4j)") connStr = flag.String("connection", "", "database connection string (or PG_CONNECTION_STRING)") iterations = flag.Int("iterations", 10, "timed iterations per scenario") output = flag.String("output", "", "markdown output file (default: stdout)") datasetDir = flag.String("dataset-dir", "integration/testdata", "path to testdata directory") localDataset = flag.String("local-dataset", "", "additional local dataset (e.g. local/phantom)") onlyDataset = flag.String("dataset", "", "run only this dataset (e.g. diamond, local/phantom)") - dbcfg = drivers.DatabaseConfiguration{} )Then at line 67:
+ dbcfg := drivers.DatabaseConfiguration{Connection: conn} - dbcfg.Connection = conn - if *driver == pg.DriverName { pool, err := pg.NewPool(dbcfg)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cmd/benchmark/main.go` at line 47, Move the non-flag variable dbcfg (type drivers.DatabaseConfiguration) out of the flag declaration block and declare it immediately before its first use in main (around where the code currently uses dbcfg, e.g., the section around line 67). Update any references to dbcfg to use the new local declaration so the flag block contains only actual flag variables; no behavior changes are needed, just relocate the declaration for clarity.drivers/config.go (1)
15-23: Consider adding validation or documentation for required fields.The struct has multiple optional fields (
Connectionvs component fields), but there's no validation to ensure a coherent configuration. Consider adding aValidate()method or documenting which field combinations are valid (e.g., eitherConnectionis set, orAddress/Database/Usernameare set).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@drivers/config.go` around lines 15 - 23, Add a Validate() method on the DatabaseConfiguration type that enforces coherent configurations: require that either Connection is non-empty OR the trio Address, Database and Username are all non-empty; return descriptive errors for missing required fields, validate MaxConcurrentSessions is positive (or set a sensible default) and if EnableRDSIAMAuth is true validate any extra constraints needed (e.g., Username present). Update the DatabaseConfiguration declaration with a short comment describing the valid combinations and call Validate() from any code that loads or uses this config (e.g., config loader/initializer) to fail fast with clear error messages.
🤖 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 33-60: RDSIAMAuthConnectionString currently logs AWS config or
token errors (from awsConfig.LoadDefaultConfig and auth.BuildAuthToken) then
silently falls back to defaultPostgreSQLConnectionString(), which hides failures
for IAM-only deployments; change the API to surface errors instead: modify
DatabaseConfiguration.RDSIAMAuthConnectionString to return (string, error) (or
add an explicit boolean/config flag) so failures during
awsConfig.LoadDefaultConfig or auth.BuildAuthToken are returned to the caller
rather than falling back, and update PostgreSQLConnectionString (or its callers)
to propagate/handle that error or decide explicitly when to fallback to
defaultPostgreSQLConnectionString.
- Around line 70-76: Neo4jConnectionString currently interpolates s.Secret
directly which breaks URIs when the secret contains special chars; update
Neo4jConnectionString to URL-escape credentials (at minimum s.Secret, and
preferably s.Username too) before building the connection string—mirror the
approach used in defaultPostgreSQLConnectionString by using Go's URL-encoding
utilities (e.g., url.UserPassword or an appropriate escape function) to produce
a valid neo4j:// userinfo component when s.Connection is empty.
- Line 47: The endpoint construction currently hardcodes ":5432" which breaks
when s.Address already contains a port or a non-standard port is required;
update the logic in RDSIAMAuthConnectionString (where endpoint, host and
strings.TrimSuffix are used) to parse s.Address for an existing port (e.g., via
net.SplitHostPort) and use that port if present, falling back to a Port field on
DatabaseConfiguration (add Port if not present) or to 5432 only as a final
default, and ensure endpoint is built using the chosen host and port rather than
concatenating a fixed ":5432".
In `@drivers/pg/pg.go`:
- Around line 74-87: The BeforeConnect handler currently uses
cfg.PostgreSQLConnectionString() which silently falls back from
RDSIAMAuthConnectionString() to defaultPostgreSQLConnectionString() on AWS
errors; change this so AWS errors are surfaced instead of allowing a potentially
empty/invalid credential string to reach pgx.ParseConfig. Concretely, modify the
config path (either make
RDSIAMAuthConnectionString()/PostgreSQLConnectionString() return an error or
expose the underlying error) and update the poolCfg.BeforeConnect (the
BeforeConnect lambda) to call the variant that can return an error and propagate
that error back (i.e., return the AWS error immediately) rather than parsing a
fallback connection string; ensure ParseConfig is only called with a validated,
non-empty connection string.
---
Nitpick comments:
In `@cmd/benchmark/main.go`:
- Line 47: Move the non-flag variable dbcfg (type drivers.DatabaseConfiguration)
out of the flag declaration block and declare it immediately before its first
use in main (around where the code currently uses dbcfg, e.g., the section
around line 67). Update any references to dbcfg to use the new local declaration
so the flag block contains only actual flag variables; no behavior changes are
needed, just relocate the declaration for clarity.
In `@drivers/config.go`:
- Around line 15-23: Add a Validate() method on the DatabaseConfiguration type
that enforces coherent configurations: require that either Connection is
non-empty OR the trio Address, Database and Username are all non-empty; return
descriptive errors for missing required fields, validate MaxConcurrentSessions
is positive (or set a sensible default) and if EnableRDSIAMAuth is true validate
any extra constraints needed (e.g., Username present). Update the
DatabaseConfiguration declaration with a short comment describing the valid
combinations and call Validate() from any code that loads or uses this config
(e.g., config loader/initializer) to fail fast with clear error messages.
In `@drivers/pg/pg.go`:
- Around line 64-66: The pool configuration hardcodes poolCfg.MinConns and
poolCfg.MaxConns; wire DatabaseConfiguration.MaxConcurrentSessions into the pool
by setting poolCfg.MaxConns = dbCfg.MaxConcurrentSessions (and keep or derive
MinConns appropriately, e.g., min 1 or a configured value) inside the function
that builds the pg pool (where poolCfg is assigned in drivers/pg/pg.go); ensure
you reference the DatabaseConfiguration instance (e.g., dbCfg or cfg) and
validate/clip MaxConcurrentSessions to sensible bounds before assigning to
poolCfg.MaxConns.
🪄 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: bb028527-5b6e-401c-ab60-2b3e49066fb4
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (5)
cmd/benchmark/main.gocmd/export/main.godrivers/config.godrivers/pg/pg.gogo.mod
| func (s DatabaseConfiguration) RDSIAMAuthConnectionString() string { | ||
| slog.Info("Loading RDS Configuration With IAM Auth") | ||
|
|
||
| if cfg, err := awsConfig.LoadDefaultConfig(context.TODO()); err != nil { | ||
| slog.Error("AWS Config Loading Error", slog.String("err", err.Error())) | ||
| } else { | ||
| host := s.Address | ||
|
|
||
| if hostCName, err := net.LookupCNAME(s.Address); err != nil { | ||
| slog.Warn("Error looking up CNAME for DB host. Using original address.", slog.String("err", err.Error())) | ||
| } else { | ||
| host = hostCName | ||
| } | ||
|
|
||
| endpoint := strings.TrimSuffix(host, ".") + ":5432" | ||
|
|
||
| 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())) | ||
| } 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 s.defaultPostgreSQLConnectionString() | ||
| } |
There was a problem hiding this comment.
Silent error fallback may mask configuration issues.
When AWS config loading or IAM token generation fails, the function logs an error but silently falls back to defaultPostgreSQLConnectionString(). For IAM-auth-only deployments (where no static password exists), this fallback will produce an invalid connection string, causing authentication failures that are harder to diagnose.
Consider returning an error or making the fallback behavior explicit/configurable.
🔧 Suggested approach to surface errors
-func (s DatabaseConfiguration) RDSIAMAuthConnectionString() string {
+func (s DatabaseConfiguration) RDSIAMAuthConnectionString() (string, error) {
slog.Info("Loading RDS Configuration With IAM Auth")
if cfg, err := awsConfig.LoadDefaultConfig(context.TODO()); err != nil {
- slog.Error("AWS Config Loading Error", slog.String("err", err.Error()))
+ return "", fmt.Errorf("failed to load AWS config: %w", err)
- } else {
+ }
// ... rest of implementation
- }
-
- return s.defaultPostgreSQLConnectionString()
+ return connString, nil
}Then update PostgreSQLConnectionString() to propagate errors or handle the fallback decision explicitly.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@drivers/config.go` around lines 33 - 60, RDSIAMAuthConnectionString currently
logs AWS config or token errors (from awsConfig.LoadDefaultConfig and
auth.BuildAuthToken) then silently falls back to
defaultPostgreSQLConnectionString(), which hides failures for IAM-only
deployments; change the API to surface errors instead: modify
DatabaseConfiguration.RDSIAMAuthConnectionString to return (string, error) (or
add an explicit boolean/config flag) so failures during
awsConfig.LoadDefaultConfig or auth.BuildAuthToken are returned to the caller
rather than falling back, and update PostgreSQLConnectionString (or its callers)
to propagate/handle that error or decide explicitly when to fallback to
defaultPostgreSQLConnectionString.
| host = hostCName | ||
| } | ||
|
|
||
| endpoint := strings.TrimSuffix(host, ".") + ":5432" |
There was a problem hiding this comment.
Hardcoded port 5432 may not match the configured address.
If s.Address contains a port (e.g., mydb.example.com:5433) or uses a non-standard port, this hardcoded :5432 will produce an incorrect endpoint. Consider parsing the port from Address or adding a Port field to DatabaseConfiguration.
🔧 Suggested approach
+func (s DatabaseConfiguration) getHostPort() (string, string) {
+ host := s.Address
+ port := "5432"
+ if h, p, err := net.SplitHostPort(s.Address); err == nil {
+ host = h
+ port = p
+ }
+ return host, port
+}Then use it in RDSIAMAuthConnectionString:
- endpoint := strings.TrimSuffix(host, ".") + ":5432"
+ _, port := s.getHostPort()
+ endpoint := strings.TrimSuffix(host, ".") + ":" + port🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@drivers/config.go` at line 47, The endpoint construction currently hardcodes
":5432" which breaks when s.Address already contains a port or a non-standard
port is required; update the logic in RDSIAMAuthConnectionString (where
endpoint, host and strings.TrimSuffix are used) to parse s.Address for an
existing port (e.g., via net.SplitHostPort) and use that port if present,
falling back to a Port field on DatabaseConfiguration (add Port if not present)
or to 5432 only as a final default, and ensure endpoint is built using the
chosen host and port rather than concatenating a fixed ":5432".
| func (s DatabaseConfiguration) Neo4jConnectionString() string { | ||
| if s.Connection == "" { | ||
| return fmt.Sprintf("neo4j://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database) | ||
| } | ||
|
|
||
| return s.Connection | ||
| } |
There was a problem hiding this comment.
Neo4jConnectionString does not URL-escape the secret.
Unlike defaultPostgreSQLConnectionString() (line 30), this method does not URL-escape s.Secret. If the secret contains special characters (e.g., @, /, #), the resulting URI will be malformed.
🔧 Proposed fix
func (s DatabaseConfiguration) Neo4jConnectionString() string {
if s.Connection == "" {
- return fmt.Sprintf("neo4j://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database)
+ return fmt.Sprintf("neo4j://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database)
}
return s.Connection
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@drivers/config.go` around lines 70 - 76, Neo4jConnectionString currently
interpolates s.Secret directly which breaks URIs when the secret contains
special chars; update Neo4jConnectionString to URL-escape credentials (at
minimum s.Secret, and preferably s.Username too) before building the connection
string—mirror the approach used in defaultPostgreSQLConnectionString by using
Go's URL-encoding utilities (e.g., url.UserPassword or an appropriate escape
function) to produce a valid neo4j:// userinfo component when s.Connection is
empty.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 36-51: The code currently calls awsConfig.LoadDefaultConfig and
net.LookupCNAME inside the per-connection BeforeConnect path; move those calls
to one-time initialization and cache their results (e.g., store awsCfg and
resolvedHost/endpoint on the driver/connector struct or package-level vars). Use
a constructor or init function (or sync.Once) to call
awsConfig.LoadDefaultConfig and net.LookupCNAME once, compute endpoint :=
strings.TrimSuffix(resolvedHost, ".")+":5432", and then in BeforeConnect reuse
the cached awsCfg, region/credentials and endpoint when calling
auth.BuildAuthToken (instead of reloading config and re-resolving CNAME on each
connection). Ensure the cache is concurrency-safe and falls back to original
behavior if initialization fails.
- Line 30: Replace the manual fmt.Sprintf URL construction that uses
url.QueryEscape on credentials with proper URL building using url.UserPassword
and url.URL: create a url.Userinfo via url.UserPassword(s.Username, s.Secret),
set the Scheme to "postgresql", Host to s.Address, and Path (or Opaque/RawPath
as appropriate) to the database name (s.Database), then call url.String() to
produce the final DSN; update the return in the function that currently uses
fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret),
s.Address, s.Database) to this url.URL-based construction so userinfo is encoded
per RFC 3986.
🪄 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: bf38cae3-d953-40f9-89de-00a017c26b6b
⛔ 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
✅ Files skipped from review due to trivial changes (1)
- go.mod
🚧 Files skipped from review as they are similar to previous changes (2)
- cmd/export/main.go
- drivers/pg/pg.go
| return s.Connection | ||
| } | ||
|
|
||
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database) |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
In Go's net/url package, is url.QueryEscape appropriate for encoding URL userinfo credentials, or should url.UserPassword/url.URL be used? Please include how spaces and reserved characters are encoded in each case.
💡 Result:
url.UserPassword/url.URL should be used for encoding URL userinfo credentials, not url.QueryEscape. url.UserPassword(username, password) creates a *Userinfo that, when used in a url.URL and stringified via u.String, properly encodes the username and password using the package's internal userinfo encoding rules (encodeUserPassword mode). This escapes only the characters that must be escaped in userinfo per RFC 3986 §3.2.1: '@', '/', '?', and ':' (to handle parsing). Other reserved characters like '&', '=', '+', ';', ',', '$' are left unescaped, and spaces are encoded as '+'. In contrast, url.QueryEscape is for query parameters (encodeQueryComponent mode). It escapes all reserved characters (including '&' as %26, '=' as %3D, etc.) and encodes spaces as '+'. Using QueryEscape for userinfo can produce incorrectly encoded URLs because it over-escapes characters safe in userinfo, potentially causing parsing issues or non-standard URLs. Official docs recommend UserPassword for constructing userinfo, noting it's for legacy sites but handles encoding correctly. Since Go 1.9.4, url.Parse requires userinfo to be percent-encoded if it contains invalid characters. Examples: For username="user name", password="pass&word?": - url.UserPassword: "user+name:pass%26word%3F" (spaces -> '+', '&'->%26, '?'->%3F; '+' and '=' unescaped if present) - url.QueryEscape on both then join: "user+name:pass%26word%3F" (same for these chars, but would escape '=' to %3D, which userinfo allows unescaped) Code snippet: package main import ( "fmt" "net/url" ) func main { u := &url.URL{ Scheme: "https", Host: "example.com", User: url.UserPassword("user name", "pass&word?"), } fmt.Println(u.String) // https://user+name:pass%26word%3F@example.com }
Citations:
- 1: https://pkg.go.dev/net/url
- 2: https://pkg.go.dev/net/url@go1.25.6
- 3: https://golang.google.cn/pkg/net/url/
- 4: https://developers.nutanix.com/api/v1/sdk/namespaces/main/files/versions/v4.0.a2/languages/golang/src/net/url/url.go.html
- 5: https://go.dev/src/net/url/url.go
- 6: https://go.googlesource.com/go/+/go1.12.8/src/net/url/url.go
- 7: https://go.googlesource.com/go/+/refs/heads/dev.boringcrypto.go1.15/src/net/url/url.go
🏁 Script executed:
fd -t f drivers/config.goRepository: SpecterOps/DAWGS
Length of output: 42
🏁 Script executed:
fd -t f "config.go" || echo "Not found with default search"Repository: SpecterOps/DAWGS
Length of output: 77
🏁 Script executed:
fd -t d "drivers" || echo "No drivers directory found"Repository: SpecterOps/DAWGS
Length of output: 68
🏁 Script executed:
cat -n drivers/config.go | head -50Repository: SpecterOps/DAWGS
Length of output: 1846
Use url.UserPassword() to properly encode URL credentials instead of url.QueryEscape().
Line 30 uses url.QueryEscape() for encoding credentials in the userinfo portion of the PostgreSQL URL, which is incorrect. url.QueryEscape is designed for query parameters and over-escapes characters that are safe in userinfo (like =, ;). Use url.UserPassword(username, password) with url.URL instead, which applies the correct userinfo encoding per RFC 3986.
🔧 Suggested fix
func (s DatabaseConfiguration) defaultPostgreSQLConnectionString() string {
if s.Connection != "" {
return s.Connection
}
- return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database)
+ u := &url.URL{
+ Scheme: "postgresql",
+ Host: s.Address,
+ Path: "/" + s.Database,
+ }
+ u.User = url.UserPassword(s.Username, s.Secret)
+ return u.String()
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@drivers/config.go` at line 30, Replace the manual fmt.Sprintf URL
construction that uses url.QueryEscape on credentials with proper URL building
using url.UserPassword and url.URL: create a url.Userinfo via
url.UserPassword(s.Username, s.Secret), set the Scheme to "postgresql", Host to
s.Address, and Path (or Opaque/RawPath as appropriate) to the database name
(s.Database), then call url.String() to produce the final DSN; update the return
in the function that currently uses fmt.Sprintf("postgresql://%s:%s@%s/%s",
s.Username, url.QueryEscape(s.Secret), s.Address, s.Database) to this
url.URL-based construction so userinfo is encoded per RFC 3986.
| if cfg, err := awsConfig.LoadDefaultConfig(context.TODO()); err != nil { | ||
| slog.Error("AWS Config Loading Error", slog.String("err", err.Error())) | ||
| } else { | ||
| host := s.Address | ||
|
|
||
| if hostCName, err := net.LookupCNAME(s.Address); err != nil { | ||
| slog.Warn("Error looking up CNAME for DB host. Using original address.", slog.String("err", err.Error())) | ||
| } else { | ||
| host = hostCName | ||
| } | ||
|
|
||
| endpoint := strings.TrimSuffix(host, ".") + ":5432" | ||
|
|
||
| slog.Info("Requesting RDS IAM Auth Token") | ||
|
|
||
| if authenticationToken, err := auth.BuildAuthToken(context.TODO(), endpoint, cfg.Region, s.Username, cfg.Credentials); err != nil { |
There was a problem hiding this comment.
Avoid repeated AWS config/DNS work on every connection attempt.
Lines 36-51 perform LoadDefaultConfig and CNAME lookup in a path that is invoked per new PG connection (via BeforeConnect). This can become a connection-latency bottleneck and create avoidable external-call load.
💡 Refactor direction
- if cfg, err := awsConfig.LoadDefaultConfig(context.TODO()); err != nil {
+ // load once (e.g., during pool initialization) and reuse:
+ // - aws.Config
+ // - normalized host/endpoint
+ // per-connection: only mint the IAM token🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@drivers/config.go` around lines 36 - 51, The code currently calls
awsConfig.LoadDefaultConfig and net.LookupCNAME inside the per-connection
BeforeConnect path; move those calls to one-time initialization and cache their
results (e.g., store awsCfg and resolvedHost/endpoint on the driver/connector
struct or package-level vars). Use a constructor or init function (or sync.Once)
to call awsConfig.LoadDefaultConfig and net.LookupCNAME once, compute endpoint
:= strings.TrimSuffix(resolvedHost, ".")+":5432", and then in BeforeConnect
reuse the cached awsCfg, region/credentials and endpoint when calling
auth.BuildAuthToken (instead of reloading config and re-resolving CNAME on each
connection). Ensure the cache is concurrency-safe and falls back to original
behavior if initialization fails.
1 participant
Description
New Features
Added support for IAM-authenticated database connections and unified configuration-driven connection strings.
Improvements
Connection pool now defaults to larger min/max connections (5–50).
Pool automatically refreshes IAM credentials before opening new connections.
Credential renewal events are logged as informational messages for observability.
Resolves: BI-1306
Type of Change
Testing
go test -tags manual_integration ./integration/...)Screenshots (if appropriate):
Driver Impact
drivers/pg)drivers/neo4j)Checklist
go.mod/go.sumare up to date if dependencies changedSummary by CodeRabbit