-
Notifications
You must be signed in to change notification settings - Fork 4
add beforeConnect to NewPool poolcfg [DRAFT] #12
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -8,6 +8,7 @@ import ( | |
|
|
||
| "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" | ||
|
|
@@ -50,15 +51,12 @@ func afterPooledConnectionRelease(conn *pgx.Conn) bool { | |
| return true | ||
| } | ||
|
|
||
| func NewPool(connectionString string) (*pgxpool.Pool, error) { | ||
| 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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 💡 Verification agent 🧩 Analysis chainPublic 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 • 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 • 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 |
||
|
|
||
| poolCtx, done := context.WithTimeout(context.Background(), poolInitConnectionTimeout) | ||
| defer done() | ||
|
|
||
| poolCfg, err := pgxpool.ParseConfig(connectionString) | ||
| poolCfg, err := pgxpool.ParseConfig(cfg.PostgreSQLConnectionString()) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
@@ -73,6 +71,17 @@ func NewPool(connectionString string) (*pgxpool.Pool, error) { | |
| poolCfg.AfterConnect = afterPooledConnectionEstablished | ||
| poolCfg.AfterRelease = afterPooledConnectionRelease | ||
|
|
||
| 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 | ||
| } | ||
|
|
||
| pool, err := pgxpool.NewWithConfig(poolCtx, poolCfg) | ||
| if err != nil { | ||
| return nil, err | ||
|
|
||
| Original file line number | Diff line number | Diff line change | ||
|---|---|---|---|---|
| @@ -1,6 +1,10 @@ | ||||
| module github.com/specterops/dawgs | ||||
|
|
||||
| go 1.23.0 | ||||
| go 1.24.4 | ||||
|
|
||||
| toolchain go1.24.6 | ||||
|
|
||||
| replace github.com/specterops/bloodhound v0.0.0 => ../bloodhound-enterprise/bhce | ||||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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:
-replace github.com/specterops/bloodhound v0.0.0 => ../bloodhound-enterprise/bhce📝 Committable suggestion
Suggested change
🤖 Prompt for AI Agents |
||||
|
|
||||
| require ( | ||||
| cuelang.org/go v0.13.2 | ||||
|
|
@@ -11,27 +15,47 @@ require ( | |||
| 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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🛠️ 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. |
||||
| github.com/stretchr/testify v1.10.0 | ||||
| ) | ||||
|
|
||||
| require ( | ||||
| github.com/aws/aws-sdk-go-v2 v1.38.0 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/config v1.31.1 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/credentials v1.18.5 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.3 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.3 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.3 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.3 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.0 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.3 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/service/sso v1.28.1 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/service/ssooidc v1.33.1 // indirect | ||||
| github.com/aws/aws-sdk-go-v2/service/sts v1.37.1 // indirect | ||||
| github.com/aws/smithy-go v1.22.5 // indirect | ||||
| github.com/bits-and-blooms/bitset v1.22.0 // indirect | ||||
| github.com/cockroachdb/apd/v3 v3.2.1 // indirect | ||||
| github.com/davecgh/go-spew v1.1.1 // indirect | ||||
| github.com/dgryski/go-metro v0.0.0-20250106013310-edb8663e5e33 // indirect | ||||
| github.com/go-ole/go-ole v1.2.6 // indirect | ||||
| github.com/jackc/pgio v1.0.0 // indirect | ||||
| github.com/jackc/pgpassfile v1.0.0 // indirect | ||||
| github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect | ||||
| github.com/jackc/puddle/v2 v2.2.2 // indirect | ||||
| github.com/kamstrup/intmap v0.5.1 // indirect | ||||
| github.com/lib/pq v1.10.9 // indirect | ||||
| github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect | ||||
| github.com/mschoch/smat v0.2.0 // indirect | ||||
| github.com/pmezard/go-difflib v1.0.0 // indirect | ||||
| github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect | ||||
| github.com/shirou/gopsutil/v3 v3.24.5 // indirect | ||||
| github.com/tklauser/go-sysconf v0.3.12 // indirect | ||||
| github.com/tklauser/numcpus v0.6.1 // indirect | ||||
| github.com/yusufpapurcu/wmi v1.2.4 // indirect | ||||
| golang.org/x/crypto v0.39.0 // indirect | ||||
| golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b // indirect | ||||
| golang.org/x/oauth2 v0.30.0 // indirect | ||||
| golang.org/x/sync v0.15.0 // indirect | ||||
| golang.org/x/sys v0.33.0 // indirect | ||||
| golang.org/x/text v0.26.0 // indirect | ||||
| golang.org/x/tools v0.34.0 // indirect | ||||
| gopkg.in/yaml.v3 v3.0.1 // indirect | ||||
| ) | ||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🛠️ 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:
Add this type near the top of the file (outside the selected range):
🤖 Prompt for AI Agents