Skip to content

feat(kernel): JWT private-key M2M auth on use_kernel=True - #921

Open
rahuls-db wants to merge 3 commits into
mainfrom
feat/kernel-jwt-private-key-m2m
Open

feat(kernel): JWT private-key M2M auth on use_kernel=True#921
rahuls-db wants to merge 3 commits into
mainfrom
feat/kernel-jwt-private-key-m2m

Conversation

@rahuls-db

Copy link
Copy Markdown
Collaborator

What

Adds OAuth machine-to-machine auth with a JWT private-key client assertion (RFC 7523) on the kernel backend (use_kernel=True). Instead of a client secret, the kernel signs a short-lived JWT with the service principal's private key and sends it as the client_assertion in the client-credentials grant; the workspace's OAuth IdP verifies it against the SP's registered public key.

Companion to the kernel-side feature (databricks-sql-kernel #249; napi token_url in #275) and the parallel databricks-sql-nodejs / databricks-sql-go changes.

How

  • auth_bridge.py — new JWT branch in build_kernel_auth_kwargs, checked before shared-secret M2M and PAT (a private-key file is unambiguous JWT M2M intent). Forwards oauth_client_id + oauth_jwt_key_file + oauth_jwt_kid (+ optional oauth_jwt_passphrase / oauth_jwt_algorithm / oauth_scopes / token_url) to the kernel's auth_type="oauth-m2m-jwt". Requires client_id + kid; mutually exclusive with oauth_client_secret / credentials_provider (both raise NotSupportedError).
  • session.py — forward the new oauth_jwt_* / token_url connect kwargs into the kernel auth options.

Usage

from databricks import sql
conn = sql.connect(
    server_hostname="adb-….azuredatabricks.net",
    http_path="/sql/1.0/warehouses/…",
    use_kernel=True,
    oauth_client_id="<sp-client-id>",
    oauth_jwt_key_file="/path/private_key.pem",
    oauth_jwt_kid="<kid>",
    token_url="https://login.microsoftonline.com/<tenant>/oauth2/v2.0/token",
    oauth_scopes=["<databricks-resource-id>/.default"],
)

Testing

  • 9 new unit tests in tests/unit/test_kernel_auth_bridge.py (routing, precedence over M2M/PAT, required-field validation, ambiguity guards); full file 56 passing.
  • Verified end-to-end against an Azure Databricks warehouse: SELECT 1[Row(n=1)], with conn.session.backend asserted to be KernelDatabricksClient (kernel path, not Thrift) and use_kernel true.

Requires databricks-sql-kernel >= 0.2.0 with JWT support.

This pull request and its description were written by Isaac.

Route JWT private-key client-assertion auth (RFC 7523) through the kernel
backend. When the caller passes `oauth_jwt_key_file` (+ `oauth_client_id`
and `oauth_jwt_kid`, optional `oauth_jwt_passphrase` / `oauth_jwt_algorithm`
/ `oauth_scopes` / `token_url`), the bridge forwards them to the kernel's
`auth_type="oauth-m2m-jwt"`, which signs a short-lived assertion with the
private key instead of sending a client secret and owns the token lifecycle.

- auth_bridge.py: new JWT branch (checked before shared-secret M2M and PAT,
  since a private-key file is unambiguous JWT M2M intent); mutually exclusive
  with oauth_client_secret / credentials_provider; requires client_id + kid.
- session.py: forward the new oauth_jwt_* / token_url kwargs into the kernel
  auth options.
- tests: 9 unit tests covering routing, precedence, validation, and
  ambiguity guards.

Verified end-to-end: `SELECT 1` via use_kernel=True against an Azure
Databricks warehouse, authenticated by Entra ID against the service
principal's registered public certificate.

Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Verdict: 1 Medium · 1 Low

Solid, well-tested addition — routing, precedence, required-field validation, and two of three ambiguity guards are covered. One medium gap: the new JWT M2M path lacks the U2M-collision guard that the shared-secret M2M path has, so oauth_jwt_key_file + auth_type="databricks-oauth" silently resolves to JWT M2M instead of failing loudly. Plus a cosmetic stale-comment numbering nit.

"(machine-to-machine). Drop oauth_client_secret for U2M, or drop "
"auth_type for M2M."
)
if has_jwt_m2m and client_secret:

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🟡 Medium — Missing ambiguity guard for JWT M2M + U2M auth_type.

The bridge explicitly rejects shared-secret M2M colliding with a U2M request (client_secret and auth_type == "databricks-oauth"NotSupportedError), on the stated principle that conflicting auth signals must "fail loudly at session-open rather than silently resolving to one flow." The new JWT branch adds guards against oauth_jwt_key_file + oauth_client_secret and oauth_jwt_key_file + credentials_provider, but there is no guard for oauth_jwt_key_file + auth_type="databricks-oauth".

Concretely, a caller who passes auth_type="databricks-oauth" (clear browser-U2M intent) while an oauth_jwt_key_file is also present (e.g. leftover ambient config) silently gets routed to oauth-m2m-jwt — the browser flow they asked for never runs, and they authenticate as the service principal instead. This is exactly the failure mode the client_secret+U2M guard was written to prevent, so the JWT path should mirror it. Consider adding:

if has_jwt_m2m and auth_type == "databricks-oauth":
    raise NotSupportedError(...)

before the JWT branch, and a corresponding unit test.

(Anchored to the nearest changed line — see the description for the exact location.)

kwargs["oauth_scopes"] = scopes
if federation_client_id:
kwargs["identity_federation_client_id"] = federation_client_id
return kwargs

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🔵 Low — Inline step-number comments weren't renumbered when the JWT branch was inserted. The docstring was correctly updated to 1=JWT, 2=M2M, 3=PAT, 4=U2M, 5=creds_provider, 6=else, but the inline # N. comments still read: # 2. OAuth M2M, # 2. PAT (duplicate 2), # 3. OAuth U2M, # 4. Custom credentials_provider, # 5. Everything else. They should be 2/3/4/5/6 to match the docstring. Purely cosmetic but the duplicated # 2 is confusing when cross-referencing the resolution order.

The JWT branch introduced an earlier untyped `kwargs =`, so mypy flagged
the M2M branch's `kwargs: Dict[str, Any]` as a redefinition. Move the
annotation to the first (JWT) assignment.

Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Verdict: 1 Medium

Solid, well-tested addition of JWT private-key M2M auth to the kernel path — routing, precedence, required-field validation, and two ambiguity guards all have unit coverage. One medium consistency gap: the JWT path lacks the U2M-conflict ambiguity guard its shared-secret sibling has, so oauth_jwt_key_file + auth_type="databricks-oauth" silently routes to M2M against the wrong principal.

"(machine-to-machine). Drop oauth_client_secret for U2M, or drop "
"auth_type for M2M."
)
if has_jwt_m2m and client_secret:

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🟡 Medium — The new JWT M2M path is missing an ambiguity guard against U2M that its shared-secret sibling has.

The shared-secret M2M path rejects oauth_client_secret combined with auth_type="databricks-oauth" (U2M browser flow) as ambiguous — see the guard at line 227-232, whose rationale is "User asked for U2M (browser) but also passed a secret (M2M). Don't silently route M2M against the wrong principal."

The JWT branch adds guards for oauth_jwt_key_file + oauth_client_secret and oauth_jwt_key_file + credentials_provider, but not for oauth_jwt_key_file + auth_type="databricks-oauth". Because the JWT branch is checked first (before the U2M branch), a caller who passes both auth_type="databricks-oauth" (asking for the browser flow) and an ambient/mistaken oauth_jwt_key_file is silently routed to oauth-m2m-jwt — authenticating as a service principal instead of the interactive user. This is precisely the silent-misroute failure the U2M/secret guard was written to prevent.

Consider adding a parallel guard, e.g.:

if has_jwt_m2m and auth_type == "databricks-oauth":
    raise NotSupportedError(
        "Ambiguous auth on use_kernel=True: auth_type='databricks-oauth' "
        "selects the U2M browser flow, but oauth_jwt_key_file (JWT "
        "private-key M2M) was also provided. Drop oauth_jwt_key_file for "
        "U2M, or drop auth_type for JWT M2M."
    )

A unit test alongside TestKernelAuthAmbiguity would also lock this in.

…ments

Address peco-review-bot review on #921:
- Medium: add the missing ambiguity guard for oauth_jwt_key_file +
  auth_type="databricks-oauth" (U2M intent), mirroring the existing
  shared-secret M2M + U2M guard. Fails loudly rather than silently
  resolving to one flow. Covered by a new unit test.
- Low: renumber the inline resolution-order comments (PAT→3, U2M→4,
  creds→5, else→6) to match the docstring after the JWT branch insert.

Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

✅ No issues identified by the review bot.

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.

1 participant