fix(export): apply SQL mutation and user impersonation in streaming CSV export

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]]:
         """

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 = (
+        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 = (
+        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 = (
+        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,44 +199,37 @@ 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 = (
+        mock_conn
     )
 
     command = StreamingCSVExportCommand(query_context)
     csv_generator_callable = command.run()
     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 = (
+        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 = (
+        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 = (
+        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 = (
+        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()