feat(worker): add AWS IAM flags to worker command for RDS IAM auth#1557
feat(worker): add AWS IAM flags to worker command for RDS IAM auth#1557tennyenoch wants to merge 10 commits into
Conversation
The worker command was missing iam.AddFlags() which provides the --aws-region, --aws-access-key-id, --aws-secret-access-key, --aws-session-token, --aws-profile and --aws-role-arn flags. Without these flags, users running the standalone worker binary (e.g. on AWS ECS Fargate without the Kubernetes operator) could not configure AWS RDS IAM authentication for the primary PostgreSQL store via CLI flags. The primary store already supports IAM auth via the --postgres-aws-enable-iam flag from connect.AddFlags, but the AWS credential/region flags were absent on the worker command. The serve command already calls both connect.AddFlags and iam.AddFlags; this commit brings the worker command to parity. Closes formancehq#1556
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe primary PostgreSQL connection path now supports AWS IAM authentication through assumed-role RDS tokens. Worker and serve commands use the shared connection-option helper, and tests verify IAM flag registration and defaults. ChangesPrimary PostgreSQL AWS IAM authentication
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant connectionOptionsFromFlags
participant AssumeRoleProvider
participant roleIAMConnector
participant PostgreSQL
CLI->>connectionOptionsFromFlags: IAM flags and PostgreSQL URI
connectionOptionsFromFlags->>AssumeRoleProvider: Configure assumed-role credentials
connectionOptionsFromFlags->>roleIAMConnector: Create connector
roleIAMConnector->>AssumeRoleProvider: Generate RDS IAM auth token
roleIAMConnector->>PostgreSQL: Connect with token as password
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
🛑 Changes requested — automated reviewThe PR adds AWS IAM flag parity to the worker command. The prior discussion shows that many earlier findings (credentials caching, DSN override, go-libs safeguards bypass, go.mod indirect entries, import ordering, trailing blank lines) were raised in resolved/outdated threads and subsequently addressed in later commits. Two concerns remain active: (1) The |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cmd/worker_test.go (2)
66-89: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider verifying the
postgres-aws-enable-iamdefault on serve too.The worker test asserts
DefValue == "false"forpostgres-aws-enable-iam, but the serve test only checks registration. Adding the same default-value assertion on serve would guard against an accidental insecure default creeping in on either command.♻️ Optional: add default-value check to serve test
func TestServeCommandHasAWSIAMFlags(t *testing.T) { t.Parallel() cmd := NewServeCommand() flags := cmd.Flags() + // Verify the IAM-enable flag defaults to false on serve as well + f := flags.Lookup(connect.PostgresAWSEnableIAMFlag) + require.NotNil(t, f, "--postgres-aws-enable-iam flag must be registered on the serve command") + assert.Equal(t, "false", f.DefValue) + iamFlags := []string{ connect.PostgresAWSEnableIAMFlag, iam.AWSRegionFlag, iam.AWSAccessKeyIDFlag, iam.AWSSecretAccessKeyFlag, iam.AWSSessionTokenFlag, iam.AWSRoleArnFlag, }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cmd/worker_test.go` around lines 66 - 89, Add an assertion in TestServeCommandHasAWSIAMFlags for connect.PostgresAWSEnableIAMFlag that verifies the registered flag’s DefValue is "false", matching the worker test and preserving the secure default while retaining the existing registration checks.
81-88: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the loop-variable shadowing. The project targets Go 1.26, so
flagName := flagNameis unnecessary here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cmd/worker_test.go` around lines 81 - 88, Remove the redundant flagName := flagName shadowing inside the iamFlags loop in the serve registration test, relying on Go 1.26 loop-variable semantics while preserving the existing parallel subtests and assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@cmd/worker_test.go`:
- Around line 66-89: Add an assertion in TestServeCommandHasAWSIAMFlags for
connect.PostgresAWSEnableIAMFlag that verifies the registered flag’s DefValue is
"false", matching the worker test and preserving the secure default while
retaining the existing registration checks.
- Around line 81-88: Remove the redundant flagName := flagName shadowing inside
the iamFlags loop in the serve registration test, relying on Go 1.26
loop-variable semantics while preserving the existing parallel subtests and
assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 31d14b3f-9c91-434f-9473-8be9c5e8b687
📒 Files selected for processing (2)
cmd/worker.gocmd/worker_test.go
- Add DefValue=false assertion for postgres-aws-enable-iam on serve command to guard against accidental insecure default, matching the worker test - Remove redundant flagName := flagName loop variable shadowing; Go 1.26 loop semantics make this unnecessary
|
Re-reviewed current head Findings:
The PR description says
Checks run:
|
- Add --aws-profile assertion to both worker and serve tests; dropping that flag would previously have gone undetected - Add comment on --aws-role-arn explaining it is registered by iam.AddFlags but not currently consumed by iam.LoadOptionFromFlags in go-libs (pre-existing on serve, newly noted here) - Remove trailing blank line at EOF (git diff --check)
|
@gfyrag - Fixed the review comments. Please review. |
NumaryBot
left a comment
NumaryBot
left a comment
There was a problem hiding this comment.
🛑 Changes requested — automated review
The PR wires AWS IAM flags onto the worker Cobra command to mirror the existing serve command pattern, and adds tests verifying flag registration and defaults. The overall approach is sound and consistent with the codebase. However, one major issue was identified: the --aws-role-arn flag is registered and exposed to operators but is never consumed by the underlying iam.LoadOptionFromFlags helper, so passing it has no effect — this is a silent no-op that can mislead users configuring RDS IAM role assumption. Two lesser issues were also noted: the new tests do not cover --aws-profile registration (meaning that flag could be silently removed without breaking tests), and there is a trailing blank line in the test file that fails whitespace checks.
Findings outside the diff
🟠 [major] --aws-role-arn flag is registered but never consumed for RDS IAM auth — cmd/worker.go
connect.ConnectionOptionsFromFlags uses iam.LoadOptionFromFlags to build AWS credentials, and the current go-libs version of that helper reads region, access key, secret, session token, and profile — but not role ARN. A worker started with --postgres-aws-enable-iam --aws-role-arn <arn> will silently ignore the ARN and fall back to the default credential chain instead of assuming the intended role. The flag is therefore a no-op that misleads operators.
Suggestion: Either (a) update iam.LoadOptionFromFlags (and the corresponding go-libs version) to support role assumption via the ARN, or (b) remove the --aws-role-arn flag from the worker command until the underlying implementation can honour it, and document the limitation.
🟡 [minor] New tests do not assert --aws-profile flag registration — cmd/worker_test.go
iam.AddFlags registers --aws-profile as part of the IAM flag set, and the PR description states that full IAM flag parity with serve is the goal. However, the new regression tests do not assert the presence of --aws-profile for either the worker or serve commands, meaning the flag could be silently dropped without causing test failures.
Suggestion: Add an assertion for --aws-profile (and any other flags registered by iam.AddFlags) alongside the existing flag checks in the worker and serve test cases.
⚪ [nit] Extra blank line at EOF causes git diff --check failure — cmd/worker_test.go
There is a trailing blank line at the end of cmd/worker_test.go that causes git diff --check to report a whitespace error.
Suggestion: Remove the extra blank line at the end of the file.
Summary: #1557 (comment)
NumaryBot
left a comment
NumaryBot
left a comment
There was a problem hiding this comment.
NumaryBot posted 1 new inline finding.
Summary: #1557 (comment)
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cmd/aws.go`:
- Around line 26-30: Update roleIAMConnector to store an
aws.CredentialsProvider, wrap the stscreds.NewAssumeRoleProvider result with
aws.NewCredentialsCache when initializing the connector, and ensure connection
creation reuses this cached provider instead of invoking the raw assumed-role
provider directly.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9fe94ea0-356e-464d-b9d8-3e4817758716
📒 Files selected for processing (4)
cmd/aws.gocmd/serve.gocmd/worker.gocmd/worker_test.go
🚧 Files skipped from review as they are similar to previous changes (2)
- cmd/worker.go
- cmd/worker_test.go
…id STS call per connection
NumaryBot
left a comment
NumaryBot
left a comment
There was a problem hiding this comment.
NumaryBot posted 3 new inline findings.
Summary: #1557 (comment)
NumaryBot
left a comment
NumaryBot
left a comment
There was a problem hiding this comment.
NumaryBot posted 5 new inline findings.
Summary: #1557 (comment)
NumaryBot
left a comment
NumaryBot
left a comment
There was a problem hiding this comment.
NumaryBot posted 5 new inline findings.
Summary: #1557 (comment)
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cmd/aws.go (1)
132-138: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winPrevent nil-pointer panic when the database URI lacks a username.
If the
--postgres-uridoes not contain user information (e.g.,postgres://db.cluster.aws.com/mydb),u.URL.Userwill benil. Callingu.URL.User.Username()directly on it will cause a process-crashing panic during connection establishment.Verify that the user information is present before extracting the username.
🐛 Proposed fix to validate the username
+ if u.URL.User == nil || u.URL.User.Username() == "" { + return nil, errors.New("aws: missing username in postgres uri for iam authentication") + } + authToken, err := auth.BuildAuthToken( ctx, u.Host, c.region, u.URL.User.Username(), c.creds, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cmd/aws.go` around lines 132 - 138, Update the auth.BuildAuthToken call in the connection-establishment flow to validate that u.URL.User is non-nil before calling Username(). Handle the missing-user case without panicking, using the surrounding error-handling path to report the invalid PostgreSQL URI and stop token construction.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@cmd/aws.go`:
- Around line 132-138: Update the auth.BuildAuthToken call in the
connection-establishment flow to validate that u.URL.User is non-nil before
calling Username(). Handle the missing-user case without panicking, using the
surrounding error-handling path to report the invalid PostgreSQL URI and stop
token construction.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 6bb3b8d2-fa20-436b-b3bc-87e944038c68
⛔ Files ignored due to path filters (1)
go.modis excluded by!**/*.mod
📒 Files selected for processing (1)
cmd/aws.go
gfyrag
left a comment
gfyrag
left a comment
There was a problem hiding this comment.
Thanks for this! The 2-line worker fix for #1556 is spot on. My main concern is about where the AssumeRole logic should live, plus a couple of concrete bugs in the reimplemented connector.
Architecture — the AssumeRole machinery belongs in go-libs, not in cmd/
cmd/aws.go copy-pastes go-libs internals: roleIAMPgxTracer ≈ connect.pgxTracer, validateReadWrite ≈ validateConnectTargetSessionAttrsReadWrite, resetReadOnly ≈ resetReadOnlySession, and roleIAMConnector ≈ connect.iamConnector. These are near-verbatim duplicates of unexported functions. They will drift from go-libs over time — and they already have (see the missing deadline below).
The role-ARN support should be added on the go-libs side (iam.LoadOptionFromFlags / connect.GetAWSIAMAuthConnector), where the connector it mirrors actually lives. That way every consumer benefits and there's a single implementation to maintain.
We're fine with a single ledger PR that bumps go-libs + keeps the 2-line worker fix, but the connector/AssumeRole logic itself must move to go-libs in all cases — it shouldn't be reimplemented here.
Scope vs. description
The PR is described as "add AWS IAM flags to the worker command", but it actually introduces a full STS AssumeRole DB connector and swaps connect.ConnectionOptionsFromFlags for a local wrapper in both serve and worker. Worth making that explicit (title/changelog), because it changes auth behavior — see next point.
Concrete issues
1. Silent behavior change for existing serve deployments. serve already registers --aws-role-arn today (via iam.AddFlags), but go-libs ignores it — iam.LoadOptionFromFlags never reads AWSRoleArnFlag. After this PR, that same flag now triggers a real STS AssumeRole. Any existing serve operator who has --aws-role-arn set (it was a no-op until now) will see DB auth switch from caller-identity to assumed-role, breaking connections if the role isn't assumable or lacks rds-db:connect. Needs to be called out in the changelog at minimum.
2. resetReadOnly has no deadline (robustness bug). go-libs wraps its show transaction_read_only probes in withPgConnDeadline (5s). The copied resetReadOnly has no deadline, yet it runs on every connection reuse via OptionResetSession. On a half-open connection this can block indefinitely and stall the goroutine returning the conn to the pool. Note validateReadWrite did keep the deadline — so the two paths already diverged. (This is exactly the drift argument for moving it to go-libs.)
3. ValidateConnect is set unconditionally. go-libs only sets it when nil (if config.ValidateConnect == nil), preserving a target_session_attrs coming from the DSN. The copy overwrites it. Low impact, but another behavioral divergence.
Verified OK
No duplicate flag registration (no pflag panic), all 6 ConnectionOptions fields repopulated, RDS token regenerated per Connect, CredentialsCache avoids per-connection STS calls, search_path DSN correctly propagated, go.mod direct/indirect changes consistent.
Suggestion
Either (a) split: land just the 2-line worker flag fix to close #1556, and move AssumeRole to a go-libs PR; or (b) single PR that bumps go-libs (with AssumeRole implemented there) + the 2-line fix. In both cases the connector logic lives in go-libs, and points 1–3 get addressed there.
NumaryBot
left a comment
NumaryBot
left a comment
There was a problem hiding this comment.
NumaryBot posted 1 new inline finding.
Summary: #1557 (comment)
NumaryBot
left a comment
NumaryBot
left a comment
There was a problem hiding this comment.
NumaryBot posted 1 new inline finding.
Summary: #1557 (comment)
NumaryBot
left a comment
NumaryBot
left a comment
There was a problem hiding this comment.
NumaryBot posted 2 new inline findings.
Summary: #1557 (comment)
| connect.AddFlags(cmd.Flags()) | ||
| metrics.AddFlags(cmd.Flags()) | ||
| traces.AddFlags(cmd.Flags()) | ||
| iam.AddFlags(cmd.Flags()) |
There was a problem hiding this comment.
🟠 [major] --aws-role-arn flag is registered but silently ignored in the Postgres IAM connection path
reported by codex, gfyrag, NumaryBot
The worker command registers --aws-role-arn (alongside --postgres-aws-enable-iam), but connect.ConnectionOptionsFromFlags loads AWS config via iam.LoadOptionFromFlags, which reads region/static keys/session/profile but not role ARN. A worker started with --postgres-aws-enable-iam --aws-role-arn will not assume the requested role; it silently falls back to the default/static credential chain. MarkHidden (if used) only hides the flag from help output — Cobra still parses and accepts it — so the no-op behaviour is invisible to operators.
Suggestion: Either (a) make the connection path consume AWSRoleArnFlag by wrapping the loaded credentials with stscreds.NewAssumeRoleProvider + aws.NewCredentialsCache when the flag is non-empty, (b) do not register --aws-role-arn on the worker command at all until the underlying library supports it and document the limitation clearly, or (c) add an explicit pre-run validation that returns a clear error when --aws-role-arn is set alongside --postgres-aws-enable-iam so operators are not silently misled.
| // the worker would fall back to the default credential chain instead | ||
| // of the intended role. Return an explicit error rather than | ||
| // misleading the operator. | ||
| if iamEnabled, _ := cmd.Flags().GetBool(connect.PostgresAWSEnableIAMFlag); iamEnabled { |
There was a problem hiding this comment.
🟠 [major] Registering --aws-role-arn breaks IRSA/web-identity environments via env-var binding
service.Execute populates every registered flag from its uppercase environment variable equivalent, so adding iam.AddFlags (which registers --aws-role-arn) causes AWS_ROLE_ARN to be bound to that flag. In IRSA or web-identity environments where AWS_ROLE_ARN is legitimately set for the AWS SDK's standard credential chain alongside POSTGRES_AWS_ENABLE_IAM=true, any validation check that rejects or errors on --aws-role-arn will abort startup even though the operator never intended to pass this flag explicitly.
Suggestion: Avoid registering --aws-role-arn on the worker command unless the connection path actually consumes it. If the flag must be registered, ensure that the presence of AWS_ROLE_ARN via environment variable does not cause an error or unexpected behaviour for operators using standard SDK credential flows.
3 participants
What
The
workercommand was missingiam.AddFlags()which provides the--aws-region,--aws-access-key-id,--aws-secret-access-key,--aws-session-token,--aws-profileand--aws-role-arnflags.Why
Without these flags, users running the standalone worker binary (e.g.
on AWS ECS Fargate without the Kubernetes operator) could not configure
AWS RDS IAM authentication for the primary PostgreSQL store via CLI flags.
The primary store already supports IAM auth via the
--postgres-aws-enable-iamflag from
connect.AddFlags, but the AWS credential/region flags were absenton the worker command.
The
servecommand already calls bothconnect.AddFlagsandiam.AddFlags;this commit brings the
workercommand to parity.Testing
Added unit tests in
cmd/worker_test.goverifying that all AWS IAM flags(
--postgres-aws-enable-iam,--aws-region,--aws-access-key-id,--aws-secret-access-key,--aws-session-token,--aws-role-arn) areregistered on both the
workerandservecommands (12 tests, all passing).Closes #1556