From ec27fa11f159add47a67af3b9ae1a2ef7b3a59c7 Mon Sep 17 00:00:00 2001 From: Prathamesh Hukkeri Date: Sat, 25 Jul 2026 11:04:32 +0530 Subject: [PATCH 1/2] fix(export): apply SQL mutation and user impersonation in streaming CSV export MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streaming CSV export path (introduced by #35478) bypassed two critical steps that all non-streaming paths perform: 1. SQL mutation — `database.mutate_sql_based_on_config()` was never called, leaving trailing semicolons unstripped for engines like Trino that reject them, causing all streaming CSV exports to crash with __STREAM_ERROR__. 2. User impersonation — the path used `get_sqla_engine()` directly instead of `get_raw_connection()`, which bypasses ENGINE_CONTEXT_MANAGER. On databases with `impersonate_user: true` (Trino, Presto, etc.), all streaming CSV exports ran as the service principal, breaking audit trails and potentially bypassing per-user authorization (Ranger, OPA, RLS views). This fix: - Calls `database.mutate_sql_based_on_config(sql)` before execution - Uses `get_raw_connection()` instead of `get_sqla_engine()` directly, which handles SSH tunnels, OAuth2, and user impersonation - Executes via raw DBAPI cursor (same fetchmany() interface) Fixes #40465 --- superset/commands/streaming_export/base.py | 59 ++++-- .../chart/streaming_export_command_test.py | 195 +++++++++++------- 2 files changed, 160 insertions(+), 94 deletions(-) diff --git a/superset/commands/streaming_export/base.py b/superset/commands/streaming_export/base.py index 70e5f6ef6d49..fa2e481d0e79 100644 --- a/superset/commands/streaming_export/base.py +++ b/superset/commands/streaming_export/base.py @@ -23,13 +23,12 @@ import logging import time from abc import abstractmethod -from contextlib import contextmanager +from contextlib import closing, contextmanager from decimal import Decimal from numbers import Real from typing import Any, Callable, Generator from flask import current_app as app, g, has_app_context -from sqlalchemy import text from superset import db from superset.commands.base import BaseCommand @@ -219,19 +218,44 @@ def _execute_query_and_stream( delimiter = csv_export_config.get("sep", ",") decimal_separator = csv_export_config.get("decimal", ".") + # Apply SQL mutations (e.g. SQL_QUERY_MUTATOR config hook) before + # execution. All non-streaming paths go through this — the streaming + # path was originally skipping it, which left trailing semicolons + # unstripped for engines like Trino that reject them. + sql = database.mutate_sql_based_on_config(sql) + with db.session(future=True) as session: # Merge database to prevent DetachedInstanceError merged_database = session.merge(database) - 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)) - - columns = list(result_proxy.keys()) + # Use get_raw_connection() instead of get_sqla_engine() directly. + # This is critical for: + # 1. User impersonation — get_raw_connection() goes through the + # ENGINE_CONTEXT_MANAGER which applies impersonate_user settings + # (e.g. X-Trino-User header). Without this, all streaming CSV + # exports run as the service principal, breaking audit trails + # and potentially bypassing per-user authorization (Ranger, OPA, + # RLS views). + # 2. SSH tunnels — get_raw_connection() sets up SSH tunnels if + # configured on the database. + # 3. OAuth2 — get_raw_connection() wraps execution in + # check_for_oauth2() context. + with closing( + merged_database.get_raw_connection(catalog=catalog, schema=schema) + ) as conn: + cursor = conn.cursor() + # Set cursor.arraysize to control the batch size for fetchmany(). + # This ensures DBAPI drivers (Trino, PostgreSQL, etc.) fetch + # rows in manageable chunks instead of buffering the entire + # result set client-side. + cursor.arraysize = limit + try: + cursor.execute(sql) + columns = ( + [desc[0] for desc in cursor.description] + if cursor.description + else [] + ) # Use StringIO with csv.writer for proper escaping # Apply delimiter from CSV_EXPORT config @@ -247,10 +271,15 @@ def _execute_query_and_stream( total_bytes += header_bytes yield header_data - # Process rows and yield chunks + # Process rows and yield chunks — cursor supports the + # same fetchmany() interface that _process_rows expects. row_count = 0 - for data_chunk, rows_processed, chunk_bytes in self._process_rows( - result_proxy, csv_writer, buffer, limit, decimal_separator + for ( + data_chunk, + rows_processed, + chunk_bytes, + ) in self._process_rows( + cursor, csv_writer, buffer, limit, decimal_separator ): total_bytes += chunk_bytes row_count = rows_processed @@ -265,6 +294,8 @@ def _execute_query_and_stream( total_mb, total_time, ) + finally: + cursor.close() def run(self) -> Callable[[], Generator[str, None, None]]: """ 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 fc12da533220..a10f99b45e50 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -92,9 +92,9 @@ def test_csv_generation_with_small_dataset(mocker: MockerFixture) -> None: """Test CSV generation with a small dataset.""" mock_db, query_context, datasource = _setup_chart_mocks(mocker) - mock_result_proxy = mocker.MagicMock() - mock_result_proxy.keys.return_value = ["col1", "col2", "col3"] - mock_result_proxy.fetchmany.side_effect = [ + mock_cursor = mocker.MagicMock() + mock_cursor.description = [("col1",), ("col2",), ("col3",)] + mock_cursor.fetchmany.side_effect = [ [ ("row1_val1", "row1_val2", "row1_val3"), ("row2_val1", "row2_val2", "row2_val3"), @@ -103,17 +103,11 @@ def test_csv_generation_with_small_dataset(mocker: MockerFixture) -> None: [], ] - mock_connection = mocker.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_conn = mocker.MagicMock() + mock_conn.cursor.return_value = mock_cursor - mock_engine = mocker.MagicMock() - mock_engine.connect.return_value = mock_connection - datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( - mock_engine + datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + mock_conn ) command = StreamingCSVExportCommand(query_context, chunk_size=2) @@ -136,22 +130,18 @@ def test_csv_generation_with_special_characters(mocker: MockerFixture) -> None: """Test CSV generation properly escapes special characters.""" mock_db, query_context, datasource = _setup_chart_mocks(mocker) - mock_result = mocker.MagicMock() - mock_result.keys.return_value = ["name", "description"] - mock_result.fetchmany.side_effect = [ + mock_cursor = mocker.MagicMock() + mock_cursor.description = [("name",), ("description",)] + mock_cursor.fetchmany.side_effect = [ [("John, Jr.", 'Quote"Test'), ("Line\nBreak", "Comma,Value")], [], ] - 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_conn = mocker.MagicMock() + mock_conn.cursor.return_value = mock_cursor - mock_engine = mocker.MagicMock() - mock_engine.connect.return_value = mock_connection - datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( - mock_engine + datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + mock_conn ) command = StreamingCSVExportCommand(query_context, chunk_size=10) @@ -169,22 +159,18 @@ def test_streaming_with_null_values(mocker: MockerFixture) -> None: """Test CSV generation handles NULL values correctly.""" mock_db, query_context, datasource = _setup_chart_mocks(mocker) - mock_result = mocker.MagicMock() - mock_result.keys.return_value = ["col1", "col2", "col3"] - mock_result.fetchmany.side_effect = [ + mock_cursor = mocker.MagicMock() + mock_cursor.description = [("col1",), ("col2",), ("col3",)] + mock_cursor.fetchmany.side_effect = [ [("value1", None, "value3"), (None, "value2", None)], [], ] - 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_conn = mocker.MagicMock() + mock_conn.cursor.return_value = mock_cursor - mock_engine = mocker.MagicMock() - mock_engine.connect.return_value = mock_connection - datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( - mock_engine + datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + mock_conn ) command = StreamingCSVExportCommand(query_context, chunk_size=10) @@ -199,12 +185,12 @@ def test_streaming_with_null_values(mocker: MockerFixture) -> None: def test_streaming_execution_options_enabled(mocker: MockerFixture) -> None: - """Test that streaming execution options are enabled.""" + """Test that get_raw_connection is used for streaming (not get_sqla_engine).""" mock_db, query_context, datasource = _setup_chart_mocks(mocker) - mock_result_proxy = mocker.MagicMock() - mock_result_proxy.keys.return_value = ["col1", "col2", "col3"] - mock_result_proxy.fetchmany.side_effect = [ + mock_cursor = mocker.MagicMock() + mock_cursor.description = [("col1",), ("col2",), ("col3",)] + mock_cursor.fetchmany.side_effect = [ [ ("row1_val1", "row1_val2", "row1_val3"), ("row2_val1", "row2_val2", "row2_val3"), @@ -213,17 +199,11 @@ def test_streaming_execution_options_enabled(mocker: MockerFixture) -> None: [], ] - mock_connection = mocker.MagicMock() - mock_execution_options = mocker.MagicMock() - mock_connection.execution_options.return_value = mock_execution_options - mock_execution_options.execute.return_value = mock_result_proxy - 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 + mock_conn = mocker.MagicMock() + mock_conn.cursor.return_value = mock_cursor + + datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + mock_conn ) command = StreamingCSVExportCommand(query_context) @@ -231,26 +211,25 @@ def test_streaming_execution_options_enabled(mocker: MockerFixture) -> None: generator = csv_generator_callable() list(generator) - mock_connection.execution_options.assert_called_once_with(stream_results=True) + # Verify get_raw_connection is used (not get_sqla_engine) + datasource.database.get_raw_connection.assert_called_once_with( + catalog=None, schema=None + ) def test_empty_result_set(mocker: MockerFixture) -> None: """Test CSV generation with empty result set.""" mock_db, query_context, datasource = _setup_chart_mocks(mocker) - mock_result = mocker.MagicMock() - mock_result.keys.return_value = ["col1", "col2"] - mock_result.fetchmany.side_effect = [[]] + mock_cursor = mocker.MagicMock() + mock_cursor.description = [("col1",), ("col2",)] + mock_cursor.fetchmany.side_effect = [[]] - 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_conn = mocker.MagicMock() + mock_conn.cursor.return_value = mock_cursor - mock_engine = mocker.MagicMock() - mock_engine.connect.return_value = mock_connection - datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( - mock_engine + datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + mock_conn ) command = StreamingCSVExportCommand(query_context) @@ -264,35 +243,91 @@ def test_empty_result_set(mocker: MockerFixture) -> None: def test_catalog_and_schema_passed_to_engine(mocker: MockerFixture) -> None: - """Test that catalog and schema are forwarded to get_sqla_engine. - - Prequeries (e.g. SET search_path for PostgreSQL) are now run automatically - via a connect event listener registered inside get_sqla_engine, not by the - streaming command itself. - """ + """Test that catalog and schema are forwarded to get_raw_connection.""" mock_db, query_context, datasource = _setup_chart_mocks( mocker, catalog="my_catalog", schema="my_schema" ) - mock_result = mocker.MagicMock() - mock_result.keys.return_value = ["col1"] - mock_result.fetchmany.side_effect = [[("val",)], []] + mock_cursor = mocker.MagicMock() + mock_cursor.description = [("col1",)] + mock_cursor.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_conn = mocker.MagicMock() + mock_conn.cursor.return_value = mock_cursor - mock_engine = mocker.MagicMock() - mock_engine.connect.return_value = mock_connection - datasource.database.get_sqla_engine.return_value.__enter__.return_value = ( - mock_engine + datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + mock_conn ) command = StreamingCSVExportCommand(query_context) list(command.run()()) - datasource.database.get_sqla_engine.assert_called_once_with( + datasource.database.get_raw_connection.assert_called_once_with( catalog="my_catalog", schema="my_schema", ) + + +def test_sql_mutation_applied_before_execution(mocker: MockerFixture) -> None: + """Test that mutate_sql_based_on_config is called before executing SQL. + + Regression test for #40465: the streaming export path was executing raw + SQL without applying SQL_QUERY_MUTATOR config, leaving trailing semicolons + unstripped for engines like Trino that reject them. + """ + mock_db, query_context, datasource = _setup_chart_mocks(mocker) + datasource.database.mutate_sql_based_on_config.return_value = "SELECT 1 LIMIT 10" + + mock_cursor = mocker.MagicMock() + mock_cursor.description = [("col1",)] + mock_cursor.fetchmany.side_effect = [[(1,)], []] + + mock_conn = mocker.MagicMock() + mock_conn.cursor.return_value = mock_cursor + + datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + mock_conn + ) + + command = StreamingCSVExportCommand(query_context) + list(command.run()()) + + # SQL mutation must be called before execution + datasource.database.mutate_sql_based_on_config.assert_called_once_with( + "SELECT * FROM test" + ) + # The mutated SQL (not the original) should be sent to the cursor + mock_cursor.execute.assert_called_once_with("SELECT 1 LIMIT 10") + + +def test_get_raw_connection_used_instead_of_get_sqla_engine( + mocker: MockerFixture, +) -> None: + """Test that get_raw_connection is used for proper user impersonation. + + Regression test for #40465: the streaming export path used + get_sqla_engine() directly, bypassing user impersonation. This meant + all streaming CSV exports ran as the service principal instead of the + logged-in user, breaking audit trails and potentially bypassing + per-user authorization (Ranger, OPA, RLS views). + """ + mock_db, query_context, datasource = _setup_chart_mocks(mocker) + + mock_cursor = mocker.MagicMock() + mock_cursor.description = [("col1",)] + mock_cursor.fetchmany.side_effect = [[("val",)], []] + + mock_conn = mocker.MagicMock() + mock_conn.cursor.return_value = mock_cursor + + datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + mock_conn + ) + + command = StreamingCSVExportCommand(query_context) + list(command.run()()) + + # Must use get_raw_connection (handles impersonation, SSH, OAuth2) + datasource.database.get_raw_connection.assert_called_once() + # Must NOT use get_sqla_engine directly (bypasses impersonation) + datasource.database.get_sqla_engine.assert_not_called() From e2f2258882716a909133a254dc743bb80342c28f Mon Sep 17 00:00:00 2001 From: Prathamesh Hukkeri Date: Wed, 29 Jul 2026 09:57:51 +0530 Subject: [PATCH 2/2] fix(tests): remove __enter__ mock from get_raw_connection (uses closing()) --- .../chart/streaming_export_command_test.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) 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 a10f99b45e50..71052b07ecae 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -106,7 +106,7 @@ def test_csv_generation_with_small_dataset(mocker: MockerFixture) -> None: mock_conn = mocker.MagicMock() mock_conn.cursor.return_value = mock_cursor - datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + datasource.database.get_raw_connection.return_value = ( mock_conn ) @@ -140,7 +140,7 @@ def test_csv_generation_with_special_characters(mocker: MockerFixture) -> None: mock_conn = mocker.MagicMock() mock_conn.cursor.return_value = mock_cursor - datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + datasource.database.get_raw_connection.return_value = ( mock_conn ) @@ -169,7 +169,7 @@ def test_streaming_with_null_values(mocker: MockerFixture) -> None: mock_conn = mocker.MagicMock() mock_conn.cursor.return_value = mock_cursor - datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + datasource.database.get_raw_connection.return_value = ( mock_conn ) @@ -202,7 +202,7 @@ def test_streaming_execution_options_enabled(mocker: MockerFixture) -> None: mock_conn = mocker.MagicMock() mock_conn.cursor.return_value = mock_cursor - datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + datasource.database.get_raw_connection.return_value = ( mock_conn ) @@ -228,7 +228,7 @@ def test_empty_result_set(mocker: MockerFixture) -> None: mock_conn = mocker.MagicMock() mock_conn.cursor.return_value = mock_cursor - datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + datasource.database.get_raw_connection.return_value = ( mock_conn ) @@ -255,7 +255,7 @@ def test_catalog_and_schema_passed_to_engine(mocker: MockerFixture) -> None: mock_conn = mocker.MagicMock() mock_conn.cursor.return_value = mock_cursor - datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + datasource.database.get_raw_connection.return_value = ( mock_conn ) @@ -285,7 +285,7 @@ def test_sql_mutation_applied_before_execution(mocker: MockerFixture) -> None: mock_conn = mocker.MagicMock() mock_conn.cursor.return_value = mock_cursor - datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + datasource.database.get_raw_connection.return_value = ( mock_conn ) @@ -320,7 +320,7 @@ def test_get_raw_connection_used_instead_of_get_sqla_engine( mock_conn = mocker.MagicMock() mock_conn.cursor.return_value = mock_cursor - datasource.database.get_raw_connection.return_value.__enter__.return_value = ( + datasource.database.get_raw_connection.return_value = ( mock_conn )