geido commented on code in PR #45952:
URL: https://github.com/apache/superset/pull/45952#discussion_r4453602668
##########
superset/commands/streaming_export/base.py:
##########
@@ -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
Review Comment:
This only reproduces with Redis lock coordination, since the default KV lock
already commits the refreshed token when it's released. The buffered export has
the same window too, and a wider one because it only commits after the query
runs, so committing right after the engine is built narrows it rather than
opening it. A real fix belongs in `superset/utils/oauth2.py`, so I'm keeping it
out of scope here.
##########
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:
Good catch, thanks! I reproduced it on a single-threaded werkzeug server
(the next anonymous request blew up with "Popped wrong request context") and
applied all three suggestions in b1455e45cc.
--
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]
