rebenitez1802 commented on code in PR #45952:
URL: https://github.com/apache/superset/pull/45952#discussion_r4448364009
##########
tests/unit_tests/commands/chart/streaming_export_command_test.py:
##########
@@ -427,15 +436,205 @@ 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()()
- # 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 _cv_request.get(None) is not view_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="[email protected]",
+ 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="[email protected]",
+ 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="[email protected]",
+ 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(()),
+ )
+
+ 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
Review Comment:
Regression test for the contextvars suggestion in `stream()`. It fails on
the current head and passes once that suggestion is applied.
```suggestion
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
```
##########
superset/commands/streaming_export/base.py:
##########
@@ -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()
Review Comment:
🟡 Medium: on a single-threaded server, an export abandoned mid-stream logs
the next visitor in as the exporter (non-blocking)
`with app_context, request_context:` keeps the request's contexts pushed on
the worker thread while the body is paused between chunks. Say something raises
outside `csv_generator`'s `try`, for example `encode_csv_stream` raising
`UnicodeEncodeError` with `CSV_EXPORT["encoding"] = "latin-1"` and a `€` in the
data. The traceback then keeps the paused generator alive, so the contexts are
never popped. On werkzeug with `threaded=False`, the next anonymous request on
that thread then got `current_user = "alice"` (the user who ran the export),
where master returned a 500. gunicorn sync, gthread and gevent were clean in
our runs, so I'm not treating this as a blocker. The leak itself predates this
PR; the PR makes the outcome worse, because the reloaded user is fully loaded.
Fix: step the inner generator inside its own `contextvars` Context, so the
request's contexts are never pushed on the worker thread itself. This goes with
the import suggestion at the top of the file and the regression test in
`streaming_export_command_test.py`. Apply all three together with "Add
suggestion to batch". Validated: ruff clean, 57 streaming tests pass, and the
new test fails on the current head with `assert <RequestContext
'http://example.com/'> is None`.
```suggestion
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()
```
##########
superset/commands/streaming_export/base.py:
##########
@@ -24,13 +24,14 @@
import logging
import time
from abc import abstractmethod
Review Comment:
Import for the contextvars suggestion further down in `stream()`.
```suggestion
from abc import abstractmethod
from contextvars import copy_context
```
--
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]
