(superset) 01/01: fix(export): keep the acting user in streaming CSV exports

[en...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

EnxDev pushed a commit to branch enxdev/fix/streaming-export-oauth2-user-context
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 0305ee44d332591c0a8becefc4f9af3b7b46398e
Author: Enzo Martellucci <en...@gmail.com>
AuthorDate: Fri Sep 18 12:06:13 2026 +0200

    fix(export): keep the acting user in streaming CSV exports
    
    CSV/Excel exports at or above CSV_STREAMING_ROW_THRESHOLD (100k rows) take
    the streaming path, which defers query execution into a generator that
    Werkzeug iterates after the request context has been popped. The generator
    replays a snapshot of flask.g inside a fresh app context.
    
    That snapshot copied g.user verbatim, but Flask-AppBuilder's before_request
    sets it to Flask-Login's `current_user` -- a LocalProxy that resolves via
    has_request_context() and so evaluates to None once the request is over.
    The generator therefore ran as no user at all: in Database._get_sqla_engine,
    `hasattr(g.user, "id")` is False, so the engine is built with
    access_token=None and the per-user OAuth2 token is dropped. The resulting
    connection failure is caught by the generator's broad except and written
    into the response body as a __STREAM_ERROR__ marker, producing a downloaded
    file containing an error and no rows. Exports below the threshold use the
    buffered path, which never leaves the request context, which is why the bug
    appeared only at the exact 100k boundary and only for databases whose
    credentials are per-user (BigQuery via OAuth2) rather than stored on the
    connection.
    
    Resolve request-bound proxies to their concrete objects while the request
    context is still available, so the generator sees the same acting user that
    every other engine-acquisition call site does.
    
    Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
---
 superset/commands/streaming_export/base.py         | 41 +++++++++++++--
 .../chart/streaming_export_command_test.py         | 61 ++++++++++++++++++++++
 2 files changed, 98 insertions(+), 4 deletions(-)

diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py
index 21ae0216cdf..fbed887ecbc 100644
--- a/superset/commands/streaming_export/base.py
+++ b/superset/commands/streaming_export/base.py
@@ -30,6 +30,7 @@ from typing import Any, Callable, Generator
 
 from flask import current_app as app, g, has_app_context
 from sqlalchemy import text
+from werkzeug.local import LocalProxy
 
 from superset import db
 from superset.commands.base import BaseCommand
@@ -38,6 +39,40 @@ from superset.utils.csv import escape_value
 logger = logging.getLogger(__name__)
 
 
+def capture_g_context() -> dict[str, Any]:
+    """
+    Snapshot ``flask.g`` so a streaming generator can replay it later.
+
+    Values held on ``g`` may be request-bound ``LocalProxy`` objects rather
+    than plain values. The most important one is ``g.user``, which
+    Flask-AppBuilder's ``before_request`` sets to Flask-Login's
+    ``current_user``: that proxy resolves through ``has_request_context()``
+    and therefore evaluates to ``None`` once the request context is gone,
+    which is exactly the situation the generator runs in.
+
+    Copying such a proxy verbatim would hand the generator a ``g.user`` that
+    silently resolves to nobody, so resolve each proxy to the concrete object
+    it currently points at while the request context is still around.
+
+    Returns:
+        Dictionary of g attributes, with request-bound proxies resolved
+    """
+    if not has_app_context():
+        return {}
+
+    captured: dict[str, Any] = {}
+    for key, value in g._get_current_object().__dict__.items():
+        if isinstance(value, LocalProxy):
+            try:
+                value = value._get_current_object()
+            except RuntimeError:
+                # Nothing is bound to the proxy, so there is no value worth
+                # carrying into the generator.
+                continue
+        captured[key] = value
+    return captured
+
+
 @contextmanager
 def preserve_g_context(
     captured_g: dict[str, Any],
@@ -49,7 +84,7 @@ def preserve_g_context(
     app context but needs access to request-scoped data from the original request.
 
     Args:
-        captured_g: Dictionary of g attributes captured before context switch
+        captured_g: Dictionary of g attributes captured by capture_g_context()
     """
     for key, value in captured_g.items():
         setattr(g, key, value)
@@ -322,9 +357,7 @@ class BaseStreamingCSVExportCommand(BaseCommand):
         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 {}
-        )
+        captured_g = capture_g_context()
 
         def csv_generator() -> Generator[str, None, None]:
             """Generator that yields CSV data chunks."""
diff --git a/tests/unit_tests/commands/chart/streaming_export_command_test.py b/tests/unit_tests/commands/chart/streaming_export_command_test.py
index 5d354e557b3..66ca9451af1 100644
--- a/tests/unit_tests/commands/chart/streaming_export_command_test.py
+++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py
@@ -16,12 +16,18 @@
 # under the License.
 """Unit tests for Chart Streaming CSV Export Command."""
 
+from types import SimpleNamespace
+from typing import Any
+
 import pytest
+from flask import g
+from flask_login import current_user
 from pytest_mock import MockerFixture
 
 from superset.commands.chart.data.streaming_export_command import (
     StreamingCSVExportCommand,
 )
+from superset.utils.core import get_user_id, get_username
 
 
 def _setup_chart_mocks(
@@ -378,3 +384,58 @@ def test_streaming_export_mutation_does_not_double_apply(
     datasource.database.mutate_sql_based_on_config.assert_called_once_with(
         "SELECT * FROM test /* mutated */", is_split=True
     )
+
+
+def test_streaming_export_keeps_acting_user_after_request_context_ends(
+    app: Any,
+    mocker: MockerFixture,
+) -> None:
+    """
+    Flask-AppBuilder's ``before_request`` sets ``g.user = current_user``, a
+    Flask-Login ``LocalProxy`` that resolves to ``None`` outside a request
+    context. Werkzeug iterates a streaming response body *after* the request
+    context is popped, so copying that proxy into the generator's fresh app
+    context yields a user-less ``g``: ``Database._get_sqla_engine`` then sees
+    ``hasattr(g.user, "id") is False`` and builds the engine with
+    ``access_token=None``, silently dropping the per-user OAuth2 token.
+    """
+    _mock_db, query_context, datasource = _setup_chart_mocks(mocker)
+
+    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
+    engine = mocker.MagicMock()
+    engine.connect.return_value = connection
+
+    # Record the acting user at the moment the engine is acquired -- the point
+    # where the real Database._get_sqla_engine resolves the OAuth2 access token.
+    seen: list[tuple[str | None, int | None]] = []
+
+    def fake_get_sqla_engine(*args: Any, **kwargs: Any) -> Any:
+        seen.append((get_username(), get_user_id()))
+        cm = mocker.MagicMock()
+        cm.__enter__.return_value = engine
+        cm.__exit__.return_value = None
+        return cm
+
+    datasource.database.get_sqla_engine.side_effect = fake_get_sqla_engine
+
+    user = SimpleNamespace(id=42, username="alice")
+
+    with app.test_request_context("/"):
+        g._login_user = user
+        g.user = current_user  # exactly what Flask-AppBuilder's before_request does
+        assert (get_username(), get_user_id()) == ("alice", 42)
+
+        command = StreamingCSVExportCommand(query_context, chunk_size=10)
+        csv_generator_callable = command.run()
+
+    # The request context is gone by the time the body is streamed.
+    list(csv_generator_callable())
+
+    assert seen == [("alice", 42)]