mirror of
https://github.com/apache/superset.git
synced 2026-07-19 13:15:49 +00:00
add databricks oauth support
This commit is contained in:
fabian_zse
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user