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.

Author: aminghadersohi

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)

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://<host>/``).
+        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

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

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,

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

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: