Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# Release History

# Unreleased
- Kernel backend (`use_kernel=True`): OAuth **M2M with a JWT private-key client assertion** (RFC 7523) is now supported. Pass `oauth_client_id` + `oauth_jwt_key_file` + `oauth_jwt_kid` (with optional `oauth_jwt_passphrase` for an encrypted PKCS#8 key, `oauth_jwt_algorithm` defaulting to `RS256`, `oauth_scopes`, and `token_url` for the IdP token endpoint) and the connector routes 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. The kernel owns the token lifecycle. A private-key file is treated as unambiguous JWT M2M intent and is mutually exclusive with `oauth_client_secret` / `credentials_provider` (both raise `NotSupportedError`). Verified end-to-end against an Azure Databricks workspace with the service principal's public certificate registered on its Entra ID app registration. Requires `databricks-sql-kernel >= 0.2.0` with JWT support.
- Kernel backend (`use_kernel=True`): OAuth U2M with `auth_type="databricks-oauth"` now forwards the connector's `databricks-sql-python` OAuth-app bundle (`client_id` + `sql offline_access` scopes + redirect port) into the kernel, so a bare U2M connection authenticates as `databricks-sql-python` — parity with the Thrift path — instead of inheriting the kernel's own `databricks-sql-connector` default. A caller-supplied `oauth_client_id` (with its coupled `oauth_redirect_port`) is honored, as is a caller-supplied `oauth_scopes`; absent one, the connector default (`sql offline_access`) is forwarded. Note: the kernel binds a single U2M redirect port, so unlike the Thrift path (which tries the full `8020..8024` range) the kernel path uses only one port and does not fall back to the next port if it is already bound — pass `oauth_redirect_port` (with `oauth_client_id`) to pick a free one on a port collision. `auth_type="azure-oauth"` (Azure AD) is not yet supported on the kernel path and raises `NotSupportedError` — use the Thrift backend for it (PECOBLR-4040; Azure tracked by PECOBLR-4120)

# 4.4.0 (2026-07-22)
Expand Down
100 changes: 89 additions & 11 deletions src/databricks/sql/backend/kernel/auth_bridge.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,11 +156,16 @@ def kernel_auth_kwargs(

(``azure-oauth`` is rejected as unsupported before these guards —
PECOBLR-4120.)
1. **OAuth M2M** — ``oauth_client_id`` + ``oauth_client_secret``
1. **OAuth M2M (JWT private key)** — ``oauth_jwt_key_file`` present →
forward the private-key + ``oauth_client_id`` + ``oauth_jwt_kid``
to the kernel's ``oauth-m2m-jwt`` (RFC 7523 client assertion). The
kernel signs the assertion and owns the token lifecycle. Checked
first because a private-key file is unambiguous JWT M2M intent.
2. **OAuth M2M** — ``oauth_client_id`` + ``oauth_client_secret``
both present → forward raw creds to the kernel's ``oauth-m2m``.
2. **PAT** — the built provider is (or wraps) an
3. **PAT** — the built provider is (or wraps) an
``AccessTokenAuthProvider`` → extract the bearer token.
3. **OAuth U2M** — ``auth_type`` is ``databricks-oauth`` → forward the
4. **OAuth U2M** — ``auth_type`` is ``databricks-oauth`` → forward the
connector's coupled ``databricks-sql-python`` bundle (``client_id``
+ ``redirect_ports`` list, defaulting scopes to ``PYSQL_OAUTH_SCOPES``
when the caller supplies none) to the kernel's ``oauth-u2m``, so a
Expand All @@ -169,9 +174,9 @@ def kernel_auth_kwargs(
``databricks-sql-connector`` default (PECOBLR-4039/4040). Unlike the
Thrift path, a caller-supplied ``oauth_scopes`` is honored here.
``azure-oauth`` is rejected as unsupported (PECOBLR-4120).
4. **Custom credentials_provider** → ``NotSupportedError`` (opaque
5. **Custom credentials_provider** → ``NotSupportedError`` (opaque
token source; no raw creds for the kernel to own).
5. Anything else → ``NotSupportedError``.
6. Anything else → ``NotSupportedError``.

M2M is checked before PAT so that a workload passing both an
access token *and* M2M creds resolves to the (refreshing) M2M path
Expand All @@ -186,7 +191,12 @@ def kernel_auth_kwargs(
client_secret = opts.get("oauth_client_secret")
federation_client_id = opts.get("identity_federation_client_id")
auth_type = opts.get("auth_type")
jwt_key_file = opts.get("oauth_jwt_key_file")
has_m2m = bool(client_id and client_secret)
# A private-key file is unambiguous JWT client-assertion M2M intent
# (RFC 7523): the kernel signs a short-lived assertion with the key
# rather than sending a client secret.
has_jwt_m2m = bool(jwt_key_file)

# azure-oauth (Azure AD U2M) is not yet supported on the kernel path.
# Reject it up front — before any M2M/U2M routing — so ANY azure-oauth
Expand Down Expand Up @@ -223,10 +233,78 @@ def kernel_auth_kwargs(
"(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.)

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.

raise NotSupportedError(
"Ambiguous auth on use_kernel=True: both oauth_jwt_key_file "
"(JWT private-key M2M) and oauth_client_secret (shared-secret "
"M2M) were provided. Pass exactly one — a private key for "
"JWT client-assertion M2M, or a client secret for shared-secret M2M."
)
if has_jwt_m2m and opts.get("credentials_provider") is not None:
raise NotSupportedError(
"Ambiguous auth on use_kernel=True: both a custom "
"credentials_provider and oauth_jwt_key_file were provided. "
"Pass exactly one — oauth_client_id + oauth_jwt_key_file for "
"kernel-managed JWT private-key M2M, or use the Thrift backend "
"(default) for credentials_provider."
)
if has_jwt_m2m and auth_type == "databricks-oauth":
raise NotSupportedError(
f"Ambiguous auth on use_kernel=True: auth_type={auth_type!r} selects "
"the U2M browser flow, but oauth_jwt_key_file was also provided "
"(JWT private-key M2M). Drop oauth_jwt_key_file for U2M, or drop "
"auth_type for JWT M2M."
)

# 1. OAuth M2M — raw client-credentials pair forwarded to the kernel.
if has_m2m:
# 1. OAuth M2M (JWT private-key client assertion) — the kernel signs a
# short-lived assertion with the private key and runs the
# client-credentials grant. Checked before shared-secret M2M and PAT
# because a private-key file is unambiguous JWT M2M intent. Requires
# oauth_client_id (the service principal / OAuth client) and
# oauth_jwt_kid (the key id the IdP uses to select the registered
# public key). Optional oauth_jwt_passphrase / oauth_jwt_algorithm /
# oauth_scopes / token_url are forwarded when present; the kernel
# fills defaults (RS256 algorithm, all-apis scope, OIDC discovery)
# for any omitted.
if has_jwt_m2m:
if not client_id:
raise ProgrammingError(
"use_kernel=True JWT private-key M2M (oauth_jwt_key_file) "
"requires oauth_client_id (the service principal / OAuth "
"client id used as the assertion issuer and subject)."
)
jwt_kid = opts.get("oauth_jwt_kid")
if not jwt_kid:
raise ProgrammingError(
"use_kernel=True JWT private-key M2M (oauth_jwt_key_file) "
"requires oauth_jwt_kid (the key id written into the JWT "
"header so the IdP can select the registered public key)."
)
kwargs: Dict[str, Any] = {
"auth_type": "oauth-m2m-jwt",
"client_id": client_id,
"jwt_key_file": jwt_key_file,
"jwt_kid": jwt_kid,
}
jwt_passphrase = opts.get("oauth_jwt_passphrase")
if jwt_passphrase:
kwargs["jwt_passphrase"] = jwt_passphrase
jwt_algorithm = opts.get("oauth_jwt_algorithm")
if jwt_algorithm:
kwargs["jwt_algorithm"] = jwt_algorithm
token_url = opts.get("token_url")
if token_url:
kwargs["token_url"] = token_url
scopes = _normalize_scopes(opts.get("oauth_scopes"))
if scopes is not None:
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.


# 2. OAuth M2M — raw client-credentials pair forwarded to the kernel.
if has_m2m:
kwargs = {
"auth_type": "oauth-m2m",
"client_id": client_id,
"client_secret": client_secret,
Expand All @@ -238,7 +316,7 @@ def kernel_auth_kwargs(
kwargs["identity_federation_client_id"] = federation_client_id
return kwargs

# 2. PAT (including TokenFederationProvider-wrapped PAT).
# 3. PAT (including TokenFederationProvider-wrapped PAT).
if _is_pat(auth_provider):
token = _extract_bearer_token(auth_provider)
if not token:
Expand All @@ -251,7 +329,7 @@ def kernel_auth_kwargs(
kwargs["identity_federation_client_id"] = federation_client_id
return kwargs

# 3. OAuth U2M — browser authorization-code flow; the kernel runs it.
# 4. OAuth U2M — browser authorization-code flow; the kernel runs it.
# Only databricks-oauth reaches here (azure-oauth rejected up front).
# Forward the connector's own databricks-sql-python bundle instead of
# the kernel's databricks-sql-connector default, for parity with the
Expand Down Expand Up @@ -283,7 +361,7 @@ def kernel_auth_kwargs(
kwargs["identity_federation_client_id"] = federation_client_id
return kwargs

# 4. Custom credentials_provider — the connector's primary M2M path
# 5. Custom credentials_provider — the connector's primary M2M path
# on Thrift/SEA, but unusable on the kernel: it's an opaque token
# source with no extractable client_id/secret, so the kernel
# can't own the token lifecycle. Point the caller at the raw
Expand All @@ -297,7 +375,7 @@ def kernel_auth_kwargs(
"credentials_provider."
)

# 5. Everything else (including no usable credentials at all —
# 6. Everything else (including no usable credentials at all —
# ``auth_provider`` is None on the kernel path when no access
# token was supplied and no OAuth kwargs resolved above).
provider_desc = (
Expand Down
9 changes: 9 additions & 0 deletions src/databricks/sql/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,15 @@ def _create_backend(
"oauth_client_secret": kwargs.get("oauth_client_secret"),
"oauth_redirect_port": kwargs.get("oauth_redirect_port"),
"oauth_scopes": kwargs.get("oauth_scopes"),
# JWT private-key M2M (RFC 7523 client assertion): the kernel
# signs a short-lived assertion with the private key instead
# of sending a client secret. token_url points the assertion
# at the workspace's OAuth IdP token endpoint (e.g. Entra ID).
"oauth_jwt_key_file": kwargs.get("oauth_jwt_key_file"),
"oauth_jwt_kid": kwargs.get("oauth_jwt_kid"),
"oauth_jwt_passphrase": kwargs.get("oauth_jwt_passphrase"),
"oauth_jwt_algorithm": kwargs.get("oauth_jwt_algorithm"),
"token_url": kwargs.get("token_url"),
"credentials_provider": kwargs.get("credentials_provider"),
"identity_federation_client_id": kwargs.get(
"identity_federation_client_id"
Expand Down
137 changes: 137 additions & 0 deletions tests/unit/test_kernel_auth_bridge.py
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,143 @@ def test_client_id_without_secret_does_not_trigger_m2m(self):
assert kwargs == {"auth_type": "pat", "access_token": "dapi-xyz"}


class TestKernelOAuthM2MJwt:
"""JWT private-key M2M (RFC 7523 client assertion) → the kernel's
``oauth-m2m-jwt``. Driven by ``oauth_jwt_key_file`` (unambiguous
private-key intent); requires ``oauth_client_id`` + ``oauth_jwt_kid``."""

def test_full_kwargs_route_to_oauth_m2m_jwt(self):
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{
"oauth_client_id": "sp-uuid",
"oauth_jwt_key_file": "/keys/jwt.pem",
"oauth_jwt_kid": "kid-1",
"oauth_jwt_passphrase": "pw",
"oauth_jwt_algorithm": "ES256",
"token_url": "https://login.microsoftonline.com/t/oauth2/v2.0/token",
"oauth_scopes": ["2ff814a6-.../.default"],
},
)
assert kwargs == {
"auth_type": "oauth-m2m-jwt",
"client_id": "sp-uuid",
"jwt_key_file": "/keys/jwt.pem",
"jwt_kid": "kid-1",
"jwt_passphrase": "pw",
"jwt_algorithm": "ES256",
"token_url": "https://login.microsoftonline.com/t/oauth2/v2.0/token",
"oauth_scopes": ["2ff814a6-.../.default"],
}

def test_minimal_kwargs_omit_optionals(self):
# Only the three required fields; the kernel fills the rest
# (RS256 algorithm, all-apis scope, OIDC discovery).
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{
"oauth_client_id": "sp-uuid",
"oauth_jwt_key_file": "/keys/jwt.pem",
"oauth_jwt_kid": "kid-1",
},
)
assert kwargs == {
"auth_type": "oauth-m2m-jwt",
"client_id": "sp-uuid",
"jwt_key_file": "/keys/jwt.pem",
"jwt_kid": "kid-1",
}

def test_normalizes_space_delimited_scopes(self):
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{
"oauth_client_id": "sp",
"oauth_jwt_key_file": "/k.pem",
"oauth_jwt_kid": "k",
"oauth_scopes": "all-apis sql",
},
)
assert kwargs["oauth_scopes"] == ["all-apis", "sql"]

def test_takes_precedence_over_pat(self):
# A private key alongside an ambient PAT resolves to the
# (refreshing) JWT M2M path, not the static token.
kwargs = kernel_auth_kwargs(
AccessTokenAuthProvider("dapi-xyz"),
{
"oauth_client_id": "sp",
"oauth_jwt_key_file": "/k.pem",
"oauth_jwt_kid": "k",
},
)
assert kwargs["auth_type"] == "oauth-m2m-jwt"

def test_missing_client_id_raises_programming_error(self):
with pytest.raises(ProgrammingError, match="oauth_client_id"):
kernel_auth_kwargs(
None,
{"oauth_jwt_key_file": "/k.pem", "oauth_jwt_kid": "k"},
)

def test_missing_kid_raises_programming_error(self):
with pytest.raises(ProgrammingError, match="oauth_jwt_kid"):
kernel_auth_kwargs(
None,
{"oauth_client_id": "sp", "oauth_jwt_key_file": "/k.pem"},
)

def test_jwt_plus_client_secret_is_rejected(self):
with pytest.raises(NotSupportedError, match="oauth_client_secret"):
kernel_auth_kwargs(
None,
{
"oauth_client_id": "sp",
"oauth_jwt_key_file": "/k.pem",
"oauth_jwt_kid": "k",
"oauth_client_secret": "shh",
},
)

def test_jwt_plus_credentials_provider_is_rejected(self):
with pytest.raises(NotSupportedError, match="credentials_provider"):
kernel_auth_kwargs(
None,
{
"oauth_client_id": "sp",
"oauth_jwt_key_file": "/k.pem",
"oauth_jwt_kid": "k",
"credentials_provider": object(),
},
)

def test_jwt_plus_databricks_oauth_auth_type_is_rejected(self):
# auth_type="databricks-oauth" signals U2M intent; a private key
# alongside it is ambiguous (mirrors the shared-secret M2M + U2M guard).
with pytest.raises(NotSupportedError, match="oauth_jwt_key_file"):
kernel_auth_kwargs(
None,
{
"oauth_client_id": "sp",
"oauth_jwt_key_file": "/k.pem",
"oauth_jwt_kid": "k",
"auth_type": "databricks-oauth",
},
)

def test_federation_client_id_forwarded(self):
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{
"oauth_client_id": "sp",
"oauth_jwt_key_file": "/k.pem",
"oauth_jwt_kid": "k",
"identity_federation_client_id": "fed",
},
)
assert kwargs["identity_federation_client_id"] == "fed"


class TestKernelOAuthU2M:
"""Only ``databricks-oauth`` U2M is supported on the kernel path.

Expand Down
Loading