EnxDev commented on code in PR #44653:
URL: https://github.com/apache/superset/pull/44653#discussion_r4121495703
##########
superset/common/utils/query_cache_manager.py:
##########
@@ -161,6 +162,12 @@ def set_query_result(
region=region,
)
except Exception as ex: # pylint: disable=broad-except
+ # `self.set` writes through the cache backend, which is a
+ # metadata-DB write under `SupersetMetastoreCache`. Roll back so a
+ # failed write doesn't poison the rest of this request — the caller
+ # carries on after this (e.g. event logging, or a second query
+ # object in the same `queries` loop).
+ db.session.rollback() # pylint: disable=consider-using-transaction
Review Comment:
`SupersetMetastoreCache.add` and `compare_and_set` already roll back on
failure, but `set` and `get` don't. That's where the poisoned session comes
from.
Could the rollback go into those two methods instead? Then every metastore
cache caller is covered, not just these three sites, and a Redis timeout on a
healthy session doesn't trigger a rollback here or at line 204.
##########
superset/utils/database.py:
##########
@@ -123,6 +127,58 @@ def warm_and_release_connection(instance: Any,
*relationships: str) -> None:
session.expire_on_commit = True
+def find_user_for_impersonation(username: str) -> User | None:
+ """
+ Resolve the login backing an impersonated database session.
+
+ ``find_user`` is a metadata-DB read, so it inherits any failed transaction
+ left behind earlier in the request and reports ``PendingRollbackError``
+ instead of the original fault — blaming this lookup for an unrelated
+ failure. Roll back and retry once so a poisoned session doesn't cost us the
+ lookup.
+
+ The resolved value becomes the identity the analytic database connects as,
+ so a lookup that still fails must not degrade to the un-resolved login:
+ that would silently query as a different principal than the one being
+ impersonated. Raise instead.
+
+ :param username: the Superset login to resolve
+ :return: the matching user, or ``None`` if no such login exists
+ :raises SupersetErrorException: if the lookup fails even after a rollback
+ """
+ # pylint: disable=import-outside-toplevel
+ from superset import db
+ from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+ from superset.exceptions import SupersetErrorException
+ from superset.extensions import security_manager
+
+ try:
+ return security_manager.find_user(username=username)
+ except SQLAlchemyError:
Review Comment:
The retry only exists for the poisoned-session case, but this catches any
`SQLAlchemyError`. That includes an `IntegrityError` from autoflushing the
caller's own pending writes, which then get rolled back without the caller
knowing.
Would `except PendingRollbackError:` be enough here? Anything else would
keep propagating the way it did before.
##########
superset/utils/rls.py:
##########
@@ -329,6 +329,12 @@ def collect_rls_predicates_for_sql(
}
)
except Exception:
+ # The block above is not only SQL parsing: `get_predicates_for_table`
+ # queries `db.session` and `get_default_catalog()` builds an engine, so
+ # a caught DB error can leave db.session in "pending rollback" state,
+ # which would poison unrelated queries later in this request.
+ db.session.rollback() # pylint: disable=consider-using-transaction
Review Comment:
None of the five new rollback sites have a test, and they're the actual fix.
The impersonation tests only cover the symptom.
A test per site that makes the guarded call raise and asserts
`db.session.rollback` was called would do it. For this one, the existing
parse-failure test in `rls_test.py` could just add that assertion.
##########
tests/unit_tests/db_engine_specs/test_starrocks.py:
##########
@@ -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
Review Comment:
Nit, take it or leave it. `get_impersonation_username()` still re-reads the
effective user and does its own lookup, and `get_prequeries` runs before
`_get_sqla_engine`, so with the flag on it's still two lookups per engine build.
The docstring says this path avoids that. Maybe just say the resolution
lives on `Database` now.
##########
superset/db_engine_specs/gsheets.py:
##########
@@ -250,9 +250,13 @@ def impersonate_user(
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():
Review Comment:
This fixes the prefix bug, but the updated test mocks
`get_impersonation_email` on a MagicMock, so nothing would catch it coming back.
Worth adding a test with a real `Database`, `IMPERSONATE_WITH_EMAIL_PREFIX`
on, login `alice` and email `[email protected]`, asserting `subject` is the
full email?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
