From 8f2a73a640cee1e1aae56fb62fcb119f0c144be1 Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Tue, 18 Aug 2026 16:04:05 -0700 Subject: [PATCH 1/7] feat(kernel): thread Azure Entra OAuth (U2M + SP M2M) through the auth bridge The kernel auth bridge rejected azure-oauth and had no azure-sp-m2m path. Route both Azure auth types onto the kernel's generic OAuth flows (the kernel needs no Azure-specific code; PR databricks/databricks-sql-kernel#263 added the token_url/scope override plumbing this relies on): - azure-oauth (Azure AD U2M) -> oauth-u2m with the Azure app client id (96eecda7-...), redirect port 8030, and the {app_id}/user_impersonation offline_access delegated scope (via AzureOAuthEndpointCollection, honoring DATABRICKS_AZURE_TENANT_ID). The kernel discovers endpoints via the workspace /oidc redirector. (PECOBLR-4120) - azure-sp-m2m (Azure service principal) -> oauth-m2m with the Azure creds, an Entra v2.0 token_url, and the {effective_app_id}/.default scope. Requires an explicit azure_tenant_id (the kernel path does not auto-discover it). The management-token header / azure_workspace_resource_id are not applied on the kernel path -- no SQL connector uses them, matching Go and Node. (PECOBLR-4141) kernel_auth_kwargs now takes hostname (for the effective Azure app id); the client passes self._server_hostname. TDD: replaced the azure-oauth NotSupportedError test with routing tests and added a TestKernelAzureSpM2M suite (routing, required tenant/creds, federation client id). 50 bridge tests pass; black clean. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 --- CHANGELOG.md | 3 +- .../sql/backend/kernel/auth_bridge.py | 134 ++++++++++++++---- src/databricks/sql/backend/kernel/client.py | 6 +- tests/unit/test_kernel_auth_bridge.py | 113 ++++++++++++--- 4 files changed, 211 insertions(+), 45 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f750fa29b..65c18f390 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,8 @@ # Release History # Unreleased -- 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) OAuth is now supported.** Both connector Azure auth types route onto the kernel's generic OAuth flows: `auth_type="azure-oauth"` (Azure AD U2M) forwards the Azure app bundle (client id `96eecda7-…`, redirect port `8030`, and the `{app_id}/user_impersonation offline_access` delegated scope, honoring `DATABRICKS_AZURE_TENANT_ID`) to the kernel's U2M browser flow; `auth_type="azure-sp-m2m"` (Azure service principal) forwards the Azure credentials with an Entra v2.0 token endpoint and the `{app_id}/.default` scope to the kernel's M2M flow. `azure-sp-m2m` requires an explicit `azure_tenant_id` on the kernel path (it does not auto-discover the tenant as the Thrift path does), and the Azure management-token header / `azure_workspace_resource_id` are not applied — the Databricks-audience token authenticates service principals that are workspace principals, matching the Go and Node SQL connectors (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/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 95374ba77..6856befbf 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -20,11 +20,24 @@ 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)** — both Azure auth types route to the + kernel's *generic* OAuth flows with Azure values as overrides (the + kernel needs no Azure-specific code): + + - ``azure-oauth`` (U2M) → ``oauth-u2m`` with the Azure app client id + (``96eecda7-…``), redirect port ``8030``, and the AAD delegated scope + ``{app_id}/user_impersonation offline_access`` (via + ``AzureOAuthEndpointCollection``, honoring ``DATABRICKS_AZURE_TENANT_ID``). + The kernel discovers endpoints via the workspace ``/oidc`` redirector, + which an Azure workspace redirects to Entra (PECOBLR-4120). + - ``azure-sp-m2m`` (M2M) → ``oauth-m2m`` with the Azure service-principal + credentials, an Entra v2.0 ``token_url``, and the + ``{effective_app_id}/.default`` scope (PECOBLR-4141). ``azure_tenant_id`` + is required (the kernel path does not auto-discover it). The Azure + management-token header and ``azure_workspace_resource_id`` are **not** + applied on the kernel path — no SQL connector uses them; an SP that is + not a workspace member (RBAC-only) is unsupported here. ``identity_federation_client_id`` is forwarded with whichever auth shape wins resolution. It selects mandatory SP-wide workload-identity token @@ -54,14 +67,23 @@ from typing import Any, Dict, Optional from databricks.sql.auth.auth import ( + PYSQL_OAUTH_AZURE_CLIENT_ID, + PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE, PYSQL_OAUTH_CLIENT_ID, PYSQL_OAUTH_REDIRECT_PORT_RANGE, PYSQL_OAUTH_SCOPES, ) from databricks.sql.auth.authenticators import AccessTokenAuthProvider, AuthProvider +from databricks.sql.auth.common import get_effective_azure_login_app_id +from databricks.sql.auth.endpoint import AzureOAuthEndpointCollection from databricks.sql.auth.token_federation import TokenFederationProvider from databricks.sql.exc import NotSupportedError, ProgrammingError +# Entra (Azure AD) v2.0 token endpoint template. The kernel's generic M2M +# provider sends the credentials as ``scope`` (v2.0), so we point it at the +# v2.0 endpoint (the connector's own SP path uses the v1.0 ``resource`` form). +_AZURE_AAD_LOGIN_HOST = "https://login.microsoftonline.com" + logger = logging.getLogger(__name__) @@ -134,6 +156,7 @@ def _extract_bearer_token(auth_provider: Optional[AuthProvider]) -> Optional[str def kernel_auth_kwargs( auth_provider: Optional[AuthProvider], auth_options: Optional[Dict[str, Any]] = None, + hostname: Optional[str] = None, ) -> Dict[str, Any]: """Build the kwargs passed to ``databricks_sql_kernel.Session(...)``. @@ -154,8 +177,9 @@ def kernel_auth_kwargs( - a U2M ``auth_type`` (``databricks-oauth``) *and* ``oauth_client_secret`` together. - (``azure-oauth`` is rejected as unsupported before these guards — - PECOBLR-4120.) + (The Azure Entra auth types — ``azure-oauth`` and ``azure-sp-m2m`` — + are routed to the kernel's generic OAuth flows up front, before these + guards; see the module docstring.) 1. **OAuth M2M** — ``oauth_client_id`` + ``oauth_client_secret`` both present → forward raw creds to the kernel's ``oauth-m2m``. 2. **PAT** — the built provider is (or wraps) an @@ -168,7 +192,6 @@ def kernel_auth_kwargs( 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. **Custom credentials_provider** → ``NotSupportedError`` (opaque token source; no raw creds for the kernel to own). 5. Anything else → ``NotSupportedError``. @@ -188,24 +211,87 @@ def kernel_auth_kwargs( auth_type = opts.get("auth_type") has_m2m = bool(client_id and client_secret) - # 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. + # 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): forward the Azure app bundle to oauth-u2m. + # The kernel runs the browser flow and discovers endpoints via the + # workspace /oidc redirector (which an Azure workspace redirects to Entra). + # The AAD delegated scope ({app_id}/user_impersonation [+ offline_access]) + # is synthesised via AzureOAuthEndpointCollection, which also honors the + # DATABRICKS_AZURE_TENANT_ID app-id override. 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." + redirect_port = opts.get("oauth_redirect_port") + caller_scopes = _normalize_scopes(opts.get("oauth_scopes")) + mapped_scopes = AzureOAuthEndpointCollection().get_scopes_mapping( + caller_scopes if caller_scopes is not None else list(PYSQL_OAUTH_SCOPES) ) + kwargs = { + "auth_type": "oauth-u2m", + "client_id": client_id or PYSQL_OAUTH_AZURE_CLIENT_ID, + "redirect_ports": ( + [_coerce_redirect_port(redirect_port)] + if client_id and redirect_port is not None + else list(PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE) + ), + "oauth_scopes": mapped_scopes, + } + if federation_client_id: + kwargs["identity_federation_client_id"] = federation_client_id + return kwargs + + # azure-sp-m2m (Azure service principal, client-credentials): forward to + # oauth-m2m with the Azure app credentials, an Entra v2.0 token endpoint, + # and the {effective_app_id}/.default scope. The kernel sends the client + # secret via HTTP Basic (which Entra v2.0 accepts) and, because a + # token_url override is set, skips workspace OIDC discovery. PECOBLR-4141. + # + # NOT applied on the kernel path: the Azure management-token header + # (X-Databricks-Azure-SP-Management-Token) and azure_workspace_resource_id. + # No SQL connector (Go, Node) uses them; the Databricks-audience token + # authenticates SPs that are workspace principals (the SQL norm). An SP with + # only an Azure RBAC role (not a workspace member) is unsupported here. + if auth_type == "azure-sp-m2m": + azure_client_id = opts.get("azure_client_id") + azure_client_secret = opts.get("azure_client_secret") + azure_tenant_id = opts.get("azure_tenant_id") + if not (azure_client_id and azure_client_secret): + raise ProgrammingError( + "auth_type='azure-sp-m2m' requires azure_client_id and " + "azure_client_secret." + ) + if not azure_tenant_id: + # The Thrift path auto-discovers the tenant from the workspace's + # /aad/auth redirect; the kernel path does not make that call, so + # require it explicitly rather than silently guessing. + raise NotSupportedError( + "use_kernel=True auth_type='azure-sp-m2m' requires an explicit " + "azure_tenant_id (the kernel path does not auto-discover the " + "Azure tenant from the workspace as the Thrift backend does)." + ) + if opts.get("azure_workspace_resource_id"): + logger.warning( + "azure_workspace_resource_id is ignored on use_kernel=True: the " + "Azure management-token flow (X-Databricks-Azure-SP-Management-" + "Token) is not applied on the kernel path. The Databricks-" + "audience token authenticates service principals that are " + "workspace principals; an RBAC-only SP is unsupported here." + ) + app_id = get_effective_azure_login_app_id(hostname or "") + kwargs = { + "auth_type": "oauth-m2m", + "client_id": azure_client_id, + "client_secret": azure_client_secret, + "token_url": f"{_AZURE_AAD_LOGIN_HOST}/{azure_tenant_id}/oauth2/v2.0/token", + "oauth_scopes": [f"{app_id}/.default"], + } + 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: diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 8df7e887d..0e5d0c61e 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -309,7 +309,11 @@ 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, + hostname=self._server_hostname, + ) # 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/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index f60943948..d9b381388 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -31,7 +31,11 @@ PYSQL_OAUTH_CLIENT_ID, PYSQL_OAUTH_SCOPES, PYSQL_OAUTH_REDIRECT_PORT_RANGE, + PYSQL_OAUTH_AZURE_CLIENT_ID, + PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE, ) +from databricks.sql.auth.common import get_effective_azure_login_app_id +from databricks.sql.auth.endpoint import AzureOAuthEndpointCollection from databricks.sql.auth.authenticators import ( AccessTokenAuthProvider, AuthProvider, @@ -257,9 +261,8 @@ 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) routes here too — see + ``test_azure_oauth_routes_to_kernel_u2m`` (PECOBLR-4120).""" def test_bare_databricks_oauth_forwards_full_python_bundle(self): # No overrides → forward the databricks-sql-python bundle in full @@ -277,23 +280,40 @@ def test_bare_databricks_oauth_forwards_full_python_bundle(self): "oauth_scopes": list(PYSQL_OAUTH_SCOPES), } - @pytest.mark.parametrize( - "opts", - [ + def test_azure_oauth_routes_to_kernel_u2m(self): + # azure-oauth (Azure AD U2M) now routes to the kernel's oauth-u2m with + # the Azure app bundle: the Azure client id, its registered port 8030, + # and the AAD delegated scope ({app_id}/user_impersonation + + # offline_access). The kernel discovers endpoints via the workspace + # /oidc redirector (which an Azure workspace redirects to Entra). + # 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_AZURE_CLIENT_ID, + "redirect_ports": list(PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE), + "oauth_scopes": AzureOAuthEndpointCollection().get_scopes_mapping( + list(PYSQL_OAUTH_SCOPES) + ), + } + # Sanity: the mapped scope is the AAD delegated form, not `sql`. + assert any(s.endswith("/user_impersonation") for s in kwargs["oauth_scopes"]) + assert "offline_access" in kwargs["oauth_scopes"] + + def test_azure_oauth_honors_custom_client_id_and_port(self): + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), + { + "auth_type": "azure-oauth", + "oauth_client_id": "custom-azure-app", + "oauth_redirect_port": 9100, + }, + ) + assert kwargs["client_id"] == "custom-azure-app" + assert kwargs["redirect_ports"] == [9100] def test_u2m_custom_client_id_port_and_scopes_honored(self): # A caller may override the coupled client_id + redirect port and the @@ -477,6 +497,61 @@ def test_u2m_auth_type_plus_client_secret_is_rejected(self): ) +class TestKernelAzureSpM2M: + """``azure-sp-m2m`` (Azure service-principal, client-credentials) routes to + the kernel's generic ``oauth-m2m`` with an Entra v2.0 token endpoint and the + ``{app_id}/.default`` scope. The management-token header is intentionally not + applied on the kernel path (no SQL connector uses it). 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_routes_to_kernel_m2m(self): + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), + dict(self._CREDS), + hostname="adb-1.azuredatabricks.net", + ) + app_id = get_effective_azure_login_app_id("adb-1.azuredatabricks.net") + assert kwargs == { + "auth_type": "oauth-m2m", + "client_id": "azure-sp", + "client_secret": "azure-secret", + "token_url": "https://login.microsoftonline.com/tenant-123/oauth2/v2.0/token", + "oauth_scopes": [f"{app_id}/.default"], + } + + def test_azure_sp_m2m_requires_tenant(self): + opts = { + "auth_type": "azure-sp-m2m", + "azure_client_id": "azure-sp", + "azure_client_secret": "azure-secret", + } + with pytest.raises(NotSupportedError, match="azure_tenant_id"): + kernel_auth_kwargs( + _FakeOAuthProvider(), opts, hostname="adb-1.azuredatabricks.net" + ) + + 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"}, + hostname="adb-1.azuredatabricks.net", + ) + + 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, hostname="adb-1.azuredatabricks.net" + ) + assert kwargs["identity_federation_client_id"] == "fed-client" + + class TestKernelScopesNormalization: def test_unknown_scope_type_raises(self): # A non-str/list/tuple oauth_scopes is a caller error; fail loudly From c2dec51f107d6a4e24fa597572df0f6917db6802 Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Tue, 18 Aug 2026 16:32:26 -0700 Subject: [PATCH 2/7] docs: update CONNECTION_PARAMETERS.md for kernel Azure OAuth support The auth table marked the azure_* fields as Kernel-unsupported and claimed azure-oauth 'still works on the kernel' (it was actually rejected). Reflect the new routing: azure-sp-m2m + azure-oauth now work on the kernel path; azure_tenant_id is required there; the management token / azure_workspace_resource_id are not applied (matching Go/Node). Co-authored-by: Isaac Signed-off-by: eric-wang-1990 --- CONNECTION_PARAMETERS.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 69a72b0f2..5f526d792 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 work on the kernel path; `azure-oauth` / `azure-sp-m2m` kernel support added in #919 (they route onto the kernel's generic OAuth flows with Azure values). | | `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. | @@ -77,7 +77,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 these route onto OAuth M2M with an Entra v2.0 token endpoint + the Databricks-resource `.default` scope (#919). **`azure_tenant_id` is required on the kernel path** — unlike Thrift, it is not auto-discovered from the workspace. | +| `azure_workspace_resource_id` | `str` | ✅ | ⚠️ | `None` | Thrift sends this with the Azure SP **management token** (`X-Databricks-Azure-SP-Management-Token`) to authorize an SP that has an Azure RBAC role but is not a workspace member. **Not applied on the kernel path** — the management-token flow is unsupported there (matching the Go and Node SQL drivers, which don't use it); add the SP as a workspace principal instead. Setting it on the kernel path logs a warning and is otherwise ignored. | | `_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`. | From 01acb39f4dac75823394ffc79a2d32c99be3bca8 Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Tue, 18 Aug 2026 18:08:29 -0700 Subject: [PATCH 3/7] =?UTF-8?q?refactor(kernel):=20thin=20azure-oauth=20?= =?UTF-8?q?=E2=80=94=20kernel=20owns=20Azure=20resolution?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel is the auth core now: for azure-oauth the bridge forwards only auth_type='azure-oauth' (+ optional client_id/redirect_port passthrough), and the kernel pins the workspace v2.0 authorize/token endpoints, the Azure app client id, port 8030, and the user_impersonation scope. Drops the connector-side endpoint/scope construction (and the AzureOAuthEndpointCollection / PYSQL_OAUTH_AZURE_* imports) from the kernel path. Live-verified end-to-end against an Azure workspace. azure-sp-m2m still routes to oauth-m2m here pending the kernel's dedicated azure-sp-m2m variant. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 --- .../sql/backend/kernel/auth_bridge.py | 38 +++++++------------ tests/unit/test_kernel_auth_bridge.py | 37 +++++++----------- 2 files changed, 28 insertions(+), 47 deletions(-) diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 6856befbf..697036f23 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -67,15 +67,12 @@ from typing import Any, Dict, Optional from databricks.sql.auth.auth import ( - PYSQL_OAUTH_AZURE_CLIENT_ID, - PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE, PYSQL_OAUTH_CLIENT_ID, PYSQL_OAUTH_REDIRECT_PORT_RANGE, PYSQL_OAUTH_SCOPES, ) from databricks.sql.auth.authenticators import AccessTokenAuthProvider, AuthProvider from databricks.sql.auth.common import get_effective_azure_login_app_id -from databricks.sql.auth.endpoint import AzureOAuthEndpointCollection from databricks.sql.auth.token_federation import TokenFederationProvider from databricks.sql.exc import NotSupportedError, ProgrammingError @@ -218,28 +215,21 @@ def kernel_auth_kwargs( # 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): forward the Azure app bundle to oauth-u2m. - # The kernel runs the browser flow and discovers endpoints via the - # workspace /oidc redirector (which an Azure workspace redirects to Entra). - # The AAD delegated scope ({app_id}/user_impersonation [+ offline_access]) - # is synthesised via AzureOAuthEndpointCollection, which also honors the - # DATABRICKS_AZURE_TENANT_ID app-id override. PECOBLR-4120. + # azure-oauth (Azure AD U2M): forward the selector; the KERNEL owns Azure + # resolution (it is the auth core). The kernel pins the workspace v2.0 + # authorize/token endpoints (`{host}/oidc/oauth2/v2.0/{authorize,token}` — + # NOT the discovered `/oidc/v1/authorize`, which the workspace redirects to + # a malformed Entra URL), the Azure app client id, port 8030, and the + # `{app_id}/user_impersonation offline_access` scope. So this binding does + # NOT construct endpoints/scopes — it just passes `auth_type='azure-oauth'` + # plus any optional client_id / redirect_port passthrough. PECOBLR-4120. if auth_type == "azure-oauth": - redirect_port = opts.get("oauth_redirect_port") - caller_scopes = _normalize_scopes(opts.get("oauth_scopes")) - mapped_scopes = AzureOAuthEndpointCollection().get_scopes_mapping( - caller_scopes if caller_scopes is not None else list(PYSQL_OAUTH_SCOPES) - ) - kwargs = { - "auth_type": "oauth-u2m", - "client_id": client_id or PYSQL_OAUTH_AZURE_CLIENT_ID, - "redirect_ports": ( - [_coerce_redirect_port(redirect_port)] - if client_id and redirect_port is not None - else list(PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE) - ), - "oauth_scopes": mapped_scopes, - } + kwargs = {"auth_type": "azure-oauth"} + if client_id: + kwargs["client_id"] = client_id + redirect_port = opts.get("oauth_redirect_port") + if redirect_port is not None: + kwargs["redirect_ports"] = [_coerce_redirect_port(redirect_port)] if federation_client_id: kwargs["identity_federation_client_id"] = federation_client_id return kwargs diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index d9b381388..33ef81eb3 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -31,11 +31,8 @@ PYSQL_OAUTH_CLIENT_ID, PYSQL_OAUTH_SCOPES, PYSQL_OAUTH_REDIRECT_PORT_RANGE, - PYSQL_OAUTH_AZURE_CLIENT_ID, - PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE, ) from databricks.sql.auth.common import get_effective_azure_login_app_id -from databricks.sql.auth.endpoint import AzureOAuthEndpointCollection from databricks.sql.auth.authenticators import ( AccessTokenAuthProvider, AuthProvider, @@ -280,30 +277,21 @@ def test_bare_databricks_oauth_forwards_full_python_bundle(self): "oauth_scopes": list(PYSQL_OAUTH_SCOPES), } - def test_azure_oauth_routes_to_kernel_u2m(self): - # azure-oauth (Azure AD U2M) now routes to the kernel's oauth-u2m with - # the Azure app bundle: the Azure client id, its registered port 8030, - # and the AAD delegated scope ({app_id}/user_impersonation + - # offline_access). The kernel discovers endpoints via the workspace - # /oidc redirector (which an Azure workspace redirects to Entra). - # PECOBLR-4120. + def test_azure_oauth_forwards_selector_kernel_owns_resolution(self): + # azure-oauth (Azure AD U2M): the bridge forwards ONLY the selector. + # The kernel owns Azure resolution — it pins the workspace v2.0 + # authorize/token endpoints, the Azure client id, port 8030, and the + # {app_id}/user_impersonation scope. So the bridge must NOT construct + # client_id / redirect_ports / oauth_scopes here. PECOBLR-4120. kwargs = kernel_auth_kwargs( _FakeOAuthProvider(), {"auth_type": "azure-oauth"}, ) - assert kwargs == { - "auth_type": "oauth-u2m", - "client_id": PYSQL_OAUTH_AZURE_CLIENT_ID, - "redirect_ports": list(PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE), - "oauth_scopes": AzureOAuthEndpointCollection().get_scopes_mapping( - list(PYSQL_OAUTH_SCOPES) - ), - } - # Sanity: the mapped scope is the AAD delegated form, not `sql`. - assert any(s.endswith("/user_impersonation") for s in kwargs["oauth_scopes"]) - assert "offline_access" in kwargs["oauth_scopes"] + assert kwargs == {"auth_type": "azure-oauth"} def test_azure_oauth_honors_custom_client_id_and_port(self): + # A caller override still passes through (client_id + its coupled port), + # but no scopes/endpoints are synthesised by the bridge. kwargs = kernel_auth_kwargs( _FakeOAuthProvider(), { @@ -312,8 +300,11 @@ def test_azure_oauth_honors_custom_client_id_and_port(self): "oauth_redirect_port": 9100, }, ) - assert kwargs["client_id"] == "custom-azure-app" - assert kwargs["redirect_ports"] == [9100] + assert kwargs == { + "auth_type": "azure-oauth", + "client_id": "custom-azure-app", + "redirect_ports": [9100], + } def test_u2m_custom_client_id_port_and_scopes_honored(self): # A caller may override the coupled client_id + redirect port and the From 311e3f07b20b927ff89606801a251c721f703e70 Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Wed, 19 Aug 2026 00:56:10 -0700 Subject: [PATCH 4/7] =?UTF-8?q?refactor(kernel):=20thin=20azure-sp-m2m=20?= =?UTF-8?q?=E2=80=94=20kernel=20owns=20Azure=20resolution=20+=20optional?= =?UTF-8?q?=20mgmt=20token?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make the `azure-sp-m2m` bridge thin, matching the kernel becoming the Azure-aware auth core. The connector now forwards `auth_type='azure-sp-m2m'` + `azure_client_id` / `azure_client_secret` (and optional `azure_tenant_id` / `azure_workspace_resource_id`) straight to the kernel Session, instead of constructing the Entra token endpoint and `{app_id}/.default` scope itself. Behavior changes on the kernel path (Thrift parity): - `azure_tenant_id` is now OPTIONAL — the kernel auto-discovers the tenant from the workspace's `/aad/auth` redirect when omitted, exactly as the Thrift backend does. (Previously the kernel path required it.) - `azure_workspace_resource_id` is now honored as an optional add-on: forward it and the kernel fetches an Azure-management token and sends the `X-Databricks-Azure-SP-Management-Token` + `X-Databricks-Azure-Workspace-Resource-Id` pair, so an SP with only an Azure RBAC role (not a workspace member) can authenticate. (Previously it was dropped with a warning.) Also thread the `azure_*` connection kwargs into `kernel_auth_options` in session.py — without this the bridge never received them and `azure-sp-m2m` failed at session-open with "requires azure_client_id". Adds a regression test for that threading, and rewrites the bridge tests for thin forwarding (tenant optional, resource id forwarded). Drops the now-unused `get_effective_azure_login_app_id` import and `_AZURE_AAD_LOGIN_HOST` constant. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 --- CHANGELOG.md | 2 +- CONNECTION_PARAMETERS.md | 6 +- .../sql/backend/kernel/auth_bridge.py | 103 ++++++++---------- src/databricks/sql/session.py | 13 +++ tests/unit/test_kernel_auth_bridge.py | 67 +++++++++--- tests/unit/test_session.py | 49 +++++++++ 6 files changed, 162 insertions(+), 78 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 65c18f390..686b2bcaa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ # Unreleased - 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) OAuth is now supported.** Both connector Azure auth types route onto the kernel's generic OAuth flows: `auth_type="azure-oauth"` (Azure AD U2M) forwards the Azure app bundle (client id `96eecda7-…`, redirect port `8030`, and the `{app_id}/user_impersonation offline_access` delegated scope, honoring `DATABRICKS_AZURE_TENANT_ID`) to the kernel's U2M browser flow; `auth_type="azure-sp-m2m"` (Azure service principal) forwards the Azure credentials with an Entra v2.0 token endpoint and the `{app_id}/.default` scope to the kernel's M2M flow. `azure-sp-m2m` requires an explicit `azure_tenant_id` on the kernel path (it does not auto-discover the tenant as the Thrift path does), and the Azure management-token header / `azure_workspace_resource_id` are not applied — the Databricks-audience token authenticates service principals that are workspace principals, matching the Go and Node SQL connectors (PECOBLR-4141; PECOBLR-4120) +- Kernel backend (`use_kernel=True`): **Azure Entra (Azure AD) OAuth is now supported.** The kernel is the Azure-aware auth core (it owns the endpoints, scopes, app ids, and tenant discovery); the connector forwards the selector and Azure credentials unchanged, so `connect()` is byte-identical between the Thrift and `use_kernel=True` paths. `auth_type="azure-oauth"` (Azure AD U2M) runs the kernel's browser flow against the workspace v2.0 authorize/token endpoints with the Azure app client id (`96eecda7-…`), redirect port `8030`, and the `{app_id}/user_impersonation offline_access` scope. `auth_type="azure-sp-m2m"` (Azure service principal) forwards `azure_client_id` / `azure_client_secret`; the kernel 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). `azure_workspace_resource_id` is an optional add-on: set it and the kernel additionally sends the Azure management-token header pair (`X-Databricks-Azure-SP-Management-Token` + `X-Databricks-Azure-Workspace-Resource-Id`) so an SP that holds only an Azure RBAC role (not a workspace member) can authenticate; omit it and the SP authenticates with the data token alone (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 5f526d792..2f04f24cd 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` (U2M), `azure-oauth` (Azure AD U2M), or `azure-sp-m2m` (Azure service-principal M2M). All three work on the kernel path; `azure-oauth` / `azure-sp-m2m` kernel support added in #919 (they route onto the kernel's generic OAuth flows with Azure values). | +| `auth_type` | `str` | ✅ | ✅ | `None` ⇒ Databricks OAuth | `databricks-oauth` (U2M), `azure-oauth` (Azure AD U2M), or `azure-sp-m2m` (Azure service-principal M2M). All three work on the kernel path; `azure-oauth` / `azure-sp-m2m` kernel support added in #919 (the connector forwards the selector + Azure credentials and the kernel owns Azure resolution). | | `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. | @@ -77,8 +77,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` | `str` | ✅ | ✅ | `None` | Azure service-principal (Entra ID M2M), selected by `auth_type="azure-sp-m2m"`. On the kernel path these route onto OAuth M2M with an Entra v2.0 token endpoint + the Databricks-resource `.default` scope (#919). **`azure_tenant_id` is required on the kernel path** — unlike Thrift, it is not auto-discovered from the workspace. | -| `azure_workspace_resource_id` | `str` | ✅ | ⚠️ | `None` | Thrift sends this with the Azure SP **management token** (`X-Databricks-Azure-SP-Management-Token`) to authorize an SP that has an Azure RBAC role but is not a workspace member. **Not applied on the kernel path** — the management-token flow is unsupported there (matching the Go and Node SQL drivers, which don't use it); add the SP as a workspace principal instead. Setting it on the kernel path logs a warning and is otherwise ignored. | +| `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` | Optional add-on for `azure-sp-m2m`. When set, the SP **management token** (`X-Databricks-Azure-SP-Management-Token` + `X-Databricks-Azure-Workspace-Resource-Id`) is sent to authorize an SP that has an Azure RBAC role but is not a workspace member. On the kernel path the connector forwards it and the kernel fetches the management token and emits the header pair; omit it and the SP authenticates with the Databricks-audience data token alone. | | `_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/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 697036f23..7e16f79cc 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -21,23 +21,25 @@ ``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 Entra (Azure AD)** — both Azure auth types route to the - kernel's *generic* OAuth flows with Azure values as overrides (the - kernel needs no Azure-specific code): - - - ``azure-oauth`` (U2M) → ``oauth-u2m`` with the Azure app client id - (``96eecda7-…``), redirect port ``8030``, and the AAD delegated scope - ``{app_id}/user_impersonation offline_access`` (via - ``AzureOAuthEndpointCollection``, honoring ``DATABRICKS_AZURE_TENANT_ID``). - The kernel discovers endpoints via the workspace ``/oidc`` redirector, - which an Azure workspace redirects to Entra (PECOBLR-4120). - - ``azure-sp-m2m`` (M2M) → ``oauth-m2m`` with the Azure service-principal - credentials, an Entra v2.0 ``token_url``, and the - ``{effective_app_id}/.default`` scope (PECOBLR-4141). ``azure_tenant_id`` - is required (the kernel path does not auto-discover it). The Azure - management-token header and ``azure_workspace_resource_id`` are **not** - applied on the kernel path — no SQL connector uses them; an SP that is - not a workspace member (RBAC-only) is unsupported here. +- **Azure Entra (Azure AD)** — both Azure auth types forward the selector + and Azure credentials to the KERNEL, which is the Azure-aware auth core + (it owns the endpoints, scopes, app ids, and tenant discovery). The + binding stays thin — it does not construct endpoints or scopes: + + - ``azure-oauth`` (U2M) → forward ``auth_type='azure-oauth'`` (plus any + optional ``oauth_client_id`` / ``oauth_redirect_port`` passthrough). The + kernel pins the workspace v2.0 authorize/token endpoints, the Azure app + client id (``96eecda7-…``), port ``8030``, and the + ``{app_id}/user_impersonation offline_access`` scope (PECOBLR-4120). + - ``azure-sp-m2m`` (M2M) → forward ``auth_type='azure-sp-m2m'`` with the + Azure service-principal ``azure_client_id`` / ``azure_client_secret``. + 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 + (Thrift parity). ``azure_workspace_resource_id`` is an optional add-on: + forward it and the kernel additionally sends the Azure management-token + header pair, so an RBAC-only SP (not a workspace member) can authenticate + (PECOBLR-4141). ``identity_federation_client_id`` is forwarded with whichever auth shape wins resolution. It selects mandatory SP-wide workload-identity token @@ -72,15 +74,9 @@ PYSQL_OAUTH_SCOPES, ) from databricks.sql.auth.authenticators import AccessTokenAuthProvider, AuthProvider -from databricks.sql.auth.common import get_effective_azure_login_app_id from databricks.sql.auth.token_federation import TokenFederationProvider from databricks.sql.exc import NotSupportedError, ProgrammingError -# Entra (Azure AD) v2.0 token endpoint template. The kernel's generic M2M -# provider sends the credentials as ``scope`` (v2.0), so we point it at the -# v2.0 endpoint (the connector's own SP path uses the v1.0 ``resource`` form). -_AZURE_AAD_LOGIN_HOST = "https://login.microsoftonline.com" - logger = logging.getLogger(__name__) @@ -175,7 +171,7 @@ def kernel_auth_kwargs( ``oauth_client_secret`` together. (The Azure Entra auth types — ``azure-oauth`` and ``azure-sp-m2m`` — - are routed to the kernel's generic OAuth flows up front, before these + are forwarded to the kernel's Azure-aware flows up front, before these guards; see the module docstring.) 1. **OAuth M2M** — ``oauth_client_id`` + ``oauth_client_secret`` both present → forward raw creds to the kernel's ``oauth-m2m``. @@ -234,51 +230,42 @@ def kernel_auth_kwargs( kwargs["identity_federation_client_id"] = federation_client_id return kwargs - # azure-sp-m2m (Azure service principal, client-credentials): forward to - # oauth-m2m with the Azure app credentials, an Entra v2.0 token endpoint, - # and the {effective_app_id}/.default scope. The kernel sends the client - # secret via HTTP Basic (which Entra v2.0 accepts) and, because a - # token_url override is set, skips workspace OIDC discovery. PECOBLR-4141. + # 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. # - # NOT applied on the kernel path: the Azure management-token header - # (X-Databricks-Azure-SP-Management-Token) and azure_workspace_resource_id. - # No SQL connector (Go, Node) uses them; the Databricks-audience token - # authenticates SPs that are workspace principals (the SQL norm). An SP with - # only an Azure RBAC role (not a workspace member) is unsupported here. + # azure_workspace_resource_id is an optional add-on: forward it and the + # kernel additionally fetches an Azure-management token and sends the + # X-Databricks-Azure-SP-Management-Token + X-Databricks-Azure-Workspace- + # Resource-Id pair, so an SP that holds only an Azure RBAC role (not a + # workspace member) can authenticate. Omit it (the common case) and the SP + # authenticates with the Databricks-audience data token alone. if auth_type == "azure-sp-m2m": azure_client_id = opts.get("azure_client_id") azure_client_secret = opts.get("azure_client_secret") - azure_tenant_id = opts.get("azure_tenant_id") if not (azure_client_id and azure_client_secret): raise ProgrammingError( "auth_type='azure-sp-m2m' requires azure_client_id and " "azure_client_secret." ) - if not azure_tenant_id: - # The Thrift path auto-discovers the tenant from the workspace's - # /aad/auth redirect; the kernel path does not make that call, so - # require it explicitly rather than silently guessing. - raise NotSupportedError( - "use_kernel=True auth_type='azure-sp-m2m' requires an explicit " - "azure_tenant_id (the kernel path does not auto-discover the " - "Azure tenant from the workspace as the Thrift backend does)." - ) - if opts.get("azure_workspace_resource_id"): - logger.warning( - "azure_workspace_resource_id is ignored on use_kernel=True: the " - "Azure management-token flow (X-Databricks-Azure-SP-Management-" - "Token) is not applied on the kernel path. The Databricks-" - "audience token authenticates service principals that are " - "workspace principals; an RBAC-only SP is unsupported here." - ) - app_id = get_effective_azure_login_app_id(hostname or "") kwargs = { - "auth_type": "oauth-m2m", - "client_id": azure_client_id, - "client_secret": azure_client_secret, - "token_url": f"{_AZURE_AAD_LOGIN_HOST}/{azure_tenant_id}/oauth2/v2.0/token", - "oauth_scopes": [f"{app_id}/.default"], + "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, and sends the data token alone when no resource id is set. + 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 diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index a83d62db1..c4cb93bed 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -177,6 +177,19 @@ 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 (kernel + # auto-discovers the tenant; the resource id gates the optional + # management token). 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 33ef81eb3..8fdd106ef 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -489,10 +489,12 @@ def test_u2m_auth_type_plus_client_secret_is_rejected(self): class TestKernelAzureSpM2M: - """``azure-sp-m2m`` (Azure service-principal, client-credentials) routes to - the kernel's generic ``oauth-m2m`` with an Entra v2.0 token endpoint and the - ``{app_id}/.default`` scope. The management-token header is intentionally not - applied on the kernel path (no SQL connector uses it). PECOBLR-4141.""" + """``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", @@ -501,31 +503,64 @@ class TestKernelAzureSpM2M: "azure_tenant_id": "tenant-123", } - def test_azure_sp_m2m_routes_to_kernel_m2m(self): + def test_azure_sp_m2m_forwards_creds_to_kernel(self): kwargs = kernel_auth_kwargs( _FakeOAuthProvider(), dict(self._CREDS), hostname="adb-1.azuredatabricks.net", ) - app_id = get_effective_azure_login_app_id("adb-1.azuredatabricks.net") + # Thin forwarding: the kernel owns endpoint/scope resolution, so no + # token_url / oauth_scopes are constructed here. assert kwargs == { - "auth_type": "oauth-m2m", - "client_id": "azure-sp", - "client_secret": "azure-secret", - "token_url": "https://login.microsoftonline.com/tenant-123/oauth2/v2.0/token", - "oauth_scopes": [f"{app_id}/.default"], + "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_requires_tenant(self): + 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", } - with pytest.raises(NotSupportedError, match="azure_tenant_id"): - kernel_auth_kwargs( - _FakeOAuthProvider(), opts, hostname="adb-1.azuredatabricks.net" - ) + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), opts, hostname="adb-1.azuredatabricks.net" + ) + 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): + # An optional add-on: when set, the kernel fetches an Azure-management + # token and emits the X-Databricks-Azure-* header pair (for an SP with + # only an Azure RBAC role, not a workspace member). The binding 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, hostname="adb-1.azuredatabricks.net" + ) + 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), + hostname="adb-1.azuredatabricks.net", + ) + 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"): 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 — From dbf2b5af49da9e40f90dc58efc98f5a32bb74c4c Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Thu, 20 Aug 2026 09:36:02 -0700 Subject: [PATCH 5/7] refactor(kernel): drop azure-sp-m2m management token on the kernel path Mirror the kernel-side simplification: the kernel's `azure-sp-m2m` no longer supports the Azure management-token flow (RBAC-only SPs), so stop forwarding `azure_workspace_resource_id` to the kernel. The SP must be a workspace member on `use_kernel=True`. `azure_workspace_resource_id` is still a valid connection parameter for the Thrift path (unchanged). On the kernel path, setting it now logs a warning (rather than silently dropping a security-relevant auth parameter and failing later with an opaque 403) and the SP authenticates with the Databricks-audience data token alone. Drops the threading through `kernel_auth_options` in session.py. RBAC-only-SP support is a documented follow-up. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 --- CHANGELOG.md | 2 +- CONNECTION_PARAMETERS.md | 2 +- .../sql/backend/kernel/auth_bridge.py | 34 +++++++++++-------- src/databricks/sql/session.py | 10 +++--- tests/unit/test_kernel_auth_bridge.py | 23 +++++++------ tests/unit/test_session.py | 2 -- 6 files changed, 38 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 686b2bcaa..5064be426 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ # Unreleased - 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) OAuth is now supported.** The kernel is the Azure-aware auth core (it owns the endpoints, scopes, app ids, and tenant discovery); the connector forwards the selector and Azure credentials unchanged, so `connect()` is byte-identical between the Thrift and `use_kernel=True` paths. `auth_type="azure-oauth"` (Azure AD U2M) runs the kernel's browser flow against the workspace v2.0 authorize/token endpoints with the Azure app client id (`96eecda7-…`), redirect port `8030`, and the `{app_id}/user_impersonation offline_access` scope. `auth_type="azure-sp-m2m"` (Azure service principal) forwards `azure_client_id` / `azure_client_secret`; the kernel 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). `azure_workspace_resource_id` is an optional add-on: set it and the kernel additionally sends the Azure management-token header pair (`X-Databricks-Azure-SP-Management-Token` + `X-Databricks-Azure-Workspace-Resource-Id`) so an SP that holds only an Azure RBAC role (not a workspace member) can authenticate; omit it and the SP authenticates with the data token alone (PECOBLR-4141; PECOBLR-4120) +- Kernel backend (`use_kernel=True`): **Azure Entra (Azure AD) OAuth is now supported.** The kernel is the Azure-aware auth core (it owns the endpoints, scopes, app ids, and tenant discovery); the connector forwards the selector and Azure credentials unchanged, so `connect()` is byte-identical between the Thrift and `use_kernel=True` paths. `auth_type="azure-oauth"` (Azure AD U2M) runs the kernel's browser flow against the workspace v2.0 authorize/token endpoints with the Azure app client id (`96eecda7-…`), redirect port `8030`, and the `{app_id}/user_impersonation offline_access` scope. `auth_type="azure-sp-m2m"` (Azure service principal) forwards `azure_client_id` / `azure_client_secret`; the kernel 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 service principal authenticates with the Databricks-audience data token, so it must be a workspace member; the kernel path does not yet support the Azure management-token flow (for an RBAC-only SP), so `azure_workspace_resource_id` is ignored with a warning on `use_kernel=True` (it still applies on the Thrift path) (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 2f04f24cd..2a0fba5d4 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -78,7 +78,7 @@ to change without notice. | `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` | `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` | Optional add-on for `azure-sp-m2m`. When set, the SP **management token** (`X-Databricks-Azure-SP-Management-Token` + `X-Databricks-Azure-Workspace-Resource-Id`) is sent to authorize an SP that has an Azure RBAC role but is not a workspace member. On the kernel path the connector forwards it and the kernel fetches the management token and emits the header pair; omit it and the SP authenticates with the Databricks-audience data token alone. | +| `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`) is sent to authorize an SP that has an Azure RBAC role but is not a workspace member. **Not yet supported on the kernel path** — the SP must be a workspace member there; setting it on `use_kernel=True` logs a warning and is otherwise ignored (RBAC-only-SP support is a planned follow-up). Applies normally on the Thrift path. | | `_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/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 7e16f79cc..b93eda005 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -36,9 +36,9 @@ 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 - (Thrift parity). ``azure_workspace_resource_id`` is an optional add-on: - forward it and the kernel additionally sends the Azure management-token - header pair, so an RBAC-only SP (not a workspace member) can authenticate + (Thrift parity). The SP must be a workspace member (the kernel path does not + yet support the Azure management-token flow, so ``azure_workspace_resource_id`` + is ignored with a warning; RBAC-only SPs are a planned follow-up) (PECOBLR-4141). ``identity_federation_client_id`` is forwarded with whichever auth shape @@ -239,12 +239,10 @@ def kernel_auth_kwargs( # (so connect() is byte-identical between Thrift and use_kernel=True). # PECOBLR-4141. # - # azure_workspace_resource_id is an optional add-on: forward it and the - # kernel additionally fetches an Azure-management token and sends the - # X-Databricks-Azure-SP-Management-Token + X-Databricks-Azure-Workspace- - # Resource-Id pair, so an SP that holds only an Azure RBAC role (not a - # workspace member) can authenticate. Omit it (the common case) and the SP - # authenticates with the Databricks-audience data token alone. + # The SP authenticates with the Databricks-audience data token, so it must be + # a workspace member. The kernel path does not yet support the Azure + # management-token flow (for an RBAC-only SP), so azure_workspace_resource_id + # is ignored with a warning here — see below. if auth_type == "azure-sp-m2m": azure_client_id = opts.get("azure_client_id") azure_client_secret = opts.get("azure_client_secret") @@ -258,14 +256,22 @@ def kernel_auth_kwargs( "azure_client_id": azure_client_id, "azure_client_secret": azure_client_secret, } - # Optional passthroughs: the kernel auto-discovers the tenant when - # absent, and sends the data token alone when no resource id is set. + # Optional passthrough: the kernel auto-discovers the tenant when absent. 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 + # azure_workspace_resource_id drives the Azure management-token flow on + # the Thrift path (for an SP with only an Azure RBAC role, not a + # workspace member). The kernel path does NOT support it yet — the SP + # must be a workspace member — so warn rather than silently dropping a + # security-relevant auth parameter and failing later with an opaque 403. + if opts.get("azure_workspace_resource_id"): + logger.warning( + "azure_workspace_resource_id is ignored on use_kernel=True: the " + "Azure SP management-token flow is not yet supported on the " + "kernel path. The service principal must be a workspace member " + "(the Databricks-audience token alone authenticates it)." + ) if federation_client_id: kwargs["identity_federation_client_id"] = federation_client_id return kwargs diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index c4cb93bed..12415f71b 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -181,15 +181,13 @@ def _create_backend( # 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 (kernel - # auto-discovers the tenant; the resource id gates the optional - # management token). Kernel-only; Thrift / SEA are unaffected. + # tenant is optional (the kernel auto-discovers it). Kernel-only; + # Thrift / SEA are unaffected. (azure_workspace_resource_id is not + # forwarded — the kernel path does not yet support the Azure + # management-token flow; the bridge warns if it is set.) "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 8fdd106ef..7225ecb75 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -537,21 +537,22 @@ def test_azure_sp_m2m_tenant_optional_kernel_autodiscovers(self): } assert "azure_tenant_id" not in kwargs - def test_azure_sp_m2m_forwards_workspace_resource_id(self): - # An optional add-on: when set, the kernel fetches an Azure-management - # token and emits the X-Databricks-Azure-* header pair (for an SP with - # only an Azure RBAC role, not a workspace member). The binding forwards - # it rather than dropping it. + def test_azure_sp_m2m_ignores_workspace_resource_id_with_warning(self, caplog): + # The kernel path does not yet support the Azure management-token flow + # (for an RBAC-only SP), so azure_workspace_resource_id is NOT forwarded + # — but we warn rather than silently dropping a security-relevant param. opts = dict( self._CREDS, azure_workspace_resource_id="/subscriptions/s/resourceGroups/rg/workspace/w", ) - kwargs = kernel_auth_kwargs( - _FakeOAuthProvider(), opts, hostname="adb-1.azuredatabricks.net" - ) - assert ( - kwargs["azure_workspace_resource_id"] - == "/subscriptions/s/resourceGroups/rg/workspace/w" + with caplog.at_level("WARNING"): + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), opts, hostname="adb-1.azuredatabricks.net" + ) + assert "azure_workspace_resource_id" not in kwargs + assert any( + "azure_workspace_resource_id is ignored" in r.message + for r in caplog.records ) def test_azure_sp_m2m_omits_workspace_resource_id_when_absent(self): diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 6fcefcade..32abbf510 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -512,7 +512,6 @@ def test_azure_sp_m2m_kwargs_threaded_into_kernel_auth_options(self): 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: @@ -522,7 +521,6 @@ def test_azure_sp_m2m_kwargs_threaded_into_kernel_auth_options(self): 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() From 9e6ffe92e3ed4cef11a7a10299b826ec2a3e3e09 Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Thu, 20 Aug 2026 10:58:33 -0700 Subject: [PATCH 6/7] refactor(kernel): reject azure-oauth on the kernel path; point to databricks-oauth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel dropped its dedicated Azure U2M flow, so `auth_type="azure-oauth"` on `use_kernel=True` is now rejected with a clear pointer to `databricks-oauth` — whose in-house OAuth browser flow works against Azure workspaces (the workspace federates login to Entra). This replaces the thin forward-the-selector branch. `azure-sp-m2m` is unchanged; `azure-oauth` on the Thrift path is unaffected. Rejecting (rather than silently remapping) keeps the user's selected flow explicit. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 --- CHANGELOG.md | 2 +- CONNECTION_PARAMETERS.md | 2 +- .../sql/backend/kernel/auth_bridge.py | 55 +++++++++---------- tests/unit/test_kernel_auth_bridge.py | 44 +++++---------- 4 files changed, 41 insertions(+), 62 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5064be426..64d0bb89d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ # Unreleased - 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) OAuth is now supported.** The kernel is the Azure-aware auth core (it owns the endpoints, scopes, app ids, and tenant discovery); the connector forwards the selector and Azure credentials unchanged, so `connect()` is byte-identical between the Thrift and `use_kernel=True` paths. `auth_type="azure-oauth"` (Azure AD U2M) runs the kernel's browser flow against the workspace v2.0 authorize/token endpoints with the Azure app client id (`96eecda7-…`), redirect port `8030`, and the `{app_id}/user_impersonation offline_access` scope. `auth_type="azure-sp-m2m"` (Azure service principal) forwards `azure_client_id` / `azure_client_secret`; the kernel 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 service principal authenticates with the Databricks-audience data token, so it must be a workspace member; the kernel path does not yet support the Azure management-token flow (for an RBAC-only SP), so `azure_workspace_resource_id` is ignored with a warning on `use_kernel=True` (it still applies on the Thrift path) (PECOBLR-4141; PECOBLR-4120) +- 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 service principal authenticates with the Databricks-audience data token, so it must be a workspace member; the kernel path does not yet support the Azure management-token flow (for an RBAC-only SP), so `azure_workspace_resource_id` is ignored with a warning on `use_kernel=True` (it still applies on the Thrift path). Azure AD **U2M** (`auth_type="azure-oauth"`) is **not** supported on the kernel path — use `auth_type="databricks-oauth"`, whose in-house OAuth browser flow works against Azure workspaces (the workspace federates login to Entra); `azure-oauth` on `use_kernel=True` is rejected with that pointer (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 2a0fba5d4..dc5206db5 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` (U2M), `azure-oauth` (Azure AD U2M), or `azure-sp-m2m` (Azure service-principal M2M). All three work on the kernel path; `azure-oauth` / `azure-sp-m2m` kernel support added in #919 (the connector forwards the selector + Azure credentials and the kernel owns Azure resolution). | +| `auth_type` | `str` | ✅ | ✅ | `None` ⇒ Databricks OAuth | `databricks-oauth` (U2M), `azure-oauth` (Azure AD U2M), or `azure-sp-m2m` (Azure service-principal M2M). On the kernel path: `databricks-oauth` and `azure-sp-m2m` are supported (#919); `azure-oauth` is **not** — use `databricks-oauth` instead (its in-house browser flow works against Azure workspaces), and `azure-oauth` on `use_kernel=True` is rejected with that pointer. 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. | diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index b93eda005..22c7e9e48 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -21,16 +21,14 @@ ``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 Entra (Azure AD)** — both Azure auth types forward the selector - and Azure credentials to the KERNEL, which is the Azure-aware auth core - (it owns the endpoints, scopes, app ids, and tenant discovery). The - binding stays thin — it does not construct endpoints or scopes: - - - ``azure-oauth`` (U2M) → forward ``auth_type='azure-oauth'`` (plus any - optional ``oauth_client_id`` / ``oauth_redirect_port`` passthrough). The - kernel pins the workspace v2.0 authorize/token endpoints, the Azure app - client id (``96eecda7-…``), port ``8030``, and the - ``{app_id}/user_impersonation offline_access`` scope (PECOBLR-4120). +- **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) → **not supported on the kernel path**; rejected with + a pointer to ``databricks-oauth``, whose in-house U2M browser flow works + against Azure workspaces (the workspace federates login to Entra). A + dedicated Azure U2M flow may return later (PECOBLR-4120). - ``azure-sp-m2m`` (M2M) → forward ``auth_type='azure-sp-m2m'`` with the Azure service-principal ``azure_client_id`` / ``azure_client_secret``. The kernel builds the Entra v2.0 token endpoint and the @@ -170,9 +168,10 @@ def kernel_auth_kwargs( - a U2M ``auth_type`` (``databricks-oauth``) *and* ``oauth_client_secret`` together. - (The Azure Entra auth types — ``azure-oauth`` and ``azure-sp-m2m`` — - are forwarded to the kernel's Azure-aware flows up front, before these - guards; see the module docstring.) + (The Azure Entra auth types are handled up front, before these guards: + ``azure-sp-m2m`` forwards to the kernel's Azure SP flow; ``azure-oauth`` + is rejected with a pointer to ``databricks-oauth``. See the module + docstring.) 1. **OAuth M2M** — ``oauth_client_id`` + ``oauth_client_secret`` both present → forward raw creds to the kernel's ``oauth-m2m``. 2. **PAT** — the built provider is (or wraps) an @@ -211,24 +210,20 @@ def kernel_auth_kwargs( # 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): forward the selector; the KERNEL owns Azure - # resolution (it is the auth core). The kernel pins the workspace v2.0 - # authorize/token endpoints (`{host}/oidc/oauth2/v2.0/{authorize,token}` — - # NOT the discovered `/oidc/v1/authorize`, which the workspace redirects to - # a malformed Entra URL), the Azure app client id, port 8030, and the - # `{app_id}/user_impersonation offline_access` scope. So this binding does - # NOT construct endpoints/scopes — it just passes `auth_type='azure-oauth'` - # plus any optional client_id / redirect_port passthrough. PECOBLR-4120. + # azure-oauth (Azure AD U2M) is NOT supported on the kernel path. The + # in-house `databricks-oauth` browser flow works against Azure workspaces + # (the workspace federates the login to Entra), so it is the U2M path on the + # kernel — reject `azure-oauth` with a clear pointer rather than silently + # changing the user's selected flow. (A dedicated Azure U2M flow may return + # later; for now Azure U2M = `databricks-oauth`.) if auth_type == "azure-oauth": - kwargs = {"auth_type": "azure-oauth"} - if client_id: - kwargs["client_id"] = client_id - redirect_port = opts.get("oauth_redirect_port") - if redirect_port is not None: - kwargs["redirect_ports"] = [_coerce_redirect_port(redirect_port)] - if federation_client_id: - kwargs["identity_federation_client_id"] = federation_client_id - return kwargs + raise NotSupportedError( + "auth_type='azure-oauth' is not supported on use_kernel=True. Use " + "auth_type='databricks-oauth' instead — the in-house OAuth U2M " + "browser flow works against Azure Databricks workspaces (the " + "workspace federates the login to Microsoft Entra). Or use the " + "Thrift backend (default) for the dedicated Azure AD U2M flow." + ) # azure-sp-m2m (Azure service principal, client-credentials): forward the # selector + Azure SP credentials; the KERNEL owns Azure resolution (it is diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index 7225ecb75..6f4e02991 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -258,8 +258,10 @@ class TestKernelOAuthU2M: may override ``oauth_scopes``; absent one, ``PYSQL_OAUTH_SCOPES`` is forwarded as the default. - ``azure-oauth`` (Azure AD U2M) routes here too — see - ``test_azure_oauth_routes_to_kernel_u2m`` (PECOBLR-4120).""" + ``azure-oauth`` (Azure AD U2M) is NOT supported on the kernel path and is + rejected with a pointer to ``databricks-oauth`` — see + ``test_azure_oauth_rejected_with_pointer_to_databricks_oauth`` + (PECOBLR-4120).""" def test_bare_databricks_oauth_forwards_full_python_bundle(self): # No overrides → forward the databricks-sql-python bundle in full @@ -277,34 +279,16 @@ def test_bare_databricks_oauth_forwards_full_python_bundle(self): "oauth_scopes": list(PYSQL_OAUTH_SCOPES), } - def test_azure_oauth_forwards_selector_kernel_owns_resolution(self): - # azure-oauth (Azure AD U2M): the bridge forwards ONLY the selector. - # The kernel owns Azure resolution — it pins the workspace v2.0 - # authorize/token endpoints, the Azure client id, port 8030, and the - # {app_id}/user_impersonation scope. So the bridge must NOT construct - # client_id / redirect_ports / oauth_scopes here. PECOBLR-4120. - kwargs = kernel_auth_kwargs( - _FakeOAuthProvider(), - {"auth_type": "azure-oauth"}, - ) - assert kwargs == {"auth_type": "azure-oauth"} - - def test_azure_oauth_honors_custom_client_id_and_port(self): - # A caller override still passes through (client_id + its coupled port), - # but no scopes/endpoints are synthesised by the bridge. - kwargs = kernel_auth_kwargs( - _FakeOAuthProvider(), - { - "auth_type": "azure-oauth", - "oauth_client_id": "custom-azure-app", - "oauth_redirect_port": 9100, - }, - ) - assert kwargs == { - "auth_type": "azure-oauth", - "client_id": "custom-azure-app", - "redirect_ports": [9100], - } + def test_azure_oauth_rejected_with_pointer_to_databricks_oauth(self): + # azure-oauth (Azure AD U2M) is not supported on the kernel path. The + # in-house databricks-oauth browser flow works against Azure workspaces, + # so the bridge rejects azure-oauth with a clear pointer rather than + # silently changing the user's selected flow. PECOBLR-4120. + with pytest.raises(NotSupportedError, match="databricks-oauth"): + kernel_auth_kwargs( + _FakeOAuthProvider(), + {"auth_type": "azure-oauth"}, + ) def test_u2m_custom_client_id_port_and_scopes_honored(self): # A caller may override the coupled client_id + redirect port and the From b121efde6c51a14d3b7bc74d4c662d1c518aeec1 Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Thu, 20 Aug 2026 12:06:17 -0700 Subject: [PATCH 7/7] feat(kernel): forward azure_workspace_resource_id for complete Azure SP M2M MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel's azure-sp-m2m now always sends the Azure SP management token and, when a workspace resource id is set, the X-Databricks-Azure-Workspace-Resource-Id header — matching the Thrift connector. So the bridge forwards azure_workspace_resource_id to the kernel (previously warned-and-ignored), and session.py threads it into kernel_auth_options. This completes Azure SP M2M parity: a service principal with an Azure RBAC role but no workspace membership now authenticates on the kernel path too. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 --- CHANGELOG.md | 2 +- CONNECTION_PARAMETERS.md | 2 +- .../sql/backend/kernel/auth_bridge.py | 47 +++++++++---------- src/databricks/sql/session.py | 11 +++-- tests/unit/test_kernel_auth_bridge.py | 22 ++++----- tests/unit/test_session.py | 2 + 6 files changed, 43 insertions(+), 43 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 64d0bb89d..e9043773b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ # Unreleased - 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 service principal authenticates with the Databricks-audience data token, so it must be a workspace member; the kernel path does not yet support the Azure management-token flow (for an RBAC-only SP), so `azure_workspace_resource_id` is ignored with a warning on `use_kernel=True` (it still applies on the Thrift path). Azure AD **U2M** (`auth_type="azure-oauth"`) is **not** supported on the kernel path — use `auth_type="databricks-oauth"`, whose in-house OAuth browser flow works against Azure workspaces (the workspace federates login to Entra); `azure-oauth` on `use_kernel=True` is rejected with that pointer (PECOBLR-4141; PECOBLR-4120) +- 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; the kernel **always** also sends the Azure SP management token (`X-Databricks-Azure-SP-Management-Token`), matching Thrift, so a service principal with an Azure RBAC role but no workspace membership can authenticate — set `azure_workspace_resource_id` to add the `X-Databricks-Azure-Workspace-Resource-Id` header too. Azure AD **U2M** (`auth_type="azure-oauth"`) is **not** supported on the kernel path — use `auth_type="databricks-oauth"`, whose in-house OAuth browser flow works against Azure workspaces (the workspace federates login to Entra); `azure-oauth` on `use_kernel=True` is rejected with that pointer (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 dc5206db5..e37b2de6d 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -78,7 +78,7 @@ to change without notice. | `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` | `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`) is sent to authorize an SP that has an Azure RBAC role but is not a workspace member. **Not yet supported on the kernel path** — the SP must be a workspace member there; setting it on `use_kernel=True` logs a warning and is otherwise ignored (RBAC-only-SP support is a planned follow-up). Applies normally on the Thrift path. | +| `azure_workspace_resource_id` | `str` | ✅ | ✅ | `None` | For `azure-sp-m2m`. The SP **management token** (`X-Databricks-Azure-SP-Management-Token`) is always sent; setting this adds the `X-Databricks-Azure-Workspace-Resource-Id` header too, to authorize an SP that has an Azure RBAC role but is not a workspace member. 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/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 22c7e9e48..06b374c22 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -30,14 +30,14 @@ against Azure workspaces (the workspace federates login to Entra). A dedicated Azure U2M flow may return later (PECOBLR-4120). - ``azure-sp-m2m`` (M2M) → forward ``auth_type='azure-sp-m2m'`` with the - Azure service-principal ``azure_client_id`` / ``azure_client_secret``. - 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 - (Thrift parity). The SP must be a workspace member (the kernel path does not - yet support the Azure management-token flow, so ``azure_workspace_resource_id`` - is ignored with a warning; RBAC-only SPs are a planned follow-up) - (PECOBLR-4141). + 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, auto-discovers the tenant from the workspace's ``/aad/auth`` redirect + when ``azure_tenant_id`` is omitted, and always sends the Azure SP + management token (adding the ``X-Databricks-Azure-Workspace-Resource-Id`` + header when ``azure_workspace_resource_id`` is set) — matching the Thrift + connector, so an RBAC-only SP can authenticate (PECOBLR-4141). ``identity_federation_client_id`` is forwarded with whichever auth shape wins resolution. It selects mandatory SP-wide workload-identity token @@ -234,10 +234,11 @@ def kernel_auth_kwargs( # (so connect() is byte-identical between Thrift and use_kernel=True). # PECOBLR-4141. # - # The SP authenticates with the Databricks-audience data token, so it must be - # a workspace member. The kernel path does not yet support the Azure - # management-token flow (for an RBAC-only SP), so azure_workspace_resource_id - # is ignored with a warning here — see below. + # The Authorization bearer is the Databricks-audience data token; the kernel + # also always sends the Azure SP management token, and adds the + # X-Databricks-Azure-Workspace-Resource-Id header when + # azure_workspace_resource_id is set — matching the Thrift connector, so an + # RBAC-only SP (Azure role, not a workspace member) can authenticate. if auth_type == "azure-sp-m2m": azure_client_id = opts.get("azure_client_id") azure_client_secret = opts.get("azure_client_secret") @@ -251,22 +252,18 @@ def kernel_auth_kwargs( "azure_client_id": azure_client_id, "azure_client_secret": azure_client_secret, } - # Optional passthrough: the kernel auto-discovers the tenant when absent. + # Optional passthroughs: the kernel auto-discovers the tenant when + # absent, and always sends the Azure SP management token. When + # azure_workspace_resource_id is set, the kernel adds the + # X-Databricks-Azure-Workspace-Resource-Id header alongside it (for an + # SP with an Azure RBAC role but no workspace membership) — matching the + # Thrift connector. azure_tenant_id = opts.get("azure_tenant_id") if azure_tenant_id: kwargs["azure_tenant_id"] = azure_tenant_id - # azure_workspace_resource_id drives the Azure management-token flow on - # the Thrift path (for an SP with only an Azure RBAC role, not a - # workspace member). The kernel path does NOT support it yet — the SP - # must be a workspace member — so warn rather than silently dropping a - # security-relevant auth parameter and failing later with an opaque 403. - if opts.get("azure_workspace_resource_id"): - logger.warning( - "azure_workspace_resource_id is ignored on use_kernel=True: the " - "Azure SP management-token flow is not yet supported on the " - "kernel path. The service principal must be a workspace member " - "(the Databricks-audience token alone authenticates it)." - ) + 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 diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 12415f71b..5c95fe24f 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -181,13 +181,16 @@ def _create_backend( # 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 is optional (the kernel auto-discovers it). Kernel-only; - # Thrift / SEA are unaffected. (azure_workspace_resource_id is not - # forwarded — the kernel path does not yet support the Azure - # management-token flow; the bridge warns if it is set.) + # 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 6f4e02991..bed6fc0be 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -521,22 +521,20 @@ def test_azure_sp_m2m_tenant_optional_kernel_autodiscovers(self): } assert "azure_tenant_id" not in kwargs - def test_azure_sp_m2m_ignores_workspace_resource_id_with_warning(self, caplog): - # The kernel path does not yet support the Azure management-token flow - # (for an RBAC-only SP), so azure_workspace_resource_id is NOT forwarded - # — but we warn rather than silently dropping a security-relevant param. + 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", ) - with caplog.at_level("WARNING"): - kwargs = kernel_auth_kwargs( - _FakeOAuthProvider(), opts, hostname="adb-1.azuredatabricks.net" - ) - assert "azure_workspace_resource_id" not in kwargs - assert any( - "azure_workspace_resource_id is ignored" in r.message - for r in caplog.records + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), opts, hostname="adb-1.azuredatabricks.net" + ) + assert ( + kwargs["azure_workspace_resource_id"] + == "/subscriptions/s/resourceGroups/rg/workspace/w" ) def test_azure_sp_m2m_omits_workspace_resource_id_when_absent(self): diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 32abbf510..6fcefcade 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -512,6 +512,7 @@ def test_azure_sp_m2m_kwargs_threaded_into_kernel_auth_options(self): 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: @@ -521,6 +522,7 @@ def test_azure_sp_m2m_kwargs_threaded_into_kernel_auth_options(self): 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()