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 4b015c55d1e..093389d106d 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 @@ -575,15 +575,31 @@ Superset automatically detects your Databricks cloud provider and uses the appro - **Azure**: Detected from hostnames containing `azure` or `azuredatabricks` - **GCP**: Detected from hostnames containing `gcp` or `googleusercontent` -You can also explicitly specify the cloud provider in your database configuration under **Advanced** → **Other** → **ENGINE PARAMETERS**: +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" + "cloud_provider": "azure", + "tenant_id": "your-azure-tenant-id" } ``` -Valid cloud provider values are: `aws`, `azure`, `gcp`. +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`. When the OAuth2 endpoints +are auto-detected, Superset substitutes this identifier into the provider's +endpoint template. If you instead supply fully-resolved `authorization_request_uri` +and `token_request_uri` values in `DATABASE_OAUTH2_CLIENTS`, those take precedence +and no `account_id`/`tenant_id` is required. ###### Usage diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py index c65dd36ab48..35b220b0690 100644 --- a/superset/db_engine_specs/databricks.py +++ b/superset/db_engine_specs/databricks.py @@ -38,7 +38,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 OAuth2RedirectError +from superset.exceptions import OAuth2Error, OAuth2RedirectError 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 @@ -325,6 +325,37 @@ class DatabricksDynamicBaseEngineSpec(BasicParametersMixin, DatabricksBaseEngine # Default to AWS for compatibility return "aws" + @classmethod + def _resolve_oauth2_endpoint( + cls, + database: Database, + provider: str, + endpoint_key: str, + ) -> str: + """ + Build a fully-resolved OAuth2 endpoint for the detected cloud provider. + + 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. + """ + template = cls._oauth2_endpoints[provider][endpoint_key] + if "{}" not in template: + return template + + extra = cls.get_extra_params(database) + account_id = extra.get("account_id") or extra.get("tenant_id") + if not account_id: + 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." + ) + return template.format(account_id) + @classmethod def impersonate_user( cls, @@ -563,26 +594,30 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec): ) -> 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. """ - from superset import db - from superset.models.core import Database + 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) - # Update config with the correct authorization URI for the cloud provider - from typing import cast + # 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._oauth2_endpoints[provider][ - "authorization_request_uri" - ] - }, - ) + 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) @@ -594,23 +629,17 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec): code_verifier: str | None = None, ) -> "OAuth2TokenResponse": """ - Exchange authorization code for refresh/access tokens with dynamic endpoint. + Exchange authorization code for refresh/access tokens. The token request URI is resolved when the OAuth2 config is built (see - `get_oauth2_config`), so it already targets the correct cloud provider and - account. Only fall back to the AWS endpoint when no URI is configured. + `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. """ - from typing import cast - if not config.get("token_request_uri"): - config = cast( - "OAuth2ClientConfig", - dict(config) - | { - "token_request_uri": cls._oauth2_endpoints["aws"][ - "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) @@ -844,26 +873,30 @@ class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec): ) -> 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. """ - from superset import db - from superset.models.core import Database + 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) - # Update config with the correct authorization URI for the cloud provider - from typing import cast + # 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._oauth2_endpoints[provider][ - "authorization_request_uri" - ] - }, - ) + 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) @@ -875,23 +908,17 @@ class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec): code_verifier: str | None = None, ) -> "OAuth2TokenResponse": """ - Exchange authorization code for refresh/access tokens with dynamic endpoint. + Exchange authorization code for refresh/access tokens. The token request URI is resolved when the OAuth2 config is built (see - `get_oauth2_config`), so it already targets the correct cloud provider and - account. Only fall back to the AWS endpoint when no URI is configured. + `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. """ - from typing import cast - if not config.get("token_request_uri"): - config = cast( - "OAuth2ClientConfig", - dict(config) - | { - "token_request_uri": cls._oauth2_endpoints["aws"][ - "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) diff --git a/tests/unit_tests/db_engine_specs/test_databricks.py b/tests/unit_tests/db_engine_specs/test_databricks.py index b6279c021e9..4b47d051850 100644 --- a/tests/unit_tests/db_engine_specs/test_databricks.py +++ b/tests/unit_tests/db_engine_specs/test_databricks.py @@ -17,19 +17,20 @@ # 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.base import OAuth2State from superset.db_engine_specs.databricks import ( DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec, ) from superset.errors import ErrorLevel, SupersetError, SupersetErrorType -from superset.exceptions import OAuth2RedirectError +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 @@ -664,6 +665,141 @@ def test_get_oauth2_fresh_token_native( ) +def _oauth2_state() -> OAuth2State: + 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( + "extra, netloc, path", + [ + ( + {"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", + ), + ], +) +def test_get_oauth2_authorization_uri_autodetects_and_resolves( + mocker: MockerFixture, + spec: Any, + extra: dict[str, Any], + netloc: str, + path: 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`. + """ + database = mocker.MagicMock() + database.extra = json.dumps(extra) + database.url_object.host = "dbc-abc.cloud.databricks.com" + 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 + + +@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 + auto-detected template, 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_account_id( + 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. + """ + database = mocker.MagicMock() + database.extra = json.dumps({"cloud_provider": "aws"}) + database.url_object.host = "dbc-abc.cloud.databricks.com" + 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,