This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 5cf8213354c fix(databricks): keep the user's OAuth2 token and extra
connect_args; re-auth on HTTP 401 (#44705)
5cf8213354c is described below
commit 5cf8213354c74cff577d9832671ff66e0440c308
Author: Amin Ghadersohi <[email protected]>
AuthorDate: Thu Oct 1 12:50:32 2026 +1000
fix(databricks): keep the user's OAuth2 token and extra connect_args;
re-auth on HTTP 401 (#44705)
---
superset/db_engine_specs/databricks.py | 91 +++++++++-
.../unit_tests/db_engine_specs/test_databricks.py | 202 ++++++++++++++++++++-
2 files changed, 289 insertions(+), 4 deletions(-)
diff --git a/superset/db_engine_specs/databricks.py
b/superset/db_engine_specs/databricks.py
index b572832fb3b..ebf4a3badb0 100644
--- a/superset/db_engine_specs/databricks.py
+++ b/superset/db_engine_specs/databricks.py
@@ -288,6 +288,25 @@ class DatabricksBaseEngineSpec(BaseEngineSpec):
return super().extract_errors(ex, context, database_name)
+# Credential connect args of databricks-sql-connector. With impersonation on,
all
+# of them are dropped from ``extra`` and the secure extra; ``impersonate_user``
+# then writes the user's OAuth2 token as ``access_token`` (when ``extra`` set
one),
+# so that token is the only credential the connector receives.
+DATABRICKS_SHARED_CREDENTIAL_CONNECT_ARGS = frozenset(
+ {
+ "access_token",
+ "auth_type",
+ "credentials_provider",
+ "oauth_client_id",
+ "oauth_client_secret",
+ "azure_client_id",
+ "azure_client_secret",
+ "azure_tenant_id",
+ "azure_workspace_resource_id",
+ }
+)
+
+
class DatabricksDynamicBaseEngineSpec(BasicParametersMixin,
DatabricksBaseEngineSpec):
default_driver = ""
encryption_parameters = {"ssl": "1"}
@@ -382,7 +401,21 @@ class
DatabricksDynamicBaseEngineSpec(BasicParametersMixin, DatabricksBaseEngine
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)
+ if any(signal in message for signal in
cls.oauth2_auth_failure_signals):
+ return True
+ # A rejected bearer token fails the request with HTTP 401 and a message
+ # such as "Credential was not sent or was of an unsupported type for
this
+ # API", which no signal above matches; the connector reports the status
+ # in ``RequestError.context["http-code"]``. 403 is left out on
purpose: it
+ # also means a missing permission, which re-authorizing cannot fix.
+ error = getattr(ex, "orig", None) or ex
+ context = getattr(error, "context", None)
+ if isinstance(context, dict):
+ try:
+ return int(context.get("http-code") or 0) == 401
+ except (TypeError, ValueError):
+ return False
+ return False
@classmethod
def get_oauth2_authorization_uri(
@@ -459,14 +492,40 @@ class
DatabricksDynamicBaseEngineSpec(BasicParametersMixin, DatabricksBaseEngine
# back to an empty string to force re-authentication when none is set.
url = url.set(password=user_token or "")
+ # Drop shared credentials from ``extra`` (a PAT, OAuth M2M client,
Azure
+ # service principal) so none of them reaches the connector next to the
+ # user's token. The secure extra is filtered the same way in
+ # ``update_params_from_encrypted_extra``.
+ connect_args = engine_kwargs.get("connect_args") or {}
+ filtered_connect_args = {
+ key: value
+ for key, value in connect_args.items()
+ if key not in DATABRICKS_SHARED_CREDENTIAL_CONNECT_ARGS
+ }
# The Python connector passes the token via ``connect_args`` instead
of the
# URL password, so keep it in sync (clearing it likewise forces
re-auth).
- connect_args = engine_kwargs.setdefault("connect_args", {})
if "access_token" in connect_args:
- connect_args["access_token"] = user_token or ""
+ filtered_connect_args["access_token"] = user_token or ""
+ engine_kwargs["connect_args"] = filtered_connect_args
return url, engine_kwargs
+ @classmethod
+ def start_oauth2_dance(cls, database: Database) -> None:
+ """
+ Start the OAuth2 dance only when the database impersonates the user.
+
+ The user's OAuth2 token only reaches the connection through
+ ``impersonate_user``. Without impersonation the connection uses the
+ shared credential, so an authorization prompt cannot fix an auth
failure
+ (e.g. a revoked shared token returning HTTP 401). Return instead, so
the
+ caller raises the original error for an admin to act on.
+ """
+ if not database.impersonate_user:
+ return
+
+ super().start_oauth2_dance(database)
+
@staticmethod
def update_params_from_encrypted_extra(
database: Database, params: dict[str, Any]
@@ -479,6 +538,14 @@ class
DatabricksDynamicBaseEngineSpec(BasicParametersMixin, DatabricksBaseEngine
consumed by ``Database.get_oauth2_config``; it is not a Databricks
driver
connection argument, so it must be stripped here to avoid poisoning the
connection when OAuth2 is configured on the database itself.
+
+ ``connect_args`` are merged key by key, so credentials kept in the
secure
+ extra (an access token, an OAuth M2M client secret) do not discard the
+ ``connect_args`` from ``extra`` or Superset's own (e.g. the
User-Agent).
+ With user impersonation enabled, shared credentials in the secure extra
+ are dropped (``impersonate_user`` drops the ones from ``extra``): the
+ user's OAuth2 token is the only credential, instead of being silently
+ replaced by a shared one.
"""
if not database.encrypted_extra:
return
@@ -488,7 +555,25 @@ class
DatabricksDynamicBaseEngineSpec(BasicParametersMixin, DatabricksBaseEngine
logger.error(ex, exc_info=True)
raise
encrypted_extra.pop("oauth2_client_info", None)
+ secure_connect_args = encrypted_extra.pop("connect_args", None)
params.update(encrypted_extra)
+ if secure_connect_args is not None and not isinstance(
+ secure_connect_args, dict
+ ):
+ params["connect_args"] = secure_connect_args
+ return
+ if secure_connect_args is None and "connect_args" not in params:
+ return
+
+ secure_connect_args = secure_connect_args or {}
+ connect_args = dict(params.get("connect_args") or {})
+ if database.impersonate_user:
+ secure_connect_args = {
+ key: value
+ for key, value in secure_connect_args.items()
+ if key not in DATABRICKS_SHARED_CREDENTIAL_CONNECT_ARGS
+ }
+ params["connect_args"] = {**connect_args, **secure_connect_args}
@staticmethod
def get_extra_params(
diff --git a/tests/unit_tests/db_engine_specs/test_databricks.py
b/tests/unit_tests/db_engine_specs/test_databricks.py
index cd1ac465d8b..fb06011bf8f 100644
--- a/tests/unit_tests/db_engine_specs/test_databricks.py
+++ b/tests/unit_tests/db_engine_specs/test_databricks.py
@@ -20,6 +20,7 @@ from datetime import datetime
from importlib import import_module
from importlib.metadata import PackageNotFoundError, requires
from typing import Any, Optional
+from unittest.mock import MagicMock
from urllib.parse import parse_qs, urlparse
import pytest
@@ -28,7 +29,7 @@ from pytest_mock import MockerFixture
from sqlalchemy import __version__ as sqlalchemy_version, create_engine, text
from sqlalchemy.engine.url import make_url
-from superset.db_engine_specs.base import OAuth2State
+from superset.db_engine_specs.base import BaseEngineSpec, OAuth2State
from superset.db_engine_specs.databricks import (
DatabricksNativeEngineSpec,
DatabricksPythonConnectorEngineSpec,
@@ -509,6 +510,48 @@ def test_needs_oauth2_matches_oauth2_redirect_error(
assert spec.needs_oauth2(ex) is True
+class _ConnectorRequestError(Exception):
+ """Stand-in for ``databricks.sql.exc.RequestError`` (message + context)."""
+
+ def __init__(self, message: str, context: dict[str, Any]) -> None:
+ super().__init__(message)
+ self.context = context
+
+
[email protected](
+ "spec",
+ [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec],
+)
[email protected](
+ "http_code, expected",
+ [(401, True), ("401", True), (403, False), (None, False)],
+)
+def test_needs_oauth2_detects_http_401_from_the_connector(
+ mocker: MockerFixture,
+ spec: Any,
+ http_code: Any,
+ expected: bool,
+) -> None:
+ """
+ A rejected bearer token fails with HTTP 401 and a message that matches no
+ signal ("Credential was not sent or was of an unsupported type for this
+ API"); the connector's structured ``http-code`` triggers the re-auth, also
+ when SQLAlchemy wraps the error. 403 (a missing permission) does not.
+ """
+ g = mocker.patch("superset.db_engine_specs.databricks.g")
+ g.user = mocker.MagicMock()
+ error = _ConnectorRequestError(
+ "Error during request to server: : Credential was not sent or was of
an "
+ "unsupported type for this API.",
+ {"method": "OpenSession", "http-code": http_code},
+ )
+ wrapped = mocker.MagicMock(orig=error)
+ wrapped.__str__ = lambda self: str(error)
+
+ assert spec.needs_oauth2(error) is expected
+ assert spec.needs_oauth2(wrapped) is expected
+
+
def test_impersonate_user_with_token(mocker: MockerFixture) -> None:
"""
Test impersonate_user method with OAuth2 token for
DatabricksNativeEngineSpec.
@@ -697,6 +740,163 @@ def test_update_params_merges_when_no_oauth2_client_info(
}
+def test_update_params_merges_connect_args(mocker: MockerFixture) -> None:
+ """
+ Secure ``connect_args`` (e.g. OAuth M2M credentials) are merged key by key
+ into the ``connect_args`` from ``extra`` instead of replacing them.
+ """
+ database: MagicMock = mocker.MagicMock()
+ database.impersonate_user = False
+ database.encrypted_extra = json.dumps(
+ {"connect_args": {"oauth_client_id": "sp", "oauth_client_secret":
"secret"}}
+ )
+ params: dict[str, Any] = {
+ "connect_args": {
+ "_user_agent_entry": "Apache Superset",
+ "session_configuration": {"TIMEZONE": "Asia/Kolkata"},
+ }
+ }
+
+ DatabricksPythonConnectorEngineSpec.update_params_from_encrypted_extra(
+ database, params
+ )
+
+ assert params == {
+ "connect_args": {
+ "_user_agent_entry": "Apache Superset",
+ "session_configuration": {"TIMEZONE": "Asia/Kolkata"},
+ "oauth_client_id": "sp",
+ "oauth_client_secret": "secret",
+ }
+ }
+
+
+def test_update_params_impersonation_keeps_only_the_user_token(
+ mocker: MockerFixture,
+) -> None:
+ """
+ With impersonation, a shared credential in the secure extra (or in
+ ``extra``) must not replace the user's OAuth2 token set by
+ ``impersonate_user``.
+ """
+ database: MagicMock = mocker.MagicMock()
+ database.impersonate_user = True
+ database.encrypted_extra = json.dumps(
+ {"connect_args": {"access_token": "shared-pat", "http_path": "/sql/1"}}
+ )
+ engine_kwargs: dict[str, Any] = {
+ "connect_args": {"access_token": "extra-pat", "oauth_client_secret":
"s"}
+ }
+
+ # Same order as ``Database._get_sqla_engine``.
+ _, engine_kwargs = DatabricksPythonConnectorEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token="user-token", # noqa: S106
+ url=make_url("databricks://token:extra-pat@host:443?http_path=/sql/1"),
+ engine_kwargs=engine_kwargs,
+ )
+ DatabricksPythonConnectorEngineSpec.update_params_from_encrypted_extra(
+ database, engine_kwargs
+ )
+
+ assert engine_kwargs == {
+ "connect_args": {"access_token": "user-token", "http_path": "/sql/1"}
+ }
+
+
+def test_impersonate_user_drops_shared_credentials_without_secure_extra(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Shared credentials in ``extra`` are dropped even when the database has no
+ secure extra (e.g. OAuth2 configured through ``DATABASE_OAUTH2_CLIENTS``),
+ so only the user's token reaches the connector.
+ """
+ database: MagicMock = mocker.MagicMock()
+ database.impersonate_user = True
+ database.encrypted_extra = None
+ engine_kwargs: dict[str, Any] = {
+ "connect_args": {
+ "_user_agent_entry": "Apache Superset",
+ "auth_type": "databricks-oauth",
+ "oauth_client_id": "sp",
+ "oauth_client_secret": "secret",
+ }
+ }
+
+ url, engine_kwargs = DatabricksPythonConnectorEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token="user-oauth-token", # noqa: S106
+ url=make_url("databricks://token:@host:443?http_path=/sql/1"),
+ engine_kwargs=engine_kwargs,
+ )
+ DatabricksPythonConnectorEngineSpec.update_params_from_encrypted_extra(
+ database, engine_kwargs
+ )
+
+ assert url.password == "user-oauth-token" # noqa: S105
+ assert engine_kwargs == {"connect_args": {"_user_agent_entry": "Apache
Superset"}}
+
+
[email protected](
+ "spec",
+ [DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec],
+)
+def test_start_oauth2_dance_requires_impersonation(
+ mocker: MockerFixture,
+ spec: Any,
+) -> None:
+ """
+ Without impersonation the user's token is never used, so an auth failure
+ (e.g. a revoked shared token) does not send the user through an
+ authorization prompt; the caller raises the original error instead.
+ """
+ base_start = mocker.patch.object(BaseEngineSpec, "start_oauth2_dance")
+ database: MagicMock = mocker.MagicMock()
+ database.impersonate_user = False
+
+ spec.start_oauth2_dance(database)
+ base_start.assert_not_called()
+
+ database.impersonate_user = True
+ spec.start_oauth2_dance(database)
+ base_start.assert_called_once_with(database)
+
+
+def test_impersonation_with_shared_token_only_in_secure_extra(
+ mocker: MockerFixture,
+) -> None:
+ """
+ With the shared token only in the secure extra, the Python connector still
+ authenticates as the user: ``impersonate_user`` puts the user's token in
the
+ URL password, which the dialect uses as ``access_token``, and the shared
+ token is dropped rather than overriding it.
+ """
+ database: MagicMock = mocker.MagicMock()
+ database.impersonate_user = True
+ database.encrypted_extra = json.dumps(
+ {"connect_args": {"access_token": "shared-pat"}}
+ )
+ engine_kwargs: dict[str, Any] = {"connect_args": {"http_path": "/sql/1"}}
+
+ # Same order as ``Database._get_sqla_engine``.
+ url, engine_kwargs = DatabricksPythonConnectorEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token="user-oauth-token", # noqa: S106
+ url=make_url("databricks://token:@host:443?http_path=/sql/1"),
+ engine_kwargs=engine_kwargs,
+ )
+ DatabricksPythonConnectorEngineSpec.update_params_from_encrypted_extra(
+ database, engine_kwargs
+ )
+
+ assert url.password == "user-oauth-token" # noqa: S105
+ assert engine_kwargs["connect_args"] == {"http_path": "/sql/1"}
+
+
def test_update_params_invalid_encrypted_extra_raises(mocker: MockerFixture)
-> None:
"""
Malformed ``encrypted_extra`` JSON raises rather than silently connecting.
