This is an automated email from the ASF dual-hosted git repository.
EnxDev pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 7c7b3adf745 fix(export): keep the acting user in streaming CSV exports
(#44425)
7c7b3adf745 is described below
commit 7c7b3adf74553aaee953f62c1a43efaa87913e20
Author: Enzo Martellucci <[email protected]>
AuthorDate: Wed Sep 23 15:40:46 2026 +0200
fix(export): keep the acting user in streaming CSV exports (#44425)
Co-authored-by: Claude Opus 5 (1M context) <[email protected]>
---
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)]
