alexandrusoare commented on code in PR #45952:
URL: https://github.com/apache/superset/pull/45952#discussion_r4455205838
##########
superset/commands/streaming_export/base.py:
##########
@@ -378,41 +369,87 @@ 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 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()
Review Comment:
`RequestContext._cv_tokens` is a plain LIFO list living on the context
object, shared across contextvars Contexts. If anything pulls a chunk before
the view's context is popped, the push inside this copied Context lands a
second token on that list, and `wsgi_app`'s `pop()` resets the wrong one:
```
ValueError: <Token var=<ContextVar name='flask.request_ctx'>> was created in
a different Context
```
Two knock-ons: `clear_request` is `len(self._cv_tokens) == 1`, so with two
tokens `teardown_request` and `request.close()` are skipped; and `_cv_request`
stays pointing at the dead context for the thread's next request — the leak
this PR is closing. Repro is just `copy_context().run(next, inner)` before
`rctx.pop()`.
Nothing in-tree triggers it (`COMPRESS_MIMETYPES` excludes `text/csv`, no
`after_request` hook touches the body, the test client doesn't pre-iterate), so
it's latent. Skipping the re-push while `_cv_request.get(None)` is already
`request_context` closes it, though the check has to hold per chunk rather than
once at generator entry.
##########
superset/commands/streaming_export/base.py:
##########
@@ -378,41 +369,87 @@ 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()
Review Comment:
The docstring promises this path ("Outside a request, consume the iterator
inside the caller's app context"), but there's no `g.user` here, so
`_reload_detached_acting_user()` no-ops and the export runs with no acting
user: no OAuth2 token, `get_effective_user()` → `None`, impersonating
connections quietly falling back to the URL's own identity. No caller does this
today — either drop the promise from the docstring or fail loudly here.
--
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]
