mirror of
https://github.com/apache/superset.git
synced 2026-07-16 19:55:39 +00:00
1f81871fb3
MUTATE_AFTER_SPLIT=True never fired the mutator for engines with run_multiple_statements_as_one=True (BigQuery, Datastore, Kusto), since those engines always build a single joined block and the per-block mutation call always passes is_split=False. Mutate each statement before joining them into that block so the flag applies consistently. Verified separately via a full `pre-commit run` (all hooks touching these files passed: auto-walrus, mypy, ruff-format, ruff, pylint) run outside the git hook, whose invocation via the system /usr/bin/python3 hits an unrelated pre-existing environment issue installing the zizmor hook (requires Python >=3.10, system python3 here is older) rather than anything in this change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1311 lines
45 KiB
Python
1311 lines
45 KiB
Python
# 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
|
|
|
|
|
|
def _passthrough_mutator(sql: str, **kwargs: Any) -> str:
|
|
"""SQL mutator stand-in that returns the SQL unchanged."""
|
|
return sql
|
|
|
|
|
|
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_mutator
|
|
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_mutator
|
|
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_mutator
|
|
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)
|
|
|
|
|
|
# =============================================================================
|
|
# 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
|
|
# =============================================================================
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"query_status,expected_result",
|
|
[
|
|
(QueryStatusEnum.STOPPED, True),
|
|
(QueryStatusEnum.RUNNING, False),
|
|
],
|
|
)
|
|
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"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"logger_configured,should_be_called",
|
|
[
|
|
(True, True),
|
|
(False, False),
|
|
],
|
|
ids=["with_logger", "no_logger"],
|
|
)
|
|
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
|