flyteorg / flyteorg/flyte

[Core feature] Add OAuth M2M and OIDC Workload Identity Federation Authentication Support to Databricks Connector

Open
#7,319 1 comment 0 reactions 0 assignees View on GitHub
flyte
Dominant language
Go
Stars
7.5k
Forks
886
Avg merge
1d 14h
Merged PRs (30d)
120

Description

### Motivation: Why do you think this is important?

Extend the Databricks connector (`flytekitplugins-spark`) with **OAuth Machine-to-Machine (M2M)** and **OIDC Workload Identity Federation** authentication, the two non-legacy authentication paths Databricks now recommends for service-to-service traffic. The connector currently only supports Personal Access Tokens (PATs).

## Motivation

### Databricks has marked PAT a legacy auth method

Databricks now classifies workspace-level PATs as **legacy** in their official documentation:

> "Where possible, Databricks recommends using OAuth instead of PATs for user account authentication because OAuth provides stronger security." Source: [Authenticate with Databricks personal access tokens (legacy)](https://docs.databricks.com/aws/en/dev-tools/auth/pat)

Operational implications of staying on PAT:

* **90-day auto-revocation**: "Databricks automatically revokes PATs that haven't been used for 90 days."
* **Per-workspace cap**: "A user can create up to 600 PATs per workspace."
* **No central rotation**: PATs can't be rotated through an identity provider; every secret is a separate operational chore.
* **Weaker auditability**: PATs are tied to a single workspace user/SP, with no IdP-level lineage on each call.

### Goal of this issue

Bring the Databricks connector into the modern Databricks auth posture (**OAuth M2M (client-credentials)** and **OIDC Workload Identity Federation (token exchange)**) while preserving everything that the prior multi-tenant PAT work in [flyteorg/flytekit#3394](https://github.com/flyteorg/flytekit/pull/3394) delivered:

* Per-namespace identity isolation (each Flyte workflow project federates as a different Databricks Service Principal, enabling per-project Unity Catalog access controls).
* Backwards-compatible fallback to existing PAT setups.
* **Zero workflow-code changes** for adoption: operators flip auth modes in connector config, workflow authors don't change a line.

### Related prior work

* [flyteorg/flyte#6911](https://github.com/flyteorg/flyte/issues/6911): Databricks Serverless Compute support (separate track, also in flight).
* [flyteorg/flytekit#3394](https://github.com/flyteorg/flytekit/pull/3394): Multi-tenant Databricks PAT via cross-namespace K8s secrets (merged Mar 10, 2026). This PR is the multi-tenancy baseline that the new auth modes preserve.
* [flyteorg/flytekit#3392](https://github.com/flyteorg/flytekit/pull/3392): Databricks Serverless support (merged earlier, related task-config surface).

## Proposed Changes

### 1. Auth strategy abstraction

Introduce a small strategy module (`databricks_auth.py`) inside the spark plugin that owns:

* Resolution of the active auth type from connector env vars or per-task config (with auto-detection when unset).
* Token acquisition for each strategy (PAT, OAuth M2M, OIDC Model 1, OIDC Model 2).
* Async-safe in-memory token cache with TTL and a pre-expiry refresh buffer.
* Exponential backoff with jitter on token-endpoint calls.

The connector continues to call a single `get_header(...)` boundary; the strategy underneath is replaceable.

### 2. Auth modes

| Auth mode | Credentials | Multi-tenancy granularity | Fallback target |
| --- | --- | --- | --- |
| **PAT** (existing) | K8s `Secret` `databricks-token` in workflow namespace | Per-namespace (per #3394) | `FLYTE_DATABRICKS_ACCESS_TOKEN` env var |
| **OAuth M2M** | K8s `Secret` `databricks-oauth` (`client_id` + `client_secret`) in workflow namespace | Per-namespace | Connector env vars `FLYTE_DATABRICKS_CLIENT_ID` / `FLYTE_DATABRICKS_CLIENT_SECRET` |
| **OIDC Model 1** (Connector-pod identity) | Connector pod's IRSA-projected JWT + connector-level `DATABRICKS_CLIENT_ID` | Single shared Databricks SP across all workflows | n/a |
| **OIDC Model 2** (Per-namespace ServiceAccount) | K8s `TokenRequest` minted for an annotated SA in the workflow namespace; SA carries `flyte.org/databricks-client-id` annotation | Per-namespace (each namespace can federate to a distinct Databricks SP) | OIDC Model 1 if no annotated SA found and Model 1 is configured; otherwise fail loudly |

### 3. Auto-detection order

If `FLYTE_DATABRICKS_AUTH_TYPE` is unset, the connector picks the strongest reachable mode at submit time, in this order:

1. **OIDC** if Model 2 SA discoverable in workflow namespace, OR Model 1 prerequisites are met.
2. **M2M** if `client_id`/`client_secret` reachable (namespace secret or connector env).
3. **PAT** as the final fallback for backwards compatibility.

If `FLYTE_DATABRICKS_AUTH_TYPE` is set explicitly, the connector uses that mode and errors loudly when its prerequisites are missing; no silent identity downgrade.

### 4. OIDC Model 2 discovery (per-namespace tenancy without per-task config)

Operators label/annotate ServiceAccounts in workflow namespaces:

```yaml
apiVersion: v1
kind: ServiceAccount
metadata:
name: my-project-databricks
namespace: my-project-namespace
labels:
flyte.org/databricks-enabled: "true"
annotations:
flyte.org/databricks-client-id: "00000000-0000-0000-0000-000000000000" # Databricks SP application ID
flyte.org/databricks-audience: "https://my-account.cloud.databricks.com/oidc/v1/token" # optional
```

At submit time, the connector lists SAs in the workflow's namespace by label, picks the one carrying the `flyte.org/databricks-client-id` annotation, and mints a JWT for it via the Kubernetes `TokenRequest` API. The result is exchanged at the Databricks token endpoint for a workspace access token. **Workflow authors write no extra config.** Different namespaces can federate to different Databricks SPs in the same connector deployment.

Connector RBAC for Model 2 (`get`/`list` on `serviceaccounts`, plus `create` on `serviceaccounts/token`) is documented in the README.

### 5. Backwards compatibility

* `DatabricksJobMetadata` and `DatabricksV2` task config get additive fields only.
* Existing PAT deployments continue to work unchanged.
* Older `DatabricksJobMetadata` payloads (without the new fields) are still consumed correctly by upgraded connectors.

## API Examples

### Operators flip auth modes via connector env vars (workflow code unchanged)

```bash
# PAT (legacy, default; unchanged behaviour from #3394)
# (no extra env vars needed; works as today)

# OAuth M2M
FLYTE_DATABRICKS_AUTH_TYPE=oauth_m2m
FLYTE_DATABRICKS_CLIENT_ID= # connector-level fallback
FLYTE_DATABRICKS_CLIENT_SECRET= # connector-level fallback
# Plus per-namespace K8s secret `databricks-oauth` for tenant overrides.

# OIDC Model 1 (connector-pod identity, shared)
FLYTE_DATABRICKS_AUTH_TYPE=oidc_federation
FLYTE_DATABRICKS_CLIENT_ID=
FLYTE_DATABRICKS_OIDC_AUDIENCE=https://...

# OIDC Model 2 (per-namespace, annotation-driven)
FLYTE_DATABRICKS_AUTH_TYPE=oidc_federation
# No connector-level client_id needed; discovered from each namespace's annotated SA.
```

### Workflow code stays identical across all four modes

```python
from flytekitplugins.spark import DatabricksV2
from flytekit import task

@task(
task_config=DatabricksV2(
databricks_conf={...},
databricks_instance="my-workspace.cloud.databricks.com",
),
container_image="my-image:tag",
)
def my_databricks_task(n: int) -> int:
...
```

## Files Changed

| File | Change |
| --- | --- |
| `plugins/flytekit-spark/flytekitplugins/spark/databricks_auth.py` | **NEW**. Strategy module: PAT / M2M / OIDC Model 1 / OIDC Model 2, async token cache, retry/backoff, auto-detection, validation |
| `plugins/flytekit-spark/flytekitplugins/spark/connector.py` | M2M/OIDC integration; `list_serviceaccounts_in_k8s` helper; persist discovered config in `DatabricksJobMetadata` |
| `plugins/flytekit-spark/flytekitplugins/spark/task.py` | New optional task-config fields surfaced for testability; auto-detection means workflow authors don't need them |
| `plugins/flytekit-spark/tests/test_databricks_auth.py` | **NEW**. 100+ tests across all four auth modes including discovery, cache, error paths |
| `plugins/flytekit-spark/tests/test_databricks_token.py` | Adjusted for new auth-resolution boundary; PAT regression coverage retained |
| `plugins/flytekit-spark/tests/test_connector.py` | Updated for the additive `DatabricksJobMetadata` fields |
| `plugins/flytekit-spark/README.md` | New "Databricks Connector Authentication" section: env var table, four-mode walkthrough, RBAC manifests, migration guide |

## Testing

### Unit tests

100+ test cases covering:

* **Auth resolution / auto-detection**: each mode selected correctly when its prerequisites are met; explicit errors when the requested mode is misconfigured (no silent downgrade).
* **PAT path**: regression suite from #3394 retained as-is.
* **M2M**: client-credentials flow, per-namespace secret read, fallback to env var, 401-driven cache invalidation on long-running jobs.
* **OIDC Model 1**: IRSA token-file read, token exchange, error paths when projected token absent.
* **OIDC Model 2**: SA discovery (single match, zero matches, multiple matches with ambiguity error), label filtering, annotation parsing, TTL discovery cache, RBAC failure surfaces.
* **Token cache**: TTL behaviour, async-safety under concurrent gets, pre-expiry refresh buffer, 401-driven refresh.
* **API resilience**: exponential backoff with jitter on the Databricks token endpoint.
* **Local dev**: lazy imports for `kubernetes` and `aiohttp` so `pyflyte run` works without K8s present.

### End-to-end (internal dev EKS cluster)

* PAT regression: passing.
* OAuth M2M: passing.
* OIDC Model 2 (per-namespace, annotation-driven): passing.
* OIDC Model 1: to be exercised before PR merge.

## Migration Path

| Today | Tomorrow |
| --- | --- |
| Workspace-scoped PAT in `FLYTE_DATABRICKS_ACCESS_TOKEN` | OAuth M2M with `client_id` / `client_secret`, OR OIDC federation via per-namespace SA annotations |
| Workspace-scoped Unity Catalog grants | Per-namespace SP, per-namespace UC grants |
| 90-day re-issuance chore | IdP-managed credentials |

Migration is opt-in and additive: existing deployments keep working until operators flip `FLYTE_DATABRICKS_AUTH_TYPE`.

## References

* Databricks PAT (legacy):
* Databricks OAuth M2M (client credentials):
* Databricks OIDC Workload Identity Federation:
* Kubernetes ServiceAccount projected tokens / `TokenRequest` API:
* Prior work, multi-tenant PAT: [flyteorg/flytekit#3394](https://github.com/flyteorg/flytekit/pull/3394)
* Tracked under (related), Databricks Serverless: [flyteorg/flyte#6911](https://github.com/flyteorg/flyte/issues/6911)

FYI: @kumare3 @pingsutw @machichima

### Goal: What should the final outcome look like, ideally?

The full feature has been implemented and tested end-to-end on an internal dev EKS cluster (PAT, M2M, OIDC Model 2 confirmed; Model 1 pending). I'm preparing a PR against `flyteorg/flytekit:master` that delivers the feature in a single coherent change. This issue is the tracking issue for that PR.

### Describe alternatives you've considered

* **Adding only OAuth M2M**: would address the legacy-PAT problem but not the auditability gap (still long-lived secrets).
* **Adding only OIDC Model 1**: secret-less, but collapses to a single shared Databricks SP, breaking the per-namespace UC tenancy goal from #6911 / #3394.
* **A per-task `databricks_oidc_service_account` field on `DatabricksV2`**: was prototyped, but it required workflow authors to know operator-level identity details and violated the zero-workflow-code-changes constraint. Replaced with the annotation-driven discovery design above.

### Propose: Link/Inline OR Additional context

Inline above. PR will land shortly with the full diff, RBAC manifests, and migration guide in the spark plugin README.

### Are you sure this issue hasn't been raised already?

* Yes. The only related issue is the multi-tenancy ask in #6911 (which #3394 partially addressed for PAT) and the serverless track (separate concern).

### Have you read the Code of Conduct?

* Yes

Contributor guide

Open the contributing guide

Research direction

Start by reading plugins/flytekit-spark/flytekitplugins/spark/connector.py and task.py, then compare the existing PAT coverage in tests/test_databricks_token.py and tests/test_connector.py. Review the proposed databricks_auth.py boundary and test_databricks_auth.py plan before making changes. Done means the four authentication modes, compatibility behavior, tests, and README/RBAC documentation are covered.

Written by the indexing model from the issue text.

Assessment

Tech stack
kubernetes, python
Domain
authentication, backend, infrastructure
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.