From 452b4264083da49fb2a3c1d9fec91ece1a0875a3 Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Wed, 19 Aug 2026 23:39:00 +0000 Subject: [PATCH 1/3] feat(kernel): support JWT private-key M2M auth on use_kernel=True 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 --- CHANGELOG.md | 1 + .../sql/backend/kernel/auth_bridge.py | 83 +++++++++++- src/databricks/sql/session.py | 9 ++ tests/unit/test_kernel_auth_bridge.py | 123 ++++++++++++++++++ 4 files changed, 210 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f750fa29b..42246f8b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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) diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 95374ba77..313090b8b 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -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 @@ -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 @@ -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 @@ -223,8 +233,69 @@ 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: + 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." + ) + + # 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 = { + "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 - # 1. OAuth M2M — raw client-credentials pair forwarded to the kernel. + # 2. OAuth M2M — raw client-credentials pair forwarded to the kernel. if has_m2m: kwargs: Dict[str, Any] = { "auth_type": "oauth-m2m", diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index a83d62db1..a62b0d081 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -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" diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index f60943948..b3933019a 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -245,6 +245,129 @@ 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_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. From 1ac6d80c41fd75f3aff98c0bf1770f7248ddc3cd Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Thu, 20 Aug 2026 02:53:27 +0000 Subject: [PATCH 2/3] fix(kernel): annotate kwargs at first assignment to satisfy mypy 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 --- src/databricks/sql/backend/kernel/auth_bridge.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 313090b8b..114802180 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -273,7 +273,7 @@ def kernel_auth_kwargs( "requires oauth_jwt_kid (the key id written into the JWT " "header so the IdP can select the registered public key)." ) - kwargs = { + kwargs: Dict[str, Any] = { "auth_type": "oauth-m2m-jwt", "client_id": client_id, "jwt_key_file": jwt_key_file, @@ -297,7 +297,7 @@ def kernel_auth_kwargs( # 2. OAuth M2M — raw client-credentials pair forwarded to the kernel. if has_m2m: - kwargs: Dict[str, Any] = { + kwargs = { "auth_type": "oauth-m2m", "client_id": client_id, "client_secret": client_secret, From 193da0802f8aa83e61317c5f59ccb55d2ce355ed Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Thu, 20 Aug 2026 02:56:11 +0000 Subject: [PATCH 3/3] fix(kernel): guard JWT M2M + databricks-oauth auth_type; renumber comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- src/databricks/sql/backend/kernel/auth_bridge.py | 15 +++++++++++---- tests/unit/test_kernel_auth_bridge.py | 14 ++++++++++++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 114802180..4ade82448 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -248,6 +248,13 @@ def kernel_auth_kwargs( "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 (JWT private-key client assertion) — the kernel signs a # short-lived assertion with the private key and runs the @@ -309,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: @@ -322,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 @@ -354,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 @@ -368,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 = ( diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index b3933019a..2bba8ceb8 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -355,6 +355,20 @@ def test_jwt_plus_credentials_provider_is_rejected(self): }, ) + 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(),