From 8cf89b63a7364de4bbb2276e96ca323ef6ce7f95 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 26 Sep 2026 12:09:57 +0000 Subject: [PATCH 1/2] Support OAuth M2M (client credentials) on the default backend; reject unknown auth_type oauth_client_id + oauth_client_secret now authenticate a Databricks service principal with the client-credentials flow against the workspace /oidc/v1/token endpoint (scope all-apis), with tokens refreshed as they expire. Previously the Thrift backend ignored oauth_client_secret and started an interactive browser login with the service principal's client ID, which blocks indefinitely on a server. This matches the credential shape the kernel backend already accepts. An auth_type the connector does not implement now raises ValueError instead of falling through to the interactive login, and a U2M auth_type combined with oauth_client_secret is rejected as ambiguous. The client-credentials token source also ignores extra fields in the token response (the workspace endpoint returns 'scope'), which made OAuthResponse(**payload) raise. --- src/databricks/sql/auth/auth.py | 30 +++++++++++ src/databricks/sql/auth/authenticators.py | 44 ++++++++++++++++ src/databricks/sql/auth/common.py | 2 + src/databricks/sql/auth/oauth.py | 8 ++- src/databricks/sql/client.py | 6 +++ tests/unit/test_auth.py | 64 +++++++++++++++++++++++ 6 files changed, 153 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py index a4d4d6f2e..92a327780 100755 --- a/src/databricks/sql/auth/auth.py +++ b/src/databricks/sql/auth/auth.py @@ -6,6 +6,7 @@ ExternalAuthProvider, DatabricksOAuthProvider, AzureServicePrincipalCredentialProvider, + DatabricksServicePrincipalCredentialProvider, ) from databricks.sql.auth.common import AuthType, ClientContext from databricks.sql.auth.token_federation import TokenFederationProvider @@ -15,8 +16,34 @@ def get_auth_provider(cfg: ClientContext, http_client): # Determine the base auth provider base_provider: Optional[AuthProvider] = None + if cfg.auth_type and cfg.auth_type not in [t.value for t in AuthType]: + # Never fall through to the interactive browser login below for an + # auth_type this connector does not implement (e.g. a typo). + raise ValueError( + f"Unsupported auth_type {cfg.auth_type!r}; supported: " + + ", ".join(t.value for t in AuthType) + ) + if cfg.credentials_provider: base_provider = ExternalAuthProvider(cfg.credentials_provider) + elif cfg.oauth_client_secret and cfg.auth_type != AuthType.AZURE_SP_M2M.value: + if cfg.auth_type in [ + AuthType.DATABRICKS_OAUTH.value, + AuthType.AZURE_OAUTH.value, + ]: + raise ValueError( + f"auth_type={cfg.auth_type!r} selects the interactive (U2M) flow, " + "but oauth_client_secret was also provided (M2M). Drop " + "oauth_client_secret for U2M, or drop auth_type for M2M." + ) + base_provider = ExternalAuthProvider( + DatabricksServicePrincipalCredentialProvider( + cfg.hostname, + cfg.oauth_client_id, + cfg.oauth_client_secret, + http_client, + ) + ) elif cfg.auth_type == AuthType.AZURE_SP_M2M.value: base_provider = ExternalAuthProvider( AzureServicePrincipalCredentialProvider( @@ -132,5 +159,8 @@ def get_python_sql_connector_auth_provider(hostname: str, http_client, **kwargs) oauth_persistence=kwargs.get("experimental_oauth_persistence"), credentials_provider=kwargs.get("credentials_provider"), identity_federation_client_id=kwargs.get("identity_federation_client_id"), + oauth_client_secret=kwargs.get("oauth_client_secret"), ) + if cfg.oauth_client_secret and not kwargs.get("oauth_client_id"): + raise ValueError("OAuth M2M needs oauth_client_id with oauth_client_secret") return get_auth_provider(cfg, http_client) diff --git a/src/databricks/sql/auth/authenticators.py b/src/databricks/sql/auth/authenticators.py index c30a287ee..5f194cd97 100644 --- a/src/databricks/sql/auth/authenticators.py +++ b/src/databricks/sql/auth/authenticators.py @@ -236,3 +236,47 @@ def header_factory() -> Dict[str, str]: return headers return header_factory + + +class DatabricksServicePrincipalCredentialProvider(CredentialsProvider): + """ + OAuth machine-to-machine (client credentials) authentication for a + Databricks service principal. + + Tokens come from the workspace's ``/oidc/v1/token`` endpoint and are + refreshed by the token source when they expire, so a long-lived connection + keeps working past the token lifetime. + + Attributes: + hostname (str): The normalized workspace URL (``https:///``). + client_id (str): The service principal's OAuth client (application) ID. + client_secret (str): The service principal's OAuth secret. + """ + + DEFAULT_SCOPE = "all-apis" + + def __init__(self, hostname, client_id, client_secret, http_client): + self.hostname = hostname + self.client_id = client_id + self.client_secret = client_secret + self._http_client = http_client + + def auth_type(self) -> str: + return "oauth-m2m" + + def __call__(self, *args, **kwargs) -> HeaderFactory: + source = ClientCredentialsTokenSource( + token_url=f"{self.hostname.rstrip('/')}/oidc/v1/token", + client_id=self.client_id, + client_secret=self.client_secret, + http_client=self._http_client, + extra_params={"scope": self.DEFAULT_SCOPE}, + ) + + def header_factory() -> Dict[str, str]: + token = source.get_token() + return { + HttpHeader.AUTHORIZATION.value: f"{token.token_type} {token.access_token}" + } + + return header_factory diff --git a/src/databricks/sql/auth/common.py b/src/databricks/sql/auth/common.py index cda30a528..0441c1f19 100644 --- a/src/databricks/sql/auth/common.py +++ b/src/databricks/sql/auth/common.py @@ -38,6 +38,7 @@ def __init__( oauth_persistence=None, credentials_provider=None, identity_federation_client_id: Optional[str] = None, + oauth_client_secret: Optional[str] = None, # HTTP client configuration parameters ssl_options=None, # SSLOptions type socket_timeout: Optional[float] = None, @@ -59,6 +60,7 @@ def __init__( self.auth_type = auth_type self.oauth_scopes = oauth_scopes self.oauth_client_id = oauth_client_id + self.oauth_client_secret = oauth_client_secret self.azure_client_id = azure_client_id self.azure_client_secret = azure_client_secret self.azure_tenant_id = azure_tenant_id diff --git a/src/databricks/sql/auth/oauth.py b/src/databricks/sql/auth/oauth.py index 231a18905..519f728a5 100644 --- a/src/databricks/sql/auth/oauth.py +++ b/src/databricks/sql/auth/oauth.py @@ -1,4 +1,5 @@ import base64 +import dataclasses import hashlib import json import logging @@ -341,7 +342,12 @@ def refresh(self) -> Token: method=HttpMethod.POST, url=self.token_url, headers=headers, body=data ) if response.status == 200: - oauth_response = OAuthResponse(**json.loads(response.data.decode("utf-8"))) + payload = json.loads(response.data.decode("utf-8")) + # Token endpoints add fields such as ``scope``; keep the known ones. + known = {f.name for f in dataclasses.fields(OAuthResponse)} + oauth_response = OAuthResponse( + **{k: v for k, v in payload.items() if k in known} + ) return Token( oauth_response.access_token, oauth_response.token_type, diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 3e06aa39b..5703ee749 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -230,6 +230,12 @@ def __init__( oauth_client_id: `str`, optional custom oauth client_id. If not specified, it will use the built-in client_id of databricks-sql-python. + oauth_client_secret: `str`, optional + OAuth secret of a service principal. Together with `oauth_client_id` + (the service principal's client ID) this selects OAuth + machine-to-machine authentication (client credentials, scope + `all-apis`), with tokens refreshed as they expire. + oauth_redirect_port: `int`, optional port of the oauth redirect uri (localhost). This is required when custom oauth client_id `oauth_client_id` is set diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py index d1b941208..75598bb4c 100644 --- a/tests/unit/test_auth.py +++ b/tests/unit/test_auth.py @@ -1,4 +1,5 @@ import unittest +from urllib.parse import parse_qs import pytest from unittest.mock import patch, MagicMock import jwt @@ -152,6 +153,69 @@ def test_get_python_sql_connector_auth_provider_access_token(self): auth_provider.add_headers(headers) self.assertEqual(headers["Authorization"], "Bearer dpi123") + def test_get_python_sql_connector_auth_provider_oauth_m2m(self): + """oauth_client_id + oauth_client_secret authenticate with client + credentials instead of starting an interactive login.""" + http_client = MagicMock() + http_client.request.return_value = MagicMock( + status=200, + data=json.dumps( + { + "access_token": "m2m-token", + "token_type": "Bearer", + "expires_in": 3600, + "scope": "all-apis", + } + ).encode(), + ) + with patch( + "databricks.sql.auth.auth.DatabricksOAuthProvider", + side_effect=AssertionError("interactive login started"), + ), patch("databricks.sql.auth.oauth.Token.is_expired", return_value=False): + auth_provider = get_python_sql_connector_auth_provider( + "example.cloud.databricks.com", + http_client, + oauth_client_id="sp-client-id", + oauth_client_secret="sp-secret", + ) + headers = {} + auth_provider.add_headers(headers) + self.assertEqual(headers["Authorization"], "Bearer m2m-token") + request = http_client.request.call_args.kwargs + self.assertEqual( + request["url"], "https://example.cloud.databricks.com/oidc/v1/token" + ) + body = parse_qs(request["body"]) + self.assertEqual(body["grant_type"], ["client_credentials"]) + self.assertEqual(body["client_id"], ["sp-client-id"]) + self.assertEqual(body["client_secret"], ["sp-secret"]) + self.assertEqual(body["scope"], ["all-apis"]) + + def test_get_python_sql_connector_auth_provider_oauth_m2m_errors(self): + with self.assertRaisesRegex(ValueError, "needs oauth_client_id"): + get_python_sql_connector_auth_provider( + "example.cloud.databricks.com", MagicMock(), oauth_client_secret="s" + ) + with self.assertRaisesRegex(ValueError, "interactive"): + get_python_sql_connector_auth_provider( + "example.cloud.databricks.com", + MagicMock(), + auth_type="databricks-oauth", + oauth_client_id="c", + oauth_client_secret="s", + ) + + def test_get_python_sql_connector_auth_provider_unknown_auth_type(self): + """An unsupported auth_type must not fall back to a browser login.""" + with patch( + "databricks.sql.auth.auth.DatabricksOAuthProvider", + side_effect=AssertionError("interactive login started"), + ): + with self.assertRaisesRegex(ValueError, "Unsupported auth_type"): + get_python_sql_connector_auth_provider( + "example.cloud.databricks.com", MagicMock(), auth_type="oauth-u2m" + ) + def test_get_python_sql_connector_auth_provider_external(self): class MyProvider(CredentialsProvider): def auth_type(self) -> str: From e1764fe7d4bbbee67efa023c7d85f1eac2e27f0d Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 26 Sep 2026 23:36:34 +0000 Subject: [PATCH 2/2] Reject credentials_provider with oauth_client_secret; document M2M options - credentials_provider together with oauth_client_secret now raises ValueError instead of silently using the provider, matching the kernel auth bridge. azure-sp-m2m keeps ignoring oauth_* values. - The missing-oauth_client_id check no longer fires for azure-sp-m2m. - oauth_redirect_port is documented as U2M-only; oauth_client_id and oauth_client_secret docs describe the M2M shape and its exclusions. - CONNECTION_PARAMETERS.md: oauth_client_secret is supported on Thrift. - Add changelog entries. Signed-off-by: Amin Ghadersohi --- CHANGELOG.md | 5 +++++ CONNECTION_PARAMETERS.md | 14 +++++++------- src/databricks/sql/auth/auth.py | 22 ++++++++++++++++++++-- src/databricks/sql/client.py | 9 ++++++--- tests/unit/test_auth.py | 18 ++++++++++++++++++ 5 files changed, 56 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bce9e16e4..d82409ee7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ # Release History +# Unreleased +- Support OAuth M2M (client credentials) for Databricks service principals on the default Thrift backend: `oauth_client_id` + `oauth_client_secret` now authenticate with tokens from the workspace `/oidc/v1/token` endpoint (scope `all-apis`), refreshed as they expire. Previously the secret was ignored and the connection started an interactive browser login. +- Reject an unsupported `auth_type` with `ValueError` instead of falling back to the interactive browser login, and reject `oauth_client_secret` together with `credentials_provider` or a U2M `auth_type`, as the kernel backend does. +- Fix: token responses with extra fields (such as `scope`) no longer fail in `ClientCredentialsTokenSource`. + # 4.6.0 (2026-09-24) - Upgrade Databricks SQL Kernel to 1.1.0; the kernel dependency is now stable and no longer experimental. - Transparently auto-recover Thrift connections to Reyden / Real-Time warehouses: when a warehouse rejects the default Thrift protocol (SQLSTATE `KP001`), the session is re-opened on the kernel backend and the warehouse is remembered so later connections skip Thrift. Applies only when no backend was chosen explicitly. diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index eaf69c57b..5570193f6 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -83,10 +83,10 @@ to change without notice. | Option | Type | Thrift | Kernel | Default Value | Note | | --------------------------------------------------- | -------------------- | :----: | :----: | ------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `access_token` (PAT) | `str` | ✅ | ✅ | `None` | Personal Access Token / bearer token. When no auth signal is supplied, Thrift falls back to Databricks OAuth U2M; Kernel requires an explicit U2M `auth_type` or another supported credential flow. | -| `auth_type` | `str` | ✅ | ✅ | `None` | `databricks-oauth` (U2M), `azure-oauth` (Azure AD U2M), or `azure-sp-m2m` (Azure service-principal M2M). All three work on both backends. Thrift treats an otherwise credential-less `None` as Databricks OAuth; Kernel does not implicitly select U2M. On Kernel, `azure-oauth` uses the same workspace-federated browser flow as `databricks-oauth`. | -| `oauth_client_id` (OAuth) | `str` | ✅ | ✅ | built-in id for U2M | Custom U2M client id on both backends. Kernel also uses it with `oauth_client_secret` or the JWT options for M2M; those M2M flows have no built-in client-id 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. | +| `auth_type` | `str` | ✅ | ✅ | `None` | `databricks-oauth` (U2M), `azure-oauth` (Azure AD U2M), or `azure-sp-m2m` (Azure service-principal M2M). All three work on both backends. Thrift rejects any other value with `ValueError`. Thrift treats an otherwise credential-less `None` as Databricks OAuth; Kernel does not implicitly select U2M. On Kernel, `azure-oauth` uses the same workspace-federated browser flow as `databricks-oauth`. | +| `oauth_client_id` (OAuth) | `str` | ✅ | ✅ | built-in id for U2M | Custom U2M client id on both backends. Both backends also use it with `oauth_client_secret` for M2M, and Kernel with the JWT options; those M2M flows have no built-in client-id default. | +| `oauth_redirect_port` (U2M) | `int` | ✅ | ✅ | `None` | Localhost redirect port for the browser (U2M) flow; not used for M2M. 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` | With `oauth_client_id`, selects OAuth M2M (client credentials) for a Databricks service principal on both backends. On Thrift the token comes from `https:///oidc/v1/token` with scope `all-apis` and is refreshed as it expires. Rejected together with `credentials_provider` or `auth_type` `databricks-oauth`/`azure-oauth`; ignored with `azure-sp-m2m`. | | `oauth_jwt_key_file` (OAuth M2M, JWT private key) | `str` | ❌ | ✅ | `None` | **Kernel-only.** Path to the PEM private key for JWT private-key M2M (RFC 7523 client assertion). Supplying it selects the JWT flow: the kernel signs a short-lived assertion with the key instead of sending a client secret. Requires `oauth_client_id` + `oauth_jwt_kid`; mutually exclusive with `oauth_client_secret` / `credentials_provider`. | | `oauth_jwt_kid` (OAuth M2M, JWT private key) | `str` | ❌ | ✅ | `None` | **Kernel-only.** Key id written into the JWT header so the IdP can select the registered public key. Required with `oauth_jwt_key_file`. (For Entra ID this is the certificate's `x5t` thumbprint.) | | `oauth_jwt_passphrase` (OAuth M2M, JWT private key) | `str` | ❌ | ✅ | `None` | **Kernel-only.** Passphrase for an encrypted PKCS#8 private key; omit for an unencrypted key. | @@ -215,9 +215,9 @@ installed kernel binding. `force_enable_telemetry` remains Python-driver-only. ### Supported on Kernel, missing / ignored on Thrift -1. Kernel-managed OAuth M2M: `oauth_client_secret`, `oauth_jwt_key_file`, - `oauth_jwt_kid`, `oauth_jwt_passphrase`, `oauth_jwt_algorithm`, and - `token_url`. +1. Kernel-managed OAuth M2M options: `oauth_jwt_key_file`, `oauth_jwt_kid`, + `oauth_jwt_passphrase`, `oauth_jwt_algorithm`, and `token_url` + (`oauth_client_secret` M2M works on both backends). 2. Custom `oauth_scopes`. 3. Kernel U2M encrypted token storage (`oauth_token_cache_enabled`). diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py index 92a327780..dfa09de35 100755 --- a/src/databricks/sql/auth/auth.py +++ b/src/databricks/sql/auth/auth.py @@ -24,9 +24,22 @@ def get_auth_provider(cfg: ClientContext, http_client): + ", ".join(t.value for t in AuthType) ) + # azure-sp-m2m is explicit and uses the azure_* credentials; oauth_* values + # are ignored for it, as on the kernel path. + oauth_m2m = ( + bool(cfg.oauth_client_secret) and cfg.auth_type != AuthType.AZURE_SP_M2M.value + ) + if oauth_m2m and cfg.credentials_provider: + # Rejected on the kernel path too; neither should silently win. + raise ValueError( + "Ambiguous auth: both a custom credentials_provider and " + "oauth_client_secret were provided. Pass oauth_client_id + " + "oauth_client_secret for OAuth M2M, or credentials_provider alone." + ) + if cfg.credentials_provider: base_provider = ExternalAuthProvider(cfg.credentials_provider) - elif cfg.oauth_client_secret and cfg.auth_type != AuthType.AZURE_SP_M2M.value: + elif oauth_m2m: if cfg.auth_type in [ AuthType.DATABRICKS_OAUTH.value, AuthType.AZURE_OAUTH.value, @@ -161,6 +174,11 @@ def get_python_sql_connector_auth_provider(hostname: str, http_client, **kwargs) identity_federation_client_id=kwargs.get("identity_federation_client_id"), oauth_client_secret=kwargs.get("oauth_client_secret"), ) - if cfg.oauth_client_secret and not kwargs.get("oauth_client_id"): + if ( + cfg.oauth_client_secret + and cfg.auth_type != AuthType.AZURE_SP_M2M.value + and not cfg.credentials_provider + and not kwargs.get("oauth_client_id") + ): raise ValueError("OAuth M2M needs oauth_client_id with oauth_client_secret") return get_auth_provider(cfg, http_client) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 5703ee749..5fe17ef5d 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -229,16 +229,19 @@ def __init__( oauth_client_id: `str`, optional custom oauth client_id. If not specified, it will use the built-in client_id of databricks-sql-python. + For OAuth M2M (with `oauth_client_secret`) this is the service principal's client ID and is required. oauth_client_secret: `str`, optional OAuth secret of a service principal. Together with `oauth_client_id` (the service principal's client ID) this selects OAuth machine-to-machine authentication (client credentials, scope - `all-apis`), with tokens refreshed as they expire. + `all-apis`), with tokens refreshed as they expire. Cannot be + combined with `credentials_provider` or with `auth_type` + `databricks-oauth` / `azure-oauth`; ignored for `azure-sp-m2m`. oauth_redirect_port: `int`, optional - port of the oauth redirect uri (localhost). This is required when custom oauth client_id - `oauth_client_id` is set + port of the oauth redirect uri (localhost) for the interactive (U2M) flow. This is required when + a custom `oauth_client_id` is used for U2M. Not used for OAuth M2M (`oauth_client_secret`). identity_federation_client_id: `str`, optional Service-principal client ID for mandatory SP-wide workload identity diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py index 75598bb4c..2dcb6b80b 100644 --- a/tests/unit/test_auth.py +++ b/tests/unit/test_auth.py @@ -205,6 +205,24 @@ def test_get_python_sql_connector_auth_provider_oauth_m2m_errors(self): oauth_client_secret="s", ) + def test_get_python_sql_connector_auth_provider_oauth_m2m_ambiguous(self): + class MyProvider(CredentialsProvider): + def auth_type(self) -> str: + return "mine" + + def __call__(self, *args, **kwargs) -> HeaderFactory: + return lambda: {"foo": "bar"} + + for client_id in ("c", None): + with self.assertRaisesRegex(ValueError, "Ambiguous auth"): + get_python_sql_connector_auth_provider( + "example.cloud.databricks.com", + MagicMock(), + credentials_provider=MyProvider(), + oauth_client_id=client_id, + oauth_client_secret="s", + ) + def test_get_python_sql_connector_auth_provider_unknown_auth_type(self): """An unsupported auth_type must not fall back to a browser login.""" with patch(