diff --git a/CHANGELOG.md b/CHANGELOG.md index 42246f8b3..439983c3f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,8 @@ # 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) +- 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 (PECOBLR-4040) +- Kernel backend (`use_kernel=True`): **Azure Entra (Azure AD) service-principal M2M is now supported.** `auth_type="azure-sp-m2m"` forwards `azure_client_id` / `azure_client_secret`; the kernel is the Azure-aware auth core — it builds the Entra v2.0 token endpoint and the `{app_id}/.default` scope, and **auto-discovers the tenant** from the workspace's `/aad/auth` redirect when `azure_tenant_id` is omitted (matching Thrift). The `Authorization` bearer is the Databricks-audience data token, which alone authenticates a workspace-member SP. Set `azure_workspace_resource_id` and the kernel also sends the Azure SP management token (`X-Databricks-Azure-SP-Management-Token`) + `X-Databricks-Azure-Workspace-Resource-Id` header (matching the JDBC driver), so a service principal with an Azure RBAC role but no workspace membership can authenticate; omit it and no ARM management-scope token is fetched. Azure AD **U2M** (`auth_type="azure-oauth"`) now routes to the kernel's OAuth U2M flow, identically to `auth_type="databricks-oauth"`: the kernel runs the in-house workspace-federated browser flow, which Azure workspaces support (the workspace federates login to Entra). It forwards the connector's `databricks-sql-python` OAuth app, not the Thrift Azure app (`96eecda7` / port 8030), which is registered for Thrift's direct-Entra flow the kernel does not perform (PECOBLR-4141; PECOBLR-4120) # 4.4.0 (2026-07-22) - Raised the minimum supported Python version to 3.10, dropping the end-of-life 3.8/3.9, to update the lockfile and clear CVE-flagged dependencies in the repo (databricks/databricks-sql-python#798) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 0fb098372..f63de23a9 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -69,7 +69,7 @@ to change without notice. | Option | Type | Thrift | Kernel | Default Value | Note | | --------------------------------------------------- | -------------------- | :----: | :----: | ------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `access_token` (PAT) | `str` | ✅ | ✅ | `None` | Personal Access Token / bearer token. The default auth mode when set; otherwise auth falls back to OAuth. | -| `auth_type` | `str` | ✅ | ✅ | `None` ⇒ Databricks OAuth | `databricks-oauth` or `azure-oauth`. | +| `auth_type` | `str` | ✅ | ✅ | `None` ⇒ Databricks OAuth | `databricks-oauth` (U2M), `azure-oauth` (Azure AD U2M), or `azure-sp-m2m` (Azure service-principal M2M). All three are supported on the kernel path (#919): `azure-oauth` routes to the kernel's OAuth U2M flow identically to `databricks-oauth` (the in-house workspace-federated browser flow, which Azure workspaces support), so it uses the `databricks-sql-python` app rather than Thrift's direct-Entra Azure app. All three work on the Thrift path. | | `oauth_client_id` (U2M) | `str` | ✅ | ✅ | built-in client id | Custom U2M client id. Forwarded on both; when absent, each path applies its own built-in default. | | `oauth_redirect_port` (U2M) | `int` | ✅ | ✅ | `None` | Localhost redirect port for the browser flow. On **both** backends it is only honored when a custom `oauth_client_id` is also supplied — then that single port becomes the redirect URI. With the built-in client id (or when omitted) the connector uses the full registered range 8020–8024 and binds the first free port, so a bare `oauth_redirect_port` has no effect. (Thrift: `auth.py` `oauth_redirect_port_range`; Kernel: same logic, forwarded as `redirect_ports`.) | | `oauth_client_secret` (OAuth M2M) | `str` | ❌ | ✅ | `None` | **Kernel-only in practice.** The Thrift auth path never reads `oauth_client_secret`; use `credentials_provider` or an Azure service principal for M2M on Thrift. | @@ -82,7 +82,8 @@ to change without notice. | `credentials_provider` | `CredentialsProvider`| ✅ | ❌ | `None` | Custom external credentials provider. **Rejected on the kernel path** (`NotSupportedError`) — it is an opaque token source, so the kernel cannot own the token lifecycle; use `oauth_client_id` + `oauth_client_secret` for M2M, or the Thrift backend. | | `identity_federation_client_id` | `str` | ✅ | ✅ | `None` | Workload identity / token-federation client id (kernel support added in #910). | | `experimental_oauth_persistence` | `OAuthPersistence` | ✅ | ❌ | `None` | **Thrift-only.** The kernel owns its own token lifecycle and does not accept a persistence store. | -| `azure_client_id` / `azure_client_secret` / `azure_tenant_id` / `azure_workspace_resource_id` | `str` | ✅ | ❌ | `None` | **Thrift-only.** The Azure service-principal (Entra ID M2M) fields are not forwarded to the kernel. (Azure *U2M* still works on the kernel via `auth_type="azure-oauth"`, the browser flow.) | +| `azure_client_id` / `azure_client_secret` / `azure_tenant_id` | `str` | ✅ | ✅ | `None` | Azure service-principal (Entra ID M2M), selected by `auth_type="azure-sp-m2m"`. On the kernel path the connector forwards these to the kernel, which owns Azure resolution (Entra v2.0 token endpoint + the Databricks-resource `.default` scope) (#919). **`azure_tenant_id` is optional on the kernel path too** — like Thrift, the kernel auto-discovers it from the workspace's `/aad/auth` redirect when omitted. | +| `azure_workspace_resource_id` | `str` | ✅ | ✅ | `None` | For `azure-sp-m2m`. When set, the SP **management token** (`X-Databricks-Azure-SP-Management-Token`) + `X-Databricks-Azure-Workspace-Resource-Id` header are sent, to authorize an SP that has an Azure RBAC role but is not a workspace member. Omit it for a workspace-member SP (the data token authenticates alone; no management token is fetched). Works on both the kernel and Thrift paths. | | `_use_cert_as_auth` (+ `_tls_client_cert_file`) | `bool` | ✅ | ❌ | `False` | Authenticate with a TLS client certificate instead of a token. Thrift-only. | | `username` / `password` | `str` | ❌ | ❌ | `None` | **Removed.** Basic auth is no longer supported; passing either raises `ValueError`. | diff --git a/KERNEL_REV b/KERNEL_REV index f74d8d55b..6cd3da53d 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -9e3dbf9c40733b176151e001c9a15202030b967a +ad78a5be3dc8bb7fc78ec574492515ab24e23d4c diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index eb9a70e8e..b0aa9ea66 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -1,7 +1,7 @@ """Translate the connector's auth configuration into ``databricks_sql_kernel`` ``Session`` auth kwargs. -Three auth shapes are supported on the kernel path: +The following auth shapes are supported on the kernel path: - **PAT** — extracted from the built ``AuthProvider`` (works for ``AccessTokenAuthProvider``, including the ``TokenFederationProvider`` @@ -15,16 +15,37 @@ connector's own OAuth provider because the kernel re-mints tokens itself and the client secret is not recoverable from a built provider. -- **OAuth U2M** — for ``auth_type`` ``databricks-oauth`` (the browser - authorization-code flow), the connector's ``databricks-sql-python`` +- **OAuth U2M** — for ``auth_type`` ``databricks-oauth`` *or* ``azure-oauth`` + (the browser authorization-code flow), the connector's ``databricks-sql-python`` app bundle (``client_id`` + ``redirect_ports`` list, with the optional ``oauth_client_id`` / ``oauth_redirect_port`` overriding it) is forwarded to the kernel's ``auth_type='oauth-u2m'`` and the kernel - runs the browser flow itself. ``azure-oauth`` (Azure AD) is **not yet - supported** on the kernel path and is rejected with - ``NotSupportedError`` — the kernel resolves OAuth endpoints only from - the workspace-native OIDC config and cannot drive the Azure AD flow - (PECOBLR-4120). + runs the browser flow itself. +- **Azure Entra (Azure AD)** — the KERNEL is the Azure-aware auth core (it + owns the endpoints, scopes, app ids, and tenant discovery); the binding + forwards the selector + credentials and does not construct endpoints: + + - ``azure-oauth`` (U2M) → routed to the kernel's ``oauth-u2m`` via the shared + U2M path above, exactly like ``databricks-oauth``. The kernel runs the + in-house workspace-federated browser flow (the workspace federates login to + Entra), which Azure Databricks workspaces support; it forwards the + ``databricks-sql-python`` bundle, **not** the Thrift Azure app + (``96eecda7`` / redirect port ``8030``) — that app is registered for + Thrift's *direct-Entra* flow (``login.microsoftonline.com`` config + + ``{tenant}/user_impersonation`` scopes), which the kernel does not perform. + So on the kernel path ``azure-oauth`` and ``databricks-oauth`` are the same + flow (PECOBLR-4120). + - ``azure-sp-m2m`` (M2M) → forward ``auth_type='azure-sp-m2m'`` with the + Azure service-principal ``azure_client_id`` / ``azure_client_secret`` (plus + optional ``azure_tenant_id`` / ``azure_workspace_resource_id``). The kernel + builds the Entra v2.0 token endpoint and the ``{effective_app_id}/.default`` + scope, and auto-discovers the tenant from the workspace's ``/aad/auth`` + redirect when ``azure_tenant_id`` is omitted. When + ``azure_workspace_resource_id`` is set, the kernel also sends the Azure SP + management token + ``X-Databricks-Azure-Workspace-Resource-Id`` header + (matching the JDBC driver), so an RBAC-only SP — one with an Azure role but + no workspace membership — can authenticate; omit it for a workspace-member + SP (PECOBLR-4141). ``identity_federation_client_id`` is forwarded with whichever auth shape wins resolution. It selects mandatory SP-wide workload-identity token @@ -151,11 +172,13 @@ def kernel_auth_kwargs( rather than silently picking one flow (and failing later as a confusing 401 against the wrong principal): - a custom ``credentials_provider`` *and* M2M kwargs together; - - a U2M ``auth_type`` (``databricks-oauth``) *and* + - a U2M ``auth_type`` (``databricks-oauth`` / ``azure-oauth``) *and* ``oauth_client_secret`` together. - (``azure-oauth`` is rejected as unsupported before these guards — - PECOBLR-4120.) + (The ``azure-sp-m2m`` Azure Entra auth type is handled up front, before + these guards, forwarding to the kernel's Azure SP flow. ``azure-oauth`` is + a U2M type and is handled by the OAuth U2M step alongside + ``databricks-oauth``. See the module docstring.) 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 @@ -165,15 +188,16 @@ def kernel_auth_kwargs( both present → forward raw creds to the kernel's ``oauth-m2m``. 3. **PAT** — the built provider is (or wraps) an ``AccessTokenAuthProvider`` → extract the bearer token. - 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 - bare U2M connection authenticates as ``databricks-sql-python`` — - forwarding the connector's own OAuth app rather than the kernel's - ``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. **OAuth U2M** — ``auth_type`` is ``databricks-oauth`` or ``azure-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 bare U2M connection authenticates as + ``databricks-sql-python`` — forwarding the connector's own OAuth app + rather than the kernel's ``databricks-sql-connector`` default + (PECOBLR-4039/4040). ``azure-oauth`` resolves identically: the kernel's + workspace-federated flow serves Azure workspaces too (PECOBLR-4120). + Unlike the Thrift path, a caller-supplied ``oauth_scopes`` is honored here. 5. **Custom credentials_provider** → ``NotSupportedError`` (opaque token source; no raw creds for the kernel to own). 6. Anything else → ``NotSupportedError``. @@ -198,24 +222,91 @@ def kernel_auth_kwargs( # 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 - # request gets a clear "not supported" error rather than being silently - # misrouted (e.g. azure-oauth + client_id + secret would otherwise look - # like M2M). The kernel resolves OAuth endpoints only from the - # workspace-native OIDC config and has no Azure AD path, so the Thrift - # azure-oauth flow (AAD token endpoint + /user_impersonation scope, see - # AzureOAuthEndpointCollection) cannot be reproduced here. Forwarding an - # azure bundle would authenticate against the wrong endpoints, so we fail - # loudly at session-open. Tracked by PECOBLR-4120. - if auth_type == "azure-oauth": - raise NotSupportedError( - "use_kernel=True does not support auth_type='azure-oauth' (Azure " - "AD U2M) yet: the kernel resolves OAuth endpoints only from the " - "workspace-native OIDC configuration and cannot drive the Azure AD " - "authorization/token flow. Use the Thrift backend (default) for " - "azure-oauth. Tracked by PECOBLR-4120." - ) + # Azure Entra (Azure AD) auth types route to the kernel's GENERIC OAuth + # flows with Azure values supplied as overrides — the kernel needs no + # Azure-specific code. Handled up front, keyed on the explicit auth_type, + # before the generic M2M/PAT/U2M routing below (azure-sp-m2m carries its + # creds in azure_* kwargs, not oauth_client_id/secret, so it would + # otherwise fall through to the final "unsupported" error). + + # azure-oauth (Azure AD U2M) routes to the kernel's oauth-u2m, handled by the + # shared U2M branch below (step 3) alongside databricks-oauth. The kernel runs + # the in-house workspace-federated browser flow — the workspace federates the + # login to Microsoft Entra — which Azure Databricks workspaces support. It uses + # the connector's databricks-sql-python app bundle, NOT the Thrift Azure app + # (96eecda7 / redirect port 8030): that app is registered for Thrift's + # direct-Entra flow (login.microsoftonline.com config + {tenant}/user_impersonation + # scopes), which the kernel does not perform. So on the kernel path azure-oauth + # and databricks-oauth are the same flow. PECOBLR-4120. + + # azure-sp-m2m (Azure service principal, client-credentials): forward the + # selector + Azure SP credentials; the KERNEL owns Azure resolution (it is + # the auth core). The kernel builds the Entra v2.0 token endpoint + # (`{login}/{tenant}/oauth2/v2.0/token`) and the `{effective_app_id}/.default` + # scope, and — when azure_tenant_id is omitted — auto-discovers the tenant + # from the workspace's /aad/auth redirect, matching the Thrift backend + # (so connect() is byte-identical between Thrift and use_kernel=True). + # PECOBLR-4141. + # + # The Authorization bearer is the Databricks-audience data token, which alone + # authenticates a workspace-member SP. When azure_workspace_resource_id is + # set, the kernel also sends the Azure SP management token + + # X-Databricks-Azure-Workspace-Resource-Id header (matching the JDBC driver), + # so an RBAC-only SP (Azure role, not a workspace member) can authenticate. + # + # This branch returns BEFORE the ambiguity guards below, so an + # azure-sp-m2m selector paired with a conflicting OAuth signal + # (oauth_client_secret / oauth_jwt_key_file / credentials_provider) does not + # fail loudly the way the other flows do. That is intentional: azure-sp-m2m + # carries its credentials in the azure_* namespace, so there is no routing + # collision, and an explicit azure-sp-m2m selector is unambiguous intent to + # use the Azure SP flow. Any oauth_*/credentials_provider values are simply + # ignored here rather than treated as ambiguous. + if auth_type == "azure-sp-m2m": + azure_client_id = opts.get("azure_client_id") + azure_client_secret = opts.get("azure_client_secret") + if not (azure_client_id and azure_client_secret): + raise ProgrammingError( + "auth_type='azure-sp-m2m' requires azure_client_id and " + "azure_client_secret." + ) + # azure-sp-m2m carries its credentials in the azure_* namespace, so a + # conflicting OAuth signal is not a routing collision and does not fail + # loudly the way the other flows do (see the routing note above). Still, + # leave a breadcrumb: a caller who paired azure-sp-m2m with an + # oauth_*/credentials_provider value likely misunderstood the routing, + # and those values are silently ignored here. + ignored_signals = [ + name + for name in ("oauth_client_secret", "oauth_jwt_key_file", "credentials_provider") + if opts.get(name) is not None + ] + if ignored_signals: + logger.warning( + "auth_type='azure-sp-m2m' selected; ignoring conflicting " + "credential signal(s) %s (Azure SP credentials in the azure_* " + "namespace take precedence on the kernel path).", + ", ".join(ignored_signals), + ) + kwargs = { + "auth_type": "azure-sp-m2m", + "azure_client_id": azure_client_id, + "azure_client_secret": azure_client_secret, + } + # Optional passthroughs: the kernel auto-discovers the tenant when + # absent. When azure_workspace_resource_id is set, the kernel sends the + # Azure SP management token + X-Databricks-Azure-Workspace-Resource-Id + # header (for an SP with an Azure RBAC role but no workspace membership) + # — matching the JDBC driver; omit it for a workspace-member SP. + azure_tenant_id = opts.get("azure_tenant_id") + if azure_tenant_id: + kwargs["azure_tenant_id"] = azure_tenant_id + azure_workspace_resource_id = opts.get("azure_workspace_resource_id") + if azure_workspace_resource_id: + kwargs["azure_workspace_resource_id"] = azure_workspace_resource_id + if federation_client_id: + kwargs["identity_federation_client_id"] = federation_client_id + return kwargs # 0. Ambiguity guards — fail before any flow is chosen. if client_secret and opts.get("credentials_provider") is not None: @@ -226,7 +317,7 @@ def kernel_auth_kwargs( "kernel-managed M2M, or use the Thrift backend (default) for " "credentials_provider." ) - if client_secret and auth_type == "databricks-oauth": + if client_secret and auth_type in ("databricks-oauth", "azure-oauth"): raise NotSupportedError( f"Ambiguous auth on use_kernel=True: auth_type={auth_type!r} selects " "the U2M browser flow, but oauth_client_secret was also provided " @@ -248,7 +339,7 @@ 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": + if has_jwt_m2m and auth_type in ("databricks-oauth", "azure-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 " @@ -336,7 +427,10 @@ def kernel_auth_kwargs( return kwargs # 4. OAuth U2M — browser authorization-code flow; the kernel runs it. - # Only databricks-oauth reaches here (azure-oauth rejected up front). + # Both databricks-oauth and azure-oauth reach here: on the kernel path + # they are the same in-house workspace-federated flow (Azure workspaces + # federate the browser login to Entra), so azure-oauth is an alias for + # databricks-oauth here — see the note near the top of this function. # Forward the connector's own databricks-sql-python bundle instead of # the kernel's databricks-sql-connector default, for parity with the # Thrift path. client_id + redirect ports are coupled per app (each @@ -348,7 +442,7 @@ def kernel_auth_kwargs( # fallback), mirroring the Thrift DatabricksOAuthProvider which retries # the next port when one is bound. A caller overriding client_id # supplies its own single registered port. - if auth_type == "databricks-oauth": + if auth_type in ("databricks-oauth", "azure-oauth"): redirect_port = opts.get("oauth_redirect_port") # Honor a caller-supplied oauth_scopes (normalized to a list of # strings); fall back to the connector default when none is given. @@ -389,8 +483,10 @@ def kernel_auth_kwargs( ) raise NotSupportedError( f"use_kernel=True requires PAT (access_token), OAuth M2M " - f"(oauth_client_id + oauth_client_secret), or OAuth U2M " - f"(auth_type='databricks-oauth'), but got " + f"(oauth_client_id + oauth_client_secret), OAuth M2M JWT " + f"(oauth_client_id + oauth_jwt_key_file + oauth_jwt_kid), OAuth U2M " + f"(auth_type='databricks-oauth' or 'azure-oauth'), or Azure SP M2M " + f"(auth_type='azure-sp-m2m'), but got " f"{provider_desc} with auth_type={auth_type!r}. Use the Thrift " "backend (default) for other auth flows." ) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index eeb496f73..01274078f 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -309,7 +309,10 @@ def open_session( auth_kwargs: Dict[str, Any] = {} tls_kwargs: Dict[str, Any] = {} try: - auth_kwargs = kernel_auth_kwargs(self._auth_provider, self._auth_options) + auth_kwargs = kernel_auth_kwargs( + self._auth_provider, + self._auth_options, + ) # Translate the connector's SSLOptions into the kernel's # ``tls_*`` Session kwargs. Empty when TLS is at defaults. tls_kwargs = _kernel_tls_kwargs(self._ssl_options) diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index a62b0d081..f35cdf525 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -186,6 +186,20 @@ def _create_backend( "identity_federation_client_id": kwargs.get( "identity_federation_client_id" ), + # Azure Entra SP credentials for the azure-sp-m2m path. The + # kernel owns Azure resolution (endpoint/scope/tenant discovery), + # so these raw kwargs are the only source; without threading them + # the bridge would fail with "requires azure_client_id". The + # tenant and workspace-resource-id are optional (the kernel + # auto-discovers the tenant; the resource id adds the + # management-token resource-id header for an RBAC-only SP). + # Kernel-only; Thrift / SEA are unaffected. + "azure_client_id": kwargs.get("azure_client_id"), + "azure_client_secret": kwargs.get("azure_client_secret"), + "azure_tenant_id": kwargs.get("azure_tenant_id"), + "azure_workspace_resource_id": kwargs.get( + "azure_workspace_resource_id" + ), } # Forward the connector's retry-tuning kwargs so the kernel's # own retry policy honours them (the kernel owns the retry diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index d2a10c879..abce5a0f4 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -8,9 +8,10 @@ look through the wrapper). - OAuth M2M (``oauth_client_id`` + ``oauth_client_secret``) routes through ``auth_type='oauth-m2m'`` with the raw creds forwarded. - - OAuth U2M (``auth_type='databricks-oauth'``) routes through - ``auth_type='oauth-u2m'``. ``azure-oauth`` (Azure AD) is not yet - supported on the kernel path and is rejected (PECOBLR-4120). + - OAuth U2M (``auth_type='databricks-oauth'`` or ``'azure-oauth'``) + routes through ``auth_type='oauth-u2m'``; on the kernel path + ``azure-oauth`` is an alias for ``databricks-oauth`` (the in-house + workspace-federated flow serves Azure workspaces too) (PECOBLR-4120). - A custom ``credentials_provider`` and any other non-PAT shape raise ``NotSupportedError`` with a clear, actionable message. """ @@ -224,7 +225,10 @@ def test_m2m_forwards_token_url(self): "token_url": "https://login.microsoftonline.com/t/oauth2/v2.0/token", }, ) - assert kwargs["token_url"] == "https://login.microsoftonline.com/t/oauth2/v2.0/token" + assert ( + kwargs["token_url"] + == "https://login.microsoftonline.com/t/oauth2/v2.0/token" + ) def test_m2m_normalizes_space_delimited_scopes(self): # DatabricksOAuthProvider stores scopes as a single @@ -368,9 +372,11 @@ 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). + @pytest.mark.parametrize("u2m_auth_type", ["databricks-oauth", "azure-oauth"]) + def test_jwt_plus_u2m_auth_type_is_rejected(self, u2m_auth_type): + # A U2M auth_type signals browser-flow intent; a private key alongside + # it is ambiguous (mirrors the shared-secret M2M + U2M guard). Both U2M + # types must trip it — azure-oauth is a U2M type on the kernel path. with pytest.raises(NotSupportedError, match="oauth_jwt_key_file"): kernel_auth_kwargs( None, @@ -378,7 +384,7 @@ def test_jwt_plus_databricks_oauth_auth_type_is_rejected(self): "oauth_client_id": "sp", "oauth_jwt_key_file": "/k.pem", "oauth_jwt_kid": "k", - "auth_type": "databricks-oauth", + "auth_type": u2m_auth_type, }, ) @@ -396,7 +402,7 @@ def test_federation_client_id_forwarded(self): class TestKernelOAuthU2M: - """Only ``databricks-oauth`` U2M is supported on the kernel path. + """``databricks-oauth`` and ``azure-oauth`` both route to the kernel's U2M. The kernel core default U2M app is ``databricks-sql-connector`` / ``sql offline_access`` / port 8030 (see PECOBLR-4039). The Python @@ -407,9 +413,10 @@ class TestKernelOAuthU2M: may override ``oauth_scopes``; absent one, ``PYSQL_OAUTH_SCOPES`` is forwarded as the default. - ``azure-oauth`` (Azure AD) is deliberately NOT handled yet — the - kernel can't drive the Azure AD authorization/token flow — so it is - rejected up front (PECOBLR-4120).""" + ``azure-oauth`` (Azure AD U2M) resolves identically to ``databricks-oauth`` + on the kernel path: the kernel runs the in-house workspace-federated flow + (which Azure workspaces support), not Thrift's direct-Entra flow — see + ``test_azure_oauth_maps_to_in_house_u2m`` (PECOBLR-4120).""" def test_bare_databricks_oauth_forwards_full_python_bundle(self): # No overrides → forward the databricks-sql-python bundle in full @@ -427,23 +434,42 @@ def test_bare_databricks_oauth_forwards_full_python_bundle(self): "oauth_scopes": list(PYSQL_OAUTH_SCOPES), } - @pytest.mark.parametrize( - "opts", - [ + def test_azure_oauth_maps_to_in_house_u2m(self): + # azure-oauth (Azure AD U2M) routes to the kernel's oauth-u2m exactly + # like databricks-oauth: the kernel runs the in-house workspace-federated + # browser flow (the workspace federates login to Entra), which Azure + # workspaces support. It forwards the same databricks-sql-python bundle + # — NOT the Thrift Azure app (96eecda7 / port 8030), which is registered + # for the direct-Entra flow the kernel does not perform. PECOBLR-4120. + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), {"auth_type": "azure-oauth"}, - {"auth_type": "azure-oauth", "oauth_client_id": "custom"}, - {"auth_type": "azure-oauth", "oauth_redirect_port": 8030}, - ], - ids=["bare", "with_client_id", "with_port"], - ) - def test_azure_oauth_not_supported(self, opts): - # azure-oauth (Azure AD U2M) can't work through the kernel yet: the - # kernel resolves OAuth endpoints only from workspace-native OIDC - # discovery and has no Azure AD path. Fail loudly at session-open - # rather than forwarding a bundle that authenticates against the - # wrong endpoints. Tracked by PECOBLR-4120. - with pytest.raises(NotSupportedError, match="azure-oauth"): - kernel_auth_kwargs(_FakeOAuthProvider(), opts) + ) + assert kwargs == { + "auth_type": "oauth-u2m", + "client_id": PYSQL_OAUTH_CLIENT_ID, + "redirect_ports": list(PYSQL_OAUTH_REDIRECT_PORT_RANGE), + "oauth_scopes": list(PYSQL_OAUTH_SCOPES), + } + + def test_azure_oauth_honors_custom_client_id_port_and_scopes(self): + # The same caller overrides available to databricks-oauth work for + # azure-oauth (they share the U2M branch). + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), + { + "auth_type": "azure-oauth", + "oauth_client_id": "custom-client", + "oauth_scopes": ["custom-scope", "offline_access"], + "oauth_redirect_port": 9999, + }, + ) + assert kwargs == { + "auth_type": "oauth-u2m", + "client_id": "custom-client", + "redirect_ports": [9999], + "oauth_scopes": ["custom-scope", "offline_access"], + } def test_u2m_custom_client_id_port_and_scopes_honored(self): # A caller may override the coupled client_id + redirect port and the @@ -612,21 +638,130 @@ def _creds_provider(): }, ) - def test_u2m_auth_type_plus_client_secret_is_rejected(self): - # User asked for U2M (browser) but also passed a secret (M2M). - # Don't silently route M2M against the wrong principal. (azure-oauth - # is rejected earlier as unsupported, so it's not exercised here.) + @pytest.mark.parametrize("u2m_auth_type", ["databricks-oauth", "azure-oauth"]) + def test_u2m_auth_type_plus_client_secret_is_rejected(self, u2m_auth_type): + # User asked for U2M (browser) but also passed a secret (M2M). Don't + # silently route M2M against the wrong principal. Both U2M auth types + # (databricks-oauth and azure-oauth) must trip the ambiguity guard — + # azure-oauth is a U2M type on the kernel path, so without the guard it + # would fall through to the M2M branch (oauth_client_id + secret). with pytest.raises(NotSupportedError, match="Ambiguous auth"): kernel_auth_kwargs( _FakeOAuthProvider(), { - "auth_type": "databricks-oauth", + "auth_type": u2m_auth_type, "oauth_client_id": "id", "oauth_client_secret": "sec", }, ) +class TestKernelAzureSpM2M: + """``azure-sp-m2m`` (Azure service-principal, client-credentials) forwards + the Azure SP credentials to the KERNEL, which owns Azure resolution: it + builds the Entra token endpoint + ``{app_id}/.default`` scope and + auto-discovers the tenant from the workspace when ``azure_tenant_id`` is + omitted (Thrift parity). The binding stays thin — it does not construct + endpoints or scopes. PECOBLR-4141.""" + + _CREDS = { + "auth_type": "azure-sp-m2m", + "azure_client_id": "azure-sp", + "azure_client_secret": "azure-secret", + "azure_tenant_id": "tenant-123", + } + + def test_azure_sp_m2m_forwards_creds_to_kernel(self): + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), + dict(self._CREDS), + ) + # Thin forwarding: the kernel owns endpoint/scope resolution, so no + # token_url / oauth_scopes are constructed here. + assert kwargs == { + "auth_type": "azure-sp-m2m", + "azure_client_id": "azure-sp", + "azure_client_secret": "azure-secret", + "azure_tenant_id": "tenant-123", + } + + def test_azure_sp_m2m_tenant_optional_kernel_autodiscovers(self): + # Unlike the earlier kernel slice, the kernel now auto-discovers the + # tenant from the workspace's /aad/auth redirect (Thrift parity), so + # omitting azure_tenant_id must NOT raise — the key is simply absent. + opts = { + "auth_type": "azure-sp-m2m", + "azure_client_id": "azure-sp", + "azure_client_secret": "azure-secret", + } + kwargs = kernel_auth_kwargs(_FakeOAuthProvider(), opts) + assert kwargs == { + "auth_type": "azure-sp-m2m", + "azure_client_id": "azure-sp", + "azure_client_secret": "azure-secret", + } + assert "azure_tenant_id" not in kwargs + + def test_azure_sp_m2m_forwards_workspace_resource_id(self): + # The kernel always sends the Azure SP management token and, when + # azure_workspace_resource_id is set, adds the resource-id header (for an + # RBAC-only SP) — so the bridge forwards it rather than dropping it. + opts = dict( + self._CREDS, + azure_workspace_resource_id="/subscriptions/s/resourceGroups/rg/workspace/w", + ) + kwargs = kernel_auth_kwargs(_FakeOAuthProvider(), opts) + assert ( + kwargs["azure_workspace_resource_id"] + == "/subscriptions/s/resourceGroups/rg/workspace/w" + ) + + def test_azure_sp_m2m_omits_workspace_resource_id_when_absent(self): + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), + dict(self._CREDS), + ) + assert "azure_workspace_resource_id" not in kwargs + + def test_azure_sp_m2m_requires_client_id_and_secret(self): + with pytest.raises(ProgrammingError, match="azure_client_id"): + kernel_auth_kwargs( + _FakeOAuthProvider(), + {"auth_type": "azure-sp-m2m", "azure_tenant_id": "t"}, + ) + + def test_azure_sp_m2m_forwards_federation_client_id(self): + opts = dict(self._CREDS, identity_federation_client_id="fed-client") + kwargs = kernel_auth_kwargs(_FakeOAuthProvider(), opts) + assert kwargs["identity_federation_client_id"] == "fed-client" + + @pytest.mark.parametrize( + "conflicting_signal", + [ + {"oauth_client_secret": "oauth-secret"}, + {"oauth_jwt_key_file": "/tmp/key.pem"}, + {"credentials_provider": object()}, + ], + ) + def test_azure_sp_m2m_ignores_conflicting_oauth_signal(self, conflicting_signal): + # Intentional asymmetry: every OTHER flow treats a conflicting credential + # signal as a hard "Ambiguous auth" error, but an explicit azure-sp-m2m + # selector carries its creds in the azure_* namespace, so a stray + # oauth_*/credentials_provider value is NOT a routing collision — it is + # silently ignored (logger.warning breadcrumb only) and the Azure SP flow + # still wins. This guards against a refactor accidentally promoting the + # ignored signal to an error (see the routing note in auth_bridge.py). + opts = dict(self._CREDS, **conflicting_signal) + kwargs = kernel_auth_kwargs(_FakeOAuthProvider(), opts) + # Does NOT raise, and routes to the Azure SP flow with only azure kwargs. + assert kwargs == { + "auth_type": "azure-sp-m2m", + "azure_client_id": "azure-sp", + "azure_client_secret": "azure-secret", + "azure_tenant_id": "tenant-123", + } + + class TestKernelScopesNormalization: def test_unknown_scope_type_raises(self): # A non-str/list/tuple oauth_scopes is a caller error; fail loudly diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index ba008b103..6fcefcade 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -477,6 +477,55 @@ def test_retry_kwargs_threaded_into_kernel_client(self): finally: conn.close() + def test_azure_sp_m2m_kwargs_threaded_into_kernel_auth_options(self): + # The Azure SP credentials a user passes to connect() must reach the + # kernel auth bridge via auth_options; without this threading the + # azure-sp-m2m path would fail at session-open with "requires + # azure_client_id". Guards the session.py -> kernel_auth_options map. + import sys + import types + + pytest.importorskip( + "pyarrow", + reason="kernel client module imports pyarrow at load", + ) + + fake = types.ModuleType("databricks_sql_kernel") + fake.KernelError = type("KernelError", (Exception,), {}) + fake.Session = MagicMock() + + with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch( + "databricks.sql.backend.kernel.client.KernelDatabricksClient" + ) as mock_kernel_client, patch( + "%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE + ): + instance = mock_kernel_client.return_value + instance.open_session.return_value = SessionId( + BackendType.SEA, "sess-id", None + ) + + conn = databricks.sql.connect( + server_hostname="foo", + http_path="/sql/1.0/warehouses/abc", + use_kernel=True, + auth_type="azure-sp-m2m", + azure_client_id="azure-sp", + azure_client_secret="azure-secret", + azure_tenant_id="tenant-123", + azure_workspace_resource_id="/subscriptions/s/rg/w", + enable_telemetry=False, + ) + try: + _, kwargs = mock_kernel_client.call_args + opts = kwargs["auth_options"] + assert opts["auth_type"] == "azure-sp-m2m" + assert opts["azure_client_id"] == "azure-sp" + assert opts["azure_client_secret"] == "azure-secret" + assert opts["azure_tenant_id"] == "tenant-123" + assert opts["azure_workspace_resource_id"] == "/subscriptions/s/rg/w" + finally: + conn.close() + class TestKernelUserAgentForwarding: """user_agent_entry must reach the kernel on the use_kernel path —