aminghadersohi commented on code in PR #44203:
URL: https://github.com/apache/superset/pull/44203#discussion_r4009567772
##########
superset/db_engine_specs/databricks.py:
##########
@@ -62,6 +62,179 @@
r"\[INSUFFICIENT_PERMISSIONS\]|\bSQLSTATE:\s*42501\b"
)
+# Per-database service-principal auth (discussion #39405). Secrets live in
+# Encrypted Extra; the engine spec injects driver args at connect time.
+AUTH_METHOD_OAUTH_M2M = "oauth-m2m"
+AUTH_METHOD_AZURE_SP_M2M = "azure-sp-m2m"
+M2M_AUTH_METHODS = frozenset({AUTH_METHOD_OAUTH_M2M, AUTH_METHOD_AZURE_SP_M2M})
+# The Databricks dialect always reads ``url.password``; a dummy token is
+# enough when M2M credentials_provider / Azure SP auth wins.
+M2M_PLACEHOLDER_TOKEN = "m2m" # noqa: S105
+
+
+def _parse_json_object(raw: Any) -> dict[str, Any]:
+ """Parse a JSON object from a string or return a dict unchanged."""
+ if isinstance(raw, dict):
+ return raw
+ if not raw or not isinstance(raw, str):
+ return {}
+ try:
+ parsed = json.loads(raw)
+ except json.JSONDecodeError:
+ return {}
+ return parsed if isinstance(parsed, dict) else {}
+
+
+def _parse_extra_for_validation(
+ extra_raw: Any,
+) -> tuple[dict[str, Any], SupersetError | None]:
+ """Parse Extra JSON for ``validate_parameters`` without raising."""
+ if isinstance(extra_raw, dict):
+ return extra_raw, None
+ if not extra_raw:
+ return {}, None
+ try:
+ parsed = json.loads(extra_raw)
+ except json.JSONDecodeError:
+ return {}, SupersetError(
+ message="The extra connection parameters are not valid JSON.",
+ error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
+ level=ErrorLevel.ERROR,
+ extra={"invalid": ["extra"]},
+ )
+ if isinstance(parsed, dict):
+ return parsed, None
+ return {}, None
+
+
+def _m2m_client_id(encrypted_extra: dict[str, Any]) -> Any:
+ return encrypted_extra.get("client_id") or
encrypted_extra.get("azure_client_id")
+
+
+def _m2m_client_secret(encrypted_extra: dict[str, Any]) -> Any:
+ return encrypted_extra.get("client_secret") or encrypted_extra.get(
+ "azure_client_secret"
+ )
+
+
+def _has_m2m_credentials(encrypted_extra: dict[str, Any]) -> bool:
+ """Return True when Encrypted Extra has a complete M2M client ID/secret.
+
+ ``mask_encrypted_extra`` redacts every top-level value (inherited ``$.*``),
+ so a masked payload still counts when ``auth_method`` is the password mask
+ and the client-id / client-secret keys (or Azure aliases) are present.
+ """
+ auth_method = encrypted_extra.get("auth_method")
+ if auth_method not in M2M_AUTH_METHODS and auth_method != PASSWORD_MASK:
+ return False
+ if not _m2m_client_id(encrypted_extra) or not
_m2m_client_secret(encrypted_extra):
+ return False
+ if auth_method == AUTH_METHOD_AZURE_SP_M2M and not encrypted_extra.get(
+ "azure_tenant_id"
+ ):
+ return False
+ return True
+
+
+def _workspace_hostname(host: str) -> str:
+ """Strip a scheme/trailing slash so Config.host is
``https://{hostname}``.""";
+ hostname = host.strip()
+ for prefix in ("https://";, "http://";):
+ if hostname.startswith(prefix):
+ hostname = hostname[len(prefix) :]
+ return hostname.rstrip("/")
+
+
+def _access_token_from_uri_password(password: str | None) -> str:
+ """Drop the dummy ``m2m`` password so the form does not treat it as a
PAT."""
+ if not password or password == M2M_PLACEHOLDER_TOKEN:
+ return ""
+ return password
+
+
+def _load_databricks_m2m_sdk() -> tuple[Any, Any, Any]:
+ """
+ Lazy-import ``databricks-sdk`` so PAT-only installs keep working.
+ """
+ try:
+ from databricks.sdk.core import Config, oauth_service_principal
+ except ImportError as ex:
+ raise ValueError(
+ "Databricks OAuth M2M requires the databricks-sdk package. "
+ "Install it with: pip install 'apache-superset[databricks]' "
+ "or pip install databricks-sdk."
+ ) from ex
+ try:
+ from databricks.sdk.core import azure_service_principal
+ except ImportError:
+ azure_service_principal = None
+ return Config, oauth_service_principal, azure_service_principal
+
+
+class _M2MCredentialsProvider:
+ """
+ Stable ``credentials_provider`` for the Databricks SQL connector.
+
+ ``Database._get_sqla_engine`` includes ``repr(engine_kwargs)`` in the
+ process-wide engine cache key. A new nested function on every build
+ would miss that cache. ``__repr__`` is identity-stable for the same
+ host / client / tenant (secret rotation evicts the cache on save).
+ """
+
+ def __init__(
+ self,
+ host: str,
+ client_id: str,
+ client_secret: str,
+ *,
+ azure: bool = False,
+ azure_tenant_id: str | None = None,
+ ) -> None:
+ hostname = _workspace_hostname(host)
+ if not hostname:
+ raise ValueError(
+ "Databricks OAuth M2M requires a workspace host on the
connection."
+ )
+ self._hostname = hostname
+ self._client_id = client_id
+ self._client_secret = client_secret
+ self._azure = azure
+ self._azure_tenant_id = azure_tenant_id
+ # Fail at engine creation, not on the first query.
+ self._config_cls, self._oauth_sp, self._azure_sp =
_load_databricks_m2m_sdk()
+ if azure and self._azure_sp is None:
+ raise ValueError(
+ "Databricks Azure service-principal M2M requires "
+ "azure_service_principal from databricks-sdk. "
+ "Install it with: pip install 'apache-superset[databricks]' "
+ "or pip install databricks-sdk."
+ )
+
+ def __call__(self) -> Any:
+ if self._azure:
+ return self._azure_sp(
+ self._config_cls(
+ host=f"https://{self._hostname}";,
+ azure_tenant_id=self._azure_tenant_id,
+ azure_client_id=self._client_id,
+ azure_client_secret=self._client_secret,
+ )
+ )
+ return self._oauth_sp(
+ self._config_cls(
+ host=f"https://{self._hostname}";,
+ client_id=self._client_id,
+ client_secret=self._client_secret,
+ )
+ )
+
+ def __repr__(self) -> str:
+ return (
+ f"_M2MCredentialsProvider(host={self._hostname!r}, "
+ f"client_id={self._client_id!r}, azure={self._azure!r}, "
+ f"azure_tenant_id={self._azure_tenant_id!r})"
+ )
Review Comment:
`repr` is the engine cache key (`models/core.py:726`), so rotating only
`client_secret` yields an identical key, and `_evict_engine_cache` is an
in-process mapper event — every other gunicorn/celery worker keeps the old
secret until restart. Needs `import hashlib`.
```suggestion
def __repr__(self) -> str:
# The secret is part of the identity (this repr is the engine cache
# key) but must never appear verbatim, so bind it as a short digest.
digest =
hashlib.sha256(self._client_secret.encode()).hexdigest()[:12]
return (
f"_M2MCredentialsProvider(host={self._hostname!r}, "
f"client_id={self._client_id!r}, azure={self._azure!r}, "
f"azure_tenant_id={self._azure_tenant_id!r}, "
f"secret_digest={digest!r})"
)
```
##########
superset/db_engine_specs/databricks.py:
##########
@@ -478,7 +666,45 @@ def update_params_from_encrypted_extra(
logger.error(ex, exc_info=True)
raise
encrypted_extra.pop("oauth2_client_info", None)
- params.update(encrypted_extra)
+
+ auth_method = encrypted_extra.pop("auth_method", None)
+ client_id = encrypted_extra.pop("client_id", None) or
encrypted_extra.pop(
+ "azure_client_id", None
+ )
+ client_secret = encrypted_extra.pop("client_secret", None) or (
+ encrypted_extra.pop("azure_client_secret", None)
+ )
+ # Always drop the Azure-named aliases so they never reach
create_engine.
+ encrypted_extra.pop("azure_client_id", None)
+ encrypted_extra.pop("azure_client_secret", None)
+ azure_tenant_id = encrypted_extra.pop("azure_tenant_id", None)
+
+ if auth_method in M2M_AUTH_METHODS:
+ if not client_id or not client_secret:
+ raise ValueError(
+ "Databricks OAuth M2M requires both client_id and "
+ "client_secret in Secure Extra."
+ )
+ is_azure = auth_method == AUTH_METHOD_AZURE_SP_M2M
+ if is_azure and not azure_tenant_id:
+ raise ValueError(
+ "Databricks Azure service-principal M2M requires "
+ "azure_tenant_id in Secure Extra."
+ )
+ host = ""
+ if getattr(database, "url_object", None) is not None:
+ host = database.url_object.host or ""
+ connect_args = params.setdefault("connect_args", {})
+ connect_args["credentials_provider"] = _M2MCredentialsProvider(
+ host,
+ client_id,
+ client_secret,
+ azure=is_azure,
+ azure_tenant_id=azure_tenant_id,
+ )
+
+ if encrypted_extra:
+ params.update(encrypted_extra)
Review Comment:
A `connect_args` key in Secure Extra replaces the dict wholesale here,
discarding the `credentials_provider` injected above; auth then silently falls
back to the `m2m` placeholder token.
```suggestion
if encrypted_extra:
# Merge without clobbering ``connect_args``: the M2M branch above
# may have injected ``credentials_provider`` into it.
extra_connect_args = encrypted_extra.pop("connect_args", None)
params.update(encrypted_extra)
if extra_connect_args:
params.setdefault("connect_args",
{}).update(extra_connect_args)
```
##########
superset/db_engine_specs/databricks.py:
##########
@@ -808,11 +1054,63 @@ class
DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec):
"?http_path={http_path}&catalog={catalog}&schema={schema}"
),
"parameters": {
- "access_token": "Personal access token from Settings > User
Settings",
+ "access_token": (
+ "Personal access token from Settings > User Settings. "
+ "Optional when OAuth M2M client ID/secret is set in Secure
Extra."
+ ),
"host": "Server hostname from cluster JDBC/ODBC settings",
"port": "Port (default 443)",
"http_path": "HTTP path from cluster JDBC/ODBC settings",
},
+ "authentication_methods": [
+ {
+ "name": "Personal Access Token",
+ "description": (
+ "Default. Use a Databricks personal access token in the "
+ "Access token field."
+ ),
+ },
+ {
+ "name": "OAuth M2M (service principal)",
+ "description": (
+ "Connect with a Databricks service principal client ID and
"
+ "client secret. No personal access token is required."
+ ),
+ "requirements": (
+ "Create a workspace service principal and OAuth secret. "
+ "Grant the principal Can Use on the SQL warehouse. "
+ "Requires databricks-sdk (included in "
+ "apache-superset[databricks])."
+ ),
+ "secure_extra": {
+ "auth_method": "oauth-m2m",
+ "client_id": "<service-principal-application-id>",
+ "client_secret": "<oauth-secret>",
+ },
+ "notes": (
+ "In the database form, leave Access token blank or use a "
+ "placeholder such as m2m. Paste the JSON into Advanced → "
+ "Security → Secure extra."
+ ),
+ },
+ {
+ "name": "Azure Entra service principal",
+ "description": (
+ "Azure Databricks only. Use an Azure Entra application "
+ "via credentials_provider (azure_service_principal)."
+ ),
+ "secure_extra": {
+ "auth_method": "azure-sp-m2m",
+ "client_id": "<azure-app-id>",
+ "client_secret": "<azure-client-secret>",
+ "azure_tenant_id": "<tenant-id>",
+ },
+ "notes": (
+ "Requires databricks-sdk. azure_tenant_id is required. "
+ "The connector does not accept auth_type=azure-sp-m2m."
+ ),
Review Comment:
databricks-sql-connector 4.5.x does accept it: `AuthType.AZURE_SP_M2M =
"azure-sp-m2m"` is handled natively in `get_auth_provider`. It is Superset that
never routes it. Same sentence in `databricks.mdx:85`.
```suggestion
"notes": (
"Requires databricks-sdk. azure_tenant_id is required. "
"Superset does not route auth_type=azure-sp-m2m to the "
"connector."
),
```
##########
tests/unit_tests/db_engine_specs/test_databricks.py:
##########
@@ -670,6 +675,800 @@ def
test_update_params_invalid_encrypted_extra_raises(mocker: MockerFixture) ->
DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, {})
+def test_encrypted_extra_sensitive_fields() -> None:
+ """
+ Inherited ``$.*`` is kept so unrelated Encrypted Extra secrets stay masked.
+ """
+ from superset.db_engine_specs.databricks import
DatabricksDynamicBaseEngineSpec
+
+ paths =
DatabricksDynamicBaseEngineSpec.encrypted_extra_sensitive_field_paths()
+ assert "$.*" in paths
+ assert "$.client_secret" in paths
+ assert "$.oauth2_client_info.secret" in paths
+
+
+def test_mask_encrypted_extra_client_secret() -> None:
+ """
+ ``$.*`` redacts every top-level Encrypted Extra value, including M2M fields
+ and unrelated driver secrets.
+ """
+ from superset.db_engine_specs.databricks import
DatabricksDynamicBaseEngineSpec
+
+ config = json.dumps(
+ {
+ "auth_method": "oauth-m2m",
+ "client_id": "sp-application-id",
+ "client_secret": "super-secret",
+ "password": "unrelated-driver-secret",
+ }
+ )
+ assert DatabricksDynamicBaseEngineSpec.mask_encrypted_extra(config) ==
json.dumps(
+ {
+ "auth_method": "XXXXXXXXXX",
+ "client_id": "XXXXXXXXXX",
+ "client_secret": "XXXXXXXXXX",
+ "password": "XXXXXXXXXX",
+ }
+ )
+
+
+def test_update_params_oauth_m2m_injects_credentials_provider(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Encrypted Extra M2M builds a ``credentials_provider`` and does not pass
+ ``client_id`` / ``client_secret`` through as driver kwargs.
+ """
+ from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+ mock_config_cls = mocker.MagicMock()
+ mock_oauth = mocker.MagicMock(return_value="oauth-provider")
+ mocker.patch(
+ "superset.db_engine_specs.databricks._load_databricks_m2m_sdk",
+ return_value=(mock_config_cls, mock_oauth, None),
+ )
+
+ database = mocker.MagicMock()
+ database.url_object.host = "dbc-abc.cloud.databricks.com"
+ database.encrypted_extra = json.dumps(
+ {
+ "auth_method": "oauth-m2m",
+ "client_id": "sp-application-id",
+ "client_secret": "super-secret",
+ "http_headers": [["X-Custom", "value"]],
+ }
+ )
+ params: dict[str, Any] = {}
+
+ DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database,
params)
+
+ assert "client_id" not in params
+ assert "client_secret" not in params
+ assert "auth_method" not in params
+ assert params["http_headers"] == [["X-Custom", "value"]]
+ provider = params["connect_args"]["credentials_provider"]
+ assert callable(provider)
+
+ assert provider() == "oauth-provider"
+ mock_config_cls.assert_called_once_with(
+ host="https://dbc-abc.cloud.databricks.com";,
+ client_id="sp-application-id",
+ client_secret="super-secret", # noqa: S106
+ )
+ mock_oauth.assert_called_once()
+
+
+def test_update_params_oauth_m2m_strips_oauth2_client_info(
+ mocker: MockerFixture,
+) -> None:
+ """
+ U2M ``oauth2_client_info`` is still stripped when M2M is also configured.
+ """
+ from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+ mocker.patch(
+ "superset.db_engine_specs.databricks._load_databricks_m2m_sdk",
+ return_value=(mocker.MagicMock(), mocker.MagicMock(), None),
+ )
+ database = mocker.MagicMock()
+ database.url_object.host = "dbc-abc.cloud.databricks.com"
+ database.encrypted_extra = json.dumps(
+ {
+ "auth_method": "oauth-m2m",
+ "client_id": "sp-application-id",
+ "client_secret": "super-secret",
+ "oauth2_client_info": {
+ "id": "u2m-client-id",
+ "secret": "u2m-client-secret",
+ },
+ }
+ )
+ params: dict[str, Any] = {}
+
+ DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database,
params)
+
+ assert "oauth2_client_info" not in params
+ assert "credentials_provider" in params["connect_args"]
+
+
+def test_update_params_oauth_m2m_missing_credentials_raises(
+ mocker: MockerFixture,
+) -> None:
+ """
+ ``auth_method=oauth-m2m`` without both credentials fails with a clear
error.
+ """
+ from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+ database = mocker.MagicMock()
+ database.encrypted_extra = json.dumps(
+ {"auth_method": "oauth-m2m", "client_id": "sp-application-id"}
+ )
+
+ with pytest.raises(ValueError, match="client_id and client_secret"):
+
DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+def test_update_params_oauth_m2m_missing_sdk_raises(mocker: MockerFixture) ->
None:
+ """
+ Missing ``databricks-sdk`` fails with an install hint, not a driver error.
+ """
+ from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+ mocker.patch(
+ "superset.db_engine_specs.databricks._load_databricks_m2m_sdk",
+ side_effect=ValueError(
+ "Databricks OAuth M2M requires the databricks-sdk package."
+ ),
+ )
+ database = mocker.MagicMock()
+ database.url_object.host = "dbc-abc.cloud.databricks.com"
+ database.encrypted_extra = json.dumps(
+ {
+ "auth_method": "oauth-m2m",
+ "client_id": "sp-application-id",
+ "client_secret": "super-secret",
+ }
+ )
+
+ with pytest.raises(ValueError, match="databricks-sdk"):
+
DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+def test_load_databricks_m2m_sdk_import_error(mocker: MockerFixture) -> None:
+ """
+ The lazy import surfaces a clear error when ``databricks-sdk`` is absent.
+ """
+ import builtins
+
+ from superset.db_engine_specs.databricks import _load_databricks_m2m_sdk
+
+ real_import = builtins.__import__
+
+ def fake_import(name: str, *args: Any, **kwargs: Any) -> Any:
+ if name == "databricks.sdk.core":
+ raise ImportError("No module named databricks.sdk")
+ return real_import(name, *args, **kwargs)
+
+ mocker.patch("builtins.__import__", side_effect=fake_import)
+ with pytest.raises(ValueError, match="databricks-sdk"):
+ _load_databricks_m2m_sdk()
+
+
+def test_load_databricks_m2m_sdk_missing_azure_helper(mocker: MockerFixture)
-> None:
+ """
+ Native M2M still loads when ``azure_service_principal`` is absent.
+ """
+ import builtins
+ from types import SimpleNamespace
+
+ from superset.db_engine_specs.databricks import _load_databricks_m2m_sdk
+
+ real_import = builtins.__import__
+
+ def fake_import(name: str, *args: Any, **kwargs: Any) -> Any:
+ if name == "databricks.sdk.core":
+ fromlist = kwargs.get("fromlist") or (args[2] if len(args) > 2
else ())
+ if fromlist and "azure_service_principal" in fromlist:
+ raise ImportError("azure_service_principal is unavailable")
+ return SimpleNamespace(Config=object,
oauth_service_principal=object)
+ return real_import(name, *args, **kwargs)
+
+ mocker.patch("builtins.__import__", side_effect=fake_import)
+ config_cls, oauth_sp, azure_sp = _load_databricks_m2m_sdk()
+ assert config_cls is object
+ assert oauth_sp is object
+ assert azure_sp is None
+
+
+def test_parse_json_object_branches() -> None:
+ from superset.db_engine_specs.databricks import _parse_json_object
+
+ payload = {"auth_method": "oauth-m2m"}
+ assert _parse_json_object(payload) is payload
+ assert _parse_json_object(None) == {}
+ assert _parse_json_object("") == {}
+ assert _parse_json_object(123) == {}
+ assert _parse_json_object("{not json") == {}
+ assert _parse_json_object("[1, 2]") == {}
+ assert _parse_json_object(json.dumps(payload)) == payload
+
+
+def test_parse_extra_for_validation_branches() -> None:
+ from superset.db_engine_specs.databricks import _parse_extra_for_validation
+
+ extra: dict[str, Any] = {"engine_params": {}}
+ assert _parse_extra_for_validation(extra) == (extra, None)
+ assert _parse_extra_for_validation(None) == ({}, None)
+ assert _parse_extra_for_validation("") == ({}, None)
+ parsed, error = _parse_extra_for_validation("[1]")
+ assert parsed == {}
+ assert error is None
+ parsed, error = _parse_extra_for_validation("{not json")
+ assert parsed == {}
+ assert error is not None
+ assert error.extra is not None
+ assert error.extra["invalid"] == ["extra"]
+
+
+def test_has_m2m_credentials_requires_complete_azure_config() -> None:
+ from superset.constants import PASSWORD_MASK
+ from superset.db_engine_specs.databricks import _has_m2m_credentials
+
+ assert _has_m2m_credentials({}) is False
+ assert _has_m2m_credentials({"auth_method": "oauth-m2m"}) is False
+ assert (
+ _has_m2m_credentials(
+ {
+ "auth_method": "azure-sp-m2m",
+ "client_id": "azure-app-id",
+ "client_secret": "azure-secret",
+ }
+ )
+ is False
+ )
+ assert (
+ _has_m2m_credentials(
+ {
+ "auth_method": PASSWORD_MASK,
+ "client_id": PASSWORD_MASK,
+ "client_secret": PASSWORD_MASK,
+ }
+ )
+ is True
+ )
+
+
+def test_workspace_hostname_strips_http_prefix() -> None:
+ from superset.db_engine_specs.databricks import _workspace_hostname
+
+ assert _workspace_hostname("http://dbc-abc.cloud.databricks.com/";) == (
+ "dbc-abc.cloud.databricks.com"
+ )
+
+
+def test_access_token_from_uri_password_keeps_real_pat() -> None:
+ from superset.db_engine_specs.databricks import
_access_token_from_uri_password
+
+ assert _access_token_from_uri_password(None) == ""
+ assert _access_token_from_uri_password("abc12345") == "abc12345"
+
+
+def test_m2m_credentials_provider_requires_host(mocker: MockerFixture) -> None:
+ from superset.db_engine_specs.databricks import _M2MCredentialsProvider
+
+ mocker.patch(
+ "superset.db_engine_specs.databricks._load_databricks_m2m_sdk",
+ return_value=(mocker.MagicMock(), mocker.MagicMock(),
mocker.MagicMock()),
+ )
+ with pytest.raises(ValueError, match="workspace host"):
+ _M2MCredentialsProvider("", "sp-application-id", "super-secret")
+
+
+def test_m2m_credentials_provider_azure_missing_helper_raises(
+ mocker: MockerFixture,
+) -> None:
+ from superset.db_engine_specs.databricks import _M2MCredentialsProvider
+
+ mocker.patch(
+ "superset.db_engine_specs.databricks._load_databricks_m2m_sdk",
+ return_value=(mocker.MagicMock(), mocker.MagicMock(), None),
+ )
+ with pytest.raises(ValueError, match="azure_service_principal"):
+ _M2MCredentialsProvider(
+ "adb-123.azuredatabricks.net",
+ "azure-app-id",
+ "azure-secret",
+ azure=True,
+ azure_tenant_id="tenant-123",
+ )
+
+
+def test_update_params_oauth_m2m_missing_host_raises(mocker: MockerFixture) ->
None:
+ from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+ database = mocker.MagicMock()
+ database.url_object = None
+ database.encrypted_extra = json.dumps(
+ {
+ "auth_method": "oauth-m2m",
+ "client_id": "sp-application-id",
+ "client_secret": "super-secret",
+ }
+ )
+ with pytest.raises(ValueError, match="workspace host"):
+
DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+def test_update_params_azure_sp_m2m(mocker: MockerFixture) -> None:
+ """
+ Azure SP M2M uses ``azure_service_principal`` via credentials_provider.
+ """
+ from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+ mock_config_cls = mocker.MagicMock()
+ mock_oauth = mocker.MagicMock()
+ mock_azure = mocker.MagicMock(return_value="azure-provider")
+ mocker.patch(
+ "superset.db_engine_specs.databricks._load_databricks_m2m_sdk",
+ return_value=(mock_config_cls, mock_oauth, mock_azure),
+ )
+ database = mocker.MagicMock()
+ database.url_object.host = "adb-123.azuredatabricks.net"
+ database.encrypted_extra = json.dumps(
+ {
+ "auth_method": "azure-sp-m2m",
+ "client_id": "azure-app-id",
+ "client_secret": "azure-secret",
+ "azure_tenant_id": "tenant-123",
+ }
+ )
+ params: dict[str, Any] = {}
+
+ DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database,
params)
+
+ assert "client_id" not in params
+ assert "client_secret" not in params
+ assert "auth_type" not in params.get("connect_args", {})
+ provider = params["connect_args"]["credentials_provider"]
+ assert provider() == "azure-provider"
+ mock_config_cls.assert_called_once_with(
+ host="https://adb-123.azuredatabricks.net";,
+ azure_tenant_id="tenant-123",
+ azure_client_id="azure-app-id",
+ azure_client_secret="azure-secret", # noqa: S106
+ )
+ mock_azure.assert_called_once()
+ mock_oauth.assert_not_called()
+
+
+def test_update_params_azure_sp_m2m_aliases(mocker: MockerFixture) -> None:
+ """
+ Azure-named client ID/secret keys are accepted and never passed to the
driver.
+ """
+ from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+ mocker.patch(
+ "superset.db_engine_specs.databricks._load_databricks_m2m_sdk",
+ return_value=(mocker.MagicMock(), mocker.MagicMock(),
mocker.MagicMock()),
+ )
+ database = mocker.MagicMock()
+ database.url_object.host = "adb-123.azuredatabricks.net"
+ database.encrypted_extra = json.dumps(
+ {
+ "auth_method": "azure-sp-m2m",
+ "azure_client_id": "azure-app-id",
+ "azure_client_secret": "azure-secret",
+ "azure_tenant_id": "tenant-123",
+ }
+ )
+ params: dict[str, Any] = {}
+
+ DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database,
params)
+
+ assert "azure_client_id" not in params
+ assert "azure_client_secret" not in params
+ assert "credentials_provider" in params["connect_args"]
+
+
+def test_update_params_azure_sp_m2m_missing_credentials_raises(
+ mocker: MockerFixture,
+) -> None:
+ from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+ database = mocker.MagicMock()
+ database.encrypted_extra = json.dumps({"auth_method": "azure-sp-m2m"})
+
+ with pytest.raises(ValueError, match="client_id and client_secret"):
+
DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+def test_update_params_azure_sp_m2m_missing_tenant_raises(
+ mocker: MockerFixture,
+) -> None:
+ from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec
+
+ database = mocker.MagicMock()
+ database.encrypted_extra = json.dumps(
+ {
+ "auth_method": "azure-sp-m2m",
+ "client_id": "azure-app-id",
+ "client_secret": "azure-secret",
+ }
+ )
+
+ with pytest.raises(ValueError, match="azure_tenant_id"):
+
DatabricksNativeEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+def test_m2m_credentials_provider_repr_is_stable(mocker: MockerFixture) ->
None:
+ """
+ ``repr`` is identity-stable (and secret-free) so the engine cache key
matches.
+ """
+ from superset.db_engine_specs.databricks import _M2MCredentialsProvider
+
+ mocker.patch(
+ "superset.db_engine_specs.databricks._load_databricks_m2m_sdk",
+ return_value=(mocker.MagicMock(), mocker.MagicMock(),
mocker.MagicMock()),
+ )
+ first = _M2MCredentialsProvider(
+ "https://dbc-abc.cloud.databricks.com/";,
+ "sp-application-id",
+ "old-secret",
+ )
+ second = _M2MCredentialsProvider(
+ "dbc-abc.cloud.databricks.com",
+ "sp-application-id",
+ "rotated-secret",
+ )
+ assert repr(first) == repr(second)
+ assert "old-secret" not in repr(first)
+ assert "rotated-secret" not in repr(second)
Review Comment:
Dropping `client_id` or `azure_tenant_id` from `__repr__` leaves all 96
tests green, so nothing binds the cache key's discriminating power.
`repr(first) == repr(second)` above also locks in the rotation collision.
```suggestion
assert "old-secret" not in repr(first)
assert "rotated-secret" not in repr(second)
# identity must still discriminate on the fields it does include
other = _M2MCredentialsProvider(
"dbc-abc.cloud.databricks.com", "other-application-id", "old-secret"
)
assert repr(first) != repr(other)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
