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..50f65090675 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 @@ -32,7 +32,7 @@ install new database drivers into your Superset configuration. ### Supported Databases and Dependencies Some of the recommended packages are shown below. Please refer to -[pyproject.toml](https://github.com/apache/superset/blob/master/pyproject.toml) for the versions that +[pyproject.toml](https://github.com/apache/superset/blob/master/pyproject.toml) for the versions thatdocs/docs/configuration/databases.mdx are compatible with Superset. |
Database
| PyPI package | Connection String | @@ -519,6 +519,63 @@ 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 +DATABASE_OAUTH2_CLIENTS = { + "Databricks (legacy)": { + "id": "your-databricks-client-id", + "secret": "your-databricks-client-secret", + "scope": "sql", + "authorization_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/{account_id}/v1/authorize", + "token_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/{account_id}/v1/token", + }, + "Databricks": { + "id": "your-databricks-client-id", + "secret": "your-databricks-client-secret", + "scope": "sql", + "authorization_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/{account_id}/v1/authorize", + "token_request_uri": "https://accounts.cloud.databricks.com/oidc/accounts/{account_id}/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 +- `{account_id}`: Your Databricks account ID (found in your workspace URL) +- `your-superset-host:port`: Your Superset instance hostname and port + +###### 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. + #### 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..24ab8018b0f 100644 --- a/superset/db_engine_specs/databricks.py +++ b/superset/db_engine_specs/databricks.py @@ -38,6 +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.utils import json from superset.utils.core import get_user_agent, QuerySource from superset.utils.network import is_hostname_valid, is_port_open @@ -277,6 +278,29 @@ class DatabricksDynamicBaseEngineSpec(BasicParametersMixin, DatabricksBaseEngine "port": "port", } + @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 OAuth2 access token for user impersonation. + """ + if user_token: + # Replace the access token in the URL with the user's OAuth2 token + url = url.set(password=user_token) + + # Also update connect_args if they contain access token + connect_args = engine_kwargs.setdefault("connect_args", {}) + if "access_token" in connect_args: + connect_args["access_token"] = user_token + + return url, engine_kwargs + @staticmethod def get_extra_params( database: Database, source: QuerySource | None = None @@ -474,6 +498,17 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec): supports_dynamic_catalog = True supports_cross_catalog_queries = True + # OAuth 2.0 support + supports_oauth2 = True + oauth2_exception = OAuth2RedirectError + oauth2_scope = "sql" + oauth2_authorization_request_uri = ( + "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/authorize" + ) + oauth2_token_request_uri = ( + "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/token" # noqa: S105 + ) + @classmethod def build_sqlalchemy_uri( # type: ignore cls, parameters: DatabricksNativeParametersType, *_ @@ -685,6 +720,17 @@ class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec): supports_dynamic_schema = supports_catalog = supports_dynamic_catalog = True + # OAuth 2.0 support + supports_oauth2 = True + oauth2_exception = OAuth2RedirectError + oauth2_scope = "sql" + oauth2_authorization_request_uri = ( + "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/authorize" + ) + oauth2_token_request_uri = ( + "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/token" # noqa: S105 + ) + @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..aa900674da8 100644 --- a/tests/unit_tests/db_engine_specs/test_databricks.py +++ b/tests/unit_tests/db_engine_specs/test_databricks.py @@ -18,13 +18,21 @@ from datetime import datetime from typing import 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.databricks import ( + DatabricksNativeEngineSpec, + DatabricksPythonConnectorEngineSpec, +) from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import 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 +299,413 @@ 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_exception is OAuth2RedirectError + assert DatabricksNativeEngineSpec.oauth2_scope == "sql" + assert ( + DatabricksNativeEngineSpec.oauth2_authorization_request_uri + == "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/authorize" + ) + assert ( + DatabricksNativeEngineSpec.oauth2_token_request_uri + == "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/token" # noqa: S105 + ) + + # Test DatabricksPythonConnectorEngineSpec + assert DatabricksPythonConnectorEngineSpec.supports_oauth2 is True + assert DatabricksPythonConnectorEngineSpec.oauth2_exception is OAuth2RedirectError + assert DatabricksPythonConnectorEngineSpec.oauth2_scope == "sql" + assert ( + DatabricksPythonConnectorEngineSpec.oauth2_authorization_request_uri + == "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/authorize" + ) + assert ( + DatabricksPythonConnectorEngineSpec.oauth2_token_request_uri + == "https://accounts.cloud.databricks.com/oidc/accounts/{}/v1/token" # noqa: S105 + ) + + +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(mocker: MockerFixture) -> None: + """ + Test impersonate_user method without OAuth2 token. + """ + 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, + ) + + # Check that nothing was changed + assert url.password == "original-token" # noqa: S105 + assert kwargs["connect_args"]["access_token"] == "original-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 + + +@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 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, + )