feat(kernel): forward identity federation client ID

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

@@ -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
 

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