feat(kernel): forward identity federation client ID

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