From 11ee222d3d877575aa6846dffe5c03c228912a5c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Geid=C5=8D?= <60598000+geido@users.noreply.github.com> Date: Mon, 2 Nov 2026 12:35:14 +0200 Subject: [PATCH 1/5] fix(export): stream CSV exports inside their request context Streaming CSV exports (100k+ rows, and SQL Lab streaming downloads) run their query from the response body generator, which the WSGI server iterates only after the view has returned. The generator replayed a snapshot of flask.g inside a fresh app context, and that breaks databases that authenticate with per-user OAuth2, Snowflake in particular, so these exports returned only the __STREAM_ERROR__ marker: - The replayed g.user is the User loaded by the request's metadata session. The event logger's commit expires it and the request teardown removes that session, so reading g.user.id to look up the OAuth2 token raised DetachedInstanceError. - The generator had no request context. Snowflake only uses the user's OAuth2 token inside a request, so even with a usable user it connected without the token ("251006: Password is empty"). The OAuth2 redirect URI and the query source passed to DB_CONNECTION_MUTATOR also come from the request. Record the request's contexts in the view and push them again only while the body is iterated, as flask.stream_with_context does from Flask 3.1.2 (earlier releases push them up front, which leaves them on the worker thread if the response never reaches the server). Reload a detached acting user by primary key, failing closed if it no longer exists, and commit right after the engine is built so a refreshed OAuth2 token is kept and no metadata transaction stays open while rows stream. This replaces the flask.g snapshot and replay. Co-Authored-By: Claude Opus 5.5 (1M context) --- superset/commands/streaming_export/base.py | 175 +++++++------ .../chart/streaming_export_command_test.py | 231 ++++++++++++++++-- .../sql_lab/streaming_export_command_test.py | 21 +- 3 files changed, 314 insertions(+), 113 deletions(-) diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py index e87065bbedbf..dac8160f0da5 100644 --- a/superset/commands/streaming_export/base.py +++ b/superset/commands/streaming_export/base.py @@ -24,13 +24,14 @@ import logging import time from abc import abstractmethod -from contextlib import contextmanager from decimal import Decimal from numbers import Real -from typing import Any, Callable, Generator, Iterable +from typing import Any, Callable, Generator, Iterable, Iterator -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 @@ -40,56 +41,39 @@ logger = logging.getLogger(__name__) -def capture_g_context() -> dict[str, Any]: +def _reload_detached_acting_user() -> None: """ - Snapshot ``flask.g`` so a streaming generator can replay it later. - - 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. - - 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 - """ - if not has_app_context(): - return {} - - 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]: + Load the acting user again if it is detached from the metadata session. + + 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 ``g.user``. Any attribute that session + expired (the event logger commits once the view returns), even + ``g.user.id``, which is how the user's OAuth2 token is looked up, would + then raise ``DetachedInstanceError``. + + The user is loaded again by primary key rather than merged, so values read + during the request are never written back by a later commit. Users that + are not ORM instances, such as embedded guest users, are left as they are. """ - Context manager that restores captured flask.g attributes. - - 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 + 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() + + state = inspect(user, raiseerr=False) + if not isinstance(state, InstanceState) or not state.detached: + return + + reloaded = db.session.get(state.class_, state.identity) + if reloaded is None: + # Fail closed: without a user, impersonation falls back to the + # connection's own credentials. + raise RuntimeError("The user running the export no longer exists") + 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 +125,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]: @@ -338,6 +321,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 +368,68 @@ 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. """ # 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 csv_generator_in_request() -> Generator[str, None, None]: + with app_context, request_context: + yield from csv_generator() + + 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..575c5ebb04fc 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,21 @@ """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 import g, has_request_context +from flask_appbuilder.security.sqla.models import User from flask_login import current_user from pytest_mock import MockerFixture +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 +from superset.security.guest_token import GuestUser from superset.utils.core import get_user_id, get_username @@ -64,7 +69,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 +390,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 +415,12 @@ 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. + 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())) + seen.append((get_username(), get_user_id(), has_request_context())) cm = mocker.MagicMock() cm.__enter__.return_value = engine cm.__exit__.return_value = None @@ -427,15 +430,203 @@ 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("/"): 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()() - # The request context is gone by the time the body is streamed. - list(csv_generator_callable()) + # ...which is streamed once the request context is gone. + assert not has_request_context() + assert "".join(stream).splitlines() == ["id", "1"] - assert seen == [("alice", 42)] + 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.__table__.create(session.get_bind()) + user = User( + first_name="Alice", + last_name="Doe", + username="alice", + email="alice@example.com", + active=True, + ) + session.add(user) + session.commit() + user_id = user.id + session.expire(user) + session.close() + 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) -> User: + """A persisted user, detached and expired as the request teardown leaves it.""" + 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() + session.expire(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 + + +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 + + 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: + """ + A body the server never iterates (the response failed after the view + returned, say) must not leave the request's contexts pushed on the worker + thread, where later requests would inherit its ``g`` and logged-in user. + """ + mocker.patch.object( + StreamingCSVExportCommand, + "_execute_query_and_stream", + return_value=iter(()), + ) + + with app.app_context(), app.test_request_context("/"): + stream = StreamingCSVExportCommand(mocker.MagicMock()).run()() + + assert not has_request_context() + del stream + assert not has_request_context() 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..a966a004dd79 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(): @@ -857,19 +856,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" From 0bcf1dab53dae87336bb17469893d82f3c700aa4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Geid=C5=8D?= <60598000+geido@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:29:34 +0200 Subject: [PATCH 2/5] test(export): check the streaming tests' own request context The two request-context tests asked whether any request context was active, so a context another test leaves pushed in the same worker would make them fail (or pass) for the wrong reason. Compare against the context of the request the test builds instead: it must be active while the body streams and gone afterwards, whatever else is on the thread. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../chart/streaming_export_command_test.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) 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 575c5ebb04fc..abfa505845b1 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -20,7 +20,8 @@ from typing import Any, Iterator import pytest -from flask import g, has_request_context +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 @@ -417,10 +418,15 @@ def test_streaming_export_runs_in_its_request_context_after_the_view_returns( # 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(), has_request_context())) + 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 @@ -433,6 +439,7 @@ def fake_get_sqla_engine(*args: Any, **kwargs: Any) -> Any: # 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) @@ -442,7 +449,7 @@ def fake_get_sqla_engine(*args: Any, **kwargs: Any) -> Any: stream = command.run()() # ...which is streamed once the request context is gone. - assert not has_request_context() + assert _cv_request.get(None) is not view_request["context"] assert "".join(stream).splitlines() == ["id", "1"] assert seen == [("alice", 42, True)] @@ -624,9 +631,10 @@ def test_streaming_export_leaves_no_request_context_behind( return_value=iter(()), ) + outer = _cv_request.get(None) with app.app_context(), app.test_request_context("/"): stream = StreamingCSVExportCommand(mocker.MagicMock()).run()() - assert not has_request_context() + assert _cv_request.get(None) is outer del stream - assert not has_request_context() + assert _cv_request.get(None) is outer From b1455e45cc0937d9053606bb41459366c30d3f78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Geid=C5=8D?= <60598000+geido@users.noreply.github.com> Date: Tue, 3 Nov 2026 12:30:40 +0200 Subject: [PATCH 3/5] fix(export): keep the streaming request's contexts off the worker thread The export pushed the request's contexts on the worker thread while the body was iterated, so they stayed pushed while it was suspended between chunks. If the server abandons the body mid-stream (e.g. encoding a chunk raises and the traceback keeps the generator alive), the thread's next request inherits them: on a single-threaded server it either gets the exporting user or fails with "Popped wrong request context" when the abandoned generator is finally closed. Step the export inside a contextvars Context of its own, so the request's contexts are only ever pushed there and never on the worker thread. Co-authored-by: Rafael Benitez <2151636+rebenitez1802@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 (1M context) --- superset/commands/streaming_export/base.py | 22 +++++++++++- .../chart/streaming_export_command_test.py | 34 +++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py index dac8160f0da5..7e9fd59c4769 100644 --- a/superset/commands/streaming_export/base.py +++ b/superset/commands/streaming_export/base.py @@ -24,6 +24,7 @@ import logging import time from abc import abstractmethod +from contextvars import copy_context from decimal import Decimal from numbers import Real from typing import Any, Callable, Generator, Iterable, Iterator @@ -426,10 +427,29 @@ def stream() -> Iterator[str]: return csv_generator() app_context = _cv_app.get() - def csv_generator_in_request() -> Generator[str, None, None]: + def in_request() -> Generator[str, None, None]: 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 abfa505845b1..7cc1559ded3d 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -638,3 +638,37 @@ def test_streaming_export_leaves_no_request_context_behind( 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 From 142c7144b7d931816292c8ed7053262c257f24ad Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Geid=C5=8D?= <60598000+geido@users.noreply.github.com> Date: Tue, 3 Nov 2026 16:08:22 +0200 Subject: [PATCH 4/5] fix(export): reload the export's database instead of merging it The export merged the Database loaded by the request into its own session and then commits right after building the engine. Merging copies the values read during the request over the row's current ones, so a change made to the connection in the meantime would be written back by that commit. Reload it by primary key, like the acting user. Also document that, outside a request, the export runs as the caller's g.user, and reuse the detached-user test helper. Co-Authored-By: Claude Opus 5.5 (1M context) --- superset/commands/streaming_export/base.py | 62 +++++++++++------ .../chart/streaming_export_command_test.py | 68 +++++++++++++------ .../sql_lab/streaming_export_command_test.py | 4 +- 3 files changed, 91 insertions(+), 43 deletions(-) diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py index 7e9fd59c4769..a3d1b50a20d4 100644 --- a/superset/commands/streaming_export/base.py +++ b/superset/commands/streaming_export/base.py @@ -27,7 +27,7 @@ from contextvars import copy_context from decimal import Decimal from numbers import Real -from typing import Any, Callable, Generator, Iterable, Iterator +from typing import Any, Callable, Generator, Iterable, Iterator, TypeVar from flask import current_app as app, g from flask.globals import _cv_app, _cv_request @@ -41,36 +41,51 @@ logger = logging.getLogger(__name__) +_T = TypeVar("_T") -def _reload_detached_acting_user() -> None: + +def _reload_if_detached(instance: _T) -> _T: """ - Load the acting user again if it is detached from the metadata session. + 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 ``g.user``. Any attribute that session - expired (the event logger commits once the view returns), even - ``g.user.id``, which is how the user's OAuth2 token is looked up, would - then raise ``DetachedInstanceError``. - - The user is loaded again by primary key rather than merged, so values read - during the request are never written back by a later commit. Users that - are not ORM instances, such as embedded guest users, are left as they are. + 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. + """ + state = inspect(instance, raiseerr=False) + if not isinstance(state, InstanceState) or not state.detached: + return instance + + 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 + + +def _reload_detached_acting_user() -> None: + """ + Point ``g.user`` at the acting user loaded again if it was detached. + + Users that are not ORM instances, such as embedded guest users, are left + as they are. """ 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() - state = inspect(user, raiseerr=False) - if not isinstance(state, InstanceState) or not state.detached: + reloaded = _reload_if_detached(user) + if reloaded is user: return - - reloaded = db.session.get(state.class_, state.identity) - if reloaded is None: - # Fail closed: without a user, impersonation falls back to the - # connection's own credentials. - raise RuntimeError("The user running the export no longer exists") g.user = reloaded if getattr(g, "_login_user", None) is user: # Keep Flask-Login's ``current_user`` on the same, reloaded user. @@ -302,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() @@ -378,7 +394,9 @@ def run(self) -> Callable[[], Iterator[str]]: 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. + 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 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 7cc1559ded3d..d3c48911b415 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -25,13 +25,17 @@ 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 +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 @@ -467,19 +471,8 @@ def test_streaming_export_reloads_a_detached_acting_user( returns). Reading ``g.user.id`` to find the user's OAuth2 token would then raise ``DetachedInstanceError`` unless the export loads the user again. """ - User.__table__.create(session.get_bind()) - user = User( - first_name="Alice", - last_name="Doe", - username="alice", - email="alice@example.com", - active=True, - ) - session.add(user) - session.commit() - user_id = user.id - session.expire(user) - session.close() + user = _detached_user(session) + user_id = inspect(user).identity[0] with pytest.raises(DetachedInstanceError): _ = user.id @@ -508,8 +501,11 @@ def fake_execute_query_and_stream( assert seen == [("alice", user_id, "alice")] -def _detached_user(session: Session) -> User: - """A persisted user, detached and expired as the request teardown leaves it.""" +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", @@ -520,7 +516,10 @@ def _detached_user(session: Session) -> User: ) session.add(user) session.commit() - session.expire(user) + if expired: + session.expire(user) + else: + session.refresh(user) session.close() return user @@ -621,9 +620,11 @@ def test_streaming_export_leaves_no_request_context_behind( mocker: MockerFixture, ) -> None: """ - A body the server never iterates (the response failed after the view - returned, say) must not leave the request's contexts pushed on the worker - thread, where later requests would inherit its ``g`` and logged-in user. + 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, @@ -672,3 +673,30 @@ def fake_execute_query_and_stream( 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_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 a966a004dd79..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 @@ -417,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" From 9ce575c15cdccee52d65989a6339bc828bf692af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Geid=C5=8D?= <60598000+geido@users.noreply.github.com> Date: Tue, 3 Nov 2026 16:27:05 +0200 Subject: [PATCH 5/5] fix(export): let the request pop when its body is read early If something read part of a streaming export's body while the request was still being handled, the export pushed the request's contexts again in its own contextvars Context. Flask's pop at the end of the request then tried to reset that token from the worker thread's Context: it raised ValueError, skipped the request's teardown and left its contexts on the thread. The contexts are already active at that point, so they are not pushed again. Co-Authored-By: Claude Opus 5.5 (1M context) --- superset/commands/streaming_export/base.py | 6 +++ .../chart/streaming_export_command_test.py | 38 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py index a3d1b50a20d4..93520e3bf63b 100644 --- a/superset/commands/streaming_export/base.py +++ b/superset/commands/streaming_export/base.py @@ -446,6 +446,12 @@ def stream() -> Iterator[str]: 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() 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 d3c48911b415..396553237a46 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -675,6 +675,44 @@ def fake_execute_query_and_stream( 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,