Skip to content
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,5 +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; scopes are fixed to match the Thrift path. `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)

# 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)
- Fix: `REMOVE` staging operations no longer require `staging_allowed_local_path` to be set, since removing a remote file does not touch the local filesystem (databricks/databricks-sql-python#726)
Expand Down
90 changes: 70 additions & 20 deletions src/databricks/sql/backend/kernel/auth_bridge.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,16 @@
connector's own OAuth provider because the kernel re-mints tokens
itself and the client secret is not recoverable from a built
provider.
- **OAuth U2M** — for ``auth_type`` ``databricks-oauth`` /
``azure-oauth`` (the browser authorization-code flow), the optional
``oauth_client_id`` / ``oauth_redirect_port`` are forwarded to the
kernel's ``auth_type='oauth-u2m'`` and the kernel runs the browser
flow itself.
- **OAuth U2M** — for ``auth_type`` ``databricks-oauth`` (the browser
authorization-code flow), the connector's ``databricks-sql-python``
app bundle (``client_id`` + ``redirect_port``, 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).

``identity_federation_client_id`` is forwarded with whichever auth shape
wins resolution. It selects mandatory SP-wide workload-identity token
Expand Down Expand Up @@ -48,6 +53,11 @@
import re
from typing import Any, Dict, Optional

from databricks.sql.auth.auth import (
PYSQL_OAUTH_CLIENT_ID,
PYSQL_OAUTH_REDIRECT_PORT_RANGE,
PYSQL_OAUTH_SCOPES,
)
from databricks.sql.auth.authenticators import AccessTokenAuthProvider, AuthProvider
from databricks.sql.auth.token_federation import TokenFederationProvider
from databricks.sql.exc import NotSupportedError, ProgrammingError
Expand Down Expand Up @@ -141,15 +151,23 @@ def kernel_auth_kwargs(
rather than silently picking one flow (and failing later as a
confusing 401 against the wrong principal):
- a custom ``credentials_provider`` *and* M2M kwargs together;
- a U2M ``auth_type`` (``databricks-oauth`` / ``azure-oauth``)
*and* ``oauth_client_secret`` together.
- a U2M ``auth_type`` (``databricks-oauth``) *and*
``oauth_client_secret`` together.

(``azure-oauth`` is rejected as unsupported before these guards —
PECOBLR-4120.)
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
``AccessTokenAuthProvider`` → extract the bearer token.
3. **OAuth U2M** — ``auth_type`` is ``databricks-oauth`` /
``azure-oauth`` → forward optional ``oauth_client_id`` /
``oauth_redirect_port`` to the kernel's ``oauth-u2m``.
3. **OAuth U2M** — ``auth_type`` is ``databricks-oauth`` → forward the
connector's coupled ``databricks-sql-python`` bundle (``client_id``
+ ``redirect_port``, with fixed ``PYSQL_OAUTH_SCOPES``) to the
kernel's ``oauth-u2m``, so a bare U2M connection authenticates as
``databricks-sql-python`` — parity with the Thrift path — rather
than the kernel's own ``databricks-sql-connector`` default
(PECOBLR-4039/4040). ``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``.
Expand All @@ -169,6 +187,25 @@ 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.
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."
)

# 0. Ambiguity guards — fail before any flow is chosen.
if client_secret and opts.get("credentials_provider") is not None:
raise NotSupportedError(
Expand All @@ -178,7 +215,7 @@ def kernel_auth_kwargs(
"kernel-managed M2M, or use the Thrift backend (default) for "
"credentials_provider."
)
if client_secret and auth_type in ("databricks-oauth", "azure-oauth"):
if client_secret and auth_type == "databricks-oauth":
raise NotSupportedError(
f"Ambiguous auth on use_kernel=True: auth_type={auth_type!r} selects "
"the U2M browser flow, but oauth_client_secret was also provided "
Expand Down Expand Up @@ -214,16 +251,29 @@ def kernel_auth_kwargs(
return kwargs

# 3. OAuth U2M — browser authorization-code flow; the kernel runs it.
if auth_type in ("databricks-oauth", "azure-oauth"):
kwargs = {"auth_type": "oauth-u2m"}
if client_id:
kwargs["client_id"] = client_id
# Only databricks-oauth reaches here (azure-oauth rejected up front).
# Forward the connector's own databricks-sql-python bundle instead of
# the kernel's databricks-sql-connector default, for parity with the
# Thrift path. client_id + redirect_port are coupled per app (each
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
# registers its own redirect URI): a caller port only overrides the
# default when an explicit client_id is also supplied. A caller may
# override oauth_scopes; absent one we forward PYSQL_OAUTH_SCOPES as
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
# the default.
if auth_type == "databricks-oauth":
redirect_port = opts.get("oauth_redirect_port")
if redirect_port is not None:
kwargs["redirect_port"] = int(redirect_port)
# Honor a caller-supplied oauth_scopes (normalized to a list of
# strings); fall back to the connector default when none is given.
scopes = _normalize_scopes(opts.get("oauth_scopes"))
if scopes is not None:
kwargs["oauth_scopes"] = scopes
kwargs = {
"auth_type": "oauth-u2m",
"client_id": client_id or PYSQL_OAUTH_CLIENT_ID,
"redirect_port": (
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

redirect_port should just get passed in, we either use the pass in one or default one, no need to judge based on client_id existense.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The current coupling (forward caller redirect_port only when an explicit client_id is also supplied) deliberately mirrors the Thrift path in src/databricks/sql/auth/auth.py:127-130, which overrides oauth_redirect_port_range only when both oauth_client_id and oauth_redirect_port are set. The rationale is correctness: the default databricks-sql-python app only registers its default port range as redirect URIs, so forwarding an arbitrary caller port without a matching custom client_id yields an unregistered redirect URI and fails the U2M flow. Rather than silently regress this parity/correctness behavior on a one-line request, I replied asking whether the kernel is intended to decouple port from app registration (differing from Thrift). If confirmed, I'll drop the client_id and guard and update test_u2m_redirect_port_ignored_without_client_id accordingly. No code change made yet.

int(redirect_port)
if client_id and redirect_port is not None
else PYSQL_OAUTH_REDIRECT_PORT_RANGE[0]
),
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
"oauth_scopes": scopes if scopes is not None else list(PYSQL_OAUTH_SCOPES),
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
}
if federation_client_id:
kwargs["identity_federation_client_id"] = federation_client_id
return kwargs
Expand Down Expand Up @@ -251,7 +301,7 @@ def kernel_auth_kwargs(
raise NotSupportedError(
f"use_kernel=True requires PAT (access_token), OAuth M2M "
f"(oauth_client_id + oauth_client_secret), or OAuth U2M "
f"(auth_type='databricks-oauth' / 'azure-oauth'), but got "
f"(auth_type='databricks-oauth'), but got "
f"{provider_desc} with auth_type={auth_type!r}. Use the Thrift "
"backend (default) for other auth flows."
)
Expand Down
146 changes: 130 additions & 16 deletions tests/unit/test_kernel_auth_bridge.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,9 @@
look through the wrapper).
- OAuth M2M (``oauth_client_id`` + ``oauth_client_secret``) routes
through ``auth_type='oauth-m2m'`` with the raw creds forwarded.
- OAuth U2M (``auth_type='databricks-oauth'`` / ``'azure-oauth'``)
routes through ``auth_type='oauth-u2m'``.
- OAuth U2M (``auth_type='databricks-oauth'``) routes through
``auth_type='oauth-u2m'``. ``azure-oauth`` (Azure AD) is not yet
supported on the kernel path and is rejected (PECOBLR-4120).
- A custom ``credentials_provider`` and any other non-PAT shape raise
``NotSupportedError`` with a clear, actionable message.
"""
Expand All @@ -26,6 +27,11 @@
# require the kernel wheel). So this test can run on the
# default-deps CI matrix without any extras. No importorskip needed.

from databricks.sql.auth.auth import (
PYSQL_OAUTH_CLIENT_ID,
PYSQL_OAUTH_SCOPES,
PYSQL_OAUTH_REDIRECT_PORT_RANGE,
)
from databricks.sql.auth.authenticators import (
AccessTokenAuthProvider,
AuthProvider,
Expand Down Expand Up @@ -240,34 +246,142 @@ def test_client_id_without_secret_does_not_trigger_m2m(self):


class TestKernelOAuthU2M:
@pytest.mark.parametrize("auth_type", ["databricks-oauth", "azure-oauth"])
def test_u2m_routes_to_kernel_u2m(self, auth_type):
"""Only ``databricks-oauth`` U2M is supported on the kernel path.

The kernel core default U2M app is ``databricks-sql-connector`` /
``sql offline_access`` / port 8030 (see PECOBLR-4039). The Python
connector is an OVERRIDE: on the kernel path it forwards its OWN
coupled ``client_id`` + ``redirect_port`` bundle so it authenticates
as ``databricks-sql-python`` rather than the kernel default. A caller
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)."""

def test_bare_databricks_oauth_forwards_full_python_bundle(self):
# No overrides → forward the databricks-sql-python bundle in full
# so the kernel does NOT fall back to its databricks-sql-connector
# default. This is the parity-with-Thrift acceptance criterion.
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{"auth_type": auth_type},
{"auth_type": "databricks-oauth"},
)
assert kwargs == {"auth_type": "oauth-u2m"}
assert kwargs == {
"auth_type": "oauth-u2m",
"client_id": PYSQL_OAUTH_CLIENT_ID,
"redirect_port": PYSQL_OAUTH_REDIRECT_PORT_RANGE[0],
"oauth_scopes": list(PYSQL_OAUTH_SCOPES),
}

def test_u2m_forwards_client_id_and_redirect_port(self):
@pytest.mark.parametrize(
"opts",
[
{"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)

def test_u2m_custom_client_id_port_and_scopes_honored(self):
# A caller may override the coupled client_id + redirect_port and
# the oauth_scopes. All three are forwarded as supplied.
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{
"auth_type": "databricks-oauth",
"oauth_client_id": "custom-client",
"oauth_redirect_port": 8030,
"oauth_scopes": ["custom-scope", "offline_access"],
"oauth_redirect_port": 9999,
},
)
assert kwargs == {
"auth_type": "oauth-u2m",
"client_id": "custom-client",
"redirect_port": 8030,
"redirect_port": 9999,
"oauth_scopes": ["custom-scope", "offline_access"],
}

@pytest.mark.parametrize("auth_type", ["databricks-oauth", "azure-oauth"])
def test_u2m_forwards_scopes(self, auth_type):
def test_u2m_custom_client_id_only_falls_back_to_connector_defaults(self):
# A custom client_id without explicit scopes/port fills the
# remaining two from the connector defaults — a custom client_id
# still uses PYSQL_OAUTH_SCOPES and the default redirect-port range.
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{
"auth_type": "databricks-oauth",
"oauth_client_id": "custom-client",
},
)
assert kwargs == {
"auth_type": "oauth-u2m",
"client_id": "custom-client",
"redirect_port": PYSQL_OAUTH_REDIRECT_PORT_RANGE[0],
"oauth_scopes": list(PYSQL_OAUTH_SCOPES),
}

def test_u2m_redirect_port_coerced_to_int(self):
# oauth_redirect_port may arrive as a string (e.g. from a DSN);
# the kernel binding wants an int. The port override is coupled to
# an explicit client_id (see the coupling test below), so supply
# one here to exercise the coercion path.
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{
"auth_type": "databricks-oauth",
"oauth_client_id": "custom-client",
"oauth_redirect_port": "8021",
},
)
assert kwargs["redirect_port"] == 8021
assert isinstance(kwargs["redirect_port"], int)

def test_u2m_redirect_port_ignored_without_client_id(self):
# A bare oauth_redirect_port (no explicit client_id) must NOT be
# forwarded: it would be paired with the default databricks-sql-python
# app, whose registered redirect URIs only cover the default port
# range, so an arbitrary port would resolve to an unregistered URI
# and fail the U2M flow. This mirrors the Thrift path's coupling,
# where oauth_redirect_port_range is only overridden when both
# oauth_client_id and oauth_redirect_port are supplied.
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{"auth_type": "databricks-oauth", "oauth_redirect_port": 9999},
)
assert kwargs["redirect_port"] == PYSQL_OAUTH_REDIRECT_PORT_RANGE[0]

def test_u2m_honors_custom_scopes(self):
# A caller-supplied oauth_scopes is forwarded to the kernel, even
# without an explicit client_id. Absent one, PYSQL_OAUTH_SCOPES is
# forwarded as the default (see the bare-bundle test above).
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{"auth_type": auth_type, "oauth_scopes": ["all-apis", "offline_access"]},
{
"auth_type": "databricks-oauth",
"oauth_scopes": ["all-apis", "offline_access"],
},
)
assert kwargs["oauth_scopes"] == ["all-apis", "offline_access"]

def test_u2m_normalizes_space_delimited_scopes(self):
# A space-delimited oauth_scopes string is normalized to a list,
# mirroring the M2M path.
kwargs = kernel_auth_kwargs(
_FakeOAuthProvider(),
{
"auth_type": "databricks-oauth",
"oauth_scopes": "all-apis offline_access",
},
)
assert kwargs["oauth_scopes"] == ["all-apis", "offline_access"]

Expand Down Expand Up @@ -331,15 +445,15 @@ def _creds_provider():
},
)

@pytest.mark.parametrize("auth_type", ["databricks-oauth", "azure-oauth"])
def test_u2m_auth_type_plus_client_secret_is_rejected(self, auth_type):
def test_u2m_auth_type_plus_client_secret_is_rejected(self):
# User asked for U2M (browser) but also passed a secret (M2M).
# Don't silently route M2M against the wrong principal.
# Don't silently route M2M against the wrong principal. (azure-oauth
# is rejected earlier as unsupported, so it's not exercised here.)
with pytest.raises(NotSupportedError, match="Ambiguous auth"):
kernel_auth_kwargs(
_FakeOAuthProvider(),
{
"auth_type": auth_type,
"auth_type": "databricks-oauth",
"oauth_client_id": "id",
"oauth_client_secret": "sec",
},
Expand Down
Loading