fix(impala): stop SQL Lab queries while their results are being fetched

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

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:

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<query_progress>[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``.

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.