diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py index 8b525a2c8226..1393cf2b193e 100644 --- a/superset/commands/streaming_export/base.py +++ b/superset/commands/streaming_export/base.py @@ -23,9 +23,10 @@ import logging import time from abc import abstractmethod +from contextlib import contextmanager from typing import Any, Callable, Generator -from flask import current_app as app +from flask import current_app as app, g, has_app_context from sqlalchemy import text from superset import db @@ -34,6 +35,24 @@ logger = logging.getLogger(__name__) +@contextmanager +def preserve_g_context( + captured_g: dict[str, Any], +) -> Generator[None, None, None]: + """ + 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 before context switch + """ + for key, value in captured_g.items(): + setattr(g, key, value) + yield + + class BaseStreamingCSVExportCommand(BaseCommand): """ Base class for streaming CSV export commands. @@ -192,23 +211,29 @@ def run(self) -> Callable[[], Generator[str, None, None]]: # to avoid DetachedInstanceError sql, database = 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 = ( + g._get_current_object().__dict__.copy() if has_app_context() else {} + ) def csv_generator() -> Generator[str, None, None]: """Generator that yields CSV data chunks.""" with self._current_app.app_context(): - try: - yield from self._execute_query_and_stream(sql, database, limit) - 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 + with preserve_g_context(captured_g): + try: + yield from self._execute_query_and_stream(sql, database, limit) + 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