feat: Databricks OAuth (M2M + Azure service principal) auth [PLATL-820] (#2781)

paulteehan · claude · pre-commit-ci[bot] · web-flow · commit 461a59489e19 · 2026-07-08T10:43:10.000+02:00
* feat(databricks): OAuth (M2M + Azure service principal) auth for soda-databricks Add an auth_type discriminator to the Databricks connector with two OAuth client-credentials modes — databricks-oauth-m2m and azure-service-principal — alongside the existing personal-access-token auth. OAuth modes build a Databricks SDK credentials_provider and strip credential fields from sql.connect kwargs; PAT stays the default when auth_type is absent. PLATL-820 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(databricks): resolve SonarCloud S5332 clear-text-HTTP finding in test _with_https used a literal 'http://' in chained startswith calls, which SonarCloud flagged as 'Using HTTP protocol is insecure' (S5332) — the sole condition failing the Quality Gate on PR #2781 (B Security Rating on New Code); also clears the paired chained-startswith smell. Replace with a scheme-agnostic '://' membership test; behavior unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(databricks): correct OAuth credentials_provider shape + fail-loud (review) Address review on PR #2781: - Blocker: credentials_provider was one level too shallow. ExternalAuthProvider calls credentials_provider() once to get the header factory, then calls that per request; passing the SDK factory directly made the per-request call raise 'dict object is not callable', so no OAuth (M2M or Azure SP) connection could be established. Wrap the SDK header factory in a zero-arg lambda. - Raise a clear ValueError when the SDK yields no header factory instead of silently degrading to a credential-less PAT connect. - Cap databricks-sdk to >=0.19,<1 (provider import paths move across majors). - Tests now drive provider()() to a headers dict (catching the shape bug) and cover the None->error path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * chore(deps): regenerate uv.lock for the databricks-sdk dependency Fixes the `uv lock --check` failure in the pre-commit & lockfile CI job. PLATL-820 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * chore(databricks): address review — sdk floor, drop provisional note PR #2781 review (Niels-b): - Bump databricks-sdk floor to >=0.117,<1 (was >=0.19): 0.19 predates the credentials-provider surface the OAuth path imports and no CI job exercises the floor; uv.lock resolves 0.119.0 and the fix was verified against 0.117. Regenerated uv.lock (uv lock --check). - Drop the 'provisional / confirm with BE' note and dangling implementation-plan.md reference on the auth_type literals — the cross-repo contract is confirmed (soda-library#759 + Cloud form); reduced to a plain contract statement. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * chore(databricks): set databricks-sdk floor to >=0.19,<1 (was 0.117) Revert the 0.117 floor. Verified against the SDK source at tag v0.19.0 that the OAuth surface we use is already present at 0.19.0 — oauth_service_principal, azure_service_principal, and Config(client_id/client_secret/azure_*). So 0.117 (a 2026-06 release) excluded ~2 years of compatible releases for no benefit and risked resolver conflicts. Keep the <1 major cap (the genuinely useful guard). Regenerated uv.lock (still resolves 0.119.0 at the top). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
diff --git a/soda-databricks/pyproject.toml b/soda-databricks/pyproject.toml
@@ -10,6 +10,7 @@ authors = [
 dependencies = [
     "soda-core==4.16.0",
     "databricks-sql-connector",
+    "databricks-sdk>=0.19,<1",
 ]
 
 [project.entry-points."soda.plugins.data_source.databricks"]
diff --git a/soda-databricks/src/soda_databricks/common/data_sources/databricks_data_source_connection.py b/soda-databricks/src/soda_databricks/common/data_sources/databricks_data_source_connection.py
@@ -13,7 +13,9 @@
     DataSourceConnectionProperties,
 )
 from soda_databricks.model.data_source.databricks_connection_properties import (
+    DatabricksAzureServicePrincipal,
     DatabricksConnectionProperties,
+    DatabricksOAuthM2M,
 )
 
 logger: logging.Logger = soda_logger
@@ -27,11 +29,77 @@ def _create_connection(
         self,
         config: DatabricksConnectionProperties,
     ):
+        connection_kwargs = config.to_connection_kwargs()
+
+        # OAuth modes: the SDK builds a credentials_provider callable (handling token
+        # acquisition + refresh). The credential fields are stripped from connection_kwargs
+        # by the properties class, so they never reach sql.connect as plain kwargs.
+        credentials_provider = self._build_credentials_provider(config, connection_kwargs)
+        if credentials_provider is not None:
+            return sql.connect(
+                user_agent_entry="Soda Core",
+                credentials_provider=credentials_provider,
+                **connection_kwargs,
+            )
+
+        # Token (PAT) auth: access_token flows through connection_kwargs unchanged.
         return sql.connect(
             user_agent_entry="Soda Core",
-            **config.to_connection_kwargs(),
+            **connection_kwargs,
+        )
+
+    @staticmethod
+    def _build_credentials_provider(config: DatabricksConnectionProperties, connection_kwargs: dict):
+        """Return an SDK credentials_provider callable for OAuth modes, else None.
+
+        The Databricks SDK derives the token endpoint from the workspace host and handles
+        refresh. ``connection_kwargs['server_hostname']`` has had any scheme stripped by the
+        properties layer, so re-add ``https://`` for the SDK ``Config``.
+        """
+        if not isinstance(config, (DatabricksOAuthM2M, DatabricksAzureServicePrincipal)):
+            return None
+
+        from databricks.sdk.core import Config
+        from databricks.sdk.credentials_provider import (
+            azure_service_principal,
+            oauth_service_principal,
         )
 
+        host = f"https://{connection_kwargs['server_hostname']}"
+
+        if isinstance(config, DatabricksOAuthM2M):
+            sdk_config = Config(
+                host=host,
+                client_id=config.client_id,
+                client_secret=config.client_secret.get_secret_value(),
+            )
+            header_factory = oauth_service_principal(sdk_config)
+            auth_desc = "OAuth (M2M)"
+        else:
+            # DatabricksAzureServicePrincipal — Entra ID (Azure AD) service principal.
+            sdk_config = Config(
+                host=host,
+                azure_client_id=config.azure_client_id,
+                azure_client_secret=config.azure_client_secret.get_secret_value(),
+                azure_tenant_id=config.azure_tenant_id,
+            )
+            header_factory = azure_service_principal(sdk_config)
+            auth_desc = "Azure service principal"
+
+        # The SDK returns None when OIDC discovery yields no token endpoint. Fail loudly here
+        # instead of returning None — otherwise the caller would silently degrade to a
+        # credential-less PAT connect for a config the user explicitly marked OAuth.
+        if header_factory is None:
+            raise ValueError(
+                f"Databricks {auth_desc} authentication setup failed: the SDK could not resolve "
+                f"an OIDC token endpoint for host '{host}'. Verify the workspace host and credentials."
+            )
+
+        # databricks-sql-connector's ExternalAuthProvider calls credentials_provider() once to
+        # obtain the header factory, then invokes that per request. So credentials_provider must
+        # be a zero-arg callable that RETURNS the SDK header factory, not the factory itself.
+        return lambda: header_factory
+
     def _fetch_session_timezone(self) -> tzinfo:
         with self.connection.cursor() as cursor:
             cursor.execute("SELECT current_timezone()")
diff --git a/soda-databricks/src/soda_databricks/model/data_source/databricks_connection_properties.py b/soda-databricks/src/soda_databricks/model/data_source/databricks_connection_properties.py
@@ -39,3 +39,47 @@ def to_connection_kwargs(self) -> dict:
 
 class DatabricksTokenAuth(DatabricksSharedConnectionProperties):
     access_token: SecretStr = Field(..., description="Personal access token")
+
+
+class DatabricksOAuthM2M(DatabricksSharedConnectionProperties):
+    """Databricks-managed OAuth machine-to-machine (service principal) auth.
+
+    The token endpoint is derived from ``host`` by the Databricks SDK, so no
+    ``token_url`` is required. The connection layer builds a ``credentials_provider``
+    from these fields via ``oauth_service_principal`` — they must NOT be emitted as
+    plain ``sql.connect`` kwargs, hence the ``to_connection_kwargs`` override below.
+    """
+
+    client_id: str = Field(..., description="Databricks OAuth service-principal client ID")
+    client_secret: SecretStr = Field(..., description="Databricks OAuth service-principal client secret")
+
+    # Consumed by the connection layer to build the credentials_provider, never passed to sql.connect.
+    _credential_fields: ClassVar[tuple] = ("client_id", "client_secret")
+
+    def to_connection_kwargs(self) -> dict:
+        connection_kwargs = super().to_connection_kwargs()
+        for field_name in self._credential_fields:
+            connection_kwargs.pop(field_name, None)
+        return connection_kwargs
+
+
+class DatabricksAzureServicePrincipal(DatabricksSharedConnectionProperties):
+    """Microsoft Entra ID (Azure AD) service-principal auth for Azure Databricks.
+
+    Distinct from Databricks-managed OAuth M2M: the SP lives in an Entra app
+    registration and the token comes from ``login.microsoftonline.com``, so a tenant
+    ID is required. The connection layer builds a ``credentials_provider`` from these
+    fields; they must NOT be emitted as plain ``sql.connect`` kwargs.
+    """
+
+    azure_client_id: str = Field(..., description="Entra ID (Azure AD) service-principal application/client ID")
+    azure_client_secret: SecretStr = Field(..., description="Entra ID (Azure AD) service-principal client secret")
+    azure_tenant_id: str = Field(..., description="Entra ID (Azure AD) directory/tenant ID")
+
+    _credential_fields: ClassVar[tuple] = ("azure_client_id", "azure_client_secret", "azure_tenant_id")
+
+    def to_connection_kwargs(self) -> dict:
+        connection_kwargs = super().to_connection_kwargs()
+        for field_name in self._credential_fields:
+            connection_kwargs.pop(field_name, None)
+        return connection_kwargs
diff --git a/soda-databricks/src/soda_databricks/model/data_source/databricks_data_source.py b/soda-databricks/src/soda_databricks/model/data_source/databricks_data_source.py
@@ -4,10 +4,24 @@
 from pydantic import Field, field_validator
 from soda_core.model.data_source.data_source import DataSourceBase
 from soda_databricks.model.data_source.databricks_connection_properties import (
+    DatabricksAzureServicePrincipal,
     DatabricksConnectionProperties,
+    DatabricksOAuthM2M,
     DatabricksTokenAuth,
 )
 
+# Explicit auth_type discriminator values (Trino/Snowflake/BigQuery style). These literals
+# are the cross-repo contract: they match soda-library#759 and the Cloud connection form.
+AUTH_TYPE_TOKEN = "personal-access-token"
+AUTH_TYPE_OAUTH_M2M = "databricks-oauth-m2m"
+AUTH_TYPE_AZURE_SP = "azure-service-principal"
+
+_AUTH_TYPE_TO_PROPERTIES = {
+    AUTH_TYPE_TOKEN: DatabricksTokenAuth,
+    AUTH_TYPE_OAUTH_M2M: DatabricksOAuthM2M,
+    AUTH_TYPE_AZURE_SP: DatabricksAzureServicePrincipal,
+}
+
 
 class DatabricksDataSource(DataSourceBase, abc.ABC):
     type: Literal["databricks"] = Field("databricks")
@@ -18,6 +32,28 @@ class DatabricksDataSource(DataSourceBase, abc.ABC):
     @field_validator("connection_properties", mode="before")
     @classmethod
     def infer_connection_type(cls, value):
+        # Already a resolved properties object (e.g. constructed in code) — pass through.
+        if isinstance(value, DatabricksConnectionProperties):
+            return value
+        if not isinstance(value, dict):
+            raise ValueError("Could not infer Databricks connection type from input")
+
+        # Copy so the discriminator can be stripped without mutating caller input; also
+        # prevents auth_type leaking into sql.connect kwargs (base model has extra='allow').
+        value = dict(value)
+        auth_type = value.pop("auth_type", None)
+
+        if auth_type is not None:
+            properties_class = _AUTH_TYPE_TO_PROPERTIES.get(auth_type)
+            if properties_class is None:
+                supported = ", ".join(sorted(_AUTH_TYPE_TO_PROPERTIES))
+                raise ValueError(f"Unknown Databricks auth_type '{auth_type}'. Supported: {supported}")
+            return properties_class(**value)
+
+        # Backward compatibility: no explicit auth_type → infer PAT from field presence.
         if "access_token" in value:
             return DatabricksTokenAuth(**value)
-        raise ValueError("Could not infer Databricks connection type from input")
+        raise ValueError(
+            "Could not infer Databricks connection type: provide 'auth_type' "
+            f"(one of {', '.join(sorted(_AUTH_TYPE_TO_PROPERTIES))}) or 'access_token'."
+        )
diff --git a/soda-databricks/tests/data_sources/test_databricks.py b/soda-databricks/tests/data_sources/test_databricks.py
@@ -19,55 +19,166 @@
 DATABRICKS_HTTP_PATH = os.getenv("DATABRICKS_HTTP_PATH")
 DATABRICKS_TOKEN = os.getenv("DATABRICKS_TOKEN")
 DATABRICKS_CATALOG = os.getenv("DATABRICKS_CATALOG", "unity_catalog")
-DATABRICKS_HOSTNAME_WITH_HTTPS = (
-    f"https://{DATABRICKS_HOST}"
-    if not (DATABRICKS_HOST.startswith("https://") or DATABRICKS_HOST.startswith("http://"))
-    else DATABRICKS_HOST
+# OAuth service-principal creds (client-credentials M2M)
+DATABRICKS_CLIENT_ID = os.getenv("DATABRICKS_CLIENT_ID")
+DATABRICKS_CLIENT_SECRET = os.getenv("DATABRICKS_CLIENT_SECRET")
+# Azure Entra ID service-principal creds
+DATABRICKS_AZURE_CLIENT_ID = os.getenv("DATABRICKS_AZURE_CLIENT_ID")
+DATABRICKS_AZURE_CLIENT_SECRET = os.getenv("DATABRICKS_AZURE_CLIENT_SECRET")
+DATABRICKS_AZURE_TENANT_ID = os.getenv("DATABRICKS_AZURE_TENANT_ID")
+
+
+def _with_https(host: str | None) -> str | None:
+    if host and "://" not in host:
+        return f"https://{host}"
+    return host
+
+
+DATABRICKS_HOSTNAME_WITH_HTTPS = _with_https(DATABRICKS_HOST)
+
+# Which live-credential sets are available. Live cases are only added when their creds are
+# present, so the module imports and the no-credential (config-validation) cases below still
+# run in any environment.
+_has_pat = all([DATABRICKS_HOST, DATABRICKS_HTTP_PATH, DATABRICKS_TOKEN])
+_has_oauth_m2m = all([DATABRICKS_HOST, DATABRICKS_HTTP_PATH, DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET])
+_has_azure_sp = all(
+    [
+        DATABRICKS_HOST,
+        DATABRICKS_HTTP_PATH,
+        DATABRICKS_AZURE_CLIENT_ID,
+        DATABRICKS_AZURE_CLIENT_SECRET,
+        DATABRICKS_AZURE_TENANT_ID,
+    ]
 )
 
 # define test cases and expected behavior (passing unless otherwise specified)
+# Config-validation cases run everywhere (no creds/live connection): they confirm auth_type
+# dispatch and required-field validation through the full YAML -> DataSourceImpl path.
 test_connections: list[TestConnection] = [
-    TestConnection(  # correct connection, should work
-        test_name="correct_connection",
-        connection_yaml_str=f"""
+    TestConnection(  # unknown discriminator must be rejected at parse time
+        test_name="unknown_auth_type_rejected",
+        connection_yaml_str="""
                 type: databricks
                 name: DATABRICKS_TEST
                 connection:
-                    host: {DATABRICKS_HOST}
-                    http_path: {DATABRICKS_HTTP_PATH}
-                    access_token: {DATABRICKS_TOKEN}
-                    catalog: {DATABRICKS_CATALOG}
+                    host: abc.cloud.databricks.com
+                    http_path: /sql/1.0/endpoints/abc
+                    auth_type: not-a-real-mode
             """,
+        valid_yaml=False,
+        expected_yaml_error="Unknown Databricks auth_type",
     ),
-    TestConnection(  # confirm session configuration is applied
-        test_name="applies_session_configuration",
-        connection_yaml_str=f"""
+    TestConnection(  # OAuth M2M without its required secret must be rejected at parse time
+        test_name="oauth_m2m_missing_client_secret_rejected",
+        connection_yaml_str="""
                 type: databricks
                 name: DATABRICKS_TEST
                 connection:
-                    host: {DATABRICKS_HOST}
-                    http_path: {DATABRICKS_HTTP_PATH}
-                    access_token: {DATABRICKS_TOKEN}
-                    catalog: {DATABRICKS_CATALOG}
-                    session_configuration: {{"foo":"bar"}}
+                    host: abc.cloud.databricks.com
+                    http_path: /sql/1.0/endpoints/abc
+                    auth_type: databricks-oauth-m2m
+                    client_id: some-client-id
             """,
-        query_should_succeed=False,
-        expected_query_error="Configuration foo is not available.",
+        valid_yaml=False,
+        expected_yaml_error="client_secret",
     ),
-    TestConnection(  # correct connection, should work
-        test_name="connection_with_https_prefix",
-        connection_yaml_str=f"""
+    TestConnection(  # Azure service principal without its tenant id must be rejected at parse time
+        test_name="azure_service_principal_missing_tenant_rejected",
+        connection_yaml_str="""
                 type: databricks
                 name: DATABRICKS_TEST
                 connection:
-                    host: {DATABRICKS_HOSTNAME_WITH_HTTPS}
-                    http_path: {DATABRICKS_HTTP_PATH}
-                    access_token: {DATABRICKS_TOKEN}
-                    catalog: {DATABRICKS_CATALOG}
+                    host: adb-123.azuredatabricks.net
+                    http_path: /sql/1.0/endpoints/abc
+                    auth_type: azure-service-principal
+                    azure_client_id: some-client-id
+                    azure_client_secret: some-secret
             """,
+        valid_yaml=False,
+        expected_yaml_error="azure_tenant_id",
     ),
 ]
 
+if _has_pat:
+    test_connections += [
+        TestConnection(  # correct connection, should work
+            test_name="correct_connection",
+            connection_yaml_str=f"""
+                    type: databricks
+                    name: DATABRICKS_TEST
+                    connection:
+                        host: {DATABRICKS_HOST}
+                        http_path: {DATABRICKS_HTTP_PATH}
+                        access_token: {DATABRICKS_TOKEN}
+                        catalog: {DATABRICKS_CATALOG}
+                """,
+        ),
+        TestConnection(  # confirm session configuration is applied
+            test_name="applies_session_configuration",
+            connection_yaml_str=f"""
+                    type: databricks
+                    name: DATABRICKS_TEST
+                    connection:
+                        host: {DATABRICKS_HOST}
+                        http_path: {DATABRICKS_HTTP_PATH}
+                        access_token: {DATABRICKS_TOKEN}
+                        catalog: {DATABRICKS_CATALOG}
+                        session_configuration: {{"foo":"bar"}}
+                """,
+            query_should_succeed=False,
+            expected_query_error="Configuration foo is not available.",
+        ),
+        TestConnection(  # correct connection, should work
+            test_name="connection_with_https_prefix",
+            connection_yaml_str=f"""
+                    type: databricks
+                    name: DATABRICKS_TEST
+                    connection:
+                        host: {DATABRICKS_HOSTNAME_WITH_HTTPS}
+                        http_path: {DATABRICKS_HTTP_PATH}
+                        access_token: {DATABRICKS_TOKEN}
+                        catalog: {DATABRICKS_CATALOG}
+                """,
+        ),
+    ]
+
+if _has_oauth_m2m:
+    test_connections.append(
+        TestConnection(  # live Databricks-managed OAuth M2M (service principal)
+            test_name="oauth_m2m_connection",
+            connection_yaml_str=f"""
+                    type: databricks
+                    name: DATABRICKS_TEST
+                    connection:
+                        host: {DATABRICKS_HOST}
+                        http_path: {DATABRICKS_HTTP_PATH}
+                        auth_type: databricks-oauth-m2m
+                        client_id: {DATABRICKS_CLIENT_ID}
+                        client_secret: {DATABRICKS_CLIENT_SECRET}
+                        catalog: {DATABRICKS_CATALOG}
+                """,
+        )
+    )
+
+if _has_azure_sp:
+    test_connections.append(
+        TestConnection(  # live Azure Entra ID service principal
+            test_name="azure_service_principal_connection",
+            connection_yaml_str=f"""
+                    type: databricks
+                    name: DATABRICKS_TEST
+                    connection:
+                        host: {DATABRICKS_HOST}
+                        http_path: {DATABRICKS_HTTP_PATH}
+                        auth_type: azure-service-principal
+                        azure_client_id: {DATABRICKS_AZURE_CLIENT_ID}
+                        azure_client_secret: {DATABRICKS_AZURE_CLIENT_SECRET}
+                        azure_tenant_id: {DATABRICKS_AZURE_TENANT_ID}
+                        catalog: {DATABRICKS_CATALOG}
+                """,
+        )
+    )
+
 
 # run tests.  parameterization means each test case will show up as an individual test
 @pytest.mark.parametrize("test_connection", test_connections, ids=[tc.test_name for tc in test_connections])
diff --git a/soda-databricks/tests/unit/test_databricks_oauth.py b/soda-databricks/tests/unit/test_databricks_oauth.py
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,7 @@ authors = [`
`10`	`10`	`dependencies = [`
`11`	`11`	`"soda-core==4.16.0",`
`12`	`12`	`"databricks-sql-connector",`
	`13`	`+ "databricks-sdk>=0.19,<1",`
`13`	`14`	`]`
`14`	`15`
`15`	`16`	`[project.entry-points."soda.plugins.data_source.databricks"]`