You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by as...@apache.org on 2026/09/21 06:21:30 UTC
(superset) 01/01: fix(database): survive a poisoned session in the impersonation user lookup
This is an automated email from the ASF dual-hosted git repository.
alexandrusoare pushed a commit to branch agorithm/fix-pendingrollback-impersonation
in repository https://gitbox.apache.org/repos/asf/superset.git
commit e9c67a1b79dbd98d57b1a78075882ad5a91a98e7
Author: Alexandru Soare <al...@ext.preset.io>
AuthorDate: Mon Sep 21 06:21:17 2026 +0000
fix(database): survive a poisoned session in the impersonation user lookup
`Database._get_sqla_engine()` resolves the email prefix for
`IMPERSONATE_WITH_EMAIL_PREFIX` with a `security_manager.find_user()`
query against the metadata session. When an earlier statement in the same
request already failed, that query raises `PendingRollbackError` in place
of the original error. Several engine-building paths acquire the engine
outside their own error handling - `ExploreMixin.query()` compiles the
query (and so builds an engine) before its broad `except` - so the error
escapes as a 500 on the chart data endpoint.
Roll back and retry the lookup once, and fall back to the unresolved
username if it still fails, mirroring the best-effort rollback in
`superset.sql_lab.get_query()`.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
---
superset/models/core.py | 54 ++++++++++++++-
tests/unit_tests/models/core_test.py | 127 +++++++++++++++++++++++++++++++++++
2 files changed, 179 insertions(+), 2 deletions(-)
diff --git a/superset/models/core.py b/superset/models/core.py
index e6953db1c45..8c2c2611617 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -55,7 +55,7 @@ from sqlalchemy import (
from sqlalchemy.engine import Connection, Dialect, Engine
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.engine.url import URL
-from sqlalchemy.exc import NoSuchModuleError
+from sqlalchemy.exc import NoSuchModuleError, SQLAlchemyError
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.orm import Mapper, relationship
from sqlalchemy.pool import NullPool
@@ -106,6 +106,7 @@ _ENGINE_CACHE: dict[tuple[int, str, str], Engine] = {}
_ENGINE_CACHE_LOCK = threading.Lock()
if TYPE_CHECKING:
+ from flask_appbuilder.security.sqla.models import User
from superset_core.queries.types import AsyncQueryHandle, QueryOptions, QueryResult
from superset.models.sql_lab import Query
@@ -171,6 +172,55 @@ def clear_bootstrap_cache(
logger.warning("Failed to clear theme bootstrap cache: %s", ex)
+def _find_user_for_impersonation(username: str) -> User | None:
+ """
+ Look up a user for email-prefix impersonation, tolerating a broken session.
+
+ The lookup is a metadata-database query, so a statement that already failed
+ earlier in the same request leaves the session needing a rollback and makes
+ this query raise ``PendingRollbackError`` in place of the original error.
+ Rolling back and retrying once both resolves the email prefix and unpoisons
+ the session, instead of turning an unrelated earlier fault into a failed
+ response on every path that builds an engine outside its own error handling
+ (the chart data endpoint among them).
+
+ The rollback is best effort, mirroring ``superset.sql_lab.get_query``: a
+ session too broken to roll back must not replace the outcome the caller
+ actually cares about. Returns ``None`` when the user can't be resolved, so
+ the caller impersonates the unresolved username — the same identity a
+ deployment without ``IMPERSONATE_WITH_EMAIL_PREFIX`` would use.
+ """
+ try:
+ return security_manager.find_user(username=username)
+ except SQLAlchemyError:
+ logger.warning(
+ "Could not look up user %s for impersonation; retrying after a "
+ "session rollback",
+ username,
+ exc_info=True,
+ )
+
+ try:
+ db.session.rollback() # pylint: disable=consider-using-transaction
+ except SQLAlchemyError:
+ logger.warning(
+ "Failed to roll back the session before retrying the impersonation "
+ "user lookup",
+ exc_info=True,
+ )
+
+ try:
+ return security_manager.find_user(username=username)
+ except SQLAlchemyError:
+ logger.warning(
+ "Could not look up user %s for impersonation; impersonating the "
+ "unresolved username instead",
+ username,
+ exc_info=True,
+ )
+ return None
+
+
class ConfigurationMethod(StrEnum):
SQLALCHEMY_FORM = "sqlalchemy_form"
DYNAMIC_FORM = "dynamic_form"
@@ -663,7 +713,7 @@ class Database(CoreDatabase, AuditMixinNullable, ImportExportMixin): # pylint:
effective_username = self.get_effective_user(sqlalchemy_url)
if effective_username and is_feature_enabled("IMPERSONATE_WITH_EMAIL_PREFIX"):
- user = security_manager.find_user(username=effective_username)
+ user = _find_user_for_impersonation(effective_username)
if user and user.email:
effective_username = user.email.split("@")[0]
diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py
index c614e44f401..65570f89e66 100644
--- a/tests/unit_tests/models/core_test.py
+++ b/tests/unit_tests/models/core_test.py
@@ -34,6 +34,7 @@ from sqlalchemy import (
)
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.engine.url import make_url
+from sqlalchemy.exc import PendingRollbackError, SQLAlchemyError
from sqlalchemy.orm.session import Session
from sqlalchemy.sql import Select
@@ -844,6 +845,132 @@ def test_get_sqla_engine_user_impersonation_email(mocker: MockerFixture) -> None
)
+@with_feature_flags(IMPERSONATE_WITH_EMAIL_PREFIX=True)
+def test_get_sqla_engine_user_impersonation_email_poisoned_session(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Test the email lookup survives a metadata session left in a failed transaction.
+
+ An earlier failed query in the same request leaves ``db.session`` needing a
+ rollback, so the lookup raises ``PendingRollbackError``. Rolling back and
+ retrying must resolve the email prefix instead of failing the whole request.
+ """
+ from superset.models.core import Database
+
+ user = mocker.MagicMock()
+ user.email = "alice.doe@example.org"
+ find_user = mocker.patch(
+ "superset.models.core.security_manager.find_user",
+ side_effect=[
+ PendingRollbackError(
+ "Can't reconnect until invalid transaction is rolled back"
+ ),
+ user,
+ ],
+ )
+ session = mocker.patch("superset.models.core.db.session")
+ mocker.patch("superset.models.core.get_username", return_value="alice")
+
+ create_engine_mock = mocker.patch(
+ "superset.models.core.create_engine",
+ return_value=create_engine("sqlite://"),
+ )
+
+ database = Database(
+ database_name="my_db",
+ sqlalchemy_uri="trino://",
+ impersonate_user=True,
+ )
+ database._get_sqla_engine(nullpool=False)
+
+ session.rollback.assert_called_once()
+ assert find_user.call_count == 2
+ create_engine_mock.assert_called_with(
+ make_url("trino:///"),
+ connect_args={"user": "alice.doe", "source": "Apache Superset"},
+ )
+
+
+@with_feature_flags(IMPERSONATE_WITH_EMAIL_PREFIX=True)
+def test_get_sqla_engine_user_impersonation_email_lookup_keeps_failing(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Test the engine is still built when the email lookup can't be recovered.
+
+ When the retry fails too the username stays unresolved, which impersonates
+ the same identity the flag-less configuration would use, rather than
+ erroring out of the request.
+ """
+ from superset.models.core import Database
+
+ find_user = mocker.patch(
+ "superset.models.core.security_manager.find_user",
+ side_effect=SQLAlchemyError("metadata database is unreachable"),
+ )
+ mocker.patch("superset.models.core.db.session")
+ mocker.patch("superset.models.core.get_username", return_value="alice")
+
+ create_engine_mock = mocker.patch(
+ "superset.models.core.create_engine",
+ return_value=create_engine("sqlite://"),
+ )
+
+ database = Database(
+ database_name="my_db",
+ sqlalchemy_uri="trino://",
+ impersonate_user=True,
+ )
+ database._get_sqla_engine(nullpool=False)
+
+ assert find_user.call_count == 2
+ create_engine_mock.assert_called_with(
+ make_url("trino:///"),
+ connect_args={"user": "alice", "source": "Apache Superset"},
+ )
+
+
+@with_feature_flags(IMPERSONATE_WITH_EMAIL_PREFIX=True)
+def test_get_sqla_engine_user_impersonation_email_rollback_fails(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Test a session too broken to roll back doesn't fail the engine creation.
+
+ The rollback is best effort: when it raises as well, the unresolved username
+ is used and the rollback error never replaces the request's own outcome.
+ """
+ from superset.models.core import Database
+
+ mocker.patch(
+ "superset.models.core.security_manager.find_user",
+ side_effect=PendingRollbackError(
+ "Can't reconnect until invalid transaction is rolled back"
+ ),
+ )
+ session = mocker.patch("superset.models.core.db.session")
+ session.rollback.side_effect = SQLAlchemyError("connection is closed")
+ mocker.patch("superset.models.core.get_username", return_value="alice")
+
+ create_engine_mock = mocker.patch(
+ "superset.models.core.create_engine",
+ return_value=create_engine("sqlite://"),
+ )
+
+ database = Database(
+ database_name="my_db",
+ sqlalchemy_uri="trino://",
+ impersonate_user=True,
+ )
+ database._get_sqla_engine(nullpool=False)
+
+ create_engine_mock.assert_called_with(
+ make_url("trino:///"),
+ connect_args={"user": "alice", "source": "Apache Superset"},
+ )
+
+
def test_get_sqla_engine_registers_prequery_event_listener(
app_context: None,
mocker: MockerFixture,