feat(kernel): forward identity federation client ID (#910)

* feat(kernel): forward identity federation client ID

* docs: preserve existing kernel comment

* test(kernel): cover federation client ID forwarding

Author: vuanhphung

KERNEL_REV

@@ -1 +1 @@
-7ffb30d533c08651ca707b8dd13894c9e01cb68e
+eff8950428f4e6cc9975c663ec919f334962f7d0

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

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

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.

src/databricks/sql/client.py

@@ -170,6 +170,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
 

src/databricks/sql/session.py

@@ -164,15 +164,19 @@ 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"),
                 "oauth_client_secret": kwargs.get("oauth_client_secret"),
                 "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

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