You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by be...@apache.org on 2025/03/03 15:16:36 UTC
(superset) 01/01: chore: simplify user impersonation
This is an automated email from the ASF dual-hosted git repository.
beto pushed a commit to branch improve-impersonate-user
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 481a1b7b932dd4bc313a375610c66ac3ddbec304
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Mon Mar 3 10:16:24 2025 -0500
chore: simplify user impersonation
---
superset/db_engine_specs/base.py | 39 +++++++++++++++++++++----
superset/db_engine_specs/drill.py | 47 +++++++++++++++---------------
superset/db_engine_specs/gsheets.py | 20 ++++++-------
superset/db_engine_specs/hive.py | 53 ++++++++-------------------------
superset/db_engine_specs/lib.py | 1 +
superset/db_engine_specs/presto.py | 35 ++++++++--------------
superset/db_engine_specs/starrocks.py | 19 ------------
superset/db_engine_specs/trino.py | 55 ++++++++---------------------------
superset/models/core.py | 41 +++++++++-----------------
9 files changed, 117 insertions(+), 193 deletions(-)
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 248f8bd0a7..7fecfe9bd3 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -22,6 +22,7 @@ import logging
import re
import warnings
from datetime import datetime
+from inspect import signature
from re import Match, Pattern
from typing import (
Any,
@@ -1414,11 +1415,6 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
should also have the attribute ``supports_dynamic_schema`` set to true, so that
Superset knows in which schema a given query is running in order to enforce
permissions (see #23385 and #23401).
-
- Currently, changing the catalog is not supported. The method accepts a catalog so
- that when catalog support is added to Superset the interface remains the same.
- This is important because DB engine specs can be installed from 3rd party
- packages, so we want to keep these methods as stable as possible.
""" # noqa: E501
return uri, {
**connect_args,
@@ -1792,6 +1788,38 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
]
@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]]:
+ """
+ Modify URL and/or engine kwargs to impersonate a different user.
+ """
+ # Update URL using old methods until 6.0.0.
+ url = cls.get_url_for_impersonation(url, True, username, user_token)
+
+ # Update engine kwargs using old methods. Note that #30674 modified the method
+ # signature, so we need to check if the method has the old signature.
+ args = [
+ engine_kwargs["connect_args"],
+ url,
+ username,
+ user_token,
+ ]
+
+ if "database" in signature(cls.update_impersonation_config).parameters:
+ args.insert(0, database)
+
+ cls.update_impersonation_config(*args)
+
+ return url, engine_kwargs
+
+ @classmethod
+ @deprecated(deprecated_in="6.0.0")
def get_url_for_impersonation(
cls,
url: URL,
@@ -1813,6 +1841,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
return url
@classmethod
+ @deprecated(deprecated_in="6.0.0")
def update_impersonation_config( # pylint: disable=too-many-arguments
cls,
database: Database,
diff --git a/superset/db_engine_specs/drill.py b/superset/db_engine_specs/drill.py
index e99d4a27f4..167ed386ea 100644
--- a/superset/db_engine_specs/drill.py
+++ b/superset/db_engine_specs/drill.py
@@ -18,7 +18,7 @@
from __future__ import annotations
from datetime import datetime
-from typing import Any
+from typing import Any, TYPE_CHECKING
from urllib import parse
from sqlalchemy import types
@@ -28,6 +28,9 @@ from superset.constants import TimeGrain
from superset.db_engine_specs.base import BaseEngineSpec
from superset.db_engine_specs.exceptions import SupersetDBAPIProgrammingError
+if TYPE_CHECKING:
+ from superset.models.core import Database
+
class DrillEngineSpec(BaseEngineSpec):
"""Engine spec for Apache Drill"""
@@ -99,31 +102,27 @@ class DrillEngineSpec(BaseEngineSpec):
return parse.unquote(sqlalchemy_uri.database).replace("/", ".")
@classmethod
- def get_url_for_impersonation(
+ def impersonate_user(
cls,
- url: URL,
- impersonate_user: bool,
+ database: Database,
username: str | None,
- access_token: str | None,
- ) -> URL:
- """
- Return a modified URL with the username set.
-
- :param url: SQLAlchemy URL object
- :param impersonate_user: Flag indicating if impersonation is enabled
- :param username: Effective username
- """
- if impersonate_user and username is not None:
- if url.drivername == "drill+odbc":
- url = url.update_query_dict({"DelegationUID": username})
- elif url.drivername in ["drill+sadrill", "drill+jdbc"]:
- url = url.update_query_dict({"impersonation_target": username})
- else:
- raise SupersetDBAPIProgrammingError(
- f"impersonation is not supported for {url.drivername}"
- )
-
- return url
+ user_token: str | None,
+ url: URL,
+ engine_kwargs: dict[str, Any],
+ ) -> tuple[URL, dict[str, Any]]:
+ if username is None:
+ return url, engine_kwargs
+
+ if url.drivername == "drill+odbc":
+ url = url.update_query_dict({"DelegationUID": username})
+ elif url.drivername in {"drill+sadrill", "drill+jdbc"}:
+ url = url.update_query_dict({"impersonation_target": username})
+ else:
+ raise SupersetDBAPIProgrammingError(
+ f"impersonation is not supported for {url.drivername}"
+ )
+
+ return url, engine_kwargs
@classmethod
def fetch_data(
diff --git a/superset/db_engine_specs/gsheets.py b/superset/db_engine_specs/gsheets.py
index 65afd9aee9..ba8d302275 100644
--- a/superset/db_engine_specs/gsheets.py
+++ b/superset/db_engine_specs/gsheets.py
@@ -118,25 +118,23 @@ class GSheetsEngineSpec(ShillelaghEngineSpec):
oauth2_exception = UnauthenticatedError
@classmethod
- def get_url_for_impersonation(
+ def impersonate_user(
cls,
- url: URL,
- impersonate_user: bool,
+ database: Database,
username: str | None,
- access_token: str | None,
- ) -> URL:
- if not impersonate_user:
- return url
-
+ user_token: str | None,
+ url: URL,
+ engine_kwargs: dict[str, Any],
+ ) -> tuple[URL, dict[str, Any]]:
if username is not None:
user = security_manager.find_user(username=username)
if user and user.email:
url = url.update_query_dict({"subject": user.email})
- if access_token:
- url = url.update_query_dict({"access_token": access_token})
+ if user_token:
+ url = url.update_query_dict({"access_token": user_token})
- return url
+ return url, engine_kwargs
@classmethod
def get_extra_table_metadata(
diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py
index 9b7afab1cc..7bc86108cb 100644
--- a/superset/db_engine_specs/hive.py
+++ b/superset/db_engine_specs/hive.py
@@ -39,7 +39,6 @@ from sqlalchemy.sql.expression import ColumnClause, Select
from superset import db
from superset.common.db_query_status import QueryStatus
from superset.constants import TimeGrain
-from superset.databases.utils import make_url_safe
from superset.db_engine_specs.base import BaseEngineSpec
from superset.db_engine_specs.presto import PrestoEngineSpec
from superset.exceptions import SupersetException
@@ -512,54 +511,26 @@ class HiveEngineSpec(PrestoEngineSpec):
)
@classmethod
- def get_url_for_impersonation(
- cls,
- url: URL,
- impersonate_user: bool,
- username: str | None,
- access_token: str | None,
- ) -> URL:
- """
- Return a modified URL with the username set.
-
- :param url: SQLAlchemy URL object
- :param impersonate_user: Flag indicating if impersonation is enabled
- :param username: Effective username
- """
- # Do nothing in the URL object since instead this should modify
- # the configuration dictionary. See get_configuration_for_impersonation
- return url
-
- @classmethod
- def update_impersonation_config( # pylint: disable=too-many-arguments
+ def impersonate_user(
cls,
database: Database,
- connect_args: dict[str, Any],
- uri: str,
username: str | None,
- access_token: str | None,
- ) -> None:
- """
- Update a configuration dictionary
- that can set the correct properties for impersonating users
- :param database: the Database Object
- :param connect_args:
- :param uri: URI string
- :param impersonate_user: Flag indicating if impersonation is enabled
- :param username: Effective username
- :return: None
- """
- url = make_url_safe(uri)
- backend_name = url.get_backend_name()
+ user_token: str | None,
+ url: URL,
+ engine_kwargs: dict[str, Any],
+ ) -> tuple[URL, dict[str, Any]]:
+ if username is None:
+ return url, engine_kwargs
- # Must be Hive connection, enable impersonation, and set optional param
- # auth=LDAP|KERBEROS
- # this will set hive.server2.proxy.user=$effective_username on connect_args['configuration'] # noqa: E501
- if backend_name == "hive" and username is not None:
+ backend_name = url.get_backend_name()
+ connect_args = engine_kwargs.setdefault("connect_args", {})
+ if backend_name == "hive":
configuration = connect_args.get("configuration", {})
configuration["hive.server2.proxy.user"] = username
connect_args["configuration"] = configuration
+ return url, engine_kwargs
+
@staticmethod
def execute( # type: ignore
cursor,
diff --git a/superset/db_engine_specs/lib.py b/superset/db_engine_specs/lib.py
index fc7a8168b2..106b9c7550 100644
--- a/superset/db_engine_specs/lib.py
+++ b/superset/db_engine_specs/lib.py
@@ -140,6 +140,7 @@ def diagnose(spec: type[BaseEngineSpec]) -> dict[str, Any]:
"user_impersonation": (
has_custom_method(spec, "update_impersonation_config")
or has_custom_method(spec, "get_url_for_impersonation")
+ or has_custom_method(spec, "impersonate_user")
),
"file_upload": spec.supports_file_upload,
"get_extra_table_metadata": has_custom_method(
diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py
index e351c5b307..feb3a72153 100644
--- a/superset/db_engine_specs/presto.py
+++ b/superset/db_engine_specs/presto.py
@@ -43,7 +43,6 @@ from sqlalchemy.sql.expression import ColumnClause, Select
from superset import cache_manager, db, is_feature_enabled
from superset.common.db_query_status import QueryStatus
from superset.constants import TimeGrain
-from superset.databases.utils import make_url_safe
from superset.db_engine_specs.base import BaseEngineSpec
from superset.errors import SupersetErrorType
from superset.exceptions import SupersetTemplateException
@@ -955,34 +954,24 @@ class PrestoEngineSpec(PrestoBaseEngineSpec):
return version is not None and Version(version) >= Version("0.319")
@classmethod
- def update_impersonation_config( # pylint: disable=too-many-arguments
+ def impersonate_user(
cls,
database: Database,
- connect_args: dict[str, Any],
- uri: str,
username: str | None,
- access_token: str | None,
- ) -> None:
- """
- Update a configuration dictionary
- that can set the correct properties for impersonating users
-
- :param connect_args: the Database object
- :param connect_args: config to be updated
- :param uri: URI string
- :param username: Effective username
- :param access_token: Personal access token for OAuth2
- :return: None
- """
- url = make_url_safe(uri)
- backend_name = url.get_backend_name()
+ user_token: str | None,
+ url: URL,
+ engine_kwargs: dict[str, Any],
+ ) -> tuple[URL, dict[str, Any]]:
+ if username is None:
+ return url, engine_kwargs
- # Must be Presto connection, enable impersonation, and set optional param
- # auth=LDAP|KERBEROS
- # Set principal_username=$effective_username
- if backend_name == "presto" and username is not None:
+ backend_name = url.get_backend_name()
+ connect_args = engine_kwargs.setdefault("connect_args", {})
+ if backend_name == "presto":
connect_args["principal_username"] = username
+ return url, engine_kwargs
+
@classmethod
def get_table_names(
cls,
diff --git a/superset/db_engine_specs/starrocks.py b/superset/db_engine_specs/starrocks.py
index 6f54329d6f..4a6563aef3 100644
--- a/superset/db_engine_specs/starrocks.py
+++ b/superset/db_engine_specs/starrocks.py
@@ -203,25 +203,6 @@ class StarRocksEngineSpec(MySQLEngineSpec):
return parse.unquote(database.split(".")[1])
- @classmethod
- def get_url_for_impersonation(
- cls,
- url: URL,
- impersonate_user: bool,
- username: Union[str, None] = None,
- access_token: Union[str, None] = None,
- ) -> URL:
- """
- Return a modified URL with the username set.
-
- :param url: SQLAlchemy URL object
- :param impersonate_user: Flag indicating if impersonation is enabled
- :param username: Effective username
- :param access_token: Personal access token
- """
- # Leave URL unchanged. We will impersonate with the pre-query below.
- return url
-
@classmethod
def get_prequeries(
cls,
diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py
index beb7c0a604..bbe6ec0f8e 100644
--- a/superset/db_engine_specs/trino.py
+++ b/superset/db_engine_specs/trino.py
@@ -30,7 +30,6 @@ from sqlalchemy.exc import NoSuchTableError
from superset import db
from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY, USER_AGENT
-from superset.databases.utils import make_url_safe
from superset.db_engine_specs.base import BaseEngineSpec, convert_inspector_columns
from superset.db_engine_specs.exceptions import (
SupersetDBAPIConnectionError,
@@ -130,56 +129,26 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
return metadata
@classmethod
- def update_impersonation_config( # pylint: disable=too-many-arguments
+ def impersonate_user(
cls,
database: Database,
- connect_args: dict[str, Any],
- uri: str,
username: str | None,
- access_token: str | None,
- ) -> None:
- """
- Update a configuration dictionary
- that can set the correct properties for impersonating users
- :param database: the Database object
- :param connect_args: config to be updated
- :param uri: URI string
- :param username: Effective username
- :param access_token: Personal access token for OAuth2
- :return: None
- """
- url = make_url_safe(uri)
- backend_name = url.get_backend_name()
+ user_token: str | None,
+ url: URL,
+ engine_kwargs: dict[str, Any],
+ ) -> tuple[URL, dict[str, Any]]:
+ if username is None:
+ return url, engine_kwargs
- # Must be Trino connection, enable impersonation, and set optional param
- # auth=LDAP|KERBEROS
- # Set principal_username=$effective_username
- if backend_name == "trino" and username is not None:
+ backend_name = url.get_backend_name()
+ connect_args = engine_kwargs.setdefault("connect_args", {})
+ if backend_name == "trino":
connect_args["user"] = username
- if access_token is not None:
+ if user_token is not None:
http_session = requests.Session()
- http_session.headers.update({"Authorization": f"Bearer {access_token}"})
+ http_session.headers.update({"Authorization": f"Bearer {user_token}"})
connect_args["http_session"] = http_session
- @classmethod
- def get_url_for_impersonation(
- cls,
- url: URL,
- impersonate_user: bool,
- username: str | None,
- access_token: str | None,
- ) -> URL:
- """
- Return a modified URL with the username set.
-
- :param access_token: Personal access token for OAuth2
- :param url: SQLAlchemy URL object
- :param impersonate_user: Flag indicating if impersonation is enabled
- :param username: Effective username
- """
- # Do nothing and let update_impersonation_config take care of impersonation
- return url
-
@classmethod
def get_allow_cost_estimate(cls, extra: dict[str, Any]) -> bool:
return True
diff --git a/superset/models/core.py b/superset/models/core.py
index 96a1953fae..55a78355fb 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -478,11 +478,12 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
self.db_engine_spec.validate_database_uri(sqlalchemy_url)
extra = self.get_extra()
- params = extra.get("engine_params", {})
+ engine_kwargs = extra.get("engine_params", {})
if nullpool:
- params["poolclass"] = NullPool
- connect_args = params.get("connect_args", {})
+ engine_kwargs["poolclass"] = NullPool
+ connect_args = engine_kwargs.setdefault("connect_args", {})
+ # modify URL/args for a specific catalog/schema
sqlalchemy_url, connect_args = self.db_engine_spec.adjust_engine_params(
uri=sqlalchemy_url,
connect_args=connect_args,
@@ -507,33 +508,19 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
if oauth2_config and hasattr(g, "user") and hasattr(g.user, "id")
else None
)
- # If using MySQL or Presto for example, will set url.username
- # If using Hive, will not do anything yet since that relies on a
- # configuration parameter instead.
- sqlalchemy_url = self.db_engine_spec.get_url_for_impersonation(
- sqlalchemy_url,
- self.impersonate_user,
- effective_username,
- access_token,
- )
-
masked_url = self.get_password_masked_url(sqlalchemy_url)
logger.debug("Database._get_sqla_engine(). Masked URL: %s", str(masked_url))
if self.impersonate_user:
- # PR #30674 changed the signature of the method to include database.
- # This ensures that the change is backwards compatible
- args = [connect_args, str(sqlalchemy_url), effective_username, access_token]
- args = self.add_database_to_signature(
- self.db_engine_spec.update_impersonation_config,
- args,
+ self.db_engine_spec.impersonate_user(
+ self,
+ effective_username,
+ access_token,
+ sqlalchemy_url,
+ engine_kwargs,
)
- self.db_engine_spec.update_impersonation_config(*args)
-
- if connect_args:
- params["connect_args"] = connect_args
- self.update_params_from_encrypted_extra(params)
+ self.update_params_from_encrypted_extra(engine_kwargs)
if DB_CONNECTION_MUTATOR:
if not source and request and request.referrer:
@@ -544,15 +531,15 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
elif "/sqllab/" in request.referrer:
source = utils.QuerySource.SQL_LAB
- sqlalchemy_url, params = DB_CONNECTION_MUTATOR(
+ sqlalchemy_url, engine_kwargs = DB_CONNECTION_MUTATOR(
sqlalchemy_url,
- params,
+ engine_kwargs,
effective_username,
security_manager,
source,
)
try:
- return create_engine(sqlalchemy_url, **params)
+ return create_engine(sqlalchemy_url, **engine_kwargs)
except Exception as ex:
raise self.db_engine_spec.get_dbapi_mapped_exception(ex) from ex