Copilot commented on code in PR #42412:
URL: https://github.com/apache/superset/pull/42412#discussion_r3659574232
##########
superset/commands/streaming_export/base.py:
##########
@@ -219,19 +218,39 @@ 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()
Review Comment:
`get_raw_connection()` in Superset is typically a context manager (your
tests also mock `__enter__`/`__exit__`). Wrapping it with `closing(...)` means
you never enter the context manager, so the ENGINE_CONTEXT_MANAGER
(impersonation/SSH/OAuth2 wrappers) may not run, and cleanup via `__exit__` may
be skipped. Use `with merged_database.get_raw_connection(catalog=catalog,
schema=schema) as conn:` and drop `closing(...)` so the context manager is
properly entered/exited.
##########
tests/unit_tests/commands/chart/streaming_export_command_test.py:
##########
@@ -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)."""
Review Comment:
The test name `test_streaming_execution_options_enabled` no longer matches
what is being asserted (it now checks `get_raw_connection` usage, not execution
options). Rename the test function to reflect the new behavior (e.g.,
`test_streaming_uses_get_raw_connection`).
##########
tests/unit_tests/commands/chart/streaming_export_command_test.py:
##########
@@ -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
)
Review Comment:
These tests configure `get_raw_connection.return_value.__enter__`, but none
of the updated tests assert that the context manager is actually entered/exited
(e.g., `__enter__` called). This can allow an implementation bug (calling
`get_raw_connection()` without a `with`) to pass silently. Add assertions like
`datasource.database.get_raw_connection.return_value.__enter__.assert_called_once()`
(and optionally `__exit__`) in at least one test that exercises the streaming
execution.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
