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/4] 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/4] 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/4] =?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/4] =?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 —