From 50ae2a852ada1851ddcc69ac223ed46ed921ca6d Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 14 Aug 2026 03:01:10 +0000 Subject: [PATCH 1/3] feat(kernel): forward identity federation client ID --- KERNEL_REV | 2 +- README.md | 3 ++- .../sql/backend/kernel/auth_bridge.py | 22 ++++++++++++++----- src/databricks/sql/backend/kernel/client.py | 3 ++- src/databricks/sql/client.py | 20 ++++++++--------- src/databricks/sql/session.py | 8 +++++-- 6 files changed, 37 insertions(+), 21 deletions(-) diff --git a/KERNEL_REV b/KERNEL_REV index 2031bb3d7..95cfce816 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -7ffb30d533c08651ca707b8dd13894c9e01cb68e +eff8950428f4e6cc9975c663ec919f334962f7d0 diff --git a/README.md b/README.md index dcd726b9b..edb4d8d4b 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,8 @@ Notes: `cp310-abi3`). On older interpreters the `[kernel]` extra installs nothing and `use_kernel=True` raises an `ImportError`. - The extra also pulls in PyArrow, which the kernel result path requires. -- Authentication supports PAT (`access_token`), OAuth M2M, and OAuth U2M. +- Authentication supports PAT (`access_token`), OAuth M2M/U2M, and SP-wide + workload identity federation (`identity_federation_client_id`). ```bash diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 5b151bddb..7402d84cf 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -21,6 +21,10 @@ kernel's ``auth_type='oauth-u2m'`` and the kernel runs the browser flow itself. +``identity_federation_client_id`` is forwarded with whichever auth shape +wins resolution. It selects mandatory SP-wide workload-identity token +exchange in the kernel; omitting it preserves BYOT / account-wide behavior. + A user-supplied custom ``credentials_provider`` is **rejected** on the kernel path with ``NotSupportedError``: it's an opaque token source with no extractable raw credentials, so the kernel can't own the @@ -125,10 +129,10 @@ def kernel_auth_kwargs( ``auth_options`` carries the raw connect() kwargs relevant to auth (``auth_type``, ``oauth_client_id``, ``oauth_client_secret``, - ``oauth_redirect_port``, ``credentials_provider``). They drive the - OAuth decisions because the OAuth secret is consumed during - ``AuthProvider`` construction and can't be read back off the built - provider. + ``oauth_redirect_port``, ``credentials_provider``, + ``identity_federation_client_id``). They drive the OAuth decisions + because the OAuth secret is consumed during ``AuthProvider`` + construction and can't be read back off the built provider. Resolution order: @@ -161,6 +165,7 @@ def kernel_auth_kwargs( client_id = opts.get("oauth_client_id") client_secret = opts.get("oauth_client_secret") + federation_client_id = opts.get("identity_federation_client_id") auth_type = opts.get("auth_type") has_m2m = bool(client_id and client_secret) @@ -191,6 +196,8 @@ def kernel_auth_kwargs( 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 # 2. PAT (including TokenFederationProvider-wrapped PAT). @@ -201,7 +208,10 @@ def kernel_auth_kwargs( "PAT auth provider did not produce a Bearer Authorization " "header; cannot route through the kernel's PAT path" ) - return {"auth_type": "pat", "access_token": token} + kwargs = {"auth_type": "pat", "access_token": token} + if federation_client_id: + kwargs["identity_federation_client_id"] = federation_client_id + return kwargs # 3. OAuth U2M — browser authorization-code flow; the kernel runs it. if auth_type in ("databricks-oauth", "azure-oauth"): @@ -214,6 +224,8 @@ def kernel_auth_kwargs( 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 # 4. Custom credentials_provider — the connector's primary M2M path diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index c6e66d661..b1a1d5b3e 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -206,7 +206,8 @@ def __init__( # Forwarded to the kernel Session in ``open_session``. self._http_headers = http_headers or [] # Raw auth-relevant connect() kwargs (auth_type, - # oauth_client_id/secret, redirect port, credentials_provider). + # oauth_client_id/secret, redirect port, credentials_provider, + # identity_federation_client_id). # The kernel auth bridge needs these to build OAuth kwargs — the # OAuth secret is consumed during ``auth_provider`` construction # and isn't recoverable from the built provider. diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index d45d51181..6a9d7abfb 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -119,17 +119,11 @@ def __init__( the Thrift backend. :param use_kernel: `bool`, optional (default is False) Route the connection through the Rust kernel - (``databricks-sql-kernel`` via PyO3). Requires the - kernel extension to be installed separately — the - wheel is not yet published on PyPI, so today the - only supported install path is a local - ``maturin develop --release`` build from the - ``databricks-sql-kernel`` repo into the same venv. - Raises ``ImportError`` if the extension is not - available. In active development — PAT auth only - today; OAuth / federation / external credentials - and native parameter binding land in follow-ups. - Mutually exclusive with ``use_sea``. + (``databricks-sql-kernel`` via PyO3). Install the + connector's ``kernel`` extra to include the extension. + Supports PAT, OAuth M2M/U2M, and workload identity + federation; custom credentials providers are not + supported. Mutually exclusive with ``use_sea``. :param use_hybrid_disposition: `bool`, optional (default is False) Use the hybrid disposition instead of the inline disposition. :param server_hostname: Databricks instance host name. @@ -170,6 +164,10 @@ def __init__( port of the oauth redirect uri (localhost). This is required when custom oauth client_id `oauth_client_id` is set + identity_federation_client_id: `str`, optional + Service-principal client ID for mandatory SP-wide workload identity + token exchange. Supported by both the default and kernel backends. + user_agent_entry: `str`, optional A custom tag to append to the User-Agent header. This is typically used by partners to identify their applications.. If not specified, it will use the default user agent PyDatabricksSqlConnector diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index c382103ea..a83d62db1 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -164,8 +164,9 @@ def _create_backend( # original credentials. On this path we intentionally did # NOT build the connector's own OAuth provider (see __init__ # above), so these raw kwargs are the only source of the - # OAuth client id/secret. These are kernel-only; the Thrift - # / SEA backends are unaffected. + # OAuth client id/secret and optional federation client id. + # These are kernel-only; the Thrift / SEA backends are + # unaffected. kernel_auth_options = { "auth_type": kwargs.get("auth_type"), "oauth_client_id": kwargs.get("oauth_client_id"), @@ -173,6 +174,9 @@ def _create_backend( "oauth_redirect_port": kwargs.get("oauth_redirect_port"), "oauth_scopes": kwargs.get("oauth_scopes"), "credentials_provider": kwargs.get("credentials_provider"), + "identity_federation_client_id": kwargs.get( + "identity_federation_client_id" + ), } # Forward the connector's retry-tuning kwargs so the kernel's # own retry policy honours them (the kernel owns the retry From ff25e44d77b9aecf82492a8856c846ae72c4894d Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 14 Aug 2026 03:06:31 +0000 Subject: [PATCH 2/3] docs: preserve existing kernel comment --- src/databricks/sql/client.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 6a9d7abfb..a8e4bab1c 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -119,11 +119,17 @@ def __init__( the Thrift backend. :param use_kernel: `bool`, optional (default is False) Route the connection through the Rust kernel - (``databricks-sql-kernel`` via PyO3). Install the - connector's ``kernel`` extra to include the extension. - Supports PAT, OAuth M2M/U2M, and workload identity - federation; custom credentials providers are not - supported. Mutually exclusive with ``use_sea``. + (``databricks-sql-kernel`` via PyO3). Requires the + kernel extension to be installed separately — the + wheel is not yet published on PyPI, so today the + only supported install path is a local + ``maturin develop --release`` build from the + ``databricks-sql-kernel`` repo into the same venv. + Raises ``ImportError`` if the extension is not + available. In active development — PAT auth only + today; OAuth / federation / external credentials + and native parameter binding land in follow-ups. + Mutually exclusive with ``use_sea``. :param use_hybrid_disposition: `bool`, optional (default is False) Use the hybrid disposition instead of the inline disposition. :param server_hostname: Databricks instance host name. From d5b27d67d34d56cbe2cdae4d1a1e698d9f327726 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 14 Aug 2026 03:08:21 +0000 Subject: [PATCH 3/3] test(kernel): cover federation client ID forwarding --- tests/unit/test_kernel_auth_bridge.py | 40 +++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index 92b0c7003..edafdf625 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -272,6 +272,46 @@ def test_u2m_forwards_scopes(self, auth_type): assert kwargs["oauth_scopes"] == ["all-apis", "offline_access"] +class TestKernelIdentityFederationClientId: + @pytest.mark.parametrize( + "auth_provider,auth_options", + [ + pytest.param(AccessTokenAuthProvider("dapi-xyz"), {}, id="pat"), + pytest.param( + _FakeOAuthProvider(), + {"oauth_client_id": "sp-uuid", "oauth_client_secret": "shh"}, + id="m2m", + ), + pytest.param( + _FakeOAuthProvider(), + {"auth_type": "databricks-oauth"}, + id="u2m", + ), + ], + ) + @pytest.mark.parametrize( + "federation_client_id", + [ + pytest.param(None, id="omitted"), + pytest.param("", id="empty"), + pytest.param("federation-client", id="supplied"), + ], + ) + def test_forwards_only_non_empty_value( + self, auth_provider, auth_options, federation_client_id + ): + options = dict(auth_options) + if federation_client_id is not None: + options["identity_federation_client_id"] = federation_client_id + + kwargs = kernel_auth_kwargs(auth_provider, options) + + if federation_client_id: + assert kwargs["identity_federation_client_id"] == federation_client_id + else: + assert "identity_federation_client_id" not in kwargs + + class TestKernelAuthAmbiguity: """Conflicting auth signals must fail loudly at session-open rather than silently resolving to one flow (which would surface later as a