rebenitez1802 commented on code in PR #42739:
URL: https://github.com/apache/superset/pull/42739#discussion_r3764418116
##########
tests/unit_tests/commands/sql_lab/streaming_export_command_test.py:
##########
@@ -771,3 +774,113 @@ 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"
+ )
Review Comment:
This mirrors the chart test's `assert_called_once_with(..., is_split=True)`
and closes the gap @sadpandajoe flagged. As written, the assertion above passes
for **both** `is_split=True` (the fix) and `is_split=False` (the bug the code
comment warns about), because `rstrip(";")` is idempotent and the flag is never
checked. Pinning the call turns this into a real regression guard:
```suggestion
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
)
```
--
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]
