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:35 UTC
(superset) branch improve-impersonate-user created (now 481a1b7b93)
This is an automated email from the ASF dual-hosted git repository.
beto pushed a change to branch improve-impersonate-user
in repository https://gitbox.apache.org/repos/asf/superset.git
at 481a1b7b93 chore: simplify user impersonation
This branch includes the following new commits:
new 481a1b7b93 chore: simplify user impersonation
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(superset) 01/01: chore: simplify user impersonation
Posted by be...@apache.org.
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