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 9cc3505c1de..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 @@ -536,25 +536,20 @@ Superset supports OAuth2 authentication for Databricks, allowing users to authen from datetime import timedelta # OAuth2 configuration for Databricks -# OAuth2 endpoints are automatically detected based on your Databricks cloud provider +# 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", - # The authorization endpoint is auto-detected from the hostname; the - # token endpoint must be set explicitly (no DB context at exchange): - # AWS: "authorization_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/{account_id}/v1/authorize", - # Azure: "authorization_request_uri": "https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/authorize", - # GCP: "authorization_request_uri": "https://accounts.gcp.databricks.com/oidc/accounts/{account_id}/v1/authorize", - # "token_request_uri": "https://", + "token_request_uri": "https://your-workspace-host/oidc/v1/token", }, "Databricks": { "id": "your-databricks-client-id", "secret": "your-databricks-client-secret", "scope": "sql", - # Authorization endpoint auto-detected from hostname; set - # "token_request_uri" explicitly for the token exchange. + "token_request_uri": "https://your-workspace-host/oidc/v1/token", }, } @@ -572,40 +567,21 @@ Replace the following placeholders: **Multi-Cloud Provider Support** -Superset automatically detects your Databricks cloud provider and uses the appropriate OAuth2 endpoints: +Databricks fronts the user-to-machine (U2M) OAuth2 flow on every workspace at +`https:///oidc/v1/authorize` and +`https:///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. -- **AWS**: Detected from hostnames containing `cloud.databricks.com` -- **Azure**: Detected from hostnames containing `azure` or `azuredatabricks` -- **GCP**: Detected from hostnames containing `gcp` or `googleusercontent` +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:///oidc/v1/token` for +your workspace. -You can also explicitly specify the cloud provider, along with the account -identifier used to build the OAuth2 endpoints, in your database configuration -under **Advanced** → **Other** → **ENGINE PARAMETERS**: - -```json -{ - "cloud_provider": "azure", - "tenant_id": "your-azure-tenant-id" -} -``` - -For AWS and GCP, supply `account_id` instead: - -```json -{ - "cloud_provider": "aws", - "account_id": "your-databricks-account-id" -} -``` - -Valid cloud provider values are: `aws`, `azure`, `gcp`. The **authorization** -endpoint is auto-detected: Superset substitutes this identifier into the -provider's authorization template. The **token** endpoint is not auto-resolved -(token exchange has no database context to detect the provider), so for the -auto-detected flow you must still supply a fully-resolved `token_request_uri` -in `DATABASE_OAUTH2_CLIENTS`. If you supply fully-resolved -`authorization_request_uri` and `token_request_uri` values, those take -precedence and no `account_id`/`tenant_id` is required. +If you supply a fully-resolved `authorization_request_uri` (and/or +`token_request_uri`), those values take precedence over the host-derived +defaults. ###### Usage diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py index d45a93d89b6..c37db319441 100644 --- a/superset/db_engine_specs/databricks.py +++ b/superset/db_engine_specs/databricks.py @@ -17,10 +17,11 @@ from __future__ import annotations 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,7 +39,7 @@ 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, OAuth2RedirectError +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 @@ -283,81 +284,111 @@ class DatabricksDynamicBaseEngineSpec(BasicParametersMixin, DatabricksBaseEngine "port": "port", } - # OAuth2 endpoints for different cloud providers - _oauth2_endpoints = { - "aws": { - "authorization_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/authorize", - "token_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/token", - }, - "azure": { - "authorization_request_uri": "https://login.microsoftonline.com/{}/oauth2/v2.0/authorize", - "token_request_uri": "https://login.microsoftonline.com/{}/oauth2/v2.0/token", - }, - "gcp": { - "authorization_request_uri": "https://accounts.gcp.databricks.com/oidc/accounts/{}/v1/authorize", - "token_request_uri": "https://accounts.gcp.databricks.com/oidc/accounts/{}/v1/token", - }, - } + # 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 = ( + "401", + "unauthorized", + "unauthenticated", + "invalid access token", + "invalid token", + "expired token", + "token expired", + ) @classmethod - def _detect_cloud_provider(cls, database: Database) -> str: + def _workspace_oauth2_endpoint(cls, database: Database, path: str) -> str: """ - Detect the cloud provider based on the database configuration. + Build a Databricks OAuth2 (U2M) endpoint from the workspace host. - Returns: - str: The cloud provider ('aws', 'azure', or 'gcp') + Databricks fronts the user-to-machine OAuth2 flow on every workspace at + ``https:///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. """ - # Check if cloud provider is explicitly configured in extra - if isinstance( - (cloud_provider := cls.get_extra_params(database).get("cloud_provider")), - str, - ): - provider = cloud_provider.lower() - if provider in cls._oauth2_endpoints: - return provider - - # Try to detect from hostname - hostname = database.url_object.host or "" - hostname = hostname.lower() - - if "azure" in hostname or "azuredatabricks" in hostname: - return "azure" - elif "gcp" in hostname or "googleusercontent" in hostname: - return "gcp" - else: - # Default to AWS for compatibility - return "aws" + 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 _resolve_oauth2_endpoint( + 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, - database: Database, - provider: str, - endpoint_key: str, + config: "OAuth2ClientConfig", + state: "OAuth2State", + code_verifier: str | None = None, ) -> str: """ - Build a fully-resolved OAuth2 endpoint for the detected cloud provider. + Return the URI for the initial OAuth2 request. - The per-provider templates carry a single ``{}`` placeholder for the - Databricks account id (or Azure tenant id), read from the database's - ``extra`` (``account_id``, or ``tenant_id`` for Azure). Raising when it - is absent keeps the flow from issuing a request to an unresolved - ``.../{}/...`` endpoint. + A fully-resolved ``authorization_request_uri`` from + ``DATABASE_OAUTH2_CLIENTS`` is preserved; otherwise the endpoint is + derived from the workspace host (``https:///oidc/v1/authorize``), + which is valid on AWS, Azure and GCP. """ - template = cls._oauth2_endpoints[provider][endpoint_key] - if "{}" not in template: - return template + if not config.get("authorization_request_uri"): + from superset import db + from superset.models.core import Database - extra = cls.get_extra_params(database) - account_id = extra.get("account_id") or extra.get("tenant_id") - if not account_id: + 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:///oidc/v1/token``) and fail fast rather than + POST to an unresolved endpoint. + """ + if not config.get("token_request_uri"): raise OAuth2Error( - "Databricks OAuth2 endpoints could not be resolved: set " - "`account_id` (or `tenant_id` for Azure) in the database's " - "engine parameters, or provide a fully-resolved " - f"`{endpoint_key}` in DATABASE_OAUTH2_CLIENTS." + "Databricks OAuth2 token endpoint is not configured: set " + "`token_request_uri` to https:///oidc/v1/token " + "in DATABASE_OAUTH2_CLIENTS." ) - return template.format(account_id) + + return super().get_oauth2_token(config, code, code_verifier) @classmethod def impersonate_user( @@ -579,73 +610,15 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec): supports_dynamic_catalog = True supports_cross_catalog_queries = True - # OAuth 2.0 support + # OAuth 2.0 support. The flow (endpoint resolution from the workspace host, + # `needs_oauth2` detection) is shared via `DatabricksDynamicBaseEngineSpec`. supports_oauth2 = True - oauth2_exception = OAuth2RedirectError oauth2_scope = "sql" - # OAuth2 endpoints are determined dynamically based on cloud provider - oauth2_authorization_request_uri = "" # Set dynamically - oauth2_token_request_uri = "" # Set dynamically - - @classmethod - def get_oauth2_authorization_uri( - cls, - config: "OAuth2ClientConfig", - state: "OAuth2State", - code_verifier: str | None = None, - ) -> str: - """ - Return URI for initial OAuth2 request with dynamic endpoint detection. - - A fully-resolved `authorization_request_uri` from `DATABASE_OAUTH2_CLIENTS` - is preserved; only fall back to the auto-detected, account-resolved - endpoint when none is configured. - """ - if not config.get("authorization_request_uri"): - from superset import db - from superset.models.core import Database - - # Get the database to detect cloud provider - database_id = state["database_id"] - if database := db.session.get(Database, database_id): - provider = cls._detect_cloud_provider(database) - from typing import cast - - config = cast( - "OAuth2ClientConfig", - dict(config) - | { - "authorization_request_uri": cls._resolve_oauth2_endpoint( - database, provider, "authorization_request_uri" - ) - }, - ) - - 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 authorization code for refresh/access tokens. - - The token request URI is resolved when the OAuth2 config is built (see - `get_oauth2_config`) and already targets the correct cloud provider and - account. There is no database context here to auto-detect it, so fail - fast rather than POST to an unresolved endpoint when it is missing. - """ - if not config.get("token_request_uri"): - raise OAuth2Error( - "Databricks OAuth2 token endpoint is not configured: provide a " - "fully-resolved `token_request_uri` in DATABASE_OAUTH2_CLIENTS." - ) - - return super().get_oauth2_token(config, code, code_verifier) + # 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 @@ -858,73 +831,15 @@ class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec): supports_dynamic_schema = supports_catalog = supports_dynamic_catalog = True - # OAuth 2.0 support + # OAuth 2.0 support. The flow (endpoint resolution from the workspace host, + # `needs_oauth2` detection) is shared via `DatabricksDynamicBaseEngineSpec`. supports_oauth2 = True - oauth2_exception = OAuth2RedirectError oauth2_scope = "sql" - # OAuth2 endpoints are determined dynamically based on cloud provider - oauth2_authorization_request_uri = "" # Set dynamically - oauth2_token_request_uri = "" # Set dynamically - - @classmethod - def get_oauth2_authorization_uri( - cls, - config: "OAuth2ClientConfig", - state: "OAuth2State", - code_verifier: str | None = None, - ) -> str: - """ - Return URI for initial OAuth2 request with dynamic endpoint detection. - - A fully-resolved `authorization_request_uri` from `DATABASE_OAUTH2_CLIENTS` - is preserved; only fall back to the auto-detected, account-resolved - endpoint when none is configured. - """ - if not config.get("authorization_request_uri"): - from superset import db - from superset.models.core import Database - - # Get the database to detect cloud provider - database_id = state["database_id"] - if database := db.session.get(Database, database_id): - provider = cls._detect_cloud_provider(database) - from typing import cast - - config = cast( - "OAuth2ClientConfig", - dict(config) - | { - "authorization_request_uri": cls._resolve_oauth2_endpoint( - database, provider, "authorization_request_uri" - ) - }, - ) - - 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 authorization code for refresh/access tokens. - - The token request URI is resolved when the OAuth2 config is built (see - `get_oauth2_config`) and already targets the correct cloud provider and - account. There is no database context here to auto-detect it, so fail - fast rather than POST to an unresolved endpoint when it is missing. - """ - if not config.get("token_request_uri"): - raise OAuth2Error( - "Databricks OAuth2 token endpoint is not configured: provide a " - "fully-resolved `token_request_uri` in DATABASE_OAUTH2_CLIENTS." - ) - - return super().get_oauth2_token(config, code, code_verifier) + # 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 diff --git a/tests/unit_tests/db_engine_specs/test_databricks.py b/tests/unit_tests/db_engine_specs/test_databricks.py index 7f2232daa69..a842e2ee85f 100644 --- a/tests/unit_tests/db_engine_specs/test_databricks.py +++ b/tests/unit_tests/db_engine_specs/test_databricks.py @@ -311,21 +311,82 @@ def test_oauth2_attributes() -> None: """ # Test DatabricksNativeEngineSpec assert DatabricksNativeEngineSpec.supports_oauth2 is True - assert DatabricksNativeEngineSpec.oauth2_exception is OAuth2RedirectError assert DatabricksNativeEngineSpec.oauth2_scope == "sql" - # OAuth2 endpoints are now dynamic and set at runtime + # 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_exception is OAuth2RedirectError assert DatabricksPythonConnectorEngineSpec.oauth2_scope == "sql" - # OAuth2 endpoints are now dynamic and set at runtime 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", + ], +) +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], +) +def test_needs_oauth2_ignores_unrelated_errors( + mocker: MockerFixture, + spec: Any, +) -> 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("Table not found")) 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. @@ -699,48 +760,33 @@ def _unresolved_oauth2_config() -> OAuth2ClientConfig: [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec], ) @pytest.mark.parametrize( - "extra, netloc, path", + "host", [ - ( - {"cloud_provider": "aws", "account_id": "acct-999"}, - "accounts.cloud.databricks.com", - "/oidc/accounts/acct-999/v1/authorize", - ), - ( - {"cloud_provider": "azure", "tenant_id": "tenant-abc"}, - "login.microsoftonline.com", - "/tenant-abc/oauth2/v2.0/authorize", - ), - ( - {"cloud_provider": "gcp", "account_id": "acct-gcp"}, - "accounts.gcp.databricks.com", - "/oidc/accounts/acct-gcp/v1/authorize", - ), + "dbc-abc.cloud.databricks.com", + "adb-123456789.12.azuredatabricks.net", + "123456789.gcp.databricks.com", ], ) -def test_get_oauth2_authorization_uri_autodetects_and_resolves( +def test_get_oauth2_authorization_uri_derives_from_workspace_host( mocker: MockerFixture, spec: Any, - extra: dict[str, Any], - netloc: str, - path: str, + host: str, ) -> None: """ - With no configured `authorization_request_uri`, the endpoint is auto-detected - per cloud provider and the `account_id`/`tenant_id` placeholder is resolved - from the database `extra`. + With no configured `authorization_request_uri`, the endpoint is derived from + the workspace host (`https:///oidc/v1/authorize`) on every cloud, with + no account/tenant identifier required. """ database = mocker.MagicMock() - database.extra = json.dumps(extra) - database.url_object.host = "dbc-abc.cloud.databricks.com" + 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 == netloc - assert parsed.path == path + assert parsed.netloc == host + assert parsed.path == "/oidc/v1/authorize" @pytest.mark.parametrize( @@ -753,7 +799,7 @@ def test_get_oauth2_authorization_uri_preserves_configured( ) -> None: """ A fully-resolved `authorization_request_uri` is never overwritten by the - auto-detected template, and no database lookup is needed. + host-derived endpoint, and no database lookup is needed. """ session_get = mocker.patch("superset.db.session.get") config = _unresolved_oauth2_config() @@ -770,17 +816,16 @@ def test_get_oauth2_authorization_uri_preserves_configured( "spec", [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec], ) -def test_get_oauth2_authorization_uri_fails_without_account_id( +def test_get_oauth2_authorization_uri_fails_without_host( mocker: MockerFixture, spec: Any, ) -> None: """ - When the endpoint must be auto-detected but no `account_id`/`tenant_id` is - available, fail fast instead of emitting an unresolved `.../{}/...` URL. + 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.extra = json.dumps({"cloud_provider": "aws"}) - database.url_object.host = "dbc-abc.cloud.databricks.com" + database.url_object.host = None mocker.patch("superset.db.session.get", return_value=database) with pytest.raises(OAuth2Error): 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 index 47e1636163b..1df0b0d127e 100644 --- a/tests/unit_tests/db_engine_specs/test_databricks_multi_cloud.py +++ b/tests/unit_tests/db_engine_specs/test_databricks_multi_cloud.py @@ -16,6 +16,7 @@ # 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 @@ -30,65 +31,27 @@ 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:///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] -@pytest.fixture -def mock_database_aws(mocker: MockerFixture) -> MagicMock: - """ - Mock database with AWS hostname. - """ - database = mocker.MagicMock() - database.url_object.host = "my-cluster.cloud.databricks.com" - database.extra = "{}" - database.id = 1 - return database - - -@pytest.fixture -def mock_database_azure(mocker: MockerFixture) -> MagicMock: - """ - Mock database with Azure hostname. - """ - database = mocker.MagicMock() - database.url_object.host = "adb-123456789.12.azuredatabricks.net" - database.extra = '{"tenant_id": "azure-tenant-id"}' - database.id = 2 - return database - - -@pytest.fixture -def mock_database_gcp(mocker: MockerFixture) -> MagicMock: - """ - Mock database with GCP hostname. - """ - database = mocker.MagicMock() - database.url_object.host = "123456789.gcp.databricks.com" - database.extra = '{"account_id": "12345"}' - database.id = 3 - return database - - -@pytest.fixture -def oauth2_config() -> OAuth2ClientConfig: - """ - Config for Databricks OAuth2 with a fully-resolved endpoint configured. - """ - 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", - } +# 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 - per-provider endpoint is auto-detected and account-resolved. + authorization endpoint is derived from the workspace host. """ return { "id": "databricks-client-id", @@ -101,100 +64,30 @@ def oauth2_config_no_uri() -> OAuth2ClientConfig: } -def test_cloud_provider_detection_aws(mock_database_aws: MagicMock) -> None: - """ - Test cloud provider detection for AWS. - """ - provider = DatabricksNativeEngineSpec._detect_cloud_provider(mock_database_aws) - assert provider == "aws" - - -def test_cloud_provider_detection_azure(mock_database_azure: MagicMock) -> None: - """ - Test cloud provider detection for Azure. - """ - provider = DatabricksNativeEngineSpec._detect_cloud_provider(mock_database_azure) - assert provider == "azure" - - -def test_cloud_provider_detection_gcp(mock_database_gcp: MagicMock) -> None: - """ - Test cloud provider detection for GCP. - """ - provider = DatabricksNativeEngineSpec._detect_cloud_provider(mock_database_gcp) - assert provider == "gcp" - - -def test_cloud_provider_detection_explicit_config(mocker: MockerFixture) -> None: - """ - Test cloud provider detection with explicit configuration. - """ +def _mock_database(mocker: MockerFixture, host: str) -> MagicMock: database = mocker.MagicMock() - database.url_object.host = "generic-host.com" - - # Mock get_extra_params to return explicit cloud provider - mocker.patch.object( - DatabricksNativeEngineSpec, - "get_extra_params", - return_value={"cloud_provider": "azure"}, - ) - - provider = DatabricksNativeEngineSpec._detect_cloud_provider(database) - assert provider == "azure" + database.url_object.host = host + database.id = 1 + return database -def test_cloud_provider_detection_invalid_config_falls_back_to_hostname( +@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: """ - An unrecognized explicit `cloud_provider` is ignored and detection falls - back to the hostname rather than raising or returning the bad value. - """ - database = mocker.MagicMock() - database.url_object.host = "adb-123456789.12.azuredatabricks.net" - - mocker.patch.object( - DatabricksNativeEngineSpec, - "get_extra_params", - return_value={"cloud_provider": "oracle"}, - ) - - provider = DatabricksNativeEngineSpec._detect_cloud_provider(database) - assert provider == "azure" - - -def test_cloud_provider_detection_non_string_falls_back_to_hostname( - mocker: MockerFixture, -) -> None: - """ - A non-string `cloud_provider` (e.g. a boolean from malformed JSON) is - ignored without raising and detection falls back to the hostname. - """ - database = mocker.MagicMock() - database.url_object.host = "adb-123456789.12.azuredatabricks.net" - - mocker.patch.object( - DatabricksNativeEngineSpec, - "get_extra_params", - return_value={"cloud_provider": True}, - ) - - provider = DatabricksNativeEngineSpec._detect_cloud_provider(database) - assert provider == "azure" - - -def test_get_oauth2_authorization_uri_aws( - mocker: MockerFixture, - oauth2_config: OAuth2ClientConfig, - mock_database_aws: MagicMock, -) -> None: - """ - Test OAuth2 authorization URI generation for AWS provider. + The authorization endpoint is the workspace host on AWS, Azure, and GCP. """ from superset.db_engine_specs.base import OAuth2State - # Mock the database query - mocker.patch("superset.extensions.db.session.get", return_value=mock_database_aws) + mocker.patch( + "superset.extensions.db.session.get", + return_value=_mock_database(mocker, host), + ) state: OAuth2State = { "database_id": 1, @@ -203,11 +96,10 @@ def test_get_oauth2_authorization_uri_aws( "tab_id": "1234", } - url = DatabricksNativeEngineSpec.get_oauth2_authorization_uri(oauth2_config, state) + url = spec.get_oauth2_authorization_uri(oauth2_config_no_uri, state) parsed = urlparse(url) - assert parsed.netloc == "accounts.cloud.databricks.com" - assert "/oidc/accounts/" in parsed.path - assert "/v1/authorize" in parsed.path + assert parsed.netloc == host + assert parsed.path == "/oidc/v1/authorize" query = parse_qs(parsed.query) assert query["scope"][0] == "sql" @@ -215,113 +107,18 @@ def test_get_oauth2_authorization_uri_aws( assert decode_oauth2_state(encoded_state) == state -def test_get_oauth2_authorization_uri_azure( +@pytest.mark.parametrize("spec", SPECS) +@pytest.mark.parametrize("host", CLOUD_HOSTS) +def test_workspace_oauth2_endpoint_builds_token_uri( mocker: MockerFixture, - oauth2_config_no_uri: OAuth2ClientConfig, - mock_database_azure: MagicMock, + spec: Any, + host: str, ) -> None: """ - Test OAuth2 authorization URI generation for Azure provider. + The helper builds the matching token endpoint from the same workspace host. """ - from superset.db_engine_specs.base import OAuth2State - - # Mock the database query - mocker.patch("superset.extensions.db.session.get", return_value=mock_database_azure) - - state: OAuth2State = { - "database_id": 2, - "user_id": 1, - "default_redirect_uri": "http://localhost:8088/api/v1/database/oauth2/", - "tab_id": "1234", - } - - url = DatabricksNativeEngineSpec.get_oauth2_authorization_uri( - oauth2_config_no_uri, state + database = _mock_database(mocker, host) + assert ( + spec._workspace_oauth2_endpoint(database, "token") + == f"https://{host}/oidc/v1/token" ) - parsed = urlparse(url) - assert parsed.netloc == "login.microsoftonline.com" - assert "/oauth2/v2.0/authorize" in parsed.path - - 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_gcp( - mocker: MockerFixture, - oauth2_config_no_uri: OAuth2ClientConfig, - mock_database_gcp: MagicMock, -) -> None: - """ - Test OAuth2 authorization URI generation for GCP provider. - """ - from superset.db_engine_specs.base import OAuth2State - - # Mock the database query - mocker.patch("superset.extensions.db.session.get", return_value=mock_database_gcp) - - state: OAuth2State = { - "database_id": 3, - "user_id": 1, - "default_redirect_uri": "http://localhost:8088/api/v1/database/oauth2/", - "tab_id": "1234", - } - - url = DatabricksNativeEngineSpec.get_oauth2_authorization_uri( - oauth2_config_no_uri, state - ) - parsed = urlparse(url) - assert parsed.netloc == "accounts.gcp.databricks.com" - assert "/oidc/accounts/" in parsed.path - assert "/v1/authorize" in parsed.path - - 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_python_connector_cloud_provider_detection_azure( - mock_database_azure: MagicMock, -) -> None: - """ - Test cloud provider detection for Python Connector with Azure. - """ - provider = DatabricksPythonConnectorEngineSpec._detect_cloud_provider( - mock_database_azure - ) - assert provider == "azure" - - -def test_python_connector_oauth2_authorization_uri_azure( - mocker: MockerFixture, - oauth2_config_no_uri: OAuth2ClientConfig, - mock_database_azure: MagicMock, -) -> None: - """ - Test OAuth2 authorization URI generation for Python Connector with Azure provider. - """ - from superset.db_engine_specs.base import OAuth2State - - # Mock the database query - mocker.patch("superset.extensions.db.session.get", return_value=mock_database_azure) - - state: OAuth2State = { - "database_id": 2, - "user_id": 1, - "default_redirect_uri": "http://localhost:8088/api/v1/database/oauth2/", - "tab_id": "1234", - } - - url = DatabricksPythonConnectorEngineSpec.get_oauth2_authorization_uri( - oauth2_config_no_uri, state - ) - parsed = urlparse(url) - assert parsed.netloc == "login.microsoftonline.com" - assert "/oauth2/v2.0/authorize" in parsed.path - - query = parse_qs(parsed.query) - assert query["scope"][0] == "sql" - encoded_state = query["state"][0].replace("%2E", ".") - assert decode_oauth2_state(encoded_state) == state