From 59959af9907a68bc70fc4a5d87a72e2f54277a4c Mon Sep 17 00:00:00 2001 From: Rafael Benitez Date: Thu, 29 Oct 2026 17:34:31 -0300 Subject: [PATCH] fix(impala): stop SQL Lab queries while their results are being fetched Impala reports a streaming SELECT as finished once its rows can be fetched, so most of such a query runs while the worker reads results, after handle_cursor stopped polling. A stop in that phase marked the query stopped while Impala kept executing it and held its admission slot. The worker now checks for a committed stop between the result RPCs of a SQL Lab fetch and cancels the operation on its own cursor, as handle_cursor does while polling, and handle_cursor checks once more when polling ends. A failed check rolls its transaction back and lets the fetch continue, while a time limit raised during a check still propagates. The async SQLExecutor task reports a fetch that failed because of a stop as the stop rather than as an unhandled error, reading only the query's status so a genuine failure keeps the statement it was executing. Co-Authored-By: Claude Opus 5.5 (1M context) --- UPDATING.md | 21 +- .../configuration/configuring-superset.mdx | 10 + superset/db_engine_specs/impala.py | 161 ++++++++- superset/sql/execution/celery_task.py | 8 + .../unit_tests/db_engine_specs/test_impala.py | 342 +++++++++++++++++- .../db_engine_specs/test_impala_pending.py | 58 ++- .../sql/execution/test_celery_task.py | 89 ++++- 7 files changed, 654 insertions(+), 35 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index fdbfd8821497..953e27189ccc 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -123,13 +123,20 @@ assists people when migrating to a new version. Lab stop on Impala is delivered to the worker that is polling the query, which cancels it with `cancel_operation()`, as Hive and Presto already do. The HTTP cancel invalidated the handle the worker was still polling, so the run failed - instead of ending as stopped. When no worker is polling, the query is marked - stopped but Impala keeps running it and holds its admission slot until it - finishes or Impala closes it: for example when the worker has died or - restarted, or when the stop arrives while results are being fetched. - Operators who relied on the HTTP cancel to free admission slots in those cases - should rely on Impala's own query and session timeouts. Chart and MCP - cancellation still use the HTTP cancel. + instead of ending as stopped. The worker also checks for a stop between result + RPCs while it fetches, so a stop still reaches a query that Impala already + reports as finished, which a streaming `SELECT` does almost at once. The + worker checks between result RPCs, at most once per `DB_POLL_INTERVAL_SECONDS` + (5 seconds for Impala by default), and Impala returns each RPC within + `FETCH_ROWS_TIMEOUT_MS` (10 seconds by default), so such a stop takes effect + within about the sum of the two. + When no worker is polling or fetching, for example because it has died or + restarted, the query is marked stopped but Impala keeps running it and holds + its admission slot until it finishes or Impala closes it. Impala checks for + cancellation between row batches, so a batch whose rows are slow to compute + runs to its end first. Operators who relied on the HTTP cancel to free + admission slots in those cases should rely on Impala's own query and session + timeouts. Chart and MCP cancellation still use the HTTP cancel. - `POST /api/v1/query/stop` returns HTTP 422 instead of 200 when the query has already finished or its worker has committed the target transaction. The response contains `message: "Could not cancel query"` and the recorded diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index 97cfb3f343e7..1aeb2fff8ce8 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -835,6 +835,16 @@ pending when result fetching begins, the additional polling wait is bounded by a timeout. Configure the driver's network timeout separately to bound individual RPC calls. +Impala reports a streaming `SELECT` as finished as soon as its rows can be fetched, +so most of such a query runs while SQL Lab reads its results. While it fetches, the +worker checks for a stop between result RPCs, at most once per +`DB_POLL_INTERVAL_SECONDS`, and cancels the query on its own connection. Each result +RPC returns within Impala's `FETCH_ROWS_TIMEOUT_MS` query option (10 seconds by +default, available from Impala 3.4) even when no rows are ready, so a stop takes +effect within about `FETCH_ROWS_TIMEOUT_MS` plus `DB_POLL_INTERVAL_SECONDS`; +lowering either makes stops take effect sooner. Impala checks for cancellation between row batches, so a batch whose +rows are slow to compute runs to its end before the query stops. + ## SQL Lab Query History Pruning SQL Lab query history is stored in the metadata database and is **not** pruned by default. To trim older rows, enable the `prune_query` Celery beat task by uncommenting it in `CELERY_BEAT_SCHEDULE` and choosing a retention window: diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index e5134e21b947..573afc3c8336 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -20,8 +20,11 @@ import logging import re import time +from collections.abc import Iterator +from contextlib import contextmanager from datetime import datetime from typing import Any, Optional, TYPE_CHECKING +from weakref import WeakKeyDictionary from celery.exceptions import SoftTimeLimitExceeded from flask import current_app as app @@ -31,6 +34,7 @@ from superset import db from superset.constants import QUERY_EARLY_CANCEL_KEY, TimeGrain from superset.db_engine_specs.base import BaseEngineSpec, DatabaseCategory +from superset.exceptions import SupersetTimeoutException from superset.models.sql_lab import Query from superset.utils.core import QueryStatus from superset.utils.network import get_ssrf_safe_requester, is_safe_host @@ -42,6 +46,16 @@ # Query 5543ffdf692b7d02:f78a944000000000: 3% Complete (17 out of 547) QUERY_PROGRESS_REGEX = re.compile(r"Query.*: (?P[0-9]+)%") +# SQL Lab queries whose result fetch watches for a stop, keyed by the driver +# operation ``handle_cursor`` polled. ``fetch_data`` reads rows from that same +# operation, also through wrapper cursors such as SQLExecutor's row budget, and +# an entry goes away with its operation. +_STOP_WATCHED_OPERATIONS: WeakKeyDictionary[Any, Query] = WeakKeyDictionary() + + +class ImpalaFetchStoppedError(Exception): + """A stop committed for a SQL Lab query ended its result fetch.""" + class ImpalaEngineSpec(BaseEngineSpec): """Engine spec for Cloudera's Impala""" @@ -135,7 +149,12 @@ def execute( @classmethod def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, ...]]: - """Wait for asynchronous operations using the public cursor API.""" + """ + Wait for asynchronous operations using the public cursor API. + + For a SQL Lab query that ``handle_cursor`` polled, a committed stop is + also checked while the results are read (see ``_stop_checked_fetch``). + """ if callable(getattr(cursor, "execute_async", None)): from impala.error import Error @@ -152,7 +171,97 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. cursor.fetchall() except Error as ex: raise cls.get_dbapi_mapped_exception(ex) from ex - return super().fetch_data(cursor, limit) + operation = getattr(cursor, "_last_operation", None) + try: + query = _STOP_WATCHED_OPERATIONS.get(operation) + except TypeError: + # The cursor exposes no driver operation that could be watched. + query = None + if query is None: + return super().fetch_data(cursor, limit) + with cls._stop_checked_fetch(cursor, operation, query): + return super().fetch_data(cursor, limit) + + @classmethod + @contextmanager + def _stop_checked_fetch( + cls, cursor: Any, operation: Any, query: Query + ) -> Iterator[None]: + """ + Check for a committed stop before each result RPC of a SQL Lab query. + + Impala reports a streaming SELECT as finished once its rows can be + fetched, so most of such a query runs while the worker reads results, + after ``handle_cursor`` stopped polling. The driver issues one fetch RPC + at a time, and each returns within the coordinator's + ``FETCH_ROWS_TIMEOUT_MS`` even when no rows are ready, so a stop goes + unseen for at most about that plus the poll interval. A stop is applied + on the worker's own cursor and thread, as in ``handle_cursor``, so no + second connection or HTTP call is involved. + """ + fetch = operation.fetch + # Read once: after a failed check the instance can be expired in a + # session that cannot load it, so reading it then would raise again. + query_id = query.id + interval = app.config["DB_POLL_INTERVAL_SECONDS"].get(cls.engine, 5) + last_check: float | None = None + checking = True + + def fetch_unless_stopped(*args: Any, **kwargs: Any) -> Any: + nonlocal last_check, checking + now = time.monotonic() + if checking and (last_check is None or now - last_check >= interval): + last_check = now + try: + stopped = cls._stop_requested(query) + except (SoftTimeLimitExceeded, SupersetTimeoutException): + raise + except Exception: # pylint: disable=broad-except + # A metadata hiccup must not fail a query nobody stopped, + # so the failed transaction must not outlive the check. + try: + db.session.rollback() # pylint: disable=consider-using-transaction + except Exception: # pylint: disable=broad-except + logger.debug("Query %s: rollback failed", query_id) + logger.warning( + "Query %s: stop check failed while fetching", query_id + ) + checking = False + stopped = False + if stopped: + cls._cancel_operation(cursor, query_id) + raise ImpalaFetchStoppedError( + f"Query {query_id} was stopped while fetching results" + ) + return fetch(*args, **kwargs) + + # Restore exactly what was there: a driver operation normally has no + # fetch of its own, so removing the wrapper re-exposes the method. + had_own_fetch = "fetch" in getattr(operation, "__dict__", {}) + operation.fetch = fetch_unless_stopped + try: + yield + finally: + if had_own_fetch: + operation.fetch = fetch + else: + del operation.fetch + + @classmethod + def _stop_requested(cls, query: Query) -> bool: + """Read whether a stop was committed for the query.""" + # End the read transaction first. Under REPEATABLE READ, refreshing + # inside one long transaction keeps returning the first snapshot, so a + # Stop committed by another request would never be seen. + db.session.commit() # pylint: disable=consider-using-transaction + db.session.refresh(query) + # Stop was requested: either before a cancel handle was published + # (early-cancel flag) or through SQL Lab's stop, which persists STOPPED + # once cancel_query() succeeds. + return bool( + query.extra.get(QUERY_EARLY_CANCEL_KEY) + or query.status == QueryStatus.STOPPED + ) @classmethod def _cancel_operation(cls, cursor: Any, query_id: int) -> None: @@ -197,20 +306,9 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: try: status = cursor.status() while status in unfinished_states: - # End the read transaction first. Under REPEATABLE READ, refreshing - # inside one long transaction keeps returning the first snapshot, - # so a Stop committed by another request would never be seen. - db.session.commit() # pylint: disable=consider-using-transaction - db.session.refresh(query) - # Stop was requested: either before a cancel handle was published - # (early-cancel flag) or through SQL Lab's stop, which persists - # STOPPED once cancel_query() succeeds. - if ( - query.extra.get(QUERY_EARLY_CANCEL_KEY) - or query.status == QueryStatus.STOPPED - ): + if cls._stop_requested(query): cls._cancel_operation(cursor, query_id) - break + return sleep_interval = app.config["DB_POLL_INTERVAL_SECONDS"].get( cls.engine, 5 @@ -234,6 +332,12 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: pending_sleep_interval *= 2 time.sleep(sleep_interval) status = cursor.status() + # The loop checks for a stop before each sleep, so a stop committed + # during the last poll interval, or before an operation that is + # finished at its first poll, is caught here, before the fetch. + if cls._stop_requested(query): + cls._cancel_operation(cursor, query_id) + return except SoftTimeLimitExceeded: # SQL Lab's own handler marks the query TIMED_OUT once this propagates; # the operation is cancelled first so it does not keep running on the @@ -243,6 +347,19 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: except Exception: # pylint: disable=broad-except logger.debug("Call to status() failed ") return + cls._watch_result_fetch(cursor, query) + + @classmethod + def _watch_result_fetch(cls, cursor: Any, query: Query) -> None: + """Have the fetch that follows the poll check for a stop of the query.""" + if (operation := getattr(cursor, "_last_operation", None)) is None: + return + try: + _STOP_WATCHED_OPERATIONS[operation] = query + except TypeError: + # The operation cannot be weakly referenced, so the fetch that + # follows reads results without stop checks. + logger.debug("Query %s: fetch not watched for a stop", query.id) @classmethod def prepare_cancel_query(cls, query: Query) -> None: @@ -260,16 +377,20 @@ def prepare_cancel_query(cls, query: Query) -> None: and persists this flag afterwards through the guarded-stop protocol. Preparation must not commit or claim the stop a second time. - Known gap: the stop only reaches Impala while a worker is polling the - operation in ``handle_cursor``. The query is marked stopped, but Impala - keeps running it and holds its admission slot until it finishes or + The worker observes the stop while it polls the operation in + ``handle_cursor`` and, for a SQL Lab query it polled, between the result + RPCs of ``fetch_data``. Known gaps: the query is marked stopped, but + Impala keeps running it and holds its admission slot until it finishes or Impala closes it on its own, when: - the worker has died or restarted; - - the stop lands after polling ended, while ``fetch_data`` is reading - results; - polling ended early because ``cursor.status()`` raised. + Impala itself checks for cancellation only between row batches, so a + batch whose rows are expensive to compute, such as a slow function in the + select list over a single batch, runs to its end before the fetch can + return. + Hive and Presto stop SQL Lab queries the same way. Chart and MCP cancellation still use the coordinator's HTTP endpoint through ``cancel_query``. diff --git a/superset/sql/execution/celery_task.py b/superset/sql/execution/celery_task.py index 0b1af28f465a..42b33891167e 100644 --- a/superset/sql/execution/celery_task.py +++ b/superset/sql/execution/celery_task.py @@ -661,6 +661,14 @@ def _execute_sql_statements( level=ErrorLevel.ERROR, ) ) from ex + except Exception: + # An engine that ends a stopped query while its results are read, + # such as Impala, fails the fetch; report the stop, not an error. + # Read only the status, so a failure keeps its pending metadata, + # such as the statement it was executing. + if not _make_check_stopped_fn(query, preserve_pending=True)(): + raise + execution_results = [] # Re-read stop intent even on autocommit connections: completed engine # effects cannot be undone, but an accepted Stop still discards results. diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index f8603e2f5410..4be71586786d 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -15,16 +15,25 @@ # specific language governing permissions and limitations # under the License. +from contextlib import nullcontext from datetime import datetime -from typing import Optional -from unittest.mock import Mock, patch +from typing import Any, Optional +from unittest.mock import MagicMock, Mock, patch, PropertyMock import pytest +from flask import current_app from pytest_mock import MockerFixture -from superset.db_engine_specs.impala import ImpalaEngineSpec as spec # noqa: N813 +from superset.constants import QUERY_CANCEL_KEY +from superset.db_engine_specs.impala import ( + ImpalaEngineSpec as spec, # noqa: N813 + ImpalaFetchStoppedError, +) from superset.models.core import Database from superset.models.sql_lab import Query +from superset.sql.execution.executor import _LimitedCursor +from superset.sql_lab import execute_query, SqlLabQueryStoppedException +from superset.utils.core import QueryStatus from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm from tests.unit_tests.fixtures.common import dttm # noqa: F401 @@ -286,6 +295,333 @@ def test_fetch_data_async_result_set(app_context: None, impala_errors: None) -> cursor.cancel_operation.assert_not_called() +class _Batch: + """A driver result batch: its rows, and whether more rows may follow.""" + + def __init__(self, rows: list[tuple[int]], expect_more_rows: bool) -> None: + """Hold the rows one result RPC returned.""" + self.rows: list[tuple[int]] = list(rows) + self.expect_more_rows: bool = expect_more_rows + + def __len__(self) -> int: + """Count the rows left in the batch.""" + return len(self.rows) + + def pop(self) -> tuple[int]: + """Hand out the rows in order.""" + return self.rows.pop(0) + + +class _Operation: + """A driver operation that answers one result RPC at a time.""" + + def __init__(self, batches: list[_Batch]) -> None: + """Script the batches the result RPCs return.""" + self.batches: list[_Batch] = batches + self.fetches: int = 0 + + def fetch(self, *args: Any, **kwargs: Any) -> _Batch: + """Return the next result RPC's batch, then empty ones at the end.""" + self.fetches += 1 + return self.batches.pop(0) if self.batches else _Batch([], False) + + +class _StreamingCursor: + """ + A finished impyla SELECT whose rows are read through one result RPC at a + time, retrying batches that come back empty while more rows are expected. + """ + + description = [("x", "INT", None, None, None, None, True)] + buffersize = 1024 + + def __init__(self, operation: _Operation) -> None: + """Start reading the operation's results.""" + self._last_operation: _Operation = operation + self._buffer: _Batch = _Batch([], True) + self._exhausted: bool = False + self.arraysize: int = 1 + self.calls: list[str] = [] + + def execute_async(self, query: str) -> None: + """Expose the asynchronous execution capability.""" + + def status(self) -> str: + """Report the finished state of a streaming SELECT.""" + return "FINISHED_STATE" + + def is_executing(self) -> bool: + """Report the operation as no longer executing.""" + return False + + def execution_failed(self) -> bool: + """Report a successful operation.""" + return False + + def _next_row(self) -> Optional[tuple[int]]: + """Read the next row, issuing result RPCs as the driver does.""" + while not self._exhausted and not len(self._buffer): + self._buffer = self._last_operation.fetch(self.description, self.buffersize) + if not len(self._buffer) and not self._buffer.expect_more_rows: + self._exhausted = True + return None if self._exhausted else self._buffer.pop() + + def fetchmany(self, size: Optional[int] = None) -> list[tuple[int]]: + """Read up to ``size`` rows.""" + rows: list[tuple[int]] = [] + while len(rows) < (size or self.arraysize): + if (row := self._next_row()) is None: + break + rows.append(row) + return rows + + def fetchall(self) -> list[tuple[int]]: + """Read every remaining row.""" + rows: list[tuple[int]] = [] + while (row := self._next_row()) is not None: + rows.append(row) + return rows + + def cancel_operation(self) -> None: + """Record the cancel RPC.""" + self.calls.append("cancel_operation") + + def close_operation(self) -> None: + """Record the close RPC.""" + self.calls.append("close_operation") + + def close(self) -> None: + """Record closing the cursor.""" + self.calls.append("close") + + +def _streaming_batches() -> list[_Batch]: + """Rows, then two empty result RPCs, then the last rows.""" + return [ + _Batch([(1,)], True), + _Batch([], True), + _Batch([], True), + _Batch([(2,)], False), + ] + + +def test_stop_while_fetching_cancels_between_result_rpcs( + app_context: None, impala_errors: None +) -> None: + """A stop committed while a SQL Lab query streams results ends its fetch.""" + operation = _Operation(_streaming_batches()) + cursor = _StreamingCursor(operation) + query = Mock(id=1, extra={}, status=QueryStatus.RUNNING) + + def refresh(refreshed: Mock) -> None: + """Another request commits the stop once the first rows have arrived.""" + if operation.fetches: + refreshed.status = QueryStatus.STOPPED + + with ( + patch.dict(current_app.config, DB_POLL_INTERVAL_SECONDS={"impala": 0}), + patch("superset.db_engine_specs.impala.db") as db, + ): + db.session.refresh.side_effect = refresh + spec.handle_cursor(cursor, query) + with pytest.raises(ImpalaFetchStoppedError): + spec.fetch_data(cursor) + + assert operation.fetches == 1 + assert cursor.calls == ["cancel_operation", "close_operation", "close"] + assert "fetch" not in vars(operation) + + +def test_watched_fetch_without_stop_reads_every_row( + app_context: None, impala_errors: None +) -> None: + """A watched fetch nobody stops returns every row, checking once per interval.""" + operation = _Operation(_streaming_batches()) + cursor = _StreamingCursor(operation) + query = Mock(id=1, extra={}, status=QueryStatus.RUNNING) + + with patch("superset.db_engine_specs.impala.db") as db: + spec.handle_cursor(cursor, query) + assert spec.fetch_data(cursor) == [(1,), (2,)] + + # One check once polling ends and one before the first result RPC; the + # default poll interval spaces out the rest. + assert db.session.refresh.call_count == 2 + # The driver confirms the end of the results with one more result RPC. + assert operation.fetches == 5 + assert cursor.calls == [] + assert "fetch" not in vars(operation) + + +def test_unwatched_fetch_reads_without_stop_checks( + app_context: None, impala_errors: None +) -> None: + """Fetches that handle_cursor never polled, as for charts, are unchanged.""" + operation = _Operation(_streaming_batches()) + cursor = _StreamingCursor(operation) + + with patch("superset.db_engine_specs.impala.db") as db: + assert spec.fetch_data(cursor) == [(1,), (2,)] + + db.session.commit.assert_not_called() + db.session.refresh.assert_not_called() + + +def test_watched_fetch_keeps_the_executor_row_budget( + app_context: None, impala_errors: None +) -> None: + """A fetch through SQLExecutor's row budget is watched and still bounded.""" + operation = _Operation([_Batch([(1,), (2,), (3,)], True)]) + cursor = _StreamingCursor(operation) + query = Mock(id=1, extra={}, status=QueryStatus.RUNNING) + + with patch("superset.db_engine_specs.impala.db") as db: + spec.handle_cursor(cursor, query) + assert spec.fetch_data(_LimitedCursor(cursor, 2)) == [(1,), (2,)] + + assert operation.fetches == 1 + # The check before the first result RPC ran through the wrapper. + assert db.session.refresh.call_count == 2 + + +def test_failed_stop_check_does_not_fail_the_fetch( + app_context: None, impala_errors: None, caplog: pytest.LogCaptureFixture +) -> None: + """A metadata error while checking for a stop leaves the results intact.""" + operation = _Operation(_streaming_batches()) + cursor = _StreamingCursor(operation) + query = Mock(id=1, extra={}, status=QueryStatus.RUNNING) + + with ( + patch.dict(current_app.config, DB_POLL_INTERVAL_SECONDS={"impala": 0}), + patch("superset.db_engine_specs.impala.db") as db, + ): + spec.handle_cursor(cursor, query) + db.session.commit.side_effect = RuntimeError("metadata unavailable") + assert spec.fetch_data(cursor) == [(1,), (2,)] + + # The failed check is not retried before every remaining result RPC. + assert db.session.commit.call_count == 2 + assert "stop check failed while fetching" in caplog.text + assert cursor.calls == [] + + +def test_failed_stop_check_keeps_the_rows_of_an_expired_query( + app_context: None, impala_errors: None +) -> None: + """A stop check that fails in a broken transaction does not fail the fetch. + + The check's commit expires the ORM query. When the refresh that follows + fails, reading the expired query again raises until the session is rolled + back, so the check must not read it. + """ + operation = _Operation(_streaming_batches()) + cursor = _StreamingCursor(operation) + query = Mock(extra={}, status=QueryStatus.RUNNING) + session_broken = False + + def read_id() -> int: + """Loading the expired query fails inside the broken transaction.""" + if session_broken: + raise RuntimeError("current transaction is aborted") + return 1 + + type(query).id = PropertyMock(side_effect=read_id) + + def failing_refresh(refreshed: Mock) -> None: + """The metadata database fails the check's refresh.""" + nonlocal session_broken + session_broken = True + raise RuntimeError("canceling statement due to statement timeout") + + def rollback() -> None: + nonlocal session_broken + session_broken = False + + with ( + patch.dict(current_app.config, DB_POLL_INTERVAL_SECONDS={"impala": 0}), + patch("superset.db_engine_specs.impala.db") as db, + ): + spec.handle_cursor(cursor, query) + db.session.refresh.side_effect = failing_refresh + db.session.rollback.side_effect = rollback + assert spec.fetch_data(cursor) == [(1,), (2,)] + + db.session.rollback.assert_called_once_with() + assert cursor.calls == [] + + +def test_timeout_during_a_stop_check_is_not_swallowed( + app_context: None, impala_errors: None +) -> None: + """SQL Lab's timeout raised while a stop check runs still ends the query.""" + from superset.errors import ErrorLevel, SupersetErrorType + from superset.exceptions import SupersetTimeoutException + + operation = _Operation(_streaming_batches()) + cursor = _StreamingCursor(operation) + query = Mock(id=1, extra={}, status=QueryStatus.RUNNING) + timeout = SupersetTimeoutException( + error_type=SupersetErrorType.SQLLAB_TIMEOUT_ERROR, + message="timed out", + level=ErrorLevel.ERROR, + ) + + with ( + patch.dict(current_app.config, DB_POLL_INTERVAL_SECONDS={"impala": 0}), + patch("superset.db_engine_specs.impala.db") as db, + ): + spec.handle_cursor(cursor, query) + db.session.refresh.side_effect = timeout + with pytest.raises(SupersetTimeoutException): + spec.fetch_data(cursor) + + assert operation.fetches == 0 + + +def test_sql_lab_stop_while_fetching_ends_stopped( + mocker: MockerFixture, app_context: None, impala_errors: None +) -> None: + """SQL Lab reports a query stopped during its fetch as stopped, not failed.""" + operation = _Operation(_streaming_batches()) + cursor = _StreamingCursor(operation) + query = MagicMock( + id=1, + status=QueryStatus.RUNNING, + progress=0, + executed_sql="SELECT x FROM t", + extra={}, + limit=None, + ) + query.database.db_engine_spec = spec + query.set_extra_json_key.side_effect = lambda key, value: query.extra.update( + {key: value} + ) + session = mocker.patch("superset.db.session") + + def refresh(refreshed: MagicMock, **kwargs: Any) -> None: + """The stop request commits STOPPED once the first rows have arrived.""" + if operation.fetches: + refreshed.status = QueryStatus.STOPPED + + session.refresh.side_effect = refresh + mocker.patch.object( + spec, "get_cancel_query_id", return_value="0123456789abcdef:fedcba9876543210" + ) + mocker.patch("superset.sql_lab.warm_and_release_connection") + mocker.patch( + "superset.sql_lab.event_logger.log_context", return_value=nullcontext() + ) + mocker.patch.dict(current_app.config, DB_POLL_INTERVAL_SECONDS={"impala": 0}) + + with pytest.raises(SqlLabQueryStoppedException): + execute_query(query, cursor) + + assert query.extra[QUERY_CANCEL_KEY] == "0123456789abcdef:fedcba9876543210" + assert operation.fetches == 1 + assert cursor.calls == ["cancel_operation", "close_operation", "close"] + + def test_the_partition_transform_default_is_impala_syntax() -> None: from superset.db_engine_specs.impala import ImpalaEngineSpec diff --git a/tests/unit_tests/db_engine_specs/test_impala_pending.py b/tests/unit_tests/db_engine_specs/test_impala_pending.py index 1e7d69b76e78..efe4d940d74c 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_pending.py +++ b/tests/unit_tests/db_engine_specs/test_impala_pending.py @@ -28,7 +28,7 @@ QUERY_DISPATCHED_KEY, QUERY_EARLY_CANCEL_KEY, ) -from superset.db_engine_specs.impala import ImpalaEngineSpec +from superset.db_engine_specs.impala import _STOP_WATCHED_OPERATIONS, ImpalaEngineSpec from superset.sql.execution.executor import SQLExecutor from superset.sql_lab import cancel_query, execute_query, SqlLabQueryStoppedException from superset.utils.core import QueryStatus @@ -218,8 +218,9 @@ def test_non_progress_logs_keep_polling(log: str) -> None: ImpalaEngineSpec.handle_cursor(cursor, query) assert cursor.status.call_count == 2 assert query.progress == 0 - # Only the poll's transaction boundary commits; no progress is written. - db.session.commit.assert_called_once_with() + # Only the stop checks commit, in the poll and before the fetch; no + # progress is written. + assert db.session.commit.call_count == 2 def test_failed_cancel_still_releases_the_operation() -> None: @@ -306,7 +307,8 @@ def test_pending_operation_polls_with_backoff() -> None: ): ImpalaEngineSpec.handle_cursor(cursor, query) assert [call.args[0] for call in sleep.call_args_list] == [0.1, 0.2, 0.3, 0.3] - assert db.session.refresh.call_count == 4 + # One stop check per poll, and one more once the operation has finished. + assert db.session.refresh.call_count == 5 db.session.query.assert_not_called() @@ -401,3 +403,51 @@ def refresh_query(refreshed_query: Mock) -> None: cursor.cancel_operation.assert_called_once_with() cursor.close_operation.assert_called_once_with() cursor.close.assert_called_once_with() + + +def test_stop_during_last_poll_interval_cancels_before_fetch() -> None: + """A stop committed while the last poll slept is caught before the fetch.""" + app: Flask = Flask(__name__) + app.config["DB_POLL_INTERVAL_SECONDS"] = {"impala": 0} + query: Mock = Mock(id=1, extra={}, progress=0, status=QueryStatus.RUNNING) + cursor: Mock = Mock() + cursor.status.side_effect = ["RUNNING_STATE", "FINISHED_STATE"] + cursor.get_log.return_value = "" + refreshes: list[Mock] = [] + + def refresh(refreshed: Mock) -> None: + """The stop is committed after the last poll's check.""" + refreshes.append(refreshed) + if len(refreshes) == 2: + refreshed.status = QueryStatus.STOPPED + + with app.app_context(), patch("superset.db_engine_specs.impala.db") as db: + db.session.refresh.side_effect = refresh + ImpalaEngineSpec.handle_cursor(cursor, query) + + cursor.cancel_operation.assert_called_once_with() + cursor.close_operation.assert_called_once_with() + cursor.close.assert_called_once_with() + assert cursor._last_operation not in _STOP_WATCHED_OPERATIONS + + +def test_stop_before_an_operation_finished_at_first_poll() -> None: + """An operation that is finished at its first poll still sees an early stop.""" + query: Mock = Mock(id=1, extra={QUERY_EARLY_CANCEL_KEY: True}, progress=0) + cursor: Mock = Mock() + cursor.status.return_value = "FINISHED_STATE" + with patch("superset.db_engine_specs.impala.db"): + ImpalaEngineSpec.handle_cursor(cursor, query) + cursor.cancel_operation.assert_called_once_with() + cursor.get_log.assert_not_called() + + +def test_finished_operation_is_watched_by_the_fetch() -> None: + """The fetch that follows a finished poll is told which query to watch.""" + query: Mock = Mock(id=1, extra={}, progress=0, status=QueryStatus.RUNNING) + cursor: Mock = Mock() + cursor.status.return_value = "FINISHED_STATE" + with patch("superset.db_engine_specs.impala.db"): + ImpalaEngineSpec.handle_cursor(cursor, query) + assert _STOP_WATCHED_OPERATIONS[cursor._last_operation] is query + cursor.cancel_operation.assert_not_called() diff --git a/tests/unit_tests/sql/execution/test_celery_task.py b/tests/unit_tests/sql/execution/test_celery_task.py index ac779184a74c..43f76994c0d0 100644 --- a/tests/unit_tests/sql/execution/test_celery_task.py +++ b/tests/unit_tests/sql/execution/test_celery_task.py @@ -18,7 +18,8 @@ """Tests for celery_task.py - async SQL execution via Celery.""" import base64 -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from decimal import Decimal from pathlib import Path from typing import Any @@ -1222,6 +1223,92 @@ def test_execute_sql_task_stopped( assert result["status"] == QueryStatusEnum.STOPPED +def test_failed_statement_keeps_its_executed_sql( + mocker: MockerFixture, app: Any, session: Any +) -> None: + """A statement that fails with no stop involved keeps the SQL it executed.""" + from superset.models.core import Database + from superset.models.sql_lab import Query + from superset.sql.execution.celery_task import execute_sql_task + + Query.metadata.create_all(session.get_bind()) # pylint: disable=no-member + database = Database(database_name="db", sqlalchemy_uri="sqlite://") + # As SQLExecutor._create_query_record leaves it: no executed_sql yet. + query = Query( + client_id="failing", + database=database, + sql="SELECT 1", + limit=100, + select_as_cta=False, + status=QueryStatusEnum.PENDING, + ) + session.add_all([database, query]) + session.commit() + query_id = query.id + + cursor = MagicMock() + cursor.execute.side_effect = RuntimeError("syntax error at or near FROM") + connection = MagicMock() + connection.cursor.return_value = cursor + + @contextmanager + def get_raw_connection(*args: Any, **kwargs: Any) -> Iterator[Any]: + yield connection + + mocker.patch.object(Database, "get_raw_connection", get_raw_connection) + + execute_sql_task(query_id, "SELECT 1") + + session.rollback() + session.expire_all() + failed = session.query(Query).filter_by(id=query_id).one() + assert failed.status == QueryStatusEnum.FAILED + assert "syntax error" in failed.error_message + assert failed.executed_sql == "SELECT\n 1" + + +@pytest.mark.parametrize("stopped", [True, False]) +def test_failed_fetch_of_a_stopped_query_is_reported_as_stopped( + mocker: MockerFixture, + app_context: None, + mock_query: MagicMock, + mock_database: MagicMock, + stopped: bool, +) -> None: + """A fetch an engine ended because of a stop is not logged as an error.""" + 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=RuntimeError("operation cancelled"), + ) + mocker.patch( + "superset.sql.execution.celery_task._make_check_stopped_fn", + return_value=lambda: stopped, + ) + mocker.patch("superset.sql.execution.celery_task.db.session") + mocker.patch("superset.sql.execution.celery_task.security_manager") + logger = mocker.patch("superset.sql.execution.celery_task.logger") + 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") + + if stopped: + assert result["status"] == QueryStatusEnum.STOPPED + logger.exception.assert_not_called() + else: + assert result["status"] == QueryStatusEnum.FAILED + logger.exception.assert_called_once() + + @pytest.mark.parametrize("stop_intent", ["executor", "early", "status"]) def test_cancelled_impala_cursor_is_not_finalized_as_success( mocker: MockerFixture,