diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py index e87065bbedbf..93520e3bf63b 100644 --- a/superset/commands/streaming_export/base.py +++ b/superset/commands/streaming_export/base.py @@ -24,13 +24,15 @@ import logging import time from abc import abstractmethod -from contextlib import contextmanager +from contextvars import copy_context from decimal import Decimal from numbers import Real -from typing import Any, Callable, Generator, Iterable +from typing import Any, Callable, Generator, Iterable, Iterator, TypeVar -from flask import current_app as app, g, has_app_context -from sqlalchemy import text +from flask import current_app as app, g +from flask.globals import _cv_app, _cv_request +from sqlalchemy import inspect, text +from sqlalchemy.orm import InstanceState from werkzeug.local import LocalProxy from superset import db @@ -39,57 +41,55 @@ logger = logging.getLogger(__name__) +_T = TypeVar("_T") -def capture_g_context() -> dict[str, Any]: + +def _reload_if_detached(instance: _T) -> _T: + """ + Return ``instance`` loaded again by primary key if it is a detached ORM + instance, otherwise ``instance`` itself. + + The WSGI server only iterates a streaming body after the view has + returned, and by then Flask has torn the request down, removing the + metadata session that loaded the instance. Any attribute that session + expired (the event logger commits once the view returns), even the + primary key, would then raise ``DetachedInstanceError``. + + Reloading rather than merging means values read during the request are + never written back over changes made since, by the commit that follows + building the engine. """ - Snapshot ``flask.g`` so a streaming generator can replay it later. + state = inspect(instance, raiseerr=False) + if not isinstance(state, InstanceState) or not state.detached: + return instance - Values held on ``g`` may be request-bound ``LocalProxy`` objects rather - than plain values. The most important one is ``g.user``, which - Flask-AppBuilder's ``before_request`` sets to Flask-Login's - ``current_user``: that proxy resolves through ``has_request_context()`` - and therefore evaluates to ``None`` once the request context is gone, - which is exactly the situation the generator runs in. + reloaded = db.session.get(state.class_, state.identity) + if reloaded is None: + # Fail closed: e.g. without the acting user, impersonation would fall + # back to the connection's own credentials. + raise RuntimeError(f"{state.class_.__name__} {state.identity} no longer exists") + return reloaded - Copying such a proxy verbatim would hand the generator a ``g.user`` that - silently resolves to nobody, so resolve each proxy to the concrete object - it currently points at while the request context is still around. - Returns: - Dictionary of g attributes, with request-bound proxies resolved +def _reload_detached_acting_user() -> None: """ - if not has_app_context(): - return {} + Point ``g.user`` at the acting user loaded again if it was detached. - captured: dict[str, Any] = {} - for key, value in g._get_current_object().__dict__.items(): - if isinstance(value, LocalProxy): - try: - value = value._get_current_object() - except RuntimeError: - # Nothing is bound to the proxy, so there is no value worth - # carrying into the generator. - continue - captured[key] = value - return captured - - -@contextmanager -def preserve_g_context( - captured_g: dict[str, Any], -) -> Generator[None, None, None]: + Users that are not ORM instances, such as embedded guest users, are left + as they are. """ - Context manager that restores captured flask.g attributes. + user = getattr(g, "user", None) + if isinstance(user, LocalProxy): + # Flask-AppBuilder sets ``g.user`` to Flask-Login's ``current_user``. + user = user._get_current_object() - This is needed for streaming responses where the generator runs in a new - app context but needs access to request-scoped data from the original request. - - Args: - captured_g: Dictionary of g attributes captured by capture_g_context() - """ - for key, value in captured_g.items(): - setattr(g, key, value) - yield + reloaded = _reload_if_detached(user) + if reloaded is user: + return + g.user = reloaded + if getattr(g, "_login_user", None) is user: + # Keep Flask-Login's ``current_user`` on the same, reloaded user. + g._login_user = reloaded def encode_csv_stream( @@ -141,7 +141,6 @@ def __init__(self, chunk_size: int = 1000): chunk_size: Number of rows to fetch per database query (default: 1000) """ self._chunk_size = chunk_size - self._current_app = app._get_current_object() @abstractmethod def _get_sql_and_database(self) -> tuple[str, Any, str | None, str | None]: @@ -318,8 +317,9 @@ def _execute_query_and_stream( decimal_separator = csv_export_config.get("decimal", ".") with db.session() as session: - # Merge database to prevent DetachedInstanceError - merged_database = session.merge(database) + # Load the database again rather than merging it, see + # _reload_if_detached. + merged_database = _reload_if_detached(database) # `is_split=True` mirrors the non-streaming download path exactly: # Database.get_df() -> _execute_sql_with_mutation_and_logging() @@ -338,6 +338,13 @@ def _execute_query_and_stream( with merged_database.get_sqla_engine( catalog=catalog, schema=schema ) as engine: + # Building the engine may have refreshed the user's OAuth2 + # token, which is only staged on the session. Commit it, as the + # buffered export's query logging does, instead of losing it + # when the session closes (a single-use refresh token would be + # spent), and so no metadata transaction is held open while the + # rows stream. + session.commit() # pylint: disable=consider-using-transaction with engine.connect() as connection: result_proxy = connection.execution_options( stream_results=True @@ -378,41 +385,95 @@ def _execute_query_and_stream( total_time, ) - def run(self) -> Callable[[], Generator[str, None, None]]: + def run(self) -> Callable[[], Iterator[str]]: """ Execute the streaming CSV export. Returns: - A callable that returns a generator yielding CSV data chunks as strings. - The callable is needed to maintain Flask app context during streaming. + A callable that returns an iterator yielding CSV data chunks as + strings. Call it while handling the request: the iterator then + runs the export inside that request's context, even though the + body is only streamed once the view has returned. Outside a + request, consume the iterator inside the caller's app context; + it then runs as the caller's ``g.user`` (set it with + ``override_user``). """ # Load all needed data while session is still active # to avoid DetachedInstanceError sql, database, catalog, schema = self._get_sql_and_database() limit = self._get_row_limit() - # Capture flask.g attributes to preserve request-scoped data - # when the streaming generator runs in a new app context. - captured_g = capture_g_context() def csv_generator() -> Generator[str, None, None]: """Generator that yields CSV data chunks.""" - with self._current_app.app_context(): - with preserve_g_context(captured_g): - try: - yield from self._execute_query_and_stream( - sql, database, limit, catalog, schema - ) - except Exception as e: - logger.error("Error in streaming CSV generator: %s", e) - import traceback - - logger.error("Traceback: %s", traceback.format_exc()) - - # Send error marker for frontend to detect - error_marker = ( - "__STREAM_ERROR__:Export failed. " - "Please try again in some time.\n" - ) - yield error_marker - - return csv_generator + try: + _reload_detached_acting_user() + yield from self._execute_query_and_stream( + sql, database, limit, catalog, schema + ) + except Exception as e: + logger.error("Error in streaming CSV generator: %s", e) + import traceback + + logger.error("Traceback: %s", traceback.format_exc()) + + # Send error marker for frontend to detect + error_marker = ( + "__STREAM_ERROR__:Export failed. Please try again in some time.\n" + ) + yield error_marker + + def stream() -> Iterator[str]: + """ + Return the CSV iterator, bound to the current request if any. + + The export runs inside the request it belongs to, as the buffered + export does, not merely inside an app context. Code on the + connection path depends on it: Snowflake only uses the user's + OAuth2 token inside a request, the OAuth2 redirect URI may be + built from it, and the query source passed to + DB_CONNECTION_MUTATOR is read from it. + + The request's contexts are recorded here and pushed again only + while the body is iterated, as ``flask.stream_with_context`` does + from Flask 3.1.2. Pushing them up front, as older releases do, + leaves them on the worker thread when the response never reaches + the server (e.g. an ``after_request`` hook fails), and later + requests on that thread would inherit this one's ``g`` and user. + """ + request_context = _cv_request.get(None) + if request_context is None: + return csv_generator() + app_context = _cv_app.get() + + def in_request() -> Generator[str, None, None]: + if _cv_request.get(None) is request_context: + # Pulled while the request is still being handled: its + # contexts are already active, and pushing them again here + # would leave a token Flask's own pop can't reset. + yield from csv_generator() + return + with app_context, request_context: + yield from csv_generator() + + def csv_generator_in_request() -> Generator[str, None, None]: + # Step the export in a contextvars Context of its own, so the + # request's contexts are only ever pushed there, never on the + # worker thread: a body the server abandons mid-stream (e.g. + # an error raised while encoding keeps it alive on a + # traceback) or closes from another thread can't leave them + # behind for the thread's next request to inherit. + context = copy_context() + inner = in_request() + try: + while True: + try: + chunk = context.run(next, inner) + except StopIteration: + return + yield chunk + finally: + context.run(inner.close) + + return csv_generator_in_request() + + return stream diff --git a/tests/unit_tests/commands/chart/streaming_export_command_test.py b/tests/unit_tests/commands/chart/streaming_export_command_test.py index 66ca9451af13..396553237a46 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -17,16 +17,26 @@ """Unit tests for Chart Streaming CSV Export Command.""" from types import SimpleNamespace -from typing import Any +from typing import Any, Iterator import pytest from flask import g +from flask.globals import _cv_request +from flask_appbuilder.security.sqla.models import User from flask_login import current_user from pytest_mock import MockerFixture +from sqlalchemy import inspect +from sqlalchemy.orm import Session +from sqlalchemy.orm.exc import DetachedInstanceError from superset.commands.chart.data.streaming_export_command import ( StreamingCSVExportCommand, ) +from superset.commands.streaming_export.base import ( + _reload_detached_acting_user, + _reload_if_detached, +) +from superset.security.guest_token import GuestUser from superset.utils.core import get_user_id, get_username @@ -64,7 +74,6 @@ def test_streaming_csv_export_command_init(mocker: MockerFixture) -> None: assert command._query_context == query_context assert command._chunk_size == 500 - assert command._current_app is not None def test_streaming_csv_export_command_default_chunk_size( @@ -386,18 +395,17 @@ def test_streaming_export_mutation_does_not_double_apply( ) -def test_streaming_export_keeps_acting_user_after_request_context_ends( +def test_streaming_export_runs_in_its_request_context_after_the_view_returns( app: Any, mocker: MockerFixture, ) -> None: """ - Flask-AppBuilder's ``before_request`` sets ``g.user = current_user``, a - Flask-Login ``LocalProxy`` that resolves to ``None`` outside a request - context. Werkzeug iterates a streaming response body *after* the request - context is popped, so copying that proxy into the generator's fresh app - context yields a user-less ``g``: ``Database._get_sqla_engine`` then sees - ``hasattr(g.user, "id") is False`` and builds the engine with - ``access_token=None``, silently dropping the per-user OAuth2 token. + The WSGI server iterates a streaming body only after the view has returned and + the request context has been popped. The export must still run as part of + that request: as the acting user, which Flask-AppBuilder exposes as + ``g.user = current_user`` (a Flask-Login proxy that resolves to ``None`` + outside a request), and with the request context active, since engine + specs such as Snowflake only use the user's OAuth2 token inside a request. """ _mock_db, query_context, datasource = _setup_chart_mocks(mocker) @@ -412,12 +420,17 @@ def test_streaming_export_keeps_acting_user_after_request_context_ends( engine = mocker.MagicMock() engine.connect.return_value = connection - # Record the acting user at the moment the engine is acquired -- the point - # where the real Database._get_sqla_engine resolves the OAuth2 access token. - seen: list[tuple[str | None, int | None]] = [] + # Record the request state at the moment the engine is acquired -- the + # point where the real Database._get_sqla_engine resolves the OAuth2 token. + # Compare against this request's own context rather than asking whether any + # request context is active, so a context another test left pushed can't + # make this pass (or fail). + view_request: dict[str, Any] = {} + seen: list[tuple[str | None, int | None, bool]] = [] def fake_get_sqla_engine(*args: Any, **kwargs: Any) -> Any: - seen.append((get_username(), get_user_id())) + in_view_request = _cv_request.get(None) is view_request["context"] + seen.append((get_username(), get_user_id(), in_view_request)) cm = mocker.MagicMock() cm.__enter__.return_value = engine cm.__exit__.return_value = None @@ -427,15 +440,301 @@ def fake_get_sqla_engine(*args: Any, **kwargs: Any) -> Any: user = SimpleNamespace(id=42, username="alice") - with app.test_request_context("/"): + # Like Flask's wsgi_app, the request pushes (and tears down) its own app + # context, so nothing it put on ``g`` outlives it by accident. + with app.app_context(), app.test_request_context("/"): + view_request["context"] = _cv_request.get() g._login_user = user g.user = current_user # exactly what Flask-AppBuilder's before_request does assert (get_username(), get_user_id()) == ("alice", 42) + # The view builds the response body... command = StreamingCSVExportCommand(query_context, chunk_size=10) - csv_generator_callable = command.run() + stream = command.run()() + + # ...which is streamed once the request context is gone. + assert _cv_request.get(None) is not view_request["context"] + assert "".join(stream).splitlines() == ["id", "1"] + + assert seen == [("alice", 42, True)] + + +def test_streaming_export_reloads_a_detached_acting_user( + app: Any, + session: Session, + mocker: MockerFixture, +) -> None: + """ + By the time the body is streamed, the metadata session that loaded the + acting user has been removed by the request teardown, and the user has + usually been expired by a commit (the event logger commits once the view + returns). Reading ``g.user.id`` to find the user's OAuth2 token would then + raise ``DetachedInstanceError`` unless the export loads the user again. + """ + user = _detached_user(session) + user_id = inspect(user).identity[0] + with pytest.raises(DetachedInstanceError): + _ = user.id + + seen: list[tuple[str | None, int | None, str]] = [] + + def fake_execute_query_and_stream( + self: StreamingCSVExportCommand, *args: Any, **kwargs: Any + ) -> Iterator[str]: + seen.append((get_username(), get_user_id(), current_user.username)) + yield "id\n" + + mocker.patch.object( + StreamingCSVExportCommand, + "_execute_query_and_stream", + fake_execute_query_and_stream, + ) + + # Like Flask's wsgi_app, the request pushes (and tears down) its own app + # context, so nothing it put on ``g`` outlives it by accident. + with app.app_context(), app.test_request_context("/"): + g._login_user = user + g.user = current_user + stream = StreamingCSVExportCommand(mocker.MagicMock()).run()() + + assert list(stream) == ["id\n"] + assert seen == [("alice", user_id, "alice")] + + +def _detached_user(session: Session, expired: bool = True) -> User: + """ + A persisted user, detached as the request teardown leaves it: expired too + when the event logger committed after the view, loaded when nothing did. + """ + User.__table__.create(session.get_bind(), checkfirst=True) + user = User( + first_name="Alice", + last_name="Doe", + username="alice", + email="alice@example.com", + active=True, + ) + session.add(user) + session.commit() + if expired: + session.expire(user) + else: + session.refresh(user) + session.close() + return user + + +def test_streaming_export_fails_closed_when_the_acting_user_is_gone( + app: Any, + session: Session, + mocker: MockerFixture, +) -> None: + """ + If the acting user was deleted before the body is streamed, the export must + not go ahead without a user: impersonating databases would then connect with + the connection's own credentials. + """ + user = _detached_user(session) + session.query(User).delete() + session.commit() + + execute = mocker.patch.object( + StreamingCSVExportCommand, "_execute_query_and_stream" + ) + + with app.app_context(), app.test_request_context("/"): + g._login_user = user + g.user = current_user + stream = StreamingCSVExportCommand(mocker.MagicMock()).run()() + + assert "".join(stream).startswith("__STREAM_ERROR__:") + execute.assert_not_called() + + +def test_reload_detached_acting_user_leaves_other_users_alone( + app: Any, + session: Session, +) -> None: + """Only a detached ORM user is reloaded; anything else on ``g.user`` stays.""" + User.__table__.create(session.get_bind(), checkfirst=True) + attached = User( + first_name="Bob", + last_name="Doe", + username="bob", + email="bob@example.com", + active=True, + ) + session.add(attached) + session.commit() + + guest = GuestUser( + token={ + "iat": 0, + "exp": 0, + "user": {"username": "guest"}, + "resources": [], + "rls_rules": [], + }, + roles=[], + ) + for user in (attached, guest, None): + g.user = user + _reload_detached_acting_user() + assert g.user is user + - # The request context is gone by the time the body is streamed. - list(csv_generator_callable()) +def test_streaming_export_commits_before_streaming_rows(mocker: MockerFixture) -> None: + """ + Building the engine can refresh the user's OAuth2 token, which is only staged + on the session. It must be committed before the rows stream: closing the + session would otherwise roll it back (spending a single-use refresh token), + and the metadata transaction would stay open for the whole export. + """ + mock_db, query_context, datasource = _setup_chart_mocks(mocker) + mock_session = mock_db.session.return_value.__enter__.return_value - assert seen == [("alice", 42)] + result_proxy = mocker.MagicMock() + result_proxy.keys.return_value = ["id"] + result_proxy.fetchmany.side_effect = [[(1,)], []] + connection = mocker.MagicMock() + connection.execution_options.return_value.execute.return_value = result_proxy + connection.__enter__.return_value = connection + connection.__exit__.return_value = None + + def connect() -> Any: + mock_session.commit.assert_called_once() + return connection + + engine = mocker.MagicMock() + engine.connect.side_effect = connect + datasource.database.get_sqla_engine.return_value.__enter__.return_value = engine + + stream = StreamingCSVExportCommand(query_context).run()() + + assert "".join(stream).splitlines() == ["id", "1"] + mock_session.commit.assert_called_once() + + +def test_streaming_export_leaves_no_request_context_behind( + app: Any, + mocker: MockerFixture, +) -> None: + """ + Building the body must not push the request's contexts: a body the server + never iterates (the response failed after the view returned, say) would + leave them on the worker thread, where later requests would inherit its + ``g`` and logged-in user. Flask < 3.1.2's ``stream_with_context`` pushes + them as soon as the body is built. + """ + mocker.patch.object( + StreamingCSVExportCommand, + "_execute_query_and_stream", + return_value=iter(()), + ) + + outer = _cv_request.get(None) + with app.app_context(), app.test_request_context("/"): + stream = StreamingCSVExportCommand(mocker.MagicMock()).run()() + + assert _cv_request.get(None) is outer + del stream + assert _cv_request.get(None) is outer + + +def test_streaming_export_does_not_leave_its_contexts_on_the_thread_mid_stream( + app: Any, + mocker: MockerFixture, +) -> None: + """ + While the body is suspended between chunks, the request's contexts must not + stay pushed on the worker thread: a body the server never finishes (an + error raised while encoding keeps it alive on a traceback) would otherwise + hand the thread's next request this one's ``g`` and logged-in user. + """ + + def fake_execute_query_and_stream( + self: StreamingCSVExportCommand, *args: Any, **kwargs: Any + ) -> Iterator[str]: + yield "id\n" + yield "1\n" + + mocker.patch.object( + StreamingCSVExportCommand, + "_execute_query_and_stream", + fake_execute_query_and_stream, + ) + + outer = _cv_request.get(None) + with app.app_context(), app.test_request_context("/"): + view_request = _cv_request.get() + stream = StreamingCSVExportCommand(mocker.MagicMock()).run()() + + assert next(stream) == "id\n" + assert _cv_request.get(None) is outer is not view_request + assert list(stream) == ["1\n"] + assert _cv_request.get(None) is outer + + +def test_streaming_export_read_in_part_before_the_request_ends( + app: Any, + mocker: MockerFixture, +) -> None: + """ + Reading part of the body while the request is still being handled (from an + ``after_request`` hook, say) must not stop Flask from popping the request's + contexts, and the rest of the body must still stream inside that request. + """ + view_request: dict[str, Any] = {} + seen: list[bool] = [] + + def fake_execute_query_and_stream( + self: StreamingCSVExportCommand, *args: Any, **kwargs: Any + ) -> Iterator[str]: + seen.append(_cv_request.get(None) is view_request["context"]) + yield "id\n" + seen.append(_cv_request.get(None) is view_request["context"]) + yield "1\n" + + mocker.patch.object( + StreamingCSVExportCommand, + "_execute_query_and_stream", + fake_execute_query_and_stream, + ) + + outer = _cv_request.get(None) + with app.app_context(), app.test_request_context("/"): + view_request["context"] = _cv_request.get() + stream = StreamingCSVExportCommand(mocker.MagicMock()).run()() + assert next(stream) == "id\n" + + assert _cv_request.get(None) is outer + assert list(stream) == ["1\n"] + assert seen == [True, True] + assert _cv_request.get(None) is outer + + +def test_reload_if_detached_keeps_changes_made_meanwhile( + app: Any, + session: Session, +) -> None: + """ + An instance the request loaded is reloaded rather than merged, so the + export's commit can't write the values it read during the request back + over a change made since (an admin editing the connection, say). + """ + user = _detached_user(session, expired=False) + assert user.email == "alice@example.com" + session.query(User).filter_by(username="alice").update( + {"email": "changed@example.com"} + ) + session.commit() + + reloaded = _reload_if_detached(user) + session.commit() + + assert reloaded is not user + assert reloaded.email == "changed@example.com" + assert ( + session.query(User.email).filter_by(username="alice").scalar() + == "changed@example.com" + ) diff --git a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py index c6aeb4ab0b05..ba690660cf22 100644 --- a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py +++ b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py @@ -92,7 +92,6 @@ def test_streaming_sql_result_export_command_init(): assert command._client_id == "client_123" assert command._chunk_size == 500 assert command._query is None - assert command._current_app is not None def test_streaming_sql_result_export_command_default_chunk_size(): @@ -418,7 +417,9 @@ def test_error_handling_yields_error_marker(mocker, mock_query): mock_db_base = mocker.patch("superset.commands.streaming_export.base.db") mock_session = MagicMock() mock_db_base.session.return_value.__enter__.return_value = mock_session - mock_session.merge.side_effect = Exception("Database connection failed") + mock_query.database.get_sqla_engine.side_effect = Exception( + "Database connection failed" + ) mock_db_sqllab = mocker.patch( "superset.commands.sql_lab.streaming_export_command.db" @@ -857,19 +858,13 @@ def test_streaming_export_preserves_impersonated_user_context( app_context, mocker, mock_query, mock_result_proxy ): """ - Checks the other half of apache/superset#40465's theory: that the - streaming path drops user impersonation because it acquires its engine - "without this context" (unlike other call sites, which the issue says - use a get_sqla_engine_with_context(user_name=...) that doesn't actually - exist anywhere in this codebase). Independently verified this does NOT - reproduce: BaseStreamingCSVExportCommand.run() captures flask.g's full - __dict__ (including g.user) before entering the deferred generator, and - csv_generator() restores it via preserve_g_context before calling - _execute_query_and_stream -- so Database.get_sqla_engine's internal - get_effective_user()/get_username() call (which every engine-acquisition - call site relies on for impersonation, streaming or not) sees the same - acting user it would anywhere else. This pins that behavior rather than - the bug the issue reports for it. + Covers the other half of apache/superset#40465: the streaming path must + impersonate the acting user. Database.get_sqla_engine resolves the + impersonated user through get_effective_user()/get_username(), which read + g.user, so the deferred generator has to see the same g.user as the code + that created it. A real, ORM-backed user is also detached by the request + teardown before the body streams; that case is covered by + test_streaming_export_reloads_a_detached_acting_user in the chart tests. """ g.user = Mock(username="alice") assert get_username() == "alice"