geido commented on code in PR #45952:
URL: https://github.com/apache/superset/pull/45952#discussion_r4458368652
##########
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:
Good catch, reproduced it: after a partial early read every later request on
that worker got a 500. Fixed by not pushing the request again when it's still
active at the first read, with a test.
##########
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:
Fair, the docstring now says it runs as the caller's `g.user` there (set
with `override_user`).
--
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]
