From 7e90184c7bb898953b8c234ccd8965e76fb2a450 Mon Sep 17 00:00:00 2001 From: Andrew Cremins Date: Mon, 22 Dec 2025 16:52:53 -0800 Subject: [PATCH 1/4] init --- pgutils/connector.go | 62 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/pgutils/connector.go b/pgutils/connector.go index 21dce91..f757874 100644 --- a/pgutils/connector.go +++ b/pgutils/connector.go @@ -6,6 +6,8 @@ import ( "fmt" "log" "net/url" + "slices" + "strings" "time" "database/sql" @@ -20,6 +22,8 @@ import ( "github.com/lib/pq" ) +const defaultPostgresPort = "5432" + type baseConnectionStringProvider interface { getBaseConnectionString(ctx context.Context) (string, error) } @@ -212,3 +216,61 @@ func MustConnectDB(conn *PostgresqlConnector) *sqlx.DB { return db } +func NewPostgresqlConnectorFromDSN(ctx context.Context, dsn string) (*PostgresqlConnector, error) { + u, err := url.Parse(dsn) + if err != nil || u.Scheme != "postgres+rds-iam" { + // Not our custom scheme: hand off to existing DSN handling. + return NewPostgresqlConnectorFromConnectionString(dsn), nil + } + + user := "" + if u.User != nil { + user = u.User.Username() + if _, hasPw := u.User.Password(); hasPw { + return nil, fmt.Errorf("postgres+rds-iam DSN must not include a password") + } + } + if user == "" { + return nil, fmt.Errorf("postgres+rds-iam DSN missing username") + } + + host := u.Hostname() + if host == "" { + return nil, fmt.Errorf("postgres+rds-iam DSN missing host") + } + + port := u.Port() + if port == "" { + port = defaultPostgresPort + } + + // Match libpq/psql defaulting: if dbname isn't specified, dbname defaults to username. + dbName := strings.TrimPrefix(u.Path, "/") + if dbName == "" { + dbName = user + } + + q := u.Query() + supportedParams := []string{"assume_role_arn", "assume_role_session_name"} + for k := range q { + if !slices.Contains(supportedParams, k) { + return nil, fmt.Errorf("postgres+rds-iam DSN has unsupported query parameter: %s", k) + } + } + + assumeRoleARN := q.Get("assume_role_arn") + assumeRoleSessionName := q.Get("assume_role_session_name") + + cfg := &IAMAuthConfig{ + RDSEndpoint: host + ":" + port, + User: user, + Database: dbName, + } + + if assumeRoleARN != "" { + cfg.AssumeRoleARN = assumeRoleARN + cfg.AssumeRoleSessionName = assumeRoleSessionName + } + + return NewPostgresqlConnectorWithIAMAuth(ctx, cfg) +} From fe2a24e573261cad057f0bc044679cb364e33fef Mon Sep 17 00:00:00 2001 From: Andrew Cremins Date: Tue, 23 Dec 2025 02:47:08 -0800 Subject: [PATCH 2/4] Error handling --- pgutils/connector.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pgutils/connector.go b/pgutils/connector.go index f757874..da1c603 100644 --- a/pgutils/connector.go +++ b/pgutils/connector.go @@ -218,7 +218,11 @@ func MustConnectDB(conn *PostgresqlConnector) *sqlx.DB { func NewPostgresqlConnectorFromDSN(ctx context.Context, dsn string) (*PostgresqlConnector, error) { u, err := url.Parse(dsn) - if err != nil || u.Scheme != "postgres+rds-iam" { + if err != nil { + return nil, fmt.Errorf("filed to parse DSN: %w", err) + } + + if u.Scheme != "postgres+rds-iam" { // Not our custom scheme: hand off to existing DSN handling. return NewPostgresqlConnectorFromConnectionString(dsn), nil } From ddb9337edb94125e82631b5b883e936bc5fb122d Mon Sep 17 00:00:00 2001 From: Andrew Cremins Date: Tue, 23 Dec 2025 03:23:10 -0800 Subject: [PATCH 3/4] Some comments --- pgutils/connector.go | 27 +++++++++++++++++++-------- 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/pgutils/connector.go b/pgutils/connector.go index da1c603..71a10c0 100644 --- a/pgutils/connector.go +++ b/pgutils/connector.go @@ -216,14 +216,27 @@ func MustConnectDB(conn *PostgresqlConnector) *sqlx.DB { return db } +// NewPostgresqlConnectorFromDSN constructs a PostgresqlConnector from either a normal +// Postgres DSN/connection string or the custom postgres+rds-iam DSN used for RDS IAM auth. +// +// IAM example 1: postgres+rds-iam://@[:]/ +// +// Optional query params (for cross-account IAM): +// - assume_role_arn: role ARN to assume. +// - assume_role_session_name: only used when assume_role_arn is set; defaults to "pgutils-rds-iam" if omitted. +// +// IAM example 2: postgres+rds-iam://@[:]/?assume_role_arn=...&assume_role_session_name=... func NewPostgresqlConnectorFromDSN(ctx context.Context, dsn string) (*PostgresqlConnector, error) { + if dsn == "" { + return nil, errors.New("DSN cannot be empty") + } + u, err := url.Parse(dsn) if err != nil { - return nil, fmt.Errorf("filed to parse DSN: %w", err) + return nil, fmt.Errorf("failed to parse DSN: %w", err) } - if u.Scheme != "postgres+rds-iam" { - // Not our custom scheme: hand off to existing DSN handling. + if u.Scheme != "postgres+rds-iam" { // Not our custom scheme: hand off to existing DSN handling. return NewPostgresqlConnectorFromConnectionString(dsn), nil } @@ -255,25 +268,23 @@ func NewPostgresqlConnectorFromDSN(ctx context.Context, dsn string) (*Postgresql } q := u.Query() - supportedParams := []string{"assume_role_arn", "assume_role_session_name"} + supportedParams := []string{"assume_role_arn", "assume_role_session_name", "assume_role_external_id", "assume_role_duration"} for k := range q { if !slices.Contains(supportedParams, k) { return nil, fmt.Errorf("postgres+rds-iam DSN has unsupported query parameter: %s", k) } } - assumeRoleARN := q.Get("assume_role_arn") - assumeRoleSessionName := q.Get("assume_role_session_name") - cfg := &IAMAuthConfig{ RDSEndpoint: host + ":" + port, User: user, Database: dbName, } + assumeRoleARN := q.Get("assume_role_arn") if assumeRoleARN != "" { cfg.AssumeRoleARN = assumeRoleARN - cfg.AssumeRoleSessionName = assumeRoleSessionName + cfg.AssumeRoleSessionName = q.Get("assume_role_session_name") } return NewPostgresqlConnectorWithIAMAuth(ctx, cfg) From 3cd2b7f8a728bf7120fecbd19677e00c1f616871 Mon Sep 17 00:00:00 2001 From: Andrew Cremins Date: Tue, 23 Dec 2025 03:30:39 -0800 Subject: [PATCH 4/4] Remove unsupported (for now) parameters --- pgutils/connector.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgutils/connector.go b/pgutils/connector.go index 71a10c0..9bdf605 100644 --- a/pgutils/connector.go +++ b/pgutils/connector.go @@ -268,7 +268,7 @@ func NewPostgresqlConnectorFromDSN(ctx context.Context, dsn string) (*Postgresql } q := u.Query() - supportedParams := []string{"assume_role_arn", "assume_role_session_name", "assume_role_external_id", "assume_role_duration"} + supportedParams := []string{"assume_role_arn", "assume_role_session_name"} for k := range q { if !slices.Contains(supportedParams, k) { return nil, fmt.Errorf("postgres+rds-iam DSN has unsupported query parameter: %s", k)