Skip to content

feat: add workload identity federation support#49

Open
c1-dev-bot[bot] wants to merge 1 commit into
mainfrom
feat/workload-identity-federation
Open

feat: add workload identity federation support#49
c1-dev-bot[bot] wants to merge 1 commit into
mainfrom
feat/workload-identity-federation

Conversation

@c1-dev-bot

@c1-dev-bot c1-dev-bot Bot commented Jul 13, 2026

Copy link
Copy Markdown

Summary

Adds a second authentication method so the connector can authenticate to Databricks using workload identity federation (Databricks "OAuth Token Federation") instead of a static service-principal client secret.

The connector presents an externally-issued JWT (e.g. a SPIFFE JWT-SVID, Kubernetes projected ServiceAccount token, or any OIDC workload-identity token) and exchanges it for a short-lived Databricks OAuth token via RFC 8693 token exchange. The existing client-secret method remains and is the default.

New CLI flags / env vars

Flag Env var Description
--databricks-token-file BATON_DATABRICKS_TOKEN_FILE Path to a file containing the external JWT. Re-read on each token refresh to support credential rotation by sidecars (e.g. SPIFFE).
--databricks-token BATON_DATABRICKS_TOKEN Inline external JWT for short-lived / CI scenarios.

Key design decisions

  • Exactly one of --databricks-client-secret, --databricks-token-file, or --databricks-token must be provided (enforced via mutual-exclusion + at-least-one config constraints)
  • --databricks-client-id remains required — it identifies the service principal for both auth methods
  • Token file is re-read on each refresh cycle, so credential rotation by a sidecar (SPIFFE agent, K8s kubelet) is supported automatically
  • Uses oauth2.ReuseTokenSource for automatic caching of the Databricks access token between API calls
  • The token exchange hits the same OIDC endpoint as the existing client-credentials flow (/oidc/accounts/{id}/v1/token)

Files changed

  • pkg/config/config.go — New fields + constraints
  • pkg/config/conf.gen.go — Regenerated struct
  • pkg/databricks/auth.goTokenFederation auth type with RFC 8693 token exchange
  • pkg/connector/connector.go — Auth method selection based on config
  • config_schema.json — Regenerated

Fixes: CXH-2028

Test plan

  • Verify build succeeds (go build ./cmd/baton-databricks/)
  • Verify --help shows the new flags with correct descriptions
  • Test with --databricks-client-secret (existing flow) still works
  • Test with --databricks-token-file pointing to a valid JWT file against a Databricks account with a federation policy
  • Test with --databricks-token providing an inline JWT
  • Verify mutual exclusion: providing both --databricks-client-secret and --databricks-token-file errors
  • Verify at-least-one: providing neither secret nor token errors
  • Test token file re-read: verify that after the Databricks token expires, the file is re-read (supports rotation)

Automated PR Notice

This PR was automatically created by c1-dev-bot as a potential implementation.

This code requires:

  • Human review of the implementation approach
  • Manual testing to verify correctness
  • Approval from the appropriate team before merging

Add a second authentication method so the connector can authenticate
to Databricks using workload identity federation (RFC 8693 token
exchange) instead of a static client secret. The connector presents
an externally-issued JWT (e.g. SPIFFE JWT-SVID, Kubernetes projected
ServiceAccount token, or any OIDC workload-identity token) and
exchanges it for a short-lived Databricks OAuth token.

New CLI flags / env vars:
- --databricks-token-file / BATON_DATABRICKS_TOKEN_FILE
  Path to a file containing the external JWT. Re-read on each token
  refresh to support credential rotation by sidecars.
- --databricks-token / BATON_DATABRICKS_TOKEN
  Inline external JWT for short-lived / CI scenarios.

The existing client-secret method remains and is the default. Exactly
one of client-secret, token-file, or token must be provided (enforced
via mutual-exclusion and at-least-one config constraints).

Fixes: CXH-2028
@c1-dev-bot
c1-dev-bot Bot requested a review from a team July 13, 2026 14:55
@linear-code

linear-code Bot commented Jul 13, 2026

Copy link
Copy Markdown

CXH-2028

Comment thread pkg/databricks/auth.go
"scope": {"all-apis"},
}

resp, err := http.PostForm(tf.tokenURL, form)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Suggestion: http.PostForm uses http.DefaultClient, which has no timeout and ignores the caller's context. A slow or hung Databricks token endpoint will block the sync indefinitely, and context cancellation during token exchange is not honored. Consider building the request with http.NewRequestWithContext(ctx, ...) on an http.Client with an explicit timeout. Note Token() currently has no ctx param — plumbing ctx from GetClient (via a bound token source) would also address R11 context propagation. (confidence: medium)

Comment thread pkg/databricks/auth.go
Comment on lines +138 to +140
if tokenResp.ExpiresIn > 0 {
token.Expiry = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Suggestion: If the token endpoint omits expires_in (ExpiresIn == 0), token.Expiry stays the zero value, which oauth2 treats as never-expiring. ReuseTokenSource would then never refresh, so the token file is never re-read — defeating the credential-rotation support described in the flag docs. Consider applying a conservative default TTL when expires_in is absent. (confidence: low)

@github-actions

Copy link
Copy Markdown
Contributor

Connector PR Review: feat: add workload identity federation support

Blocking Issues: 0 | Suggestions: 2 | Threads Resolved: 0
Criteria: Criteria status: loaded .claude/skills/ci-review.md from trusted base e6d89a38400f.
Review mode: full
View review run

Review Summary

Scanned the full PR diff for security and correctness. This adds a second auth method (TokenFederation) implementing RFC 8693 token exchange, gated behind --databricks-token/--databricks-token-file, with mutual-exclusion plus at-least-one config constraints so exactly one credential is required. The change is backward compatible: the existing client-secret flow remains the default and existing configs keep working (the now-optional databricks-client-secret is covered by the at-least-one constraint). SDK constraint helpers verified present in the vendored SDK; no go.mod/go.sum changes. No blocking issues found, only two non-blocking robustness suggestions on the token-exchange path.

Security Issues

None found.

Correctness Issues

None found.

Suggestions

  • pkg/databricks/auth.go:110http.PostForm uses http.DefaultClient (no timeout, ignores context); a hung token endpoint blocks the sync and cancellation is not honored. (confidence: medium)
  • pkg/databricks/auth.go:138 — if the endpoint omits expires_in, token.Expiry is zero and oauth2 treats it as never-expiring, so ReuseTokenSource never refreshes and the token file is never re-read, defeating the documented rotation support. (confidence: low)
Prompt for AI agents
Verify each finding against the current code and only fix it if needed.

## Suggestions

In pkg/databricks/auth.go:
- Around line 110: TokenFederation.Token() calls http.PostForm(tf.tokenURL, form), which uses http.DefaultClient (no timeout) and does not propagate a context. Build the request with http.NewRequestWithContext(ctx, http.MethodPost, tf.tokenURL, strings.NewReader(form.Encode())), set the Content-Type application/x-www-form-urlencoded header, and execute it on an http.Client with an explicit Timeout. Since Token() has no ctx parameter, thread a context through from GetClient(ctx) so cancellation and timeouts are honored during token exchange.
- Around line 138-140: token.Expiry is only set when tokenResp.ExpiresIn is greater than 0. If the endpoint omits expires_in, Expiry stays the zero value and oauth2 treats the token as never expiring, so ReuseTokenSource never refreshes and the token file is never re-read, breaking the rotation behavior promised by --databricks-token-file. Apply a conservative default expiry when ExpiresIn is 0 or less.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No blocking issues found.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

0 participants