You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ru...@apache.org on 2026/08/12 12:56:08 UTC
(superset) branch master updated: fix(sql-lab): apply SQL_QUERY_MUTATOR to streaming exports (#40465) (#42739)
This is an automated email from the ASF dual-hosted git repository.
rusackas 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 4354b37b963 fix(sql-lab): apply SQL_QUERY_MUTATOR to streaming exports (#40465) (#42739)
4354b37b963 is described below
commit 4354b37b96309029f1bd49a4249c3f68605d2e78
Author: Evan Rusackas <ev...@preset.io>
AuthorDate: Wed Aug 12 05:56:01 2026 -0700
fix(sql-lab): apply SQL_QUERY_MUTATOR to streaming exports (#40465) (#42739)
Co-authored-by: Claude Code <no...@anthropic.com>
---
superset/commands/streaming_export/base.py | 24 ++++-
.../chart/streaming_export_command_test.py | 46 ++++++++
.../sql_lab/streaming_export_command_test.py | 120 +++++++++++++++++++++
3 files changed, 189 insertions(+), 1 deletion(-)
diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py
index 70e5f6ef6d4..1113c2eff99 100644
--- a/superset/commands/streaming_export/base.py
+++ b/superset/commands/streaming_export/base.py
@@ -85,6 +85,14 @@ class BaseStreamingCSVExportCommand(BaseCommand):
"""
Get the SQL query, database, catalog, and schema for execution.
+ The returned SQL is expected to already carry any
+ ``is_split=False`` mutation applied upstream (e.g. by
+ ``get_query_str_extended`` for charts, or SQL Lab's stored
+ ``executed_sql``); ``run()`` applies the complementary
+ ``is_split=True`` mutation exactly once before execution, so the
+ mutator fires once total under either ``MUTATE_AFTER_SPLIT``
+ setting.
+
Returns:
Tuple of (sql_query, database_object, catalog, schema)
"""
@@ -223,13 +231,27 @@ class BaseStreamingCSVExportCommand(BaseCommand):
# Merge database to prevent DetachedInstanceError
merged_database = session.merge(database)
+ # `is_split=True` mirrors the non-streaming download path exactly:
+ # Database.get_df() -> _execute_sql_with_mutation_and_logging()
+ # always calls mutate_sql_based_on_config(..., is_split=True) on
+ # SQL Lab's stored select_sql/executed_sql, and the chart query
+ # path applies its own mutation upstream (in
+ # get_query_str_extended, with is_split=False) before landing
+ # here. Since `is_split` and `MUTATE_AFTER_SPLIT` are compared
+ # for equality, `is_split=True` is the complement of that
+ # upstream chart mutation -- together they mutate the SQL
+ # exactly once for either MUTATE_AFTER_SPLIT setting, instead of
+ # double-mutating when it's False and never mutating when it's
+ # True.
+ mutated_sql = merged_database.mutate_sql_based_on_config(sql, is_split=True)
+
with merged_database.get_sqla_engine(
catalog=catalog, schema=schema
) as engine:
with engine.connect() as connection:
result_proxy = connection.execution_options(
stream_results=True
- ).execute(text(sql))
+ ).execute(text(mutated_sql))
columns = list(result_proxy.keys())
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 fc12da53322..5f9246eea53 100644
--- a/tests/unit_tests/commands/chart/streaming_export_command_test.py
+++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py
@@ -39,6 +39,9 @@ def _setup_chart_mocks(
datasource = mocker.MagicMock()
datasource.get_query_str.return_value = sql
datasource.database = mocker.MagicMock()
+ datasource.database.mutate_sql_based_on_config.side_effect = (
+ lambda sql_, **kwargs: sql_
+ )
datasource.catalog = catalog
datasource.schema = schema
query_context.datasource = datasource
@@ -296,3 +299,46 @@ def test_catalog_and_schema_passed_to_engine(mocker: MockerFixture) -> None:
catalog="my_catalog",
schema="my_schema",
)
+
+
+def test_streaming_export_mutation_does_not_double_apply(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Regression test for a bug caught in review of apache/superset#40465's fix:
+ ``datasource.get_query_str()`` (used to build the SQL for chart streaming
+ exports) already runs it through
+ ``Database.mutate_sql_based_on_config(sql)`` -- with ``is_split=False`` --
+ upstream in ``get_query_str_extended``. If
+ ``_execute_query_and_stream`` mutated again with the same default
+ ``is_split=False``, a ``SQL_QUERY_MUTATOR`` would run twice under the
+ default ``MUTATE_AFTER_SPLIT=False`` config, and never run at all under
+ ``MUTATE_AFTER_SPLIT=True`` (since neither call would use
+ ``is_split=True``). Passing ``is_split=True`` here is the complement of
+ the upstream call, so exactly one of the two fires for either setting.
+ """
+ mock_db, query_context, datasource = _setup_chart_mocks(
+ mocker, sql="SELECT * FROM test /* mutated */"
+ )
+
+ mock_result = mocker.MagicMock()
+ mock_result.keys.return_value = ["col1"]
+ mock_result.fetchmany.side_effect = [[("val",)], []]
+
+ mock_connection = mocker.MagicMock()
+ mock_connection.execution_options.return_value.execute.return_value = mock_result
+ mock_connection.__enter__.return_value = mock_connection
+ mock_connection.__exit__.return_value = None
+
+ mock_engine = mocker.MagicMock()
+ mock_engine.connect.return_value = mock_connection
+ datasource.database.get_sqla_engine.return_value.__enter__.return_value = (
+ mock_engine
+ )
+
+ command = StreamingCSVExportCommand(query_context)
+ list(command.run()())
+
+ datasource.database.mutate_sql_based_on_config.assert_called_once_with(
+ "SELECT * FROM test /* mutated */", is_split=True
+ )
diff --git a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py
index 12406a51f6d..3e33f6c4ac1 100644
--- a/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py
+++ b/tests/unit_tests/commands/sql_lab/streaming_export_command_test.py
@@ -20,6 +20,7 @@ from decimal import Decimal
from unittest.mock import MagicMock, Mock, patch
import pytest
+from flask import g
from pytest_mock import MockerFixture
from superset.commands.sql_lab.streaming_export_command import (
@@ -28,6 +29,7 @@ from superset.commands.sql_lab.streaming_export_command import (
from superset.errors import SupersetErrorType
from superset.exceptions import SupersetErrorException, SupersetSecurityException
from superset.sqllab.limiting_factor import LimitingFactor
+from superset.utils.core import get_username
def _setup_sqllab_mocks(
@@ -61,6 +63,7 @@ def mock_query():
query.database = MagicMock()
query.database.db_engine_spec = MagicMock()
query.database.db_engine_spec.engine = "postgresql"
+ query.database.mutate_sql_based_on_config.side_effect = lambda sql_, **kwargs: sql_
query.raise_for_access = MagicMock()
return query
@@ -771,3 +774,120 @@ def test_csv_export_config_custom_decimal_for_decimal_type(mocker, mock_query) -
assert "2;56,78" in csv_data
assert "12.34" not in csv_data
assert "56.78" not in csv_data
+
+
+def test_streaming_export_applies_sql_mutator(mocker, mock_query, mock_result_proxy):
+ """
+ Regression for apache/superset#40465: the non-streaming SQL execution
+ path (Database._execute_sql_with_mutation_and_logging, used by get_df())
+ calls Database.mutate_sql_based_on_config on every statement before
+ executing it -- that's the hook a deployment's SQL_QUERY_MUTATOR runs
+ through (e.g. to strip a trailing semicolon some engines' HTTP endpoints
+ reject, per the issue's Trino repro). _execute_query_and_stream sends
+ the raw SQL straight to engine.connect().execute(text(sql)) instead,
+ bypassing it entirely.
+ """
+ mock_query.select_sql = None
+ mock_query.executed_sql = "SELECT * FROM test_table;"
+ mock_query.limiting_factor = LimitingFactor.NOT_LIMITED
+
+ mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query)
+ mock_query.database.mutate_sql_based_on_config.side_effect = (
+ lambda sql_, is_split=False: sql_.rstrip(";")
+ )
+
+ mock_connection = MagicMock()
+ mock_connection.execution_options.return_value.execute.return_value = (
+ mock_result_proxy
+ )
+ mock_connection.__enter__.return_value = mock_connection
+ mock_connection.__exit__.return_value = None
+
+ mock_engine = MagicMock()
+ mock_engine.connect.return_value = mock_connection
+ mock_query.database.get_sqla_engine.return_value.__enter__.return_value = (
+ mock_engine
+ )
+
+ command = StreamingSqlResultExportCommand("test_client_123", chunk_size=10)
+ command.validate()
+
+ csv_generator_callable = command.run()
+ list(csv_generator_callable())
+
+ executed_sql = str(
+ mock_connection.execution_options.return_value.execute.call_args[0][0]
+ )
+ assert not executed_sql.rstrip().endswith(";"), (
+ "streaming export sent the raw, unmutated SQL to the engine -- "
+ "Database.mutate_sql_based_on_config was never called on it"
+ )
+ # Pin the contract the fix relies on: mutate_sql_based_on_config is
+ # called exactly once, with is_split=True -- the complement of the
+ # is_split=False mutation applied upstream. Without this, the assertion
+ # above still passes with is_split=False (rstrip is idempotent).
+ mock_query.database.mutate_sql_based_on_config.assert_called_once_with(
+ "SELECT * FROM test_table;", is_split=True
+ )
+
+
+def test_streaming_export_preserves_impersonated_user_context(
+ app_context, mocker, mock_query, mock_result_proxy
+):
+ """
+ Checks the other half of apache/superset#40465's theory: that the
+ streaming path drops user impersonation because it acquires its engine
+ "without this context" (unlike other call sites, which the issue says
+ use a get_sqla_engine_with_context(user_name=...) that doesn't actually
+ exist anywhere in this codebase). Independently verified this does NOT
+ reproduce: BaseStreamingCSVExportCommand.run() captures flask.g's full
+ __dict__ (including g.user) before entering the deferred generator, and
+ csv_generator() restores it via preserve_g_context before calling
+ _execute_query_and_stream -- so Database.get_sqla_engine's internal
+ get_effective_user()/get_username() call (which every engine-acquisition
+ call site relies on for impersonation, streaming or not) sees the same
+ acting user it would anywhere else. This pins that behavior rather than
+ the bug the issue reports for it.
+ """
+ g.user = Mock(username="alice")
+ assert get_username() == "alice"
+
+ mock_query.select_sql = None
+ mock_query.executed_sql = "SELECT * FROM test_table"
+ mock_query.limiting_factor = LimitingFactor.NOT_LIMITED
+
+ mock_db, mock_session = _setup_sqllab_mocks(mocker, mock_query)
+
+ mock_connection = MagicMock()
+ mock_connection.execution_options.return_value.execute.return_value = (
+ mock_result_proxy
+ )
+ mock_connection.__enter__.return_value = mock_connection
+ mock_connection.__exit__.return_value = None
+
+ mock_engine = MagicMock()
+ mock_engine.connect.return_value = mock_connection
+
+ seen_usernames = []
+
+ def fake_get_sqla_engine(*args, **kwargs):
+ # get_sqla_engine is where Database._get_sqla_engine would call
+ # get_effective_user() -> get_username() -> g.user.username in the
+ # real (unmocked) implementation; capture what g.user resolves to
+ # at the moment this is invoked, from inside the generator's
+ # restored app/g context.
+ seen_usernames.append(get_username())
+ cm = MagicMock()
+ cm.__enter__.return_value = mock_engine
+ cm.__exit__.return_value = None
+ return cm
+
+ mock_query.database.get_sqla_engine.side_effect = fake_get_sqla_engine
+
+ command = StreamingSqlResultExportCommand("test_client_123", chunk_size=10)
+ command.validate()
+
+ csv_generator_callable = command.run()
+ list(csv_generator_callable())
+
+ assert seen_usernames == ["alice"]