feat(kernel): route Azure Entra auth (azure-sp-m2m + azure-oauth) through the auth bridge - #919
feat(kernel): route Azure Entra auth (azure-sp-m2m + azure-oauth) through the auth bridge#919eric-wang-1990 wants to merge 17 commits into
Conversation
…h 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 <e.wang@databricks.com>
There was a problem hiding this comment.
Pull request overview
This PR extends the use_kernel=True auth bridge to support Azure Entra (Azure AD) OAuth by routing the connector’s Azure auth types (azure-oauth U2M and azure-sp-m2m M2M) onto the kernel’s generic OAuth flows with Azure-specific overrides.
Changes:
- Route
auth_type="azure-oauth"to kerneloauth-u2musing the Azure OAuth app bundle and Azure scope mapping. - Route
auth_type="azure-sp-m2m"to kerneloauth-m2musing Entra v2.0 token URL and{effective_app_id}/.defaultscope; requireazure_tenant_id. - Thread
hostnameintokernel_auth_kwargs(viaKernelDatabricksClient.open_session) and add/adjust unit tests + changelog entry.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/databricks/sql/backend/kernel/auth_bridge.py |
Adds Azure auth-type routing and related override plumbing for kernel session kwargs. |
src/databricks/sql/backend/kernel/client.py |
Passes the server hostname into kernel_auth_kwargs during session open. |
tests/unit/test_kernel_auth_bridge.py |
Updates U2M tests for azure-oauth routing and adds azure-sp-m2m routing/validation tests. |
CHANGELOG.md |
Documents Azure Entra OAuth support on the kernel backend. |
Suppressed comments (1)
src/databricks/sql/backend/kernel/auth_bridge.py:262
auth_type='azure-sp-m2m'usesazure_client_id/azure_client_secret, but because this branch returns before the generic ambiguity guards, a caller can also supplyoauth_client_id/oauth_client_secretand have them silently ignored. This makes misconfigurations hard to diagnose.
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):
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| 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")) |
There was a problem hiding this comment.
Already resolved in the current code. The early-return azure-oauth branch the comment targeted was refactored away: azure-oauth now falls through to the shared U2M branch (auth_bridge.py:419), which runs AFTER the ambiguity guards. The guard at auth_bridge.py:294-300 explicitly rejects auth_type in ('databricks-oauth','azure-oauth') together with oauth_client_secret (and the JWT-M2M guard at 316-322 covers oauth_jwt_key_file), so a stray M2M secret is now a loud NotSupportedError rather than silently ignored. Covered by TestKernelAuthAmbiguity::test_u2m_auth_type_plus_client_secret_is_rejected, parametrized over both databricks-oauth and azure-oauth — passing.
| app_id = get_effective_azure_login_app_id(hostname or "") | ||
| kwargs = { |
There was a problem hiding this comment.
This concern no longer applies to the current branch. The azure-sp-m2m branch (src/databricks/sql/backend/kernel/auth_bridge.py:256–283) was refactored: it no longer calls get_effective_azure_login_app_id(hostname or "") and no longer computes oauth_scopes in Python. Azure endpoint/app-id/scope resolution (including the {effective_app_id}/.default scope and Entra v2.0 token endpoint) now happens inside the kernel, and the bridge only forwards the azure-sp-m2m selector plus the Azure SP credentials and optional passthroughs. There is therefore no PROD-app-id fallback and no scope-defaulting path that could go wrong when hostname is omitted, so no code change is needed.
There was a problem hiding this comment.
Verdict: 1 Medium
Routing logic is clean and well-tested overall. One medium concern: the new Azure branches return before the ambiguity guards, so azure-oauth + oauth_client_secret (or + credentials_provider) is silently accepted as U2M whereas the parallel databricks-oauth case raises NotSupportedError. Nit (not filed inline): the step-3 comment # Only databricks-oauth reaches here (azure-oauth rejected up front) in the U2M block is now stale — azure-oauth is routed up front, not rejected.
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 <e.wang@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — clean, well-documented routing with strong test coverage for both new Azure flows (routing, required tenant/creds, federation client id). One low-severity consistency note: the azure-oauth branch returns before the ambiguity guards, so oauth_client_secret / credentials_provider are silently ignored there, unlike the databricks-oauth U2M path which rejects them.
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 <e.wang@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Medium · 2 Low
Overall solid — the azure-sp-m2m v2.0 .default scope + v2.0 token endpoint correctly mirror the Thrift path's v1.0 resource= form, and the empty-hostname fallback is safe. Main issue is a stale module docstring (F1, medium) that claims the azure-oauth bridge synthesises the client id/port/scope via AzureOAuthEndpointCollection when the code actually just forwards the azure-oauth selector and lets the kernel own resolution — the inline comment already contradicts it. Plus two low doc/consistency nits.
…optional mgmt token
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 <e.wang@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Medium · 1 Low
Looks good overall — the Azure routing logic is thin, well-tested, and correct. Two cleanup items: an unused hostname parameter now threaded through client.py for no effect (medium), and a dead test import (low). Nit: the comment at auth_bridge.py:318 ("azure-oauth rejected up front") is now stale — azure-oauth is routed up front, not rejected.
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — clean, thin forwarding of both Azure Entra flows with thorough unit coverage (routing, ambiguity guards extended to azure-oauth, required-creds, optional passthroughs, session threading). One low-severity consistency note: the azure-sp-m2m early return bypasses the module's fail-loud ambiguity guards, so conflicting cross-namespace signals are silently ignored rather than rejected.
Addresses: - #3828326882 at src/databricks/sql/backend/kernel/auth_bridge.py:200 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — clean, thin binding that correctly forwards the Azure selector + creds and lets the kernel own resolution, with thorough unit coverage (bridge + session threading + ambiguity guards parametrized over both U2M types). One low-severity note about the intentional ambiguity-guard exemption on the azure-sp-m2m branch diverging from the module's loud-failure philosophy.
Addresses: - #3828367785 at src/databricks/sql/backend/kernel/auth_bridge.py:265 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
| | --------------------------------------------------- | -------------------- | :----: | :----: | ------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | | ||
| | `access_token` (PAT) | `str` | ✅ | ✅ | `None` | Personal Access Token / bearer token. The default auth mode when set; otherwise auth falls back to OAuth. | | ||
| | `auth_type` | `str` | ✅ | ✅ | `None` ⇒ Databricks OAuth | `databricks-oauth` or `azure-oauth`. | | ||
| | `auth_type` | `str` | ✅ | ✅ | `None` ⇒ Databricks OAuth | `databricks-oauth` (U2M), `azure-oauth` (Azure AD U2M), or `azure-sp-m2m` (Azure service-principal M2M). All three are supported on the kernel path (#919): `azure-oauth` routes to the kernel's OAuth U2M flow identically to `databricks-oauth` (the in-house workspace-federated browser flow, which Azure workspaces support), so it uses the `databricks-sql-python` app rather than Thrift's direct-Entra Azure app. All three work on the Thrift path. | |
There was a problem hiding this comment.
azure-oauth for Kernel mode is same with databricks-oauth, which is inHouse OAuth U2M
There was a problem hiding this comment.
Agreed — and the doc already states exactly this. CONNECTION_PARAMETERS.md:72 says azure-oauth on the kernel path "routes to the kernel's OAuth U2M flow identically to databricks-oauth (the in-house workspace-federated browser flow, which Azure workspaces support)." That matches your point that kernel-mode azure-oauth is the same in-house OAuth U2M as databricks-oauth, so no change is needed.
There was a problem hiding this comment.
Verdict: 1 Low · 1 Nit
Looks good — a clean, thin binding that forwards the Azure selector + credentials to the kernel, with the azure-sp-m2m branch correctly placed before the ambiguity guards and session.py threading the azure_* kwargs only into the kernel path. Test coverage is solid; just one low (untested silent-ignore asymmetry for conflicting signals on azure-sp-m2m) and one nit (stale docstring count).
Other findings
- ⚪ Nit — Module docstring says "Three auth shapes are supported on the kernel path" but the bullet list now enumerates four top-level categories (PAT, OAuth M2M, OAuth U2M, and Azure Entra) plus JWT private-key M2M. The count was already slightly stale before this PR (JWT M2M), and this PR's new Azure Entra section makes it more so. Consider dropping the fixed count (e.g. "The following auth shapes are supported on the kernel path:") so it doesn't need updating each time a flow is added.
Addresses: - #3828404026 at src/databricks/sql/backend/kernel/auth_bridge.py:279 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — clean, well-tested addition routing azure-sp-m2m and azure-oauth through the kernel auth bridge, with matching session threading and thorough unit coverage (branch ordering, ambiguity guards, and the intentional azure-sp-m2m asymmetry are all exercised). One low note on log visibility for the silently-ignored conflicting-credential path.
Addresses: - #3828445075 at src/databricks/sql/backend/kernel/auth_bridge.py:285 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
…urface Point kernel-e2e's KERNEL_REV at the kernel commit that adds the pyo3 azure-sp-m2m surface (databricks-sql-kernel#263), so the connector's kernel-e2e builds a kernel wheel that can accept auth_type='azure-sp-m2m'. Temporary pin to the unmerged #263 branch tip; re-point to a kernel main SHA once #263 merges and a release is cut. Customer-facing pin (pyproject databricks-sql-kernel ^0.2.0) still needs a release bump to a published kernel with #263. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — the Azure Entra routing (azure-sp-m2m up-front handling, azure-oauth folded into the shared U2M branch, both ambiguity guards extended) is correct, and session.py threads the azure_* kwargs only on the kernel path. Tests cover every new branch. Only one minor docstring-count inaccuracy noted inline.
Addresses: - #3828603835 at src/databricks/sql/backend/kernel/auth_bridge.py:18 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
What
Teaches the
use_kernel=Trueauth bridge to route both Azure Entra auth flows to the kernel (databricks/databricks-sql-kernel#263):auth_type="azure-sp-m2m") — forward the SP credentials to the kernel, which owns Azure resolution.auth_type="azure-oauth") — route to the kernel's OAuth U2M flow, identically todatabricks-oauth.How
src/databricks/sql/backend/kernel/auth_bridge.py:azure-sp-m2m→ kernelazure-sp-m2m— forwardazure_client_id/azure_client_secret(+ optionalazure_tenant_id,azure_workspace_resource_id). The kernel builds the Entra v2.0 token endpoint and the{effective_app_id}/.defaultscope, and auto-discovers the tenant from the workspace's/aad/authredirect whenazure_tenant_idis omitted (Thrift parity). TheAuthorizationbearer (data token) alone authenticates a workspace-member SP; whenazure_workspace_resource_idis set, the kernel also sends the Azure SP management token (X-Databricks-Azure-SP-Management-Token) +X-Databricks-Azure-Workspace-Resource-Idheader (matching the JDBC driver), so an SP with an Azure RBAC role but no workspace membership authenticates. PECOBLR-4141.azure-oauth→ kerneloauth-u2m— routed through the shared U2M branch, exactly likedatabricks-oauth. The kernel runs the in-house workspace-federated browser flow, which Azure Databricks workspaces support (the workspace federates login to Entra). It forwards the connector'sdatabricks-sql-pythonapp bundle, not Thrift's Azure app (96eecda7/ redirect port8030): that app is registered for Thrift's direct-Entra flow (login.microsoftonline.comconfig +{tenant}/user_impersonationscopes), which the kernel does not perform. So on the kernel pathazure-oauthanddatabricks-oauthare the same flow. The U2M-vs-M2M ambiguity guard now also coversazure-oauth, soazure-oauth+oauth_client_secretis rejected as ambiguous rather than silently falling through to M2M. PECOBLR-4120.src/databricks/sql/session.pythreads theazure_*connection kwargs intokernel_auth_options(the kernel owns Azure resolution, so these raw kwargs are the only source).The binding is thin — it does not construct endpoints or scopes; the kernel does.
Tests
TestKernelAzureSpM2M: forwards creds; tenant optional (kernel auto-discovers);azure_workspace_resource_idforwarded; required-creds error; federation client id.TestKernelOAuthU2M:azure-oauthmaps tooauth-u2mwith thedatabricks-sql-pythonbundle (identical todatabricks-oauth); customclient_id/port/scopes honored; the ambiguity guard (U2M auth_type+oauth_client_secret) is parametrized over bothdatabricks-oauthandazure-oauth.azure_*kwargs reachauth_options.Bridge + session + auth unit tests pass;
blackclean. (test_kernel_client.pyskips without the kernel wheel, as before.)E2E verification (live Azure workspace)
Ran the full connector path (
sql.connect(use_kernel=True, ...)) against a live Azure Databricks workspace, kernel wheel built from databricks/databricks-sql-kernel#263. All three Azure kernel-path flows round-tripSELECT 1:auth_type="azure-oauth") — full interactive browser authorization-code flow (cache cleared to force it). Log confirms the design: authorize URL is the workspace's own OIDC endpoint ({host}/oidc/v1/authorize),client_id=databricks-sql-python,redirect_uri=http://localhost:8020,scope=sql offline_access— i.e. the in-house workspace-federated flow, not Thrift's direct-Entra app (96eecda7/ port 8030 /user_impersonation). Authenticated interactively and returnedSELECT 1. ✅auth_type="azure-sp-m2m", noazure_workspace_resource_id) — workspace-member SP,Row(one=1, who='<sp-app-id>'). ✅azure-sp-m2m+azure_workspace_resource_id) — the kernel mints the ARM management token via the Entra v1.0resource=grant (audhttps://management.core.windows.net/, trailing slash) and sends theX-Databricks-Azure-SP-Management-Token+-Workspace-Resource-Idheader pair; round-trip succeeds. ✅Negative controls confirm the pipeline is real, not a no-op: a mismatched SP secret surfaces Entra
AADSTS7000215, and a service principal with no workspace access returnsHTTP 403 User not authorized(a real control-plane authz decision, not a bridge error).Note: this required a kernel wheel with the
azure-sp-m2mpyo3 surface — the currently-publisheddatabricks-sql-kernel0.2.0 wheel lacks it (Session.__new__() got an unexpected keyword argument 'azure_client_id'), so merging this depends on a kernel release that includes #263. (Azure U2M works on the old wheel —oauth-u2mpredates it.) Not yet covered: the RBAC-only authorization case (a non-member SP authorized purely by an Azure role via the management token), which needs an Azure role assignment still pending.Related
This pull request and its description were written by Isaac.