You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ru...@apache.org on 2026/07/02 16:26:16 UTC
(superset) branch master updated: feat(database): add databricks oauth support (#41421)
This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 2992d7b4c87 feat(database): add databricks oauth support (#41421)
2992d7b4c87 is described below
commit 2992d7b4c87807d181169709a92df154c3d5c060
Author: Evan Rusackas <ev...@preset.io>
AuthorDate: Thu Jul 2 09:26:08 2026 -0700
feat(database): add databricks oauth support (#41421)
Co-authored-by: fabian_zse <fa...@zalando.de>
Co-authored-by: Claude Opus 4.8 <no...@anthropic.com>
---
.../version-6.0.0/configuration/databases.mdx | 74 +++
superset/db_engine_specs/databricks.py | 195 +++++-
.../unit_tests/db_engine_specs/test_databricks.py | 731 ++++++++++++++++++++-
.../db_engine_specs/test_databricks_multi_cloud.py | 127 ++++
4 files changed, 1124 insertions(+), 3 deletions(-)
diff --git a/docs/user_docs_versioned_docs/version-6.0.0/configuration/databases.mdx b/docs/user_docs_versioned_docs/version-6.0.0/configuration/databases.mdx
index 779c33b05d1..3c4816f5826 100644
--- a/docs/user_docs_versioned_docs/version-6.0.0/configuration/databases.mdx
+++ b/docs/user_docs_versioned_docs/version-6.0.0/configuration/databases.mdx
@@ -519,6 +519,80 @@ For a connection to a SQL endpoint you need to use the HTTP path from the endpoi
{"connect_args": {"http_path": "/sql/1.0/endpoints/****", "driver_path": "/path/to/odbc/driver"}}
```
+##### OAuth2 Authentication
+
+Superset supports OAuth2 authentication for Databricks, allowing users to authenticate with their personal Databricks accounts instead of using shared access tokens. This provides better security and audit capabilities.
+
+###### Prerequisites
+
+1. Create an OAuth2 application in your Databricks account:
+ - Go to your Databricks account console
+ - Navigate to **Settings** → **Developer** → **OAuth apps**
+ - Create a new OAuth app with the redirect URI: `http://your-superset-host:port/api/v1/database/oauth2/`
+
+2. Configure OAuth2 in your `superset_config.py`:
+
+```python
+from datetime import timedelta
+
+# OAuth2 configuration for Databricks
+# The authorization endpoint is derived from your Databricks workspace host; the
+# token endpoint must be set explicitly (see notes below).
+DATABASE_OAUTH2_CLIENTS = {
+ "Databricks (legacy)": {
+ "id": "your-databricks-client-id",
+ "secret": "your-databricks-client-secret",
+ "scope": "sql",
+ "token_request_uri": "https://your-workspace-host/oidc/v1/token",
+ },
+ "Databricks": {
+ "id": "your-databricks-client-id",
+ "secret": "your-databricks-client-secret",
+ "scope": "sql",
+ "token_request_uri": "https://your-workspace-host/oidc/v1/token",
+ },
+}
+
+# OAuth2 redirect URI (adjust hostname/port for your setup)
+DATABASE_OAUTH2_REDIRECT_URI = "http://your-superset-host:port/api/v1/database/oauth2/"
+
+# Optional: OAuth2 timeout
+DATABASE_OAUTH2_TIMEOUT = timedelta(seconds=30)
+```
+
+Replace the following placeholders:
+- `your-databricks-client-id`: Your Databricks OAuth2 application client ID
+- `your-databricks-client-secret`: Your Databricks OAuth2 application client secret
+- `your-superset-host:port`: Your Superset instance hostname and port
+
+**Multi-Cloud Provider Support**
+
+Databricks fronts the user-to-machine (U2M) OAuth2 flow on every workspace at
+`https://<workspace-host>/oidc/v1/authorize` and
+`https://<workspace-host>/oidc/v1/token`, regardless of whether the workspace
+runs on AWS, Azure, or GCP. Superset derives the **authorization** endpoint
+directly from your connection's host, so no cloud provider or account/tenant
+identifier needs to be configured.
+
+The **token** endpoint cannot be auto-derived (token exchange has no database
+context to read the host), so you must supply `token_request_uri` in
+`DATABASE_OAUTH2_CLIENTS`, set to `https://<workspace-host>/oidc/v1/token` for
+your workspace.
+
+If you supply a fully-resolved `authorization_request_uri` (and/or
+`token_request_uri`), those values take precedence over the host-derived
+defaults.
+
+###### Usage
+
+Once configured, users can:
+
+1. Connect to Databricks databases normally using access tokens
+2. When querying data, Superset will automatically redirect users to authenticate with Databricks if needed
+3. User-specific OAuth2 tokens will be used for database connections, providing better security and audit trails
+
+This feature works with both "Databricks (legacy)" and "Databricks" engine types and automatically supports all major cloud providers (AWS, Azure, GCP).
+
#### Denodo
The recommended connector library for Denodo is
diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py
index 1393608468e..db3d382f2ad 100644
--- a/superset/db_engine_specs/databricks.py
+++ b/superset/db_engine_specs/databricks.py
@@ -16,11 +16,13 @@
# under the License.
from __future__ import annotations
+import logging
from datetime import datetime
-from typing import Any, Callable, TYPE_CHECKING, TypedDict, Union
+from typing import Any, Callable, cast, TYPE_CHECKING, TypedDict, Union
from apispec import APISpec
from apispec.ext.marshmallow import MarshmallowPlugin
+from flask import g
from flask_babel import gettext as __
from marshmallow import fields, Schema
from marshmallow.validate import Range
@@ -38,12 +40,21 @@ from superset.db_engine_specs.base import (
)
from superset.db_engine_specs.hive import HiveEngineSpec
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import OAuth2Error
from superset.utils import json
from superset.utils.core import get_user_agent, QuerySource
from superset.utils.network import is_hostname_valid, is_port_open
if TYPE_CHECKING:
from superset.models.core import Database
+ from superset.superset_typing import (
+ OAuth2ClientConfig,
+ OAuth2State,
+ OAuth2TokenResponse,
+ )
+
+
+logger = logging.getLogger(__name__)
try:
@@ -277,6 +288,168 @@ class DatabricksDynamicBaseEngineSpec(BasicParametersMixin, DatabricksBaseEngine
"port": "port",
}
+ # The Databricks SQL driver has no dedicated authentication exception, so an
+ # expired or missing token surfaces as a generic driver error. These case-
+ # insensitive substrings flag the errors that should bootstrap a re-auth.
+ oauth2_auth_failure_signals = (
+ "http 401",
+ "unauthorized",
+ "unauthenticated",
+ "invalid access token",
+ "invalid token",
+ "expired token",
+ "token expired",
+ # Raised by the databricks-sql-connector when no usable credentials are
+ # present (e.g. an OAuth2 token that has been cleared/expired).
+ "no valid authentication settings",
+ )
+
+ @classmethod
+ def _workspace_oauth2_endpoint(cls, database: Database, path: str) -> str:
+ """
+ Build a Databricks OAuth2 (U2M) endpoint from the workspace host.
+
+ Databricks fronts the user-to-machine OAuth2 flow on every workspace at
+ ``https://<workspace-host>/oidc/v1/{authorize,token}`` across AWS, Azure
+ and GCP, so the endpoints derive directly from the connection host and
+ need no account or tenant identifier.
+ """
+ host = database.url_object.host
+ if not host:
+ raise OAuth2Error(
+ "Databricks OAuth2 endpoint could not be resolved: the database "
+ "connection has no host."
+ )
+ return f"https://{host}/oidc/v1/{path}"
+
+ @classmethod
+ def needs_oauth2(cls, ex: Exception) -> bool:
+ """
+ Identify driver errors that should trigger the OAuth2 dance.
+
+ Unlike Trino (``TrinoAuthError``) or GSheets (``UnauthenticatedError``),
+ the Databricks driver raises no dedicated auth exception, so in addition
+ to the base ``isinstance`` check we match the auth signals above on the
+ error message (mirrors ``GSheetsEngineSpec.needs_oauth2``).
+ """
+ if not (g and hasattr(g, "user")):
+ return False
+ if isinstance(ex, cls.oauth2_exception):
+ return True
+ message = str(ex).lower()
+ return any(signal in message for signal in cls.oauth2_auth_failure_signals)
+
+ @classmethod
+ def get_oauth2_authorization_uri(
+ cls,
+ config: "OAuth2ClientConfig",
+ state: "OAuth2State",
+ code_verifier: str | None = None,
+ ) -> str:
+ """
+ Return the URI for the initial OAuth2 request.
+
+ A fully-resolved ``authorization_request_uri`` from
+ ``DATABASE_OAUTH2_CLIENTS`` is preserved; otherwise the endpoint is
+ derived from the workspace host (``https://<host>/oidc/v1/authorize``),
+ which is valid on AWS, Azure and GCP.
+ """
+ if not config.get("authorization_request_uri"):
+ from superset import db
+ from superset.models.core import Database
+
+ database_id = state["database_id"]
+ if database := db.session.get(Database, database_id):
+ config = cast(
+ "OAuth2ClientConfig",
+ dict(config)
+ | {
+ "authorization_request_uri": cls._workspace_oauth2_endpoint(
+ database, "authorize"
+ )
+ },
+ )
+
+ return super().get_oauth2_authorization_uri(config, state, code_verifier)
+
+ @classmethod
+ def get_oauth2_token(
+ cls,
+ config: "OAuth2ClientConfig",
+ code: str,
+ code_verifier: str | None = None,
+ ) -> "OAuth2TokenResponse":
+ """
+ Exchange the authorization code for refresh/access tokens.
+
+ Token exchange runs in a separate request with no database context, so
+ the workspace host is not available to derive the endpoint here. Require
+ a configured ``token_request_uri``
+ (``https://<workspace-host>/oidc/v1/token``) and fail fast rather than
+ POST to an unresolved endpoint.
+ """
+ if not config.get("token_request_uri"):
+ raise OAuth2Error(
+ "Databricks OAuth2 token endpoint is not configured: set "
+ "`token_request_uri` to https://<workspace-host>/oidc/v1/token "
+ "in DATABASE_OAUTH2_CLIENTS."
+ )
+
+ return super().get_oauth2_token(config, code, code_verifier)
+
+ @classmethod
+ def impersonate_user(
+ cls,
+ database: Database,
+ username: str | None,
+ user_token: str | None,
+ url: URL,
+ engine_kwargs: dict[str, Any],
+ ) -> tuple[URL, dict[str, Any]]:
+ """
+ Update connection with the user's OAuth2 access token for impersonation.
+
+ When impersonation is enabled but no user token is available yet (e.g. the
+ first connection, before the OAuth2 dance has run), the stored credential
+ is cleared rather than left in place. The driver then raises a "no valid
+ authentication settings" error, which ``needs_oauth2`` catches to bootstrap
+ the OAuth2 flow instead of silently connecting with a stale credential.
+ """
+ # Replace the credential in the URL with the user's OAuth2 token, falling
+ # back to an empty string to force re-authentication when none is set.
+ url = url.set(password=user_token or "")
+
+ # The Python connector passes the token via ``connect_args`` instead of the
+ # URL password, so keep it in sync (clearing it likewise forces re-auth).
+ connect_args = engine_kwargs.setdefault("connect_args", {})
+ if "access_token" in connect_args:
+ connect_args["access_token"] = user_token or ""
+
+ return url, engine_kwargs
+
+ @staticmethod
+ def update_params_from_encrypted_extra(
+ database: Database, params: dict[str, Any]
+ ) -> None:
+ """
+ Merge ``encrypted_extra`` into the connection params, dropping the
+ ``oauth2_client_info`` block.
+
+ ``oauth2_client_info`` holds the per-database OAuth2 client configuration
+ consumed by ``Database.get_oauth2_config``; it is not a Databricks driver
+ connection argument, so it must be stripped here to avoid poisoning the
+ connection when OAuth2 is configured on the database itself.
+ """
+ if not database.encrypted_extra:
+ return
+ try:
+ encrypted_extra = json.loads(database.encrypted_extra)
+ except json.JSONDecodeError as ex:
+ logger.error(ex, exc_info=True)
+ raise
+ encrypted_extra.pop("oauth2_client_info", None)
+ params.update(encrypted_extra)
+
@staticmethod
def get_extra_params(
database: Database, source: QuerySource | None = None
@@ -474,6 +647,16 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec):
supports_dynamic_catalog = True
supports_cross_catalog_queries = True
+ # OAuth 2.0 support. The flow (endpoint resolution from the workspace host,
+ # `needs_oauth2` detection) is shared via `DatabricksDynamicBaseEngineSpec`.
+ supports_oauth2 = True
+ oauth2_scope = "sql"
+
+ # Authorization endpoint is derived from the workspace host at runtime; the
+ # token endpoint must be configured (no DB context at exchange time).
+ oauth2_authorization_request_uri = ""
+ oauth2_token_request_uri = ""
+
@classmethod
def build_sqlalchemy_uri( # type: ignore
cls, parameters: DatabricksNativeParametersType, *_
@@ -685,6 +868,16 @@ class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec):
supports_dynamic_schema = supports_catalog = supports_dynamic_catalog = True
+ # OAuth 2.0 support. The flow (endpoint resolution from the workspace host,
+ # `needs_oauth2` detection) is shared via `DatabricksDynamicBaseEngineSpec`.
+ supports_oauth2 = True
+ oauth2_scope = "sql"
+
+ # Authorization endpoint is derived from the workspace host at runtime; the
+ # token endpoint must be configured (no DB context at exchange time).
+ oauth2_authorization_request_uri = ""
+ oauth2_token_request_uri = ""
+
@classmethod
def build_sqlalchemy_uri( # type: ignore
cls, parameters: DatabricksPythonConnectorParametersType, *_
diff --git a/tests/unit_tests/db_engine_specs/test_databricks.py b/tests/unit_tests/db_engine_specs/test_databricks.py
index 61366136e6a..f7cfaa5c663 100644
--- a/tests/unit_tests/db_engine_specs/test_databricks.py
+++ b/tests/unit_tests/db_engine_specs/test_databricks.py
@@ -17,14 +17,23 @@
# pylint: disable=unused-argument, import-outside-toplevel, protected-access
from datetime import datetime
-from typing import Optional
+from typing import Any, Optional
+from urllib.parse import parse_qs, urlparse
import pytest
from pytest_mock import MockerFixture
+from sqlalchemy.engine.url import make_url
-from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+from superset.db_engine_specs.base import OAuth2State
+from superset.db_engine_specs.databricks import (
+ DatabricksNativeEngineSpec,
+ DatabricksPythonConnectorEngineSpec,
+)
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import OAuth2Error, OAuth2RedirectError
+from superset.superset_typing import OAuth2ClientConfig
from superset.utils import json
+from superset.utils.oauth2 import decode_oauth2_state
from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
from tests.unit_tests.fixtures.common import dttm # noqa: F401
@@ -291,3 +300,721 @@ def test_get_prequeries(mocker: MockerFixture) -> None:
"USE CATALOG `evil`` USE CATALOG bad`",
"USE SCHEMA `evil`` USE SCHEMA bad`",
]
+
+
+# OAuth2 Tests
+
+
+def test_oauth2_attributes() -> None:
+ """
+ Test that OAuth2 attributes are properly set for both engine specs.
+ """
+ # Test DatabricksNativeEngineSpec
+ assert DatabricksNativeEngineSpec.supports_oauth2 is True
+ assert DatabricksNativeEngineSpec.oauth2_scope == "sql"
+ # The authorization endpoint is derived from the workspace host at runtime;
+ # the token endpoint must be configured explicitly.
+ assert DatabricksNativeEngineSpec.oauth2_authorization_request_uri == ""
+ assert DatabricksNativeEngineSpec.oauth2_token_request_uri == ""
+
+ # Test DatabricksPythonConnectorEngineSpec
+ assert DatabricksPythonConnectorEngineSpec.supports_oauth2 is True
+ assert DatabricksPythonConnectorEngineSpec.oauth2_scope == "sql"
+ assert DatabricksPythonConnectorEngineSpec.oauth2_authorization_request_uri == ""
+ assert DatabricksPythonConnectorEngineSpec.oauth2_token_request_uri == ""
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec],
+)
+@pytest.mark.parametrize(
+ "message",
+ [
+ "Error during request to server: HTTP 401 Unauthorized",
+ "Invalid access token",
+ "The access token expired",
+ "UNAUTHENTICATED: token is no longer valid",
+ "RuntimeError: No valid authentication settings!",
+ ],
+)
+def test_needs_oauth2_detects_auth_failure_from_message(
+ mocker: MockerFixture,
+ spec: Any,
+ message: str,
+) -> None:
+ """
+ The Databricks driver has no dedicated auth exception, so `needs_oauth2`
+ matches auth-failure signals in the error message to bootstrap a re-auth.
+ """
+ g = mocker.patch("superset.db_engine_specs.databricks.g")
+ g.user = mocker.MagicMock()
+
+ assert spec.needs_oauth2(Exception(message)) is True
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec],
+)
+@pytest.mark.parametrize(
+ "message",
+ [
+ "Table not found",
+ # A bare 401 in an unrelated position must not look like an auth failure.
+ "Query failed at line 401: syntax error",
+ ],
+)
+def test_needs_oauth2_ignores_unrelated_errors(
+ mocker: MockerFixture,
+ spec: Any,
+ message: str,
+) -> None:
+ """
+ A non-auth driver error must not trigger the OAuth2 dance.
+ """
+ g = mocker.patch("superset.db_engine_specs.databricks.g")
+ g.user = mocker.MagicMock()
+
+ assert spec.needs_oauth2(Exception(message)) is False
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec],
+)
+def test_needs_oauth2_matches_oauth2_redirect_error(
+ mocker: MockerFixture,
+ spec: Any,
+) -> None:
+ """
+ The inherited `isinstance` check against `oauth2_exception` still holds.
+ """
+ g = mocker.patch("superset.db_engine_specs.databricks.g")
+ g.user = mocker.MagicMock()
+
+ ex = OAuth2RedirectError("https://example/authorize", "tab", "redirect")
+ assert spec.needs_oauth2(ex) is True
+
+
+def test_impersonate_user_with_token(mocker: MockerFixture) -> None:
+ """
+ Test impersonate_user method with OAuth2 token for DatabricksNativeEngineSpec.
+ """
+ database = mocker.MagicMock()
+ original_url = make_url(
+ "databricks+connector://token:original-token@host:443/database"
+ )
+ engine_kwargs = {"connect_args": {"access_token": "original-token"}}
+
+ # Test with user token
+ url, kwargs = DatabricksNativeEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token="user-oauth-token", # noqa: S106
+ url=original_url,
+ engine_kwargs=engine_kwargs,
+ )
+
+ # Check that the password (token) was updated in the URL
+ assert url.password == "user-oauth-token" # noqa: S105
+ # Check that access_token was updated in connect_args
+ assert kwargs["connect_args"]["access_token"] == "user-oauth-token" # noqa: S105
+
+
+def test_impersonate_user_without_token_forces_oauth2(mocker: MockerFixture) -> None:
+ """
+ Without a user token, impersonate_user clears any stored credential so the
+ driver fails authentication and the OAuth2 dance is bootstrapped, rather than
+ silently connecting with a stale credential.
+ """
+ database = mocker.MagicMock()
+ original_url = make_url(
+ "databricks+connector://token:original-token@host:443/database"
+ )
+ engine_kwargs = {"connect_args": {"access_token": "original-token"}}
+
+ # Test without user token
+ url, kwargs = DatabricksNativeEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token=None,
+ url=original_url,
+ engine_kwargs=engine_kwargs,
+ )
+
+ # The stored credential is cleared in both the URL and connect_args
+ assert url.password == "" # noqa: S105
+ assert kwargs["connect_args"]["access_token"] == "" # noqa: S105
+
+
+def test_impersonate_user_python_connector(mocker: MockerFixture) -> None:
+ """
+ Test impersonate_user method for DatabricksPythonConnectorEngineSpec.
+ """
+ database = mocker.MagicMock()
+ original_url = make_url(
+ "databricks://token:original-token@host:443?http_path=path&catalog=main&schema=default"
+ )
+ engine_kwargs = {"connect_args": {"access_token": "original-token"}}
+
+ # Test with user token
+ url, kwargs = DatabricksPythonConnectorEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token="user-oauth-token", # noqa: S106
+ url=original_url,
+ engine_kwargs=engine_kwargs,
+ )
+
+ # Check that the password (token) was updated in the URL
+ assert url.password == "user-oauth-token" # noqa: S105
+ # Check that access_token was updated in connect_args
+ assert kwargs["connect_args"]["access_token"] == "user-oauth-token" # noqa: S105
+
+
+def test_impersonate_user_python_connector_without_token_forces_oauth2(
+ mocker: MockerFixture,
+) -> None:
+ """
+ The Python connector path also clears the credential when no user token is
+ available, so per-user OAuth2 is enforced rather than falling back to a stale
+ connection-level token.
+ """
+ database = mocker.MagicMock()
+ original_url = make_url(
+ "databricks://token:original-token@host:443?http_path=path&catalog=main&schema=default"
+ )
+ engine_kwargs = {"connect_args": {"access_token": "original-token"}}
+
+ url, kwargs = DatabricksPythonConnectorEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token=None,
+ url=original_url,
+ engine_kwargs=engine_kwargs,
+ )
+
+ assert url.password == "" # noqa: S105
+ assert kwargs["connect_args"]["access_token"] == "" # noqa: S105
+
+
+def test_impersonate_user_without_connect_args_token(mocker: MockerFixture) -> None:
+ """
+ When ``connect_args`` carries no ``access_token`` (the URL-only auth path), the
+ token is applied to the URL password and no spurious ``access_token`` key is
+ introduced into ``connect_args``.
+ """
+ database = mocker.MagicMock()
+ original_url = make_url(
+ "databricks+connector://token:original-token@host:443/database"
+ )
+ engine_kwargs: dict[str, Any] = {"connect_args": {"http_path": "/sql/path"}}
+
+ url, kwargs = DatabricksNativeEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token="user-oauth-token", # noqa: S106
+ url=original_url,
+ engine_kwargs=engine_kwargs,
+ )
+
+ assert url.password == "user-oauth-token" # noqa: S105
+ assert "access_token" not in kwargs["connect_args"]
+ # Unrelated connect_args are preserved
+ assert kwargs["connect_args"]["http_path"] == "/sql/path"
+
+
+def test_update_params_strips_oauth2_client_info(mocker: MockerFixture) -> None:
+ """
+ ``oauth2_client_info`` is the per-database OAuth2 client config consumed by
+ ``Database.get_oauth2_config``; it must not leak into the driver connection
+ params, while the rest of ``encrypted_extra`` is still merged.
+ """
+ database = mocker.MagicMock()
+ database.encrypted_extra = json.dumps(
+ {
+ "oauth2_client_info": {
+ "id": "client-id",
+ "secret": "client-secret",
+ "authorization_request_uri": "https://host/oidc/v1/authorize",
+ "token_request_uri": "https://host/oidc/v1/token",
+ },
+ "http_headers": [["X-Custom", "value"]],
+ }
+ )
+ params: dict[str, Any] = {}
+
+ DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, params)
+
+ assert "oauth2_client_info" not in params
+ assert params == {"http_headers": [["X-Custom", "value"]]}
+
+
+def test_update_params_no_encrypted_extra(mocker: MockerFixture) -> None:
+ """
+ A database without ``encrypted_extra`` leaves the params untouched.
+ """
+ database = mocker.MagicMock()
+ database.encrypted_extra = None
+ params: dict[str, Any] = {"existing": "value"}
+
+ DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, params)
+
+ assert params == {"existing": "value"}
+
+
+def test_update_params_merges_when_no_oauth2_client_info(
+ mocker: MockerFixture,
+) -> None:
+ """
+ ``encrypted_extra`` without an ``oauth2_client_info`` block is merged in full,
+ so the strip never removes legitimate driver params.
+ """
+ database = mocker.MagicMock()
+ database.encrypted_extra = json.dumps(
+ {"http_headers": [["X-Custom", "value"]], "_tls_verify_hostname": True}
+ )
+ params: dict[str, Any] = {}
+
+ DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, params)
+
+ assert params == {
+ "http_headers": [["X-Custom", "value"]],
+ "_tls_verify_hostname": True,
+ }
+
+
+def test_update_params_invalid_encrypted_extra_raises(mocker: MockerFixture) -> None:
+ """
+ Malformed ``encrypted_extra`` JSON raises rather than silently connecting.
+ """
+ database = mocker.MagicMock()
+ database.encrypted_extra = "{not valid json"
+
+ with pytest.raises(json.JSONDecodeError):
+ DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+@pytest.fixture
+def oauth2_config_native() -> OAuth2ClientConfig:
+ """
+ Config for Databricks Native OAuth2.
+ """
+ return {
+ "id": "databricks-client-id",
+ "secret": "databricks-client-secret",
+ "scope": "sql",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "authorization_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/authorize",
+ "token_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token",
+ "request_content_type": "json",
+ }
+
+
+@pytest.fixture
+def oauth2_config_python() -> OAuth2ClientConfig:
+ """
+ Config for Databricks Python Connector OAuth2.
+ """
+ return {
+ "id": "databricks-client-id",
+ "secret": "databricks-client-secret",
+ "scope": "sql",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "authorization_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/authorize",
+ "token_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token",
+ "request_content_type": "json",
+ }
+
+
+def test_is_oauth2_enabled_no_config_native(mocker: MockerFixture) -> None:
+ """
+ Test `is_oauth2_enabled` when OAuth2 is not configured for Native engine.
+ """
+ mocker.patch(
+ "flask.current_app.config",
+ new={"DATABASE_OAUTH2_CLIENTS": {}},
+ )
+
+ assert DatabricksNativeEngineSpec.is_oauth2_enabled() is False
+
+
+def test_is_oauth2_enabled_config_native(mocker: MockerFixture) -> None:
+ """
+ Test `is_oauth2_enabled` when OAuth2 is configured for Native engine.
+ """
+ mocker.patch(
+ "flask.current_app.config",
+ new={
+ "DATABASE_OAUTH2_CLIENTS": {
+ "Databricks (legacy)": {
+ "id": "client-id",
+ "secret": "client-secret",
+ },
+ }
+ },
+ )
+
+ assert DatabricksNativeEngineSpec.is_oauth2_enabled() is True
+
+
+def test_is_oauth2_enabled_no_config_python(mocker: MockerFixture) -> None:
+ """
+ Test `is_oauth2_enabled` when OAuth2 is not configured for Python Connector engine.
+ """
+ mocker.patch(
+ "flask.current_app.config",
+ new={"DATABASE_OAUTH2_CLIENTS": {}},
+ )
+
+ assert DatabricksPythonConnectorEngineSpec.is_oauth2_enabled() is False
+
+
+def test_is_oauth2_enabled_config_python(mocker: MockerFixture) -> None:
+ """
+ Test `is_oauth2_enabled` when OAuth2 is configured for Python Connector engine.
+ """
+ mocker.patch(
+ "flask.current_app.config",
+ new={
+ "DATABASE_OAUTH2_CLIENTS": {
+ "Databricks": {
+ "id": "client-id",
+ "secret": "client-secret",
+ },
+ }
+ },
+ )
+
+ assert DatabricksPythonConnectorEngineSpec.is_oauth2_enabled() is True
+
+
+def test_get_oauth2_authorization_uri_native(
+ mocker: MockerFixture,
+ oauth2_config_native: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_authorization_uri` for Native engine.
+ """
+ from superset.db_engine_specs.base import OAuth2State
+
+ state: OAuth2State = {
+ "database_id": 1,
+ "user_id": 1,
+ "default_redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "tab_id": "1234",
+ }
+
+ url = DatabricksNativeEngineSpec.get_oauth2_authorization_uri(
+ oauth2_config_native, state
+ )
+ parsed = urlparse(url)
+ assert parsed.netloc == "accounts.cloud.databricks.com"
+ assert parsed.path == "/oidc/accounts/12345/v1/authorize"
+
+ query = parse_qs(parsed.query)
+ assert query["scope"][0] == "sql"
+ encoded_state = query["state"][0].replace("%2E", ".")
+ assert decode_oauth2_state(encoded_state) == state
+
+
+def test_get_oauth2_authorization_uri_python(
+ mocker: MockerFixture,
+ oauth2_config_python: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_authorization_uri` for Python Connector engine.
+ """
+ from superset.db_engine_specs.base import OAuth2State
+
+ state: OAuth2State = {
+ "database_id": 1,
+ "user_id": 1,
+ "default_redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "tab_id": "1234",
+ }
+
+ url = DatabricksPythonConnectorEngineSpec.get_oauth2_authorization_uri(
+ oauth2_config_python, state
+ )
+ parsed = urlparse(url)
+ assert parsed.netloc == "accounts.cloud.databricks.com"
+ assert parsed.path == "/oidc/accounts/12345/v1/authorize"
+
+ query = parse_qs(parsed.query)
+ assert query["scope"][0] == "sql"
+ encoded_state = query["state"][0].replace("%2E", ".")
+ assert decode_oauth2_state(encoded_state) == state
+
+
+def test_get_oauth2_token_native(
+ mocker: MockerFixture,
+ oauth2_config_native: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_token` for Native engine.
+ """
+ requests = mocker.patch("superset.db_engine_specs.base.requests")
+ requests.post().json.return_value = {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+
+ assert DatabricksNativeEngineSpec.get_oauth2_token(
+ oauth2_config_native, "authorization-code"
+ ) == {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+ requests.post.assert_called_with(
+ "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token",
+ json={
+ "code": "authorization-code",
+ "client_id": "databricks-client-id",
+ "client_secret": "databricks-client-secret",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "grant_type": "authorization_code",
+ },
+ timeout=30.0,
+ )
+
+
+def test_get_oauth2_token_python(
+ mocker: MockerFixture,
+ oauth2_config_python: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_token` for Python Connector engine.
+ """
+ requests = mocker.patch("superset.db_engine_specs.base.requests")
+ requests.post().json.return_value = {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+
+ assert DatabricksPythonConnectorEngineSpec.get_oauth2_token(
+ oauth2_config_python, "authorization-code"
+ ) == {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+ requests.post.assert_called_with(
+ "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token",
+ json={
+ "code": "authorization-code",
+ "client_id": "databricks-client-id",
+ "client_secret": "databricks-client-secret",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "grant_type": "authorization_code",
+ },
+ timeout=30.0,
+ )
+
+
+def test_get_oauth2_fresh_token_native(
+ mocker: MockerFixture,
+ oauth2_config_native: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_fresh_token` for Native engine.
+ """
+ requests = mocker.patch("superset.db_engine_specs.base.requests")
+ requests.post().json.return_value = {
+ "access_token": "new-access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "new-refresh-token",
+ }
+
+ assert DatabricksNativeEngineSpec.get_oauth2_fresh_token(
+ oauth2_config_native, "old-refresh-token"
+ ) == {
+ "access_token": "new-access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "new-refresh-token",
+ }
+ requests.post.assert_called_with(
+ "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token",
+ json={
+ "client_id": "databricks-client-id",
+ "client_secret": "databricks-client-secret",
+ "refresh_token": "old-refresh-token",
+ "grant_type": "refresh_token",
+ },
+ timeout=30.0,
+ )
+
+
+def _oauth2_state() -> OAuth2State:
+ """
+ Build the default OAuth2 state shared by the OAuth2 tests.
+ """
+ state: OAuth2State = {
+ "database_id": 1,
+ "user_id": 1,
+ "default_redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "tab_id": "1234",
+ }
+ return state
+
+
+def _unresolved_oauth2_config() -> OAuth2ClientConfig:
+ """
+ Config as built by `get_oauth2_config` when no endpoints are overridden:
+ the URIs default to the spec's empty class attributes.
+ """
+ return {
+ "id": "databricks-client-id",
+ "secret": "databricks-client-secret",
+ "scope": "sql",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "authorization_request_uri": "",
+ "token_request_uri": "",
+ "request_content_type": "json",
+ }
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec],
+)
+@pytest.mark.parametrize(
+ "host",
+ [
+ "dbc-abc.cloud.databricks.com",
+ "adb-123456789.12.azuredatabricks.net",
+ "123456789.gcp.databricks.com",
+ ],
+)
+def test_get_oauth2_authorization_uri_derives_from_workspace_host(
+ mocker: MockerFixture,
+ spec: Any,
+ host: str,
+) -> None:
+ """
+ With no configured `authorization_request_uri`, the endpoint is derived from
+ the workspace host (`https://<host>/oidc/v1/authorize`) on every cloud, with
+ no account/tenant identifier required.
+ """
+ database = mocker.MagicMock()
+ database.url_object.host = host
+ mocker.patch("superset.db.session.get", return_value=database)
+
+ url = spec.get_oauth2_authorization_uri(
+ _unresolved_oauth2_config(), _oauth2_state()
+ )
+ parsed = urlparse(url)
+ assert parsed.netloc == host
+ assert parsed.path == "/oidc/v1/authorize"
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec],
+)
+def test_get_oauth2_authorization_uri_preserves_configured(
+ mocker: MockerFixture,
+ spec: Any,
+) -> None:
+ """
+ A fully-resolved `authorization_request_uri` is never overwritten by the
+ host-derived endpoint, and no database lookup is needed.
+ """
+ session_get = mocker.patch("superset.db.session.get")
+ config = _unresolved_oauth2_config()
+ config["authorization_request_uri"] = (
+ "https://accounts.cloud.databricks.com/oidc/accounts/override/v1/authorize"
+ )
+
+ url = spec.get_oauth2_authorization_uri(config, _oauth2_state())
+ assert urlparse(url).path == "/oidc/accounts/override/v1/authorize"
+ session_get.assert_not_called()
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec],
+)
+def test_get_oauth2_authorization_uri_fails_without_host(
+ mocker: MockerFixture,
+ spec: Any,
+) -> None:
+ """
+ When the endpoint must be derived but the connection has no host, fail fast
+ instead of emitting an invalid `https:///oidc/v1/authorize` URL.
+ """
+ database = mocker.MagicMock()
+ database.url_object.host = None
+ mocker.patch("superset.db.session.get", return_value=database)
+
+ with pytest.raises(OAuth2Error):
+ spec.get_oauth2_authorization_uri(_unresolved_oauth2_config(), _oauth2_state())
+
+
+@pytest.mark.parametrize(
+ "spec",
+ [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec],
+)
+def test_get_oauth2_token_fails_without_uri(
+ mocker: MockerFixture,
+ spec: Any,
+) -> None:
+ """
+ Token exchange has no database context to auto-detect the endpoint, so a
+ missing `token_request_uri` fails fast rather than POSTing to `.../{}/...`.
+ """
+ with pytest.raises(OAuth2Error):
+ spec.get_oauth2_token(_unresolved_oauth2_config(), "authorization-code")
+
+
+def test_get_oauth2_fresh_token_python(
+ mocker: MockerFixture,
+ oauth2_config_python: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_fresh_token` for Python Connector engine.
+ """
+ requests = mocker.patch("superset.db_engine_specs.base.requests")
+ requests.post().json.return_value = {
+ "access_token": "new-access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "new-refresh-token",
+ }
+
+ assert DatabricksPythonConnectorEngineSpec.get_oauth2_fresh_token(
+ oauth2_config_python, "old-refresh-token"
+ ) == {
+ "access_token": "new-access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "new-refresh-token",
+ }
+ requests.post.assert_called_with(
+ "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token",
+ json={
+ "client_id": "databricks-client-id",
+ "client_secret": "databricks-client-secret",
+ "refresh_token": "old-refresh-token",
+ "grant_type": "refresh_token",
+ },
+ timeout=30.0,
+ )
diff --git a/tests/unit_tests/db_engine_specs/test_databricks_multi_cloud.py b/tests/unit_tests/db_engine_specs/test_databricks_multi_cloud.py
new file mode 100644
index 00000000000..65f8218e3a1
--- /dev/null
+++ b/tests/unit_tests/db_engine_specs/test_databricks_multi_cloud.py
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-argument, import-outside-toplevel, protected-access
+
+from typing import Any
+from unittest.mock import MagicMock
+from urllib.parse import parse_qs, urlparse
+
+import pytest
+from pytest_mock import MockerFixture
+
+from superset.db_engine_specs.databricks import (
+ DatabricksNativeEngineSpec,
+ DatabricksPythonConnectorEngineSpec,
+)
+from superset.superset_typing import OAuth2ClientConfig
+from superset.utils.oauth2 import decode_oauth2_state
+
+# Multi-Cloud Provider Tests
+#
+# Databricks fronts the user-to-machine OAuth2 flow on every workspace at
+# `https://<workspace-host>/oidc/v1/{authorize,token}`, regardless of cloud, so
+# the authorization endpoint derives from the connection host with no per-cloud
+# account/tenant identifier.
+
+SPECS = [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec]
+
+# Representative workspace hosts for each cloud provider.
+CLOUD_HOSTS = [
+ "my-cluster.cloud.databricks.com", # AWS
+ "adb-123456789.12.azuredatabricks.net", # Azure
+ "123456789.gcp.databricks.com", # GCP
+]
+
+
+@pytest.fixture
+def oauth2_config_no_uri() -> OAuth2ClientConfig:
+ """
+ Config for Databricks OAuth2 without a pre-configured endpoint, so the
+ authorization endpoint is derived from the workspace host.
+ """
+ return {
+ "id": "databricks-client-id",
+ "secret": "databricks-client-secret",
+ "scope": "sql",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "authorization_request_uri": "",
+ "token_request_uri": "",
+ "request_content_type": "json",
+ }
+
+
+def _mock_database(mocker: MockerFixture, host: str) -> MagicMock:
+ """
+ Build a mock database whose URL resolves to the given workspace host.
+ """
+ database = mocker.MagicMock()
+ database.url_object.host = host
+ database.id = 1
+ return database
+
+
+@pytest.mark.parametrize("spec", SPECS)
+@pytest.mark.parametrize("host", CLOUD_HOSTS)
+def test_get_oauth2_authorization_uri_uses_workspace_host(
+ mocker: MockerFixture,
+ spec: Any,
+ host: str,
+ oauth2_config_no_uri: OAuth2ClientConfig,
+) -> None:
+ """
+ The authorization endpoint is the workspace host on AWS, Azure, and GCP.
+ """
+ from superset.db_engine_specs.base import OAuth2State
+
+ mocker.patch(
+ "superset.db.session.get",
+ return_value=_mock_database(mocker, host),
+ )
+
+ state: OAuth2State = {
+ "database_id": 1,
+ "user_id": 1,
+ "default_redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
+ "tab_id": "1234",
+ }
+
+ url = spec.get_oauth2_authorization_uri(oauth2_config_no_uri, state)
+ parsed = urlparse(url)
+ assert parsed.netloc == host
+ assert parsed.path == "/oidc/v1/authorize"
+
+ query = parse_qs(parsed.query)
+ assert query["scope"][0] == "sql"
+ encoded_state = query["state"][0].replace("%2E", ".")
+ assert decode_oauth2_state(encoded_state) == state
+
+
+@pytest.mark.parametrize("spec", SPECS)
+@pytest.mark.parametrize("host", CLOUD_HOSTS)
+def test_workspace_oauth2_endpoint_builds_token_uri(
+ mocker: MockerFixture,
+ spec: Any,
+ host: str,
+) -> None:
+ """
+ The helper builds the matching token endpoint from the same workspace host.
+ """
+ database = _mock_database(mocker, host)
+ assert (
+ spec._workspace_oauth2_endpoint(database, "token")
+ == f"https://{host}/oidc/v1/token"
+ )