rusackas commented on code in PR #42739:
URL: https://github.com/apache/superset/pull/42739#discussion_r3765287127
##########
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:
Applied as suggested, and added a docstring note on
`_get_sql_and_database()` per the optional follow-up.
--
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]
