add beforeConnect to NewPool poolcfg [DRAFT] #12
Conversation
WalkthroughReworks PostgreSQL pool initialization to accept a configuration object and refresh credentials on each new connection via a BeforeConnect hook. Updates go.mod to Go 1.24.x with a toolchain directive, adds a local replace for bloodhound, introduces numerous indirect dependencies, and removes some prior indirect modules. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant C as Caller
participant PG as NewPool (drivers/pg)
participant CFG as DatabaseConfiguration
participant P as pgxpool
participant DB as PostgreSQL
C->>PG: NewPool(cfg)
PG->>CFG: PostgreSQLConnectionString()
PG->>P: pgxpool.ParseConfig(connStr)<br/>MinConns=5, MaxConns=50
note over PG,P: Register hooks
rect rgba(220,240,255,0.5)
note right of P: BeforeConnect (credential refresh)
P->>CFG: PostgreSQLConnectionString()
P->>P: Reparse DSN, update Password
end
P->>DB: Establish connection
note over P,DB: AfterConnect (composite types)
DB-->>P: Ready
C-->>P: Pool returned
note over P: AfterRelease (validation) on release
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
MitchellAmundsen
changed the title
![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 (4)
drivers/pg/pg.go (3)
59-59: Parse from provider OK; consider context-aware retrievalUsing cfg.PostgreSQLConnectionString() is fine, but it ignores the ctx you already have (poolCtx). If the provider needs AWS/IMDS/ST S calls, a ctx-aware method helps cancellations/timeouts.
Possible evolution (if/when feasible in your config layer):
type PGConnStringProvider interface { PostgreSQLConnectionStringWithContext(ctx context.Context) (string, error) }Then:
cs, err := cfg.PostgreSQLConnectionStringWithContext(poolCtx) if err != nil { return nil, err } poolCfg, err := pgxpool.ParseConfig(cs)
74-83: BeforeConnect: fix log typo + lower log level + avoid shadowing + minor cleanup
- Typo: “credentional” → “credentials”.
- Prefer Debug to avoid noisy logs on every physical connect.
- Avoid fmt.Sprint for a static string.
- Rename param to avoid shadowing outer poolCfg.
- poolCfg.BeforeConnect = func(ctx context.Context, poolCfg *pgx.ConnConfig) error { - slog.Info(fmt.Sprint("RDS credentional beforeConnect(), creating new IAM credentials")) - refreshConnectionString := cfg.PostgreSQLConnectionString() - newPoolCfg, err := pgxpool.ParseConfig(refreshConnectionString) - if err != nil { - return err - } - poolCfg.Password = newPoolCfg.ConnConfig.Password - return nil - } + poolCfg.BeforeConnect = func(ctx context.Context, connCfg *pgx.ConnConfig) error { + slog.Debug("Refreshing RDS IAM credentials before establishing a new connection") + refreshConnectionString := cfg.PostgreSQLConnectionString() + newPoolCfg, err := pgxpool.ParseConfig(refreshConnectionString) + if err != nil { + return err + } + connCfg.Password = newPoolCfg.ConnConfig.Password + return nil + }Optional: if the provider can change host/port/user, consider updating those too instead of only Password.
65-66: Pool sizing should be externally configurableHard-coding MinConns=5 / MaxConns=50 makes tuning in different deployments difficult (and may slow startup if eager connections are made). Expose these via configuration with sensible defaults; you already left a TODO.
go.mod (1)
23-36: Trim indirect AWS SDK dependencies by decoupling the PG driverIt looks like lines 23–36 in go.mod list a cluster of AWS SDK for Go v2 modules (all marked “// indirect”) that are being pulled in transitively via
github.com/specterops/bloodhound. To keep the coredawgslibrary lean, consider extracting the PostgreSQL driver (and its Bloodhound dependency) into its own module:
- Move all pg-related code (e.g.
drivers/pg,cypher/models/pgsql) into a separate repo or submodule.- Remove the
github.com/specterops/bloodhoundrequire and its localreplacewhen consuming only the core library.- Verify that only truly needed dependencies remain in the core.
After doing so, rerun:
go mod tidy rg -nP 'require\s*\(' -n go.modYou should see the AWS SDK v2 entries drop out of the require block.
📜 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 ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (2)
drivers/pg/pg.go(3 hunks)go.mod(2 hunks)
🔇 Additional comments (1)
go.mod (1)
3-6: Verify CI Go toolchain supportLocal development environment confirms alignment with the module’s directives:
- File: go.mod (lines 3–6):
•go 1.24.4
•toolchain go1.24.6- Installed Go binary:
/usr/local/go/bin/go→go version go1.24.6 linux/amd64(fetched viatoolchaindirective)- GOTOOLCHAIN setting:
auto(ensures thetoolchaindirective is applied)Please verify that your CI runners are using Go 1.21 or newer (so the
toolchaindirective is supported) and either pre-install Go 1.24.6 or allow the directive to download it. This will ensure consistent toolchain usage across local and CI environments.
|
|
||
| "github.com/jackc/pgx/v5" | ||
| "github.com/jackc/pgx/v5/pgxpool" | ||
| "github.com/specterops/bloodhound/cmd/api/src/config" |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Decouple DAWGS from BloodHound config package
Importing github.com/specterops/bloodhound/cmd/api/src/config couples this library to a specific app and forces the go.mod replace. Prefer a tiny local interface or a callback to supply/refresh the connection string so DAWGS remains app-agnostic.
Apply this focused change to drop the hard import and accept a minimal provider:
-import (
+import (
"context"
"fmt"
"log/slog"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
- "github.com/specterops/bloodhound/cmd/api/src/config"
"github.com/specterops/dawgs"
"github.com/specterops/dawgs/cypher/models/pgsql"
"github.com/specterops/dawgs/graph"
)Add this type near the top of the file (outside the selected range):
// PGConnStringProvider abstracts how a PostgreSQL connection string is produced.
type PGConnStringProvider interface {
PostgreSQLConnectionString() string
}🤖 Prompt for AI Agents
In drivers/pg/pg.go around line 11, the file currently imports
github.com/specterops/bloodhound/cmd/api/src/config which couples this library
to a specific app; replace that hard import by introducing a small local
interface and accepting it instead. Add a PGConnStringProvider interface near
the top of the file (outside the current range) with method
PostgreSQLConnectionString() string, remove the direct config import, and change
any constructors/functions that currently rely on the config package to take
this provider (or a callback) and call PostgreSQLConnectionString() to obtain
the DSN; update usages and tests to pass an implementation from the caller.
| if connectionString == "" { | ||
| return nil, fmt.Errorf("graph connection requires a connection url to be set") | ||
| } | ||
| func NewPool(cfg config.DatabaseConfiguration) (*pgxpool.Pool, error) { |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Public API break: NewPool now depends on app-specific type
Changing NewPool to take config.DatabaseConfiguration breaks existing callers and ties DAWGS to BloodHound. Keep the new capability but avoid the hard type by using a local interface (see prior comment) and/or provide a backward-compatible helper.
Option A (preferred): switch to the local interface.
-func NewPool(cfg config.DatabaseConfiguration) (*pgxpool.Pool, error) {
+func NewPool(cfg PGConnStringProvider) (*pgxpool.Pool, error) {Option B: add a compatibility constructor so existing users with raw DSNs don’t break:
// NewPoolFromConnString keeps backward-compat for existing callers.
func NewPoolFromConnString(connString string) (*pgxpool.Pool, error) {
type staticProvider string
func (s staticProvider) PostgreSQLConnectionString() string { return string(s) }
return NewPool(staticProvider(connString))
}To find any impacted call sites, run:
🏁 Script executed:
#!/bin/bash
# Find NewPool call sites that may need updating
rg -nP -C2 '\bNewPool\s*\('Length of output: 292
Public API break: avoid tying NewPool to app-specific config
No internal usages of NewPool were found within the codebase (the only match is its own declaration), but this function is part of your public driver API and changing its signature will break downstream callers. To preserve existing integrations while still gaining the new flexibility, please choose one of the following:
• Option A (preferred): introduce a minimal local interface
// Define in drivers/pg/pg.go
type PGConnStringProvider interface {
PostgreSQLConnectionString() string
}
// Change signature to depend on the interface
-func NewPool(cfg config.DatabaseConfiguration) (*pgxpool.Pool, error) {
+func NewPool(cfg PGConnStringProvider) (*pgxpool.Pool, error) {
// existing implementation
}This lets you implement PGConnStringProvider for config.DatabaseConfiguration (and any other future types) without breaking callers.
• Option B: add a backward-compatible constructor for raw DSNs
// NewPoolFromConnString preserves the old behavior for string callers.
func NewPoolFromConnString(connString string) (*pgxpool.Pool, error) {
type staticProvider string
func (s staticProvider) PostgreSQLConnectionString() string { return string(s) }
return NewPool(staticProvider(connString))
}Callers that previously did:
pool, err := pg.NewPool(myConfig.DatabaseConnectionString)can continue doing so by switching to:
pool, err := pg.NewPoolFromConnString(myConfig.DatabaseConnectionString)Please implement one of these to ensure the driver remains backward-compatible.
🤖 Prompt for AI Agents
In drivers/pg/pg.go around line 54, NewPool currently accepts an app-specific
config type which breaks the public driver API; fix by introducing a minimal
local interface and updating NewPool to depend on it (preferred): declare type
PGConnStringProvider with method PostgreSQLConnectionString() string, change
NewPool(cfg config.DatabaseConfiguration) to NewPool(cfg PGConnStringProvider),
and implement this interface on config.DatabaseConfiguration (returning its
connection string) so existing callers continue to compile; alternatively,
implement a backward-compatible wrapper NewPoolFromConnString(connString string)
that adapts a plain string to the interface and forwards to NewPool so callers
using raw DSNs remain supported.
|
|
||
| toolchain go1.24.6 | ||
|
|
||
| replace github.com/specterops/bloodhound v0.0.0 => ../bloodhound-enterprise/bhce |
There was a problem hiding this comment.
Do not commit a local replace to an enterprise path
replace github.com/specterops/bloodhound v0.0.0 => ../bloodhound-enterprise/bhce will break external consumers and CI outside your workstation. This is a release blocker for a shared library.
Recommended:
- Remove the replace from go.mod; keep it in a dev-only go.work or a private branch.
- If DAWGS must integrate with BloodHound, depend on a tagged module version in a public or private module proxy, not a relative path.
-replace github.com/specterops/bloodhound v0.0.0 => ../bloodhound-enterprise/bhce📝 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.
| replace github.com/specterops/bloodhound v0.0.0 => ../bloodhound-enterprise/bhce |
🤖 Prompt for AI Agents
In go.mod around line 7, the local replace directive pointing to a relative
enterprise path (replace github.com/specterops/bloodhound v0.0.0 =>
../bloodhound-enterprise/bhce) must be removed because it breaks external
consumers and CI; delete that replace line from go.mod, run go mod tidy to
update module files, and instead keep any local development overrides in a
dev-only go.work or private branch; if the project needs to depend on BloodHound
for integration, switch the dependency to a tagged module version hosted in a
public or private module proxy (or document the dev-only replacement in
contributor docs) and update CI to use the published module version.
| github.com/jackc/pgtype v1.14.4 | ||
| github.com/jackc/pgx/v5 v5.7.5 | ||
| github.com/neo4j/neo4j-go-driver/v5 v5.28.1 | ||
| github.com/specterops/bloodhound v0.0.0 |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Direct dependency on bloodhound introduces tight coupling
Adding github.com/specterops/bloodhound as a direct require drags in a large dependency tree and couples DAWGS to an application. The pg package should rely on a small interface (or callback) instead.
If you adopt the PGConnStringProvider interface in drivers/pg/pg.go, this direct require can be dropped.
1 participant
adds beforeConnect function to poolCfg. enables RDS iam token auth by requesting a new PostgresConnectionString prior to each new connection
Summary by CodeRabbit