This is an automated email from the ASF dual-hosted git repository.
alexandrusoare pushed a commit to branch
alexandrusoare/fix/pending-rollback-chart-data-impersonation
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to
refs/heads/alexandrusoare/fix/pending-rollback-chart-data-impersonation by this
push:
new 90f059418e7 fixing database jamming issues
90f059418e7 is described below
commit 90f059418e7d1e65bd19ca20b1277547830d1820
Author: alexandrusoare <[email protected]>
AuthorDate: Fri Sep 25 15:33:52 2026 +0300
fixing database jamming issues
---
superset/db_engine_specs/gsheets.py | 12 ++--
superset/db_engine_specs/snowflake.py | 8 +--
superset/db_engine_specs/starrocks.py | 12 +---
superset/models/core.py | 49 ++++++++++++--
tests/unit_tests/db_engine_specs/test_gsheets.py | 7 +-
tests/unit_tests/db_engine_specs/test_starrocks.py | 74 +++++----------------
tests/unit_tests/models/core_test.py | 72 ++++++++++++++++++++
tests/unit_tests/utils/test_database.py | 76 +++++++++++++++++++++-
8 files changed, 220 insertions(+), 90 deletions(-)
diff --git a/superset/db_engine_specs/gsheets.py
b/superset/db_engine_specs/gsheets.py
index a75c46eac71..6b9642b283e 100644
--- a/superset/db_engine_specs/gsheets.py
+++ b/superset/db_engine_specs/gsheets.py
@@ -38,7 +38,7 @@ from sqlalchemy.engine import create_engine
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.engine.url import URL
-from superset import db, security_manager
+from superset import db
from superset.databases.schemas import encrypted_field_properties,
EncryptedString
from superset.db_engine_specs.base import DatabaseCategory
from superset.db_engine_specs.shillelagh import ShillelaghEngineSpec
@@ -249,9 +249,13 @@ class GSheetsEngineSpec(ShillelaghEngineSpec):
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})
+ # Resolved from the database rather than from ``username``: with
+ # ``IMPERSONATE_WITH_EMAIL_PREFIX`` enabled the caller has already
+ # substituted the email prefix into ``username``, so looking it up
+ # here as if it were still the login finds nothing whenever the two
+ # differ, silently leaving the subject unset.
+ if email := database.get_impersonation_email():
+ url = url.update_query_dict({"subject": email})
if user_token:
url = url.update_query_dict({"access_token": user_token})
diff --git a/superset/db_engine_specs/snowflake.py
b/superset/db_engine_specs/snowflake.py
index 63ae8886228..53fdec2b85a 100644
--- a/superset/db_engine_specs/snowflake.py
+++ b/superset/db_engine_specs/snowflake.py
@@ -37,7 +37,7 @@ from sqlalchemy.exc import DatabaseError as
SqlalchemyDatabaseError
from sqlalchemy.sql import quoted_name
from sqlalchemy.sql.elements import ColumnElement
-from superset import is_feature_enabled, security_manager
+from superset import is_feature_enabled
from superset.constants import TimeGrain
from superset.databases.utils import make_url_safe
from superset.db_engine_specs.base import (
@@ -353,10 +353,8 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
# leaving the default/service-account username paired
# with this user's OAuth token. Use it as given.
url = url.set(username=username)
- else:
- user = security_manager.find_user(username=username)
- if user and user.email:
- url = url.set(username=user.email)
+ elif email := database.get_impersonation_email():
+ url = url.set(username=email)
url = url.update_query_dict({"token": user_token})
diff --git a/superset/db_engine_specs/starrocks.py
b/superset/db_engine_specs/starrocks.py
index 9fcaab308d9..4ad6eba429f 100644
--- a/superset/db_engine_specs/starrocks.py
+++ b/superset/db_engine_specs/starrocks.py
@@ -28,11 +28,9 @@ from sqlalchemy.engine.url import URL
from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.sql.type_api import TypeEngine
-from superset import is_feature_enabled
from superset.db_engine_specs.base import DatabaseCategory
from superset.db_engine_specs.mysql import MySQLEngineSpec
from superset.errors import SupersetErrorType
-from superset.extensions import security_manager
from superset.models.core import Database
from superset.utils.core import GenericDataType
@@ -426,15 +424,9 @@ class StarRocksEngineSpec(MySQLEngineSpec):
For StarRocks with user impersonation enabled, returns an EXECUTE AS
statement.
"""
if database.impersonate_user:
- username = database.get_effective_user(database.url_object)
-
- if username:
- effective_username = username
- if is_feature_enabled("IMPERSONATE_WITH_EMAIL_PREFIX"):
- user = security_manager.find_user(username=username)
- if user and user.email:
- effective_username = user.email.split("@", 1)[0]
+ effective_username = database.get_impersonation_username()
+ if effective_username:
escaped = effective_username.replace('"', '""')
return [f'EXECUTE AS "{escaped}" WITH NO REVERT;']
diff --git a/superset/models/core.py b/superset/models/core.py
index 24ebfd5c735..65811d6794b 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -537,6 +537,49 @@ class Database(CoreDatabase, AuditMixinNullable,
ImportExportMixin): # pylint:
else None
)
+ def get_impersonation_email(self, object_url: URL | None = None) -> str |
None:
+ """
+ Get the email address of the user being impersonated.
+
+ Resolves the effective login against the metadata database. DB engine
+ specs that need the email (or a part of it) to build a connection must
+ call this rather than looking the login up themselves: the lookup is a
+ metadata-DB read that can inherit a failed transaction from earlier in
+ the request, and centralising it keeps that handling in one place.
+
+ :param object_url: URL to read the login from; defaults to this
+ database's own URL
+ :return: The impersonated user's email, or ``None`` if there is no
+ effective user or the login has no email on record
+ """
+ username = self.get_effective_user(object_url or self.url_object)
+ if not username:
+ return None
+
+ user = find_user_for_impersonation(username)
+ return user.email if user and user.email else None
+
+ def get_impersonation_username(self, object_url: URL | None = None) -> str
| None:
+ """
+ Get the username to impersonate on the analytic database.
+
+ With ``IMPERSONATE_WITH_EMAIL_PREFIX`` enabled this is the local part
of
+ the user's email address; otherwise it is the effective login. Falls
+ back to the login when the flag is on but the user has no email on
+ record, matching the behaviour of a connection made without the flag.
+
+ :param object_url: URL to read the login from; defaults to this
+ database's own URL
+ :return: The username to connect as, or ``None`` if there is no
+ effective user
+ """
+ username = self.get_effective_user(object_url or self.url_object)
+ if not username or not
is_feature_enabled("IMPERSONATE_WITH_EMAIL_PREFIX"):
+ return username
+
+ email = self.get_impersonation_email(object_url)
+ return email.split("@")[0] if email else username
+
@contextmanager
def get_sqla_engine( # pylint: disable=too-many-arguments
self,
@@ -662,11 +705,7 @@ class Database(CoreDatabase, AuditMixinNullable,
ImportExportMixin): # pylint:
)
engine_kwargs["connect_args"] = connect_args
- effective_username = self.get_effective_user(sqlalchemy_url)
- if effective_username and
is_feature_enabled("IMPERSONATE_WITH_EMAIL_PREFIX"):
- user = find_user_for_impersonation(effective_username)
- if user and user.email:
- effective_username = user.email.split("@")[0]
+ effective_username = self.get_impersonation_username(sqlalchemy_url)
oauth2_config = self.get_oauth2_config()
access_token = (
diff --git a/tests/unit_tests/db_engine_specs/test_gsheets.py
b/tests/unit_tests/db_engine_specs/test_gsheets.py
index 5981caa85c6..5fadc258f27 100644
--- a/tests/unit_tests/db_engine_specs/test_gsheets.py
+++ b/tests/unit_tests/db_engine_specs/test_gsheets.py
@@ -571,13 +571,8 @@ def test_impersonate_user_username(mocker: MockerFixture)
-> None:
"""
from superset.db_engine_specs.gsheets import GSheetsEngineSpec
- user = mocker.MagicMock()
- user.email = "[email protected]"
- mocker.patch(
- "superset.db_engine_specs.gsheets.security_manager.find_user",
- return_value=user,
- )
database = mocker.MagicMock()
+ database.get_impersonation_email.return_value = "[email protected]"
assert GSheetsEngineSpec.impersonate_user(
database,
diff --git a/tests/unit_tests/db_engine_specs/test_starrocks.py
b/tests/unit_tests/db_engine_specs/test_starrocks.py
index 5590c9d70d7..5aa7b9ca8aa 100644
--- a/tests/unit_tests/db_engine_specs/test_starrocks.py
+++ b/tests/unit_tests/db_engine_specs/test_starrocks.py
@@ -35,7 +35,6 @@ from superset.db_engine_specs.starrocks import (
TINYINT,
)
from superset.utils.core import GenericDataType
-from tests.unit_tests.conftest import with_feature_flags
from tests.unit_tests.db_engine_specs.utils import assert_column_spec
@@ -158,7 +157,7 @@ def test_impersonation_username(mocker: MockerFixture) ->
None:
database = mocker.MagicMock()
database.impersonate_user = True
- database.get_effective_user.return_value = "alice"
+ database.get_impersonation_username.return_value = "alice"
assert StarRocksEngineSpec.impersonate_user(
database,
@@ -172,7 +171,9 @@ def test_impersonation_username(mocker: MockerFixture) ->
None:
'EXECUTE AS "alice" WITH NO REVERT;'
]
- database.get_effective_user.return_value = 'evil" WITH NO REVERT; DROP
TABLE x--'
+ database.get_impersonation_username.return_value = (
+ 'evil" WITH NO REVERT; DROP TABLE x--'
+ )
assert StarRocksEngineSpec.get_prequeries(database) == [
'EXECUTE AS "evil"" WITH NO REVERT; DROP TABLE x--" WITH NO REVERT;'
]
@@ -297,74 +298,29 @@ def test_adjust_engine_params_with_catalog(
assert returned_url.database == expected_database
-@with_feature_flags(IMPERSONATE_WITH_EMAIL_PREFIX=True)
-def test_get_prequeries_with_email_prefix(mocker: MockerFixture) -> None:
- """Test that get_prequeries uses email prefix when
IMPERSONATE_WITH_EMAIL_PREFIX"""
- from superset.db_engine_specs.starrocks import StarRocksEngineSpec
-
- user = mocker.MagicMock()
- user.email = "[email protected]"
- mocker.patch(
- "superset.db_engine_specs.starrocks.security_manager.find_user",
- return_value=user,
- )
-
- database = mocker.MagicMock()
- database.impersonate_user = True
- database.url_object = make_url("starrocks://localhost:9030/")
- database.get_effective_user.return_value = "[email protected]"
-
- assert StarRocksEngineSpec.get_prequeries(database) == [
- 'EXECUTE AS "alice" WITH NO REVERT;'
- ]
-
-
-@with_feature_flags(IMPERSONATE_WITH_EMAIL_PREFIX=True)
-def test_get_prequeries_with_email_prefix_dotted_local_part(
+def test_get_prequeries_defers_impersonation_resolution(
mocker: MockerFixture,
) -> None:
- """Test that get_prequeries uses email prefix when
IMPERSONATE_WITH_EMAIL_PREFIX"""
- from superset.db_engine_specs.starrocks import StarRocksEngineSpec
-
- user = mocker.MagicMock()
- user.email = "[email protected]"
- mocker.patch(
- "superset.db_engine_specs.starrocks.security_manager.find_user",
- return_value=user,
- )
-
- database = mocker.MagicMock()
- database.impersonate_user = True
- database.url_object = make_url("starrocks://localhost:9030/")
- database.get_effective_user.return_value = "[email protected]"
-
- assert StarRocksEngineSpec.get_prequeries(database) == [
- 'EXECUTE AS "alice.doe" WITH NO REVERT;'
- ]
-
+ """
+ Test that `get_prequeries` impersonates whoever `Database` resolves.
-@with_feature_flags(IMPERSONATE_WITH_EMAIL_PREFIX=True)
-def
test_get_prequeries_with_email_prefix_from_user_email_when_effective_user_differs(
- mocker: MockerFixture,
-) -> None:
- """Use looked-up user.email local part when effective username is
different."""
+ Deriving the name here instead (re-reading the effective user and looking
up
+ its email) would duplicate work `Database._get_sqla_engine` has already
done
+ for the same connection, and would put a second unguarded metadata-DB read
+ on the query path. The prefix substitution itself is covered by the
+ `Database.get_impersonation_username` tests.
+ """
from superset.db_engine_specs.starrocks import StarRocksEngineSpec
- user = mocker.MagicMock()
- user.email = "[email protected]"
- mocker.patch(
- "superset.db_engine_specs.starrocks.security_manager.find_user",
- return_value=user,
- )
-
database = mocker.MagicMock()
database.impersonate_user = True
database.url_object = make_url("starrocks://localhost:9030/")
- database.get_effective_user.return_value = "alice"
+ database.get_impersonation_username.return_value = "alice.doe"
assert StarRocksEngineSpec.get_prequeries(database) == [
'EXECUTE AS "alice.doe" WITH NO REVERT;'
]
+ database.get_effective_user.assert_not_called()
def test_time_grain_expressions_inherit_mysql() -> None:
diff --git a/tests/unit_tests/models/core_test.py
b/tests/unit_tests/models/core_test.py
index c614e44f401..6a8a7e8150b 100644
--- a/tests/unit_tests/models/core_test.py
+++ b/tests/unit_tests/models/core_test.py
@@ -844,6 +844,78 @@ def test_get_sqla_engine_user_impersonation_email(mocker:
MockerFixture) -> None
)
+@with_feature_flags(IMPERSONATE_WITH_EMAIL_PREFIX=True)
+def test_get_impersonation_username_uses_email_prefix(mocker: MockerFixture)
-> None:
+ """
+ Test that the impersonated username is the local part of the user's email.
+
+ The login and the email prefix commonly differ, so the lookup result is
+ used rather than the login it was resolved from.
+ """
+ user = mocker.MagicMock()
+ user.email = "[email protected]"
+ mocker.patch(
+ "superset.models.core.find_user_for_impersonation",
+ return_value=user,
+ )
+ mocker.patch("superset.models.core.get_username", return_value="alice")
+
+ database = Database(
+ database_name="my_db",
+ sqlalchemy_uri="trino://",
+ impersonate_user=True,
+ )
+
+ assert database.get_impersonation_username() == "alice.doe"
+ assert database.get_impersonation_email() == "[email protected]"
+
+
+@with_feature_flags(IMPERSONATE_WITH_EMAIL_PREFIX=True)
+def test_get_impersonation_username_without_email(mocker: MockerFixture) ->
None:
+ """
+ Test that a user with no email on record falls back to the login.
+
+ This matches how the connection would be made with the flag off, rather
+ than impersonating nobody.
+ """
+ user = mocker.MagicMock()
+ user.email = None
+ mocker.patch(
+ "superset.models.core.find_user_for_impersonation",
+ return_value=user,
+ )
+ mocker.patch("superset.models.core.get_username", return_value="alice")
+
+ database = Database(
+ database_name="my_db",
+ sqlalchemy_uri="trino://",
+ impersonate_user=True,
+ )
+
+ assert database.get_impersonation_username() == "alice"
+ assert database.get_impersonation_email() is None
+
+
+def test_get_impersonation_username_without_flag(mocker: MockerFixture) ->
None:
+ """
+ Test that the login is used verbatim when the flag is off.
+
+ No lookup should happen at all: it would be a metadata-DB read on the query
+ path whose result is then discarded.
+ """
+ find_user =
mocker.patch("superset.models.core.find_user_for_impersonation")
+ mocker.patch("superset.models.core.get_username", return_value="alice")
+
+ database = Database(
+ database_name="my_db",
+ sqlalchemy_uri="trino://",
+ impersonate_user=True,
+ )
+
+ assert database.get_impersonation_username() == "alice"
+ find_user.assert_not_called()
+
+
def test_get_sqla_engine_registers_prequery_event_listener(
app_context: None,
mocker: MockerFixture,
diff --git a/tests/unit_tests/utils/test_database.py
b/tests/unit_tests/utils/test_database.py
index 22bea17765e..3743938a417 100644
--- a/tests/unit_tests/utils/test_database.py
+++ b/tests/unit_tests/utils/test_database.py
@@ -17,12 +17,15 @@
"""Tests for superset.utils.database module."""
import pytest
+from pytest_mock import MockerFixture
from sqlalchemy import Sequence
from sqlalchemy.dialects import mysql, postgresql
+from sqlalchemy.exc import PendingRollbackError
from sqlalchemy.schema import CreateSequence
from sqlalchemy.sql.compiler import DDLCompiler
-from superset.utils.database import apply_mariadb_ddl_fix
+from superset.exceptions import SupersetErrorException
+from superset.utils.database import apply_mariadb_ddl_fix,
find_user_for_impersonation
@pytest.fixture(scope="module", autouse=True)
@@ -51,3 +54,74 @@ def test_nocycle_fix_not_applied_for_postgresql():
result = compiler.visit_create_sequence(CreateSequence(seq))
assert "NO CYCLE" in result
+
+
+def test_find_user_for_impersonation(mocker: MockerFixture) -> None:
+ """Test that a healthy session resolves the user without rolling back."""
+ user = mocker.MagicMock()
+ find_user = mocker.patch(
+ "superset.extensions.security_manager.find_user",
+ return_value=user,
+ )
+ session = mocker.patch("superset.db.session")
+
+ assert find_user_for_impersonation("alice") is user
+ find_user.assert_called_once_with(username="alice")
+ session.rollback.assert_not_called()
+
+
+def test_find_user_for_impersonation_retries_after_rollback(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Test that a poisoned session is rolled back and the lookup retried.
+
+ An earlier failure in the same request leaves ``db.session`` in a failed
+ transaction, so this lookup reports ``PendingRollbackError`` for a fault
+ that has nothing to do with it. Rolling back makes the session usable
+ again, and the retry returns the user the caller asked for.
+ """
+ user = mocker.MagicMock()
+ find_user = mocker.patch(
+ "superset.extensions.security_manager.find_user",
+ side_effect=[PendingRollbackError("poisoned"), user],
+ )
+ session = mocker.patch("superset.db.session")
+
+ assert find_user_for_impersonation("alice") is user
+ session.rollback.assert_called_once()
+ assert find_user.call_count == 2
+
+
+def test_find_user_for_impersonation_raises_when_retry_fails(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Test that a lookup failing past the retry raises rather than falling back.
+
+ The resolved name becomes the identity the analytic database connects as,
+ so degrading to the un-resolved login would silently run the query as a
+ different principal. Failing loudly is the safe outcome.
+ """
+ mocker.patch(
+ "superset.extensions.security_manager.find_user",
+ side_effect=PendingRollbackError("still poisoned"),
+ )
+ session = mocker.patch("superset.db.session")
+
+ with pytest.raises(SupersetErrorException):
+ find_user_for_impersonation("alice")
+
+ session.rollback.assert_called_once()
+
+
+def test_find_user_for_impersonation_unknown_login(mocker: MockerFixture) ->
None:
+ """Test that an unknown login is reported as absent, not as a failure."""
+ find_user = mocker.patch(
+ "superset.extensions.security_manager.find_user",
+ return_value=None,
+ )
+ mocker.patch("superset.db.session")
+
+ assert find_user_for_impersonation("nobody") is None
+ find_user.assert_called_once_with(username="nobody")
