Spaces:
Paused
Paused
| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| """Tests for celery_task.py - async SQL execution via Celery.""" | |
| from typing import Any | |
| from unittest.mock import MagicMock | |
| import msgpack | |
| import pytest | |
| from celery.exceptions import SoftTimeLimitExceeded | |
| from flask import current_app | |
| from pytest_mock import MockerFixture | |
| from superset.common.db_query_status import QueryStatus as QueryStatusEnum | |
| from superset.errors import ErrorLevel, SupersetError, SupersetErrorType | |
| from superset.exceptions import SupersetErrorException, SupersetErrorsException | |
| # Note: mock_query, mock_database, mock_result_set, and mock_db_session | |
| # fixtures are imported from conftest.py | |
| from .conftest import _passthrough_mutate_sql_based_on_config | |
| def _prefixing_mutator(sql: str, **kwargs: Any) -> str: | |
| """SQL mutator stand-in that prepends a marker comment.""" | |
| return f"-- mutated\n{sql}" | |
| # ============================================================================= | |
| # Query Retrieval Tests | |
| # ============================================================================= | |
| def test_get_query_success( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test successful query retrieval.""" | |
| from superset.sql.execution.celery_task import _get_query | |
| mock_session = mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mock_session.query.return_value.filter_by.return_value.one.return_value = mock_query | |
| result = _get_query(123) | |
| assert result == mock_query | |
| # ============================================================================= | |
| # Error Handling Tests | |
| # ============================================================================= | |
| def test_handle_query_error_basic( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test basic error handling.""" | |
| from superset.sql.execution.celery_task import _handle_query_error | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| ex = Exception("Something went wrong") | |
| payload = _handle_query_error(ex, mock_query) | |
| assert payload["status"] == QueryStatusEnum.FAILED | |
| assert "Something went wrong" in payload["error"] | |
| def test_handle_query_error_with_end_time_set( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test error handling when end_time is already set (line 116->120).""" | |
| from superset.sql.execution.celery_task import _handle_query_error | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| # Set end_time to trigger the branch skip | |
| mock_query.end_time = 12345.0 | |
| ex = Exception("Error with end_time set") | |
| payload = _handle_query_error(ex, mock_query) | |
| assert payload["status"] == QueryStatusEnum.FAILED | |
| # end_time should not be modified | |
| assert mock_query.end_time == 12345.0 | |
| def test_handle_query_error_sets_end_time( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test error handling sets end_time when not set.""" | |
| from superset.sql.execution.celery_task import _handle_query_error | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.now_as_float", return_value=99999.0 | |
| ) | |
| # end_time is None | |
| mock_query.end_time = None | |
| ex = Exception("Error") | |
| _handle_query_error(ex, mock_query) | |
| # Should set end_time | |
| assert mock_query.end_time == 99999.0 | |
| def test_handle_query_error_superset_error_exception( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test error handling with SupersetErrorException.""" | |
| from superset.sql.execution.celery_task import _handle_query_error | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| error = SupersetError( | |
| message="Test error", | |
| error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR, | |
| level=ErrorLevel.ERROR, | |
| ) | |
| ex = SupersetErrorException(error) | |
| payload = _handle_query_error(ex, mock_query) | |
| assert len(payload["errors"]) == 1 | |
| assert payload["errors"][0]["message"] == "Test error" | |
| def test_handle_query_error_superset_errors_exception( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test error handling with SupersetErrorsException.""" | |
| from superset.sql.execution.celery_task import _handle_query_error | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| errors = [ | |
| SupersetError( | |
| message="Error 1", | |
| error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR, | |
| level=ErrorLevel.ERROR, | |
| ), | |
| SupersetError( | |
| message="Error 2", | |
| error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR, | |
| level=ErrorLevel.ERROR, | |
| ), | |
| ] | |
| ex = SupersetErrorsException(errors) | |
| payload = _handle_query_error(ex, mock_query) | |
| assert len(payload["errors"]) == 2 | |
| def test_handle_query_error_with_troubleshooting_link( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test error handling includes troubleshooting link.""" | |
| from superset.sql.execution.celery_task import _handle_query_error | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch.dict( | |
| current_app.config, {"TROUBLESHOOTING_LINK": "https://help.example.com"} | |
| ) | |
| ex = Exception("Error") | |
| payload = _handle_query_error(ex, mock_query) | |
| assert payload["link"] == "https://help.example.com" | |
| def test_handle_query_error_no_stacktrace_when_format_exc_empty( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test that stacktrace key is omitted when traceback.format_exc returns empty.""" | |
| from superset.sql.execution.celery_task import _handle_query_error | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.traceback.format_exc", return_value="" | |
| ) | |
| ex = Exception("Error") | |
| payload = _handle_query_error(ex, mock_query) | |
| assert "stacktrace" not in payload | |
| def test_handle_query_error_includes_stacktrace_when_show_stacktrace_enabled( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test stacktrace key is included when SHOW_STACKTRACE=True. | |
| Covers lines 107-108 in celery_task._handle_query_error. | |
| """ | |
| from superset.sql.execution.celery_task import _handle_query_error | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch.dict(current_app.config, {"SHOW_STACKTRACE": True}) | |
| try: | |
| raise RuntimeError("boom") | |
| except RuntimeError as ex: | |
| payload = _handle_query_error(ex, mock_query) | |
| assert "stacktrace" in payload | |
| assert payload["stacktrace"] | |
| def test_handle_query_error_omits_stacktrace_when_show_stacktrace_enabled_but_empty( | |
| mocker: MockerFixture, app_context: None, mock_query: MagicMock | |
| ) -> None: | |
| """Test stacktrace key is omitted when SHOW_STACKTRACE=True but format_exc is empty. | |
| Covers the branch at line 107 where SHOW_STACKTRACE is True but | |
| traceback.format_exc() returns an empty string (no active exception context). | |
| """ | |
| from superset.sql.execution.celery_task import _handle_query_error | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch.dict(current_app.config, {"SHOW_STACKTRACE": True}) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.traceback.format_exc", return_value="" | |
| ) | |
| ex = Exception("Error") | |
| payload = _handle_query_error(ex, mock_query) | |
| assert "stacktrace" not in payload | |
| # ============================================================================= | |
| # Serialization Tests | |
| # ============================================================================= | |
| def test_serialize_payload_json(mocker: MockerFixture, app_context: None) -> None: | |
| """Test JSON serialization when msgpack config is False.""" | |
| from superset.sql.execution.celery_task import _serialize_payload | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| payload = {"status": "success", "data": [1, 2, 3]} | |
| result = _serialize_payload(payload) | |
| # Now always returns bytes (encoded UTF-8) | |
| assert isinstance(result, bytes) | |
| assert b"success" in result | |
| def test_serialize_payload_msgpack(mocker: MockerFixture, app_context: None) -> None: | |
| """Test msgpack serialization when msgpack config is True.""" | |
| from superset.sql.execution.celery_task import _serialize_payload | |
| mocker.patch("superset.results_backend_use_msgpack", True) | |
| payload = {"status": "success", "data": [1, 2, 3]} | |
| result = _serialize_payload(payload) | |
| assert isinstance(result, bytes) | |
| unpacked = msgpack.loads(result) | |
| assert unpacked["status"] == "success" | |
| # ============================================================================= | |
| # Statement Preparation Tests | |
| # ============================================================================= | |
| def test_prepare_statement_blocks_single_statement( | |
| app_context: None, mock_database: MagicMock | |
| ) -> None: | |
| """Test statement block preparation for single statement.""" | |
| from superset.sql.execution.celery_task import _prepare_statement_blocks | |
| mock_database.mutate_sql_based_on_config = _passthrough_mutate_sql_based_on_config | |
| sql = "SELECT * FROM users" | |
| script, blocks = _prepare_statement_blocks( | |
| sql, mock_database.db_engine_spec, mock_database | |
| ) | |
| assert len(blocks) == 1 | |
| def test_prepare_statement_blocks_multiple_statements( | |
| app_context: None, mock_database: MagicMock | |
| ) -> None: | |
| """Test statement block preparation for multiple statements.""" | |
| from superset.sql.execution.celery_task import _prepare_statement_blocks | |
| mock_database.mutate_sql_based_on_config = _passthrough_mutate_sql_based_on_config | |
| sql = "SELECT * FROM users; SELECT * FROM orders;" | |
| script, blocks = _prepare_statement_blocks( | |
| sql, mock_database.db_engine_spec, mock_database | |
| ) | |
| assert len(blocks) == 2 | |
| def test_prepare_statement_blocks_run_as_one( | |
| app_context: None, mock_database: MagicMock | |
| ) -> None: | |
| """Test statement block preparation when engine runs multiple as one.""" | |
| from superset.sql.execution.celery_task import _prepare_statement_blocks | |
| mock_database.db_engine_spec.run_multiple_statements_as_one = True | |
| mock_database.mutate_sql_based_on_config = _passthrough_mutate_sql_based_on_config | |
| sql = "SELECT * FROM users; SELECT * FROM orders;" | |
| script, blocks = _prepare_statement_blocks( | |
| sql, mock_database.db_engine_spec, mock_database | |
| ) | |
| assert len(blocks) == 1 | |
| def test_prepare_statement_blocks_mutates_before_split_when_configured( | |
| app_context: None, mock_database: MagicMock, mocker: MockerFixture | |
| ) -> None: | |
| """ | |
| `MUTATE_AFTER_SPLIT=False` should mutate the whole, un-split query before | |
| it gets broken into per-statement blocks, for engines that execute | |
| statements individually. Regression guard for issue #30169. | |
| """ | |
| from superset.sql.execution.celery_task import _prepare_statement_blocks | |
| mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": False}) | |
| mutate_mock = mocker.patch.object( | |
| mock_database, | |
| "mutate_sql_based_on_config", | |
| side_effect=_prefixing_mutator, | |
| ) | |
| sql = "SELECT * FROM users; SELECT * FROM orders;" | |
| _, blocks = _prepare_statement_blocks( | |
| sql, mock_database.db_engine_spec, mock_database | |
| ) | |
| assert len(blocks) == 2 | |
| assert "mutated" in blocks[0] | |
| mutate_mock.assert_called_once_with(mocker.ANY, is_split=False) | |
| def test_prepare_statement_blocks_skips_pre_split_mutation_when_configured( | |
| app_context: None, mock_database: MagicMock, mocker: MockerFixture | |
| ) -> None: | |
| """ | |
| `MUTATE_AFTER_SPLIT=True` means the mutator should see each already-split | |
| statement instead, so `_prepare_statement_blocks` must not mutate the | |
| whole, un-split query up front. | |
| """ | |
| from superset.sql.execution.celery_task import _prepare_statement_blocks | |
| mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": True}) | |
| mutate_mock = mocker.patch.object( | |
| mock_database, | |
| "mutate_sql_based_on_config", | |
| side_effect=_prefixing_mutator, | |
| ) | |
| sql = "SELECT * FROM users; SELECT * FROM orders;" | |
| _, blocks = _prepare_statement_blocks( | |
| sql, mock_database.db_engine_spec, mock_database | |
| ) | |
| assert len(blocks) == 2 | |
| assert "mutated" not in blocks[0] | |
| mutate_mock.assert_not_called() | |
| def test_prepare_statement_blocks_mutates_per_statement_when_run_as_one( | |
| app_context: None, mock_database: MagicMock, mocker: MockerFixture | |
| ) -> None: | |
| """ | |
| Engines that always run statements as a single block (e.g. BigQuery, Kusto) | |
| never see `is_split=True` in the per-block mutation call, so with | |
| `MUTATE_AFTER_SPLIT=True` the mutator must instead be applied to each | |
| statement here, before they're joined into that single block. | |
| """ | |
| from superset.sql.execution.celery_task import _prepare_statement_blocks | |
| mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": True}) | |
| mock_database.db_engine_spec.run_multiple_statements_as_one = True | |
| mutate_mock = mocker.patch.object( | |
| mock_database, | |
| "mutate_sql_based_on_config", | |
| side_effect=_prefixing_mutator, | |
| ) | |
| sql = "SELECT * FROM users; SELECT * FROM orders;" | |
| _, blocks = _prepare_statement_blocks( | |
| sql, mock_database.db_engine_spec, mock_database | |
| ) | |
| assert len(blocks) == 1 | |
| assert blocks[0].count("mutated") == 2 | |
| is_split_values = [ | |
| call.kwargs.get("is_split") for call in mutate_mock.call_args_list | |
| ] | |
| assert is_split_values == [True, True] | |
| def test_prepare_statement_blocks_raises_when_mutator_strips_all_statements( | |
| app_context: None, mock_database: MagicMock, mocker: MockerFixture | |
| ) -> None: | |
| """ | |
| A `SQL_QUERY_MUTATOR` that strips a query down to nothing (e.g. only | |
| comments/whitespace) must raise a clean error instead of silently | |
| producing an empty block list. | |
| """ | |
| from superset.sql.execution.celery_task import _prepare_statement_blocks | |
| mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": False}) | |
| mocker.patch.object( | |
| mock_database, | |
| "mutate_sql_based_on_config", | |
| side_effect=lambda sql, **kw: "-- just a comment", | |
| ) | |
| sql = "SELECT * FROM users" | |
| with pytest.raises(SupersetErrorException): | |
| _prepare_statement_blocks(sql, mock_database.db_engine_spec, mock_database) | |
| def test_prepare_statement_blocks_raises_when_mutator_strips_single_block( | |
| app_context: None, mock_database: MagicMock, mocker: MockerFixture | |
| ) -> None: | |
| """ | |
| The empty-statement guard must also cover engines that run all statements | |
| as one block: with `MUTATE_AFTER_SPLIT=True` the per-statement mutator | |
| outputs are joined into a single block, and a comment-only/empty result | |
| must raise a clean error instead of reaching execution as an empty block. | |
| """ | |
| from superset.sql.execution.celery_task import _prepare_statement_blocks | |
| mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": True}) | |
| mock_database.db_engine_spec.run_multiple_statements_as_one = True | |
| mocker.patch.object( | |
| mock_database, | |
| "mutate_sql_based_on_config", | |
| side_effect=lambda sql, **kw: "-- just a comment", | |
| ) | |
| sql = "SELECT * FROM users; SELECT * FROM orders;" | |
| with pytest.raises(SupersetErrorException): | |
| _prepare_statement_blocks(sql, mock_database.db_engine_spec, mock_database) | |
| # ============================================================================= | |
| # Result Finalization Tests | |
| # ============================================================================= | |
| def test_finalize_successful_query( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_result_set: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test successful query finalization.""" | |
| from superset.sql.execution.celery_task import _finalize_successful_query | |
| from superset.sql.parse import SQLScript | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch("superset.dataframe.df_to_records", return_value=[{"id": 1}]) | |
| payload: dict[str, Any] = {} | |
| # Create original script | |
| original_script = SQLScript( | |
| "SELECT * FROM users", mock_database.db_engine_spec.engine | |
| ) | |
| # New signature: (query, original_script, execution_results, payload, | |
| # total_execution_time_ms). execution_results is a list of tuples: | |
| # (executed_sql, result_set, exec_time, rowcount) | |
| execution_results = [ | |
| ("SELECT * FROM users WHERE rls_filter", mock_result_set, 10.5, 2) | |
| ] | |
| _finalize_successful_query( | |
| mock_query, | |
| original_script, | |
| execution_results, # type: ignore[arg-type] | |
| payload, | |
| 10.5, | |
| ) | |
| assert mock_query.rows == 2 | |
| assert mock_query.progress == 100 | |
| assert payload["status"] == QueryStatusEnum.SUCCESS | |
| assert "statements" in payload | |
| # SQL is formatted by SQLScript, so we can't compare exact whitespace | |
| assert "SELECT" in payload["statements"][0]["original_sql"] | |
| assert "FROM users" in payload["statements"][0]["original_sql"] | |
| assert ( | |
| payload["statements"][0]["executed_sql"] | |
| == "SELECT * FROM users WHERE rls_filter" | |
| ) | |
| def test_finalize_successful_query_statement_count_mismatch( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_result_set: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """ | |
| A `SQL_QUERY_MUTATOR` that changes the number of statements (e.g. by | |
| prepending a statement when run on the whole, un-split query) must not | |
| crash finalization when `original_script`'s statement count no longer | |
| matches `execution_results`. | |
| """ | |
| from superset.sql.execution.celery_task import _finalize_successful_query | |
| from superset.sql.parse import SQLScript | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch("superset.dataframe.df_to_records", return_value=[{"id": 1}]) | |
| payload: dict[str, Any] = {} | |
| # Only one statement in the un-mutated original... | |
| original_script = SQLScript( | |
| "SELECT * FROM users", mock_database.db_engine_spec.engine | |
| ) | |
| # ...but the mutator turned it into two. | |
| execution_results = [ | |
| ("SET ROLE 'bob'", None, 1.0, 0), | |
| ("SELECT * FROM users", mock_result_set, 10.5, 2), | |
| ] | |
| _finalize_successful_query( | |
| mock_query, | |
| original_script, | |
| execution_results, # type: ignore[arg-type] | |
| payload, | |
| 11.5, | |
| ) | |
| assert mock_query.rows == 2 | |
| assert payload["status"] == QueryStatusEnum.SUCCESS | |
| assert len(payload["statements"]) == 2 | |
| # Falls back to labeling each result with its own executed SQL. | |
| assert payload["statements"][0]["original_sql"] == "SET ROLE 'bob'" | |
| assert payload["statements"][1]["original_sql"] == "SELECT * FROM users" | |
| def test_finalize_successful_query_with_msgpack( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_result_set: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test successful query finalization with Arrow/msgpack.""" | |
| from superset.sql.execution.celery_task import _finalize_successful_query | |
| from superset.sql.parse import SQLScript | |
| mocker.patch("superset.results_backend_use_msgpack", True) | |
| mock_buffer = MagicMock() | |
| mock_buffer.to_pybytes.return_value = b"arrow_data" | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.write_ipc_buffer", return_value=mock_buffer | |
| ) | |
| # Mock stats_logger to cover the stats timing branch | |
| mock_stats = MagicMock() | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": mock_stats}) | |
| payload: dict[str, Any] = {} | |
| # Create original script | |
| original_script = SQLScript( | |
| "SELECT * FROM users", mock_database.db_engine_spec.engine | |
| ) | |
| execution_results = [ | |
| ("SELECT * FROM users WHERE rls_filter", mock_result_set, 10.5, 2) | |
| ] | |
| _finalize_successful_query( | |
| mock_query, | |
| original_script, | |
| execution_results, # type: ignore[arg-type] | |
| payload, | |
| 10.5, | |
| ) | |
| assert payload["statements"][0]["data"] == b"arrow_data" | |
| def test_finalize_successful_query_msgpack_no_stats( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_result_set: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test finalization with msgpack when has_app_context() is False.""" | |
| from superset.sql.execution.celery_task import _finalize_successful_query | |
| from superset.sql.parse import SQLScript | |
| mocker.patch("superset.results_backend_use_msgpack", True) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.has_app_context", return_value=False | |
| ) | |
| mock_buffer = MagicMock() | |
| mock_buffer.to_pybytes.return_value = b"arrow_data" | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.write_ipc_buffer", return_value=mock_buffer | |
| ) | |
| payload: dict[str, Any] = {} | |
| # Create original script | |
| original_script = SQLScript( | |
| "SELECT * FROM users", mock_database.db_engine_spec.engine | |
| ) | |
| execution_results = [ | |
| ("SELECT * FROM users WHERE rls_filter", mock_result_set, 10.5, 2) | |
| ] | |
| _finalize_successful_query( | |
| mock_query, | |
| original_script, | |
| execution_results, # type: ignore[arg-type] | |
| payload, | |
| 10.5, | |
| ) | |
| assert payload["statements"][0]["data"] == b"arrow_data" | |
| def test_finalize_successful_query_with_dml( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test successful query finalization with DML statement (no result_set).""" | |
| from superset.sql.execution.celery_task import _finalize_successful_query | |
| from superset.sql.parse import SQLScript | |
| payload: dict[str, Any] = {} | |
| # Create original script | |
| original_script = SQLScript( | |
| "INSERT INTO users VALUES (1)", mock_database.db_engine_spec.engine | |
| ) | |
| # DML statement: result_set is None, rowcount indicates affected rows | |
| execution_results: list[tuple[str, None, float, int]] = [ | |
| ("INSERT INTO users VALUES (1)", None, 5.0, 1) | |
| ] | |
| _finalize_successful_query( | |
| mock_query, | |
| original_script, | |
| execution_results, # type: ignore[arg-type] | |
| payload, | |
| 5.0, | |
| ) | |
| assert mock_query.rows == 0 # No rows returned for DML | |
| assert mock_query.progress == 100 | |
| assert payload["status"] == QueryStatusEnum.SUCCESS | |
| assert payload["statements"][0]["data"] is None | |
| assert payload["statements"][0]["row_count"] == 1 | |
| assert payload["statements"][0]["columns"] == [] | |
| # ============================================================================= | |
| # Results Storage Tests | |
| # ============================================================================= | |
| def test_store_results_in_backend_success( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test successful results storage.""" | |
| from superset.sql.execution.celery_task import _store_results_in_backend | |
| mock_results_backend = MagicMock() | |
| mock_results_backend.set.return_value = True | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.results_backend", mock_results_backend | |
| ) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.zlib_compress", return_value=b"compressed" | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| payload = {"status": "success", "data": [], "query": {}} | |
| _store_results_in_backend(mock_query, payload, mock_database) | |
| assert mock_query.results_key is not None | |
| mock_results_backend.set.assert_called_once() | |
| def test_store_results_in_backend_with_size_check( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test results storage with payload size check (covers lines 232-247).""" | |
| from superset.sql.execution.celery_task import _store_results_in_backend | |
| mock_results_backend = MagicMock() | |
| mock_results_backend.set.return_value = True | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.results_backend", mock_results_backend | |
| ) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.zlib_compress", return_value=b"compressed" | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| # Set a high payload max to pass the size check | |
| mocker.patch.dict(current_app.config, {"SQLLAB_PAYLOAD_MAX_MB": 100}) | |
| payload = {"status": "success", "data": [], "query": {}} | |
| _store_results_in_backend(mock_query, payload, mock_database) | |
| mock_results_backend.set.assert_called_once() | |
| def test_store_results_in_backend_payload_too_large( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test results storage with payload exceeding size limit.""" | |
| from superset.sql.execution.celery_task import _store_results_in_backend | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| # Set very low limit | |
| mocker.patch.dict(current_app.config, {"SQLLAB_PAYLOAD_MAX_MB": 0.000001}) | |
| large_payload = {"data": "x" * 1000, "query": {}} | |
| with pytest.raises(SupersetErrorException) as exc_info: | |
| _store_results_in_backend(mock_query, large_payload, mock_database) | |
| assert exc_info.value.error.error_type == SupersetErrorType.RESULT_TOO_LARGE_ERROR | |
| def test_store_results_in_backend_default_cache_timeout( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test storage uses default cache timeout when database timeout is None.""" | |
| from superset.sql.execution.celery_task import _store_results_in_backend | |
| mock_results_backend = MagicMock() | |
| mock_results_backend.set.return_value = True | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.results_backend", mock_results_backend | |
| ) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.zlib_compress", return_value=b"compressed" | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| # Set database cache_timeout to None | |
| mock_database.cache_timeout = None | |
| payload = {"status": "success", "data": [], "query": {}} | |
| _store_results_in_backend(mock_query, payload, mock_database) | |
| mock_results_backend.set.assert_called_once() | |
| def test_store_results_in_backend_write_failure( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test results storage write failure.""" | |
| from superset.sql.execution.celery_task import _store_results_in_backend | |
| mock_results_backend = MagicMock() | |
| mock_results_backend.set.return_value = False | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.results_backend", mock_results_backend | |
| ) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.zlib_compress", return_value=b"compressed" | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| payload = {"status": "success", "data": [], "query": {}} | |
| with pytest.raises(SupersetErrorException) as exc_info: | |
| _store_results_in_backend(mock_query, payload, mock_database) | |
| assert exc_info.value.error.error_type == SupersetErrorType.RESULTS_BACKEND_ERROR | |
| # ============================================================================= | |
| # Data Serialization Tests | |
| # ============================================================================= | |
| def test_serialize_result_set_msgpack( | |
| mocker: MockerFixture, app_context: None, mock_result_set: MagicMock | |
| ) -> None: | |
| """Test result set serialization with msgpack/Arrow when config is True.""" | |
| from superset.sql.execution.celery_task import _serialize_result_set | |
| mocker.patch("superset.results_backend_use_msgpack", True) | |
| mock_buffer = MagicMock() | |
| mock_buffer.to_pybytes.return_value = b"arrow_data" | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.write_ipc_buffer", return_value=mock_buffer | |
| ) | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| data, columns = _serialize_result_set(mock_result_set) | |
| assert data == b"arrow_data" | |
| assert columns == mock_result_set.columns | |
| def test_serialize_result_set_json( | |
| mocker: MockerFixture, app_context: None, mock_result_set: MagicMock | |
| ) -> None: | |
| """Test result set serialization with JSON when msgpack config is False.""" | |
| from superset.sql.execution.celery_task import _serialize_result_set | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch( | |
| "superset.dataframe.df_to_records", | |
| return_value=[{"id": 1, "name": "Alice"}], | |
| ) | |
| data, columns = _serialize_result_set(mock_result_set) | |
| assert data == [{"id": 1, "name": "Alice"}] | |
| assert columns == mock_result_set.columns | |
| # ============================================================================= | |
| # Helper Function Tests | |
| # ============================================================================= | |
| def test_make_check_stopped_fn( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| query_status: QueryStatusEnum, | |
| expected_result: bool, | |
| ) -> None: | |
| """Test check_stopped function returns correct value based on query status.""" | |
| from superset.sql.execution.celery_task import _make_check_stopped_fn | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mock_query.status = query_status | |
| check_stopped = _make_check_stopped_fn(mock_query) | |
| result = check_stopped() | |
| assert result is expected_result | |
| def test_make_execute_fn( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test execute function creation.""" | |
| from superset.sql.execution.celery_task import _make_execute_fn | |
| mock_cursor = MagicMock() | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| execute_fn = _make_execute_fn(mock_query, mock_database.db_engine_spec) | |
| execute_fn(mock_cursor, "SELECT * FROM users") | |
| assert mock_query.executed_sql == "SELECT * FROM users" | |
| def test_make_log_query_fn( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_database: MagicMock, | |
| logger_configured: bool, | |
| should_be_called: bool, | |
| ) -> None: | |
| """Test log query function with and without logger configured.""" | |
| from superset.sql.execution.celery_task import _make_log_query_fn | |
| mock_logger = MagicMock() if logger_configured else None | |
| mocker.patch.dict(current_app.config, {"QUERY_LOGGER": mock_logger}) | |
| mocker.patch("superset.sql.execution.celery_task.security_manager", MagicMock()) | |
| log_fn = _make_log_query_fn(mock_database) | |
| log_fn("SELECT * FROM users", "public") | |
| if should_be_called: | |
| assert mock_logger is not None | |
| mock_logger.assert_called_once() | |
| # If no logger, the function should complete without error | |
| # ============================================================================= | |
| # Main Task Execution Tests | |
| # ============================================================================= | |
| def test_execute_sql_task_success( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| mock_result_set: MagicMock, | |
| ) -> None: | |
| """Test successful execute_sql_task (covers lines 339-352, 400-473).""" | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| from .conftest import setup_mock_raw_connection | |
| mock_query.database = mock_database | |
| mock_query.status = QueryStatusEnum.PENDING | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.execute_sql_with_cursor", | |
| return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.results_backend", None) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.dataframe.df_to_records", return_value=[]) | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| setup_mock_raw_connection(mock_database) | |
| result = execute_sql_task(123, "SELECT * FROM users", username="admin") | |
| assert result is not None # Success returns payload | |
| assert result["status"] == QueryStatusEnum.SUCCESS | |
| assert mock_query.status == QueryStatusEnum.SUCCESS | |
| def test_execute_sql_task_with_start_time( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| mock_result_set: MagicMock, | |
| ) -> None: | |
| """Test execute_sql_task accepts start_time parameter (covers line 400-402).""" | |
| import time | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| mock_query.database = mock_database | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.execute_sql_with_cursor", | |
| return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.results_backend", None) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.dataframe.df_to_records", return_value=[]) | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| from .conftest import setup_mock_raw_connection | |
| setup_mock_raw_connection(mock_database) | |
| start_time = time.time() - 1.0 | |
| result = execute_sql_task(123, "SELECT * FROM users", start_time=start_time) | |
| # Verify task completes successfully with start_time | |
| assert result is not None | |
| assert result["status"] == QueryStatusEnum.SUCCESS | |
| def test_execute_sql_task_with_cancel_query_id( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| mock_result_set: MagicMock, | |
| ) -> None: | |
| """Test execute_sql_task sets cancel_query_id when available.""" | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| mock_query.database = mock_database | |
| mock_database.db_engine_spec.get_cancel_query_id.return_value = "cancel_123" | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.execute_sql_with_cursor", | |
| return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.results_backend", None) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.dataframe.df_to_records", return_value=[]) | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| from .conftest import setup_mock_raw_connection | |
| setup_mock_raw_connection(mock_database) | |
| execute_sql_task(123, "SELECT * FROM users") | |
| # Verify cancel_query_id was set via db_engine_spec | |
| mock_database.db_engine_spec.get_cancel_query_id.assert_called_once() | |
| mock_query.set_extra_json_key.assert_any_call("cancel_query", "cancel_123") | |
| def test_execute_sql_task_stopped( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test execute_sql_task when query is stopped (covers lines 456-458).""" | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| mock_query.database = mock_database | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| # Empty list indicates stopped (check_stopped_fn returned True mid-execution) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.execute_sql_with_cursor", | |
| return_value=[], | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| from .conftest import setup_mock_raw_connection | |
| setup_mock_raw_connection(mock_database) | |
| result = execute_sql_task(123, "SELECT * FROM users") | |
| assert result["status"] == QueryStatusEnum.STOPPED | |
| def test_execute_sql_task_with_mutation( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| mock_result_set: MagicMock, | |
| ) -> None: | |
| """Test execute_sql_task commits for mutations (covers lines 461-462).""" | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| mock_query.database = mock_database | |
| mock_query.select_as_cta = True # Trigger mutation commit | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.execute_sql_with_cursor", | |
| return_value=[("INSERT INTO users VALUES (1)", mock_result_set, 5.0, 1)], | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.results_backend", None) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.dataframe.df_to_records", return_value=[]) | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| from .conftest import setup_mock_raw_connection | |
| mock_conn = setup_mock_raw_connection(mock_database) | |
| execute_sql_task(123, "INSERT INTO users VALUES (1)") | |
| mock_conn.commit.assert_called() | |
| def test_execute_sql_task_with_results_backend( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| mock_result_set: MagicMock, | |
| ) -> None: | |
| """Test execute_sql_task stores results in backend (covers lines 466-467).""" | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| mock_query.database = mock_database | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.execute_sql_with_cursor", | |
| return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], | |
| ) | |
| mock_results_backend = MagicMock() | |
| mock_results_backend.set.return_value = True | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.results_backend", mock_results_backend | |
| ) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.zlib_compress", return_value=b"data" | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.dataframe.df_to_records", return_value=[]) | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| from .conftest import setup_mock_raw_connection | |
| setup_mock_raw_connection(mock_database) | |
| result = execute_sql_task(123, "SELECT * FROM users") | |
| mock_results_backend.set.assert_called_once() | |
| assert result is not None | |
| assert result["status"] == QueryStatusEnum.SUCCESS | |
| def test_execute_sql_task_timeout( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| ) -> None: | |
| """Test execute_sql_task handles timeout (covers lines 438-453).""" | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| mock_query.database = mock_database | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.execute_sql_with_cursor", | |
| side_effect=SoftTimeLimitExceeded(), | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict( | |
| current_app.config, | |
| {"STATS_LOGGER": MagicMock(), "SQLLAB_ASYNC_TIME_LIMIT_SEC": 300}, | |
| ) | |
| from .conftest import setup_mock_raw_connection | |
| setup_mock_raw_connection(mock_database) | |
| result = execute_sql_task(123, "SELECT * FROM users") | |
| # TIMED_OUT status is preserved (not overwritten to FAILED) | |
| assert result["status"] == QueryStatusEnum.TIMED_OUT.value | |
| def test_execute_sql_task_unhandled_exception( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| ) -> None: | |
| """Test execute_sql_task handles unhandled exceptions (covers lines 347-352).""" | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| # Mock _get_query to succeed first time (for override_user), then return mock_query | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._execute_sql_statements", | |
| side_effect=Exception("Unexpected error"), | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| result = execute_sql_task(123, "SELECT * FROM users") | |
| assert result["status"] == QueryStatusEnum.FAILED | |
| def test_execute_sql_task_success_final_commit( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| mock_result_set: MagicMock, | |
| ) -> None: | |
| """Test execute_sql_task final success path (covers lines 469-473).""" | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| mock_query.database = mock_database | |
| mock_query.status = QueryStatusEnum.RUNNING # Will be changed to SUCCESS | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.execute_sql_with_cursor", | |
| return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.results_backend", None) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mock_session = mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.dataframe.df_to_records", return_value=[]) | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| from .conftest import setup_mock_raw_connection | |
| setup_mock_raw_connection(mock_database) | |
| result = execute_sql_task(123, "SELECT * FROM users") | |
| assert result is not None | |
| assert result["status"] == QueryStatusEnum.SUCCESS | |
| assert mock_query.status == QueryStatusEnum.SUCCESS | |
| mock_session.commit.assert_called() | |
| def test_execute_sql_task_with_failed_status_before_final_commit( | |
| mocker: MockerFixture, | |
| app_context: None, | |
| mock_query: MagicMock, | |
| mock_database: MagicMock, | |
| mock_result_set: MagicMock, | |
| ) -> None: | |
| """Test execute_sql_task final commit when query.status is already FAILED.""" | |
| from superset.sql.execution.celery_task import execute_sql_task | |
| mock_query.database = mock_database | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._get_query", return_value=mock_query | |
| ) | |
| # execute_sql_with_cursor returns (exec_sql, result_set, time, rowcount) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task.execute_sql_with_cursor", | |
| return_value=[("SELECT * FROM users", mock_result_set, 10.5, 2)], | |
| ) | |
| # Mock _store_results_in_backend to set status to FAILED without raising | |
| def mock_store_results(query, payload, database): | |
| query.status = QueryStatusEnum.FAILED | |
| mocker.patch("superset.sql.execution.celery_task.results_backend", MagicMock()) | |
| mocker.patch("superset.results_backend_use_msgpack", False) | |
| mocker.patch( | |
| "superset.sql.execution.celery_task._store_results_in_backend", | |
| side_effect=mock_store_results, | |
| ) | |
| mocker.patch("superset.sql.execution.celery_task.db.session") | |
| mocker.patch("superset.dataframe.df_to_records", return_value=[]) | |
| mocker.patch("superset.sql.execution.celery_task.security_manager") | |
| mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()}) | |
| from .conftest import setup_mock_raw_connection | |
| setup_mock_raw_connection(mock_database) | |
| result = execute_sql_task(123, "SELECT * FROM users") | |
| # Verify query status remains FAILED and is not changed to SUCCESS | |
| assert result is not None | |
| assert mock_query.status == QueryStatusEnum.FAILED | |