This is an automated email from the ASF dual-hosted git repository.
sadpandajoe pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 026ca216cb2 fix(sqllab): implement query cancellation for ClickHouse
(Stop button) (#45070)
026ca216cb2 is described below
commit 026ca216cb24fa6c80890ba2bb510aedc1048d99
Author: Joe Li <[email protected]>
AuthorDate: Fri Oct 23 10:19:00 2026 -0700
fix(sqllab): implement query cancellation for ClickHouse (Stop button)
(#45070)
Co-authored-by: Claude Sonnet 5 <[email protected]>
Co-authored-by: Evan Rusackas <[email protected]>
Co-authored-by: Evan Rusackas <[email protected]>
---
superset/db_engine_specs/clickhouse.py | 223 +++++++++++++
.../db_engine_specs/test_clickhouse.py | 118 +++++++
tests/unit_tests/dao/test_query_stop_race.py | 95 ++++++
.../unit_tests/db_engine_specs/test_clickhouse.py | 344 ++++++++++++++++++++-
tests/unit_tests/sql/execution/test_celery_task.py | 47 +++
tests/unit_tests/sql_lab_test.py | 38 +++
6 files changed, 864 insertions(+), 1 deletion(-)
diff --git a/superset/db_engine_specs/clickhouse.py
b/superset/db_engine_specs/clickhouse.py
index d748cd0973b..433ea8b2e61 100644
--- a/superset/db_engine_specs/clickhouse.py
+++ b/superset/db_engine_specs/clickhouse.py
@@ -18,6 +18,7 @@ from __future__ import annotations
import logging
import re
+import uuid
from datetime import datetime, timezone
from typing import Any, cast, TYPE_CHECKING
from urllib import parse
@@ -31,6 +32,9 @@ from sqlalchemy.engine.url import URL
from sqlalchemy.sql.expression import ColumnElement
from urllib3.exceptions import NewConnectionError
+from superset import db
+from superset.common.db_query_status import QueryStatus
+from superset.constants import QUERY_CANCEL_KEY
from superset.databases.utils import make_url_safe
from superset.db_engine_specs.base import (
BaseEngineSpec,
@@ -47,10 +51,27 @@ from superset.utils.network import is_hostname_valid,
is_port_open
if TYPE_CHECKING:
from superset.models.core import Database
+ from superset.models.sql_lab import Query
logger = logging.getLogger(__name__)
+class ClickHouseQueryStoppedBeforeDispatchError(Exception):
+ """
+ Raised internally, right before a statement would be sent to
+ ClickHouse, when ``query.status`` is already STOPPED. Any exception
+ type works for this purpose: ``execute_query()``
+ (``superset/sql_lab.py``) wraps ``execute_with_cursor`` in its own
+ generic ``except Exception`` handler that refreshes ``query`` and
+ converts whatever lands here into the standard
+ ``SqlLabQueryStoppedException`` flow if the status is STOPPED. A
+ dedicated type exists only so the intent reads clearly at the raise
+ site and in stack traces -- importing ``SqlLabQueryStoppedException``
+ itself would be circular, since ``sql_lab.py`` already imports from
+ ``superset.db_engine_specs``.
+ """
+
+
class ClickHouseBaseEngineSpec(BaseEngineSpec):
"""Shared engine spec for ClickHouse."""
@@ -729,3 +750,205 @@ class ClickHouseConnectEngineSpec(BasicParametersMixin,
ClickHouseEngineSpec):
# bare outer SELECT to satisfy it without altering or dropping any of
# the mutator's comments.
return f"SELECT * FROM (\n{sql}\n) AS __superset_type_probe LIMIT 0"
# noqa: S608
+
+ # ClickHouse never hands back a server-assigned identifier Superset could
+ # read and reuse later -- but its HTTP interface accepts a client-chosen
+ # `query_id`, and `KILL QUERY WHERE query_id = ...` accepts that same id
+ # from a second, independent connection. Minting it here needs no round
+ # trip: `has_query_id_before_execute` (the base default) runs this before
+ # the statement is sent, so the id is threaded into `execute_with_cursor`
+ # below in time to be passed to the driver.
+ #
+ # Only `execute_with_cursor` does that threading. The GTF chart-data path
+ # (`Database.get_df`) dispatches through the base `execute()` instead,
+ # which never forwards a `settings` kwarg, so an id minted for that path
+ # would never reach the driver -- `cancel_query`'s later `KILL QUERY
+ # WHERE query_id = ...` would have nothing real to kill, and its "nothing
+ # to confirm" fallback would then report a cancel that never happened.
+ # `superset.tasks.query_cancel._CancellationQuery` stubs `query.id` as
+ # `None` for exactly this case so engines can decline gracefully; do
+ # that here rather than minting an id this engine can't actually use.
+ @classmethod
+ def get_cancel_query_id( # pylint: disable=unused-argument
+ cls,
+ cursor: Any,
+ query: Query,
+ ) -> str | None:
+ if query.id is None:
+ return None
+ return str(uuid.uuid4())
+
+ @classmethod
+ def execute_with_cursor(
+ cls,
+ cursor: Any,
+ sql: str,
+ query: Query,
+ ) -> None:
+ """
+ Forward the id `get_cancel_query_id` recorded into `query.extra` on
+ to `execute` as a `settings` kwarg -- the base `execute_with_cursor`/
+ `execute` pair never threads per-query kwargs through to
+ `cursor.execute()`, so cancellation needs this override to reach the
+ driver at all.
+ """
+ logger.debug("Query %d: Running query: %s", query.id, sql)
+
+ # Narrows (see the exception class's own docstring for how this is
+ # converted, and RCA.md for why this is closable for ClickHouse but
+ # not a given for every engine) the window between the per-block
+ # loop's own STOPPED check (sql_lab.py, which already ran once
+ # before this call) and actually sending the statement below.
+ # Postgres/MySQL are accidentally safe in this exact window because
+ # their cancel_query() kills the whole underlying connection, so a
+ # cursor.execute() landing after that fails outright and gets
+ # caught by execute_query()'s generic exception handler. ClickHouse
+ # KILL QUERY has no such side effect when nothing has matched yet
+ # -- the connection stays healthy and an unguarded cursor.execute()
+ # here would run the statement to completion even though Superset
+ # already told the user it stopped.
+ try:
+ # flush() first: refresh() does NOT autoflush -- without this,
+ # a pending, uncommitted attribute set just before this call
+ # (e.g. the async/Celery path's own execute_fn sets
+ # query.executed_sql = sql immediately before calling
+ # execute_with_cursor, with no commit in between) would be
+ # silently discarded and reloaded back to its previous
+ # committed value instead of surviving to this query's later
+ # commits. Same idiom and same reasoning as the existing
+ # flush()-before-refresh() pairs already in sql_lab.py.
+ db.session.flush()
+ db.session.refresh(query)
+ already_stopped = query.status == QueryStatus.STOPPED
+ except Exception: # pylint: disable=broad-except
+ # Best-effort, same as cancel_query()'s own refresh below: if
+ # this fails for any reason, fall back to dispatching the
+ # statement rather than letting a diagnostic check become a
+ # new failure mode of query execution itself.
+ already_stopped = False
+
+ if already_stopped:
+ raise ClickHouseQueryStoppedBeforeDispatchError()
+
+ cancel_query_id = query.extra.get(QUERY_CANCEL_KEY)
+ settings = {"query_id": cancel_query_id} if cancel_query_id else None
+ cls.execute(cursor, sql, query.database, settings=settings)
+ logger.debug("Query %d: Handling cursor", query.id)
+ cls.handle_cursor(cursor, query)
+
+ @classmethod
+ def execute( # pylint: disable=unused-argument
+ cls,
+ cursor: Any,
+ query: str,
+ database: Database,
+ **kwargs: Any,
+ ) -> None:
+ if cls.arraysize:
+ cursor.arraysize = cls.arraysize
+ try:
+ cursor.execute(query, settings=kwargs.get("settings"))
+ except Exception as ex:
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+
+ # ClickHouse's own enum for the `kill_status` column `KILL QUERY SYNC`
+ # returns per matched process (https://clickhouse.com/docs/sql-reference/
+ # statements/kill): "finished" is the only value confirming the query
+ # was actually terminated. The other four are documented failure/
+ # indeterminate states -- a real process matched, but termination wasn't
+ # confirmed -- and must not be reported as a successful cancel.
+ _KILL_STATUS_CONFIRMED = "finished"
+ _KILL_STATUS_KNOWN_NOT_CONFIRMED = frozenset(
+ {"waiting", "cant_cancel", "pending", "unknown_status"}
+ )
+
+ @classmethod
+ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) ->
bool:
+ """
+ :param cancel_query_id: client-chosen `query_id` minted by
+ `get_cancel_query_id` and passed to the driver at execute time
+ :return: True if query cancelled successfully, False otherwise
+ """
+ # UUID4-shaped: defense-in-depth against SQL injection, mirroring
+ # every other engine's own id-shaped pattern for this check (e.g.
+ # Postgres/MySQL's `^\\d+$` for their integer ids).
+ if not cls.validate_cancel_query_id(
+ cancel_query_id,
+ r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-"
+ r"[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$",
+ ):
+ return False
+
+ try:
+ # SYNC blocks until ClickHouse confirms the outcome. The default
+ # ASYNC mode returns immediately with kill_status="waiting" --
+ # before termination is confirmed -- which would make this
+ # method's return value meaningless.
+ cursor.execute( # noqa: S608
+ f"KILL QUERY WHERE query_id = '{cancel_query_id}' SYNC"
+ )
+ rows = cursor.fetchall()
+ except Exception: # pylint: disable=broad-except
+ return False
+
+ kill_status = rows[0][0] if rows and rows[0] else None
+ if kill_status == cls._KILL_STATUS_CONFIRMED:
+ return True
+ if kill_status in cls._KILL_STATUS_KNOWN_NOT_CONFIRMED:
+ # A real process matched our query_id, but ClickHouse itself
+ # couldn't confirm termination -- there's positive evidence the
+ # query may still be running, so don't report success.
+ return False
+
+ # No recognizable kill_status landed in the first column. Verified
+ # empirically against a real server, this covers two cases that
+ # can't be told apart from this response alone:
+ # 1. No process matched `query_id` at all -- most likely the query
+ # had already finished on its own between the stop request and
+ # this KILL QUERY reaching the server: ClickHouse returns a
+ # genuinely empty HTTP body for a non-matching `KILL QUERY`
+ # (0 result_rows per its own X-ClickHouse-Summary header). The
+ # desired end state -- query not running -- already holds.
+ # 2. clickhouse-connect==1.9.0's DB-API cursor synthesizes a
+ # 1-row "stats" result from that same X-ClickHouse-Summary
+ # header whenever the real HTTP body is empty, so `rows` is
+ # never actually `[]` here even when nothing matched; its
+ # first column is an int (a row count), never one of the
+ # known kill_status strings above.
+ # Both report "nothing to confirm", not "confirmed failure" --
+ # treating this as success avoids stranding the SQL Lab UI on a
+ # query that has already finished on its own. See RCA.md for the
+ # full reasoning and empirical verification.
+ #
+ # Narrow, don't ignore, the one way "nothing to confirm" can still
+ # be wrong: the query may have finished -- with a real SUCCESS/
+ # FAILED/TIMED_OUT result already committed -- *during* the
+ # KILL QUERY SYNC round-trip above, after QueryDAO.stop_query()'s
+ # own already-complete guard ran (it only checks once, before
+ # calling into this method, and never re-checks after). Returning
+ # True here would make that caller unconditionally overwrite the
+ # just-committed status with STOPPED, destroying a legitimate
+ # result. `query` is the same session-bound ORM object
+ # `QueryDAO.stop_query()` already loaded, so re-reading it is a
+ # plain metadata-DB read, not a change to sql_lab.py/daos/query.py.
+ # This narrows the race to the (much smaller) gap between this
+ # refresh and that caller's own status assignment -- it cannot be
+ # closed from here without a lock that code doesn't have either.
+ # See RCA.md for the full reasoning.
+ try:
+ db.session.refresh(query)
+ already_terminal = query.status in (
+ QueryStatus.SUCCESS,
+ QueryStatus.FAILED,
+ QueryStatus.TIMED_OUT,
+ )
+ except Exception: # pylint: disable=broad-except
+ # Best-effort: if the refresh itself fails for any reason,
+ # fall back to the pre-existing behavior rather than letting a
+ # diagnostic check become a new failure mode of its own.
+ already_terminal = False
+
+ if already_terminal:
+ return False
+
+ return True
diff --git a/tests/testcontainers/db_engine_specs/test_clickhouse.py
b/tests/testcontainers/db_engine_specs/test_clickhouse.py
index d66d800114a..1c0a4e9a528 100644
--- a/tests/testcontainers/db_engine_specs/test_clickhouse.py
+++ b/tests/testcontainers/db_engine_specs/test_clickhouse.py
@@ -41,6 +41,8 @@ the real app instance `tests/conftest.py` already builds for
the rest of the
test run rather than constructing a second one.
"""
+import threading
+import time
from collections.abc import Iterator
import pytest
@@ -171,3 +173,119 @@ def test_group_by_all_groups_correctly(engine: Engine) ->
None:
).fetchall()
assert [tuple(row) for row in rows] == [("a", 15), ("b", 7)]
+
+
+def _processes_count_for(engine: Engine, query_id: str) -> int:
+ conn = engine.raw_connection()
+ cursor = conn.cursor()
+ cursor.execute(
+ "SELECT count() FROM system.processes WHERE query_id = %(query_id)s",
+ {"query_id": query_id},
+ )
+ return cursor.fetchone()[0]
+
+
+def test_cancel_query_actually_stops_the_query_server_side(engine: Engine) ->
None:
+ """
+ SQL Lab's Stop button must make the query actually stop running on the
+ ClickHouse server -- not just make Superset give up waiting on it. Drives
+ the real `get_cancel_query_id` -> `execute_with_cursor` -> `cancel_query`
+ sequence `superset/sql_lab.py` uses, against a real server, with a second,
+ independent connection/cursor for the cancellation itself (mirroring
+ `superset/sql_lab.py::cancel_query()`'s own `engine.raw_connection()`).
+ A regression that merely raises client-side without reaching a real
+ `KILL QUERY` on the server would leave the query in `system.processes`
+ and this test would catch that even though the client-side call returned.
+
+ Two things this test deliberately gets right that an earlier version of
+ it got wrong (caught by review, see RCA.md):
+
+ - `sleep(N)` runs once PER BLOCK, not once per row. `numbers(20)` alone
+ fits in a single block under ClickHouse's default `max_block_size`, so
+ `SELECT sleep(1) FROM numbers(20)` finishes in ~1s regardless of
+ whether cancellation works at all -- a no-op `KILL QUERY` would look
+ identical to a real one. `SETTINGS max_block_size=1` forces 20
+ separate blocks (and therefore 20 separate ~1s `sleep()` calls),
+ guaranteeing the query is still running seconds in, verified
+ empirically against a real server.
+ - The test confirms the query is ACTUALLY still running (queries
+ `system.processes` directly) before cancelling, rather than just
+ sleeping an arbitrary amount of time and hoping that's enough.
+ """
+ from tests.integration_tests.test_app import app
+
+ with app.app_context():
+ from superset.db_engine_specs.clickhouse import
ClickHouseConnectEngineSpec
+
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ query.id = 1
+
+ worker_conn = engine.raw_connection()
+ worker_cursor = worker_conn.cursor()
+
+ cancel_query_id = ClickHouseConnectEngineSpec.get_cancel_query_id(
+ worker_cursor, query
+ )
+ assert cancel_query_id is not None
+ query.set_extra_json_key("cancel_query", cancel_query_id)
+
+ # Forces 20 separate ~1s `sleep()` calls (one per block), guaranteeing
+ # genuine multi-second runtime regardless of ClickHouse's default block
+ # sizing -- see the docstring above.
+ sql = "SELECT sleep(1) FROM numbers(20) SETTINGS max_block_size=1"
+ elapsed_seconds: list[float] = []
+ errors: list[BaseException] = []
+ start = time.monotonic()
+
+ def run_query() -> None:
+ try:
+ ClickHouseConnectEngineSpec.execute_with_cursor(worker_cursor,
sql, query)
+ except Exception as ex: # pylint: disable=broad-except
+ # A genuine mid-flight kill makes clickhouse-connect raise
+ # QUERY_WAS_CANCELLED on the worker's own cursor.execute() call
+ # -- verified empirically against a real server. This is
+ # expected and not itself a failure; the real proof of
+ # cancellation is `system.processes` below, not whether this
+ # particular call happened to raise.
+ errors.append(ex)
+ finally:
+ elapsed_seconds.append(time.monotonic() - start)
+
+ worker_thread = threading.Thread(target=run_query)
+ worker_thread.start()
+
+ # Confirm the query is ACTUALLY still running server-side before
+ # cancelling it -- not just "we slept a bit and assumed so". Polls
+ # briefly since the query needs a moment to register in
+ # system.processes after the thread starts.
+ deadline = time.monotonic() + 10
+ while time.monotonic() < deadline:
+ if _processes_count_for(engine, cancel_query_id) == 1:
+ break
+ time.sleep(0.2)
+ else:
+ raise AssertionError(
+ "query never appeared in system.processes -- can't prove "
+ "cancellation works against a query that was never confirmed "
+ "running in the first place"
+ )
+
+ cancel_conn = engine.raw_connection()
+ cancel_cursor = cancel_conn.cursor()
+ assert (
+ ClickHouseConnectEngineSpec.cancel_query(cancel_cursor, query,
cancel_query_id)
+ is True
+ )
+
+ worker_thread.join(timeout=25)
+ assert not worker_thread.is_alive(), (
+ "worker query did not stop after cancel_query() returned True"
+ )
+ # The real assertion: fewer than the full ~20s elapsed, proving the
+ # query was actually killed server-side rather than running to
+ # completion while Superset's cancel_query() call merely returned True.
+ assert elapsed_seconds[0] < 15
+
+ assert _processes_count_for(engine, cancel_query_id) == 0
diff --git a/tests/unit_tests/dao/test_query_stop_race.py
b/tests/unit_tests/dao/test_query_stop_race.py
index 324fc1d7def..62f5a21c0b8 100644
--- a/tests/unit_tests/dao/test_query_stop_race.py
+++ b/tests/unit_tests/dao/test_query_stop_race.py
@@ -1224,3 +1224,98 @@ def
test_handle_query_error_status_refresh_does_not_autoflush_other_state(
f"expected no write before the targeted status SELECT, but found "
f"one: {preceding}"
)
+
+
+def test_execute_with_cursor_refresh_does_not_discard_pending_executed_sql(
+ app: Any, session: Session
+) -> None:
+ """
+ Round 9 review: `ClickHouseConnectEngineSpec.execute_with_cursor()`'s own
+ pre-dispatch STOPPED check does `db.session.refresh(query)` -- a *full*
+ refresh (every column, not the targeted single-attribute one
+ `handle_query_error()` uses), since nothing it does locally beforehand
+ needs protecting from an autoflush the way `handle_query_error()`'s
+ exception-handler callers do. But `refresh()` does NOT autoflush on its
+ own either: the async/Celery path's own `_make_execute_fn().
+ execute_with_stats()` (superset/sql/execution/celery_task.py) sets
+ `query.executed_sql = sql` immediately before calling
+ `execute_with_cursor()`, with no commit or flush in between. Without a
+ `flush()` first, that pending, uncommitted `executed_sql` would be
+ silently discarded and reloaded back to its previous committed value --
+ leaving the persisted Query row describing the wrong executed SQL,
+ verified here by actually reproducing the sequence and reading the
+ column back from a fresh query, not just asserting the in-memory
+ attribute.
+
+ Getting this reproduction right required care: a commit expires every
+ attribute by default (`expire_on_commit=True`), and accessing an
+ *expired* attribute (e.g. this method's own `query.id` in its first
+ debug-log line) triggers an incidental autoflush+reload as a side
+ effect of the lazy-load SELECT -- which would flush the pending
+ `executed_sql` "for free" and mask the bug entirely. The real call
+ sequence doesn't hit that accidental protection: `execute_sql_with_cursor`
+ (superset/sql/execution/executor.py) calls `check_stopped_fn()` --
+ which itself does a plain `db.session.refresh(query)` -- immediately
+ before `execute_fn()` in the *same* loop iteration, so every attribute
+ is already loaded (not expired) by the time `execute_with_cursor` sets
+ `executed_sql` and runs. This test replicates that ordering explicitly
+ (an explicit refresh to simulate `check_stopped_fn()`, then the dirty
+ write, then the call) rather than relying on a bare post-commit object,
+ which would accidentally dodge the bug it's meant to catch.
+ """
+ from unittest.mock import MagicMock
+
+ from superset import db
+ from superset.common.db_query_status import QueryStatus
+ from superset.models.core import Database
+ from superset.models.sql_lab import Query
+
+ engine = db.session.get_bind()
+ Query.metadata.create_all(engine) # pylint: disable=no-member
+
+ database = Database(database_name="my_database",
sqlalchemy_uri="sqlite://")
+ query_obj = Query(
+ client_id="pending-executed-sql",
+ database=database,
+ tab_name="test_tab",
+ sql_editor_id="test_editor_id",
+ sql="select * from bar",
+ select_sql="select * from bar",
+ executed_sql="select * from bar",
+ limit=100,
+ select_as_cta=False,
+ status=QueryStatus.RUNNING,
+ )
+ db.session.add(database)
+ db.session.add(query_obj)
+ db.session.commit()
+
+ with app.app_context():
+ from superset.db_engine_specs.clickhouse import
ClickHouseConnectEngineSpec
+
+ # Simulates execute_sql_with_cursor()'s own check_stopped_fn() call,
+ # which always runs immediately before execute_fn() in the same loop
+ # iteration -- this is what leaves every attribute already loaded (not
+ # expired) by the time the dirty write below happens, matching the real
+ # sequence instead of a freshly-expired, accidentally-protected object.
+ db.session.refresh(query_obj)
+
+ new_sql = "select * from bar where id = 42"
+ # Mirrors celery_task.py's execute_with_stats(): set the attribute, then
+ # call execute_with_cursor() with no flush/commit in between.
+ query_obj.executed_sql = new_sql
+
+ cursor_mock = MagicMock()
+ ClickHouseConnectEngineSpec.execute_with_cursor(cursor_mock, new_sql,
query_obj)
+
+ cursor_mock.execute.assert_called_once_with(new_sql, settings=None)
+
+ db.session.expire_all()
+ refreshed = (
+
db.session.query(Query).filter_by(client_id="pending-executed-sql").one()
+ )
+ assert refreshed.executed_sql == new_sql, (
+ f"expected the pending executed_sql set just before "
+ f"execute_with_cursor() to survive its internal refresh(), got "
+ f"{refreshed.executed_sql!r}"
+ )
diff --git a/tests/unit_tests/db_engine_specs/test_clickhouse.py
b/tests/unit_tests/db_engine_specs/test_clickhouse.py
index ac22b949cef..3b78fcb24f7 100644
--- a/tests/unit_tests/db_engine_specs/test_clickhouse.py
+++ b/tests/unit_tests/db_engine_specs/test_clickhouse.py
@@ -17,7 +17,7 @@
from datetime import datetime, timedelta, timezone
from typing import Any, Optional
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
import pytest
from sqlalchemy.engine.url import make_url
@@ -765,3 +765,345 @@ def test_multivalue_contains_any_numeric_coercion_sql()
-> None:
# element type) and confirm the emitted array literal is numeric.
expr = spec.array_contains_any(column("scores"), [5, 6])
assert _compile(expr) == "hasAny(scores, array(5, 6))"
+
+
+def test_get_cancel_query_id_mints_a_uuid_without_a_round_trip() -> None:
+ """
+ No SQL should be issued to obtain the id: ClickHouse never hands one
+ back, the id is entirely client-chosen (verified against a real server
+ in tests/testcontainers/db_engine_specs/test_clickhouse.py).
+ """
+ import uuid
+
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ query.id = 1
+ cursor_mock = Mock()
+ cancel_query_id = spec.get_cancel_query_id(cursor_mock, query)
+
+ assert cancel_query_id is not None
+ assert uuid.UUID(cancel_query_id).version == 4
+ cursor_mock.execute.assert_not_called()
+
+
+def test_get_cancel_query_id_declines_for_the_gtf_chart_data_stub() -> None:
+ """
+ `superset.tasks.query_cancel._CancellationQuery` stubs `query.id` as
+ `None` for the GTF chart-data path, which dispatches through the base
+ `execute()` rather than `execute_with_cursor` and so never forwards a
+ `settings` kwarg. An id minted for that path would never reach the
+ driver, so `cancel_query`'s `KILL QUERY WHERE query_id = ...` would have
+ nothing real to kill and would still report success via its "nothing to
+ confirm" fallback -- worse than the pre-existing, honest non-cancellable
+ behavior. Decline by returning `None` instead.
+ """
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+ query = Query()
+ assert query.id is None
+
+ assert spec.get_cancel_query_id(cursor_mock, query) is None
+ cursor_mock.execute.assert_not_called()
+
+
+def test_execute_with_cursor_forwards_recorded_cancel_id_as_query_id_setting()
-> None:
+ """
+ `execute_with_cursor` is the only place that can thread the id recorded
+ by `get_cancel_query_id` into the driver -- the base `execute`/
+ `execute_with_cursor` pair (superset/db_engine_specs/base.py) never
+ forwards per-query kwargs to `cursor.execute()`. Without this, the id
+ Superset records in `query.extra` would never reach the server, and
+ `KILL QUERY WHERE query_id = ...` would have nothing real to kill.
+ """
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ query.id = 1
+ query.set_extra_json_key("cancel_query",
"d14b9e5b-47ec-478e-bc7c-f8b56be98576")
+ cursor_mock = Mock()
+
+ spec.execute_with_cursor(cursor_mock, "SELECT 1", query)
+
+ cursor_mock.execute.assert_called_once_with(
+ "SELECT 1",
+ settings={"query_id": "d14b9e5b-47ec-478e-bc7c-f8b56be98576"},
+ )
+
+
+def test_execute_with_cursor_omits_settings_when_no_cancel_id_recorded() ->
None:
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ query.id = 1
+ cursor_mock = Mock()
+
+ spec.execute_with_cursor(cursor_mock, "SELECT 1", query)
+
+ cursor_mock.execute.assert_called_once_with("SELECT 1", settings=None)
+
+
+def test_execute_with_cursor_refuses_to_dispatch_an_already_stopped_query() ->
None:
+ """
+ The window this closes: the per-block loop's own STOPPED check
+ (sql_lab.py, which already ran once before this call) can pass, and
+ then Stop can be clicked and committed before the statement below is
+ actually sent. Unlike Postgres/MySQL (whose cancel_query() kills the
+ whole connection, so a cursor.execute() landing in this window fails
+ outright and gets caught upstream), ClickHouse's KILL QUERY has no such
+ side effect when nothing has matched yet -- without this check, the
+ statement would be sent and run to completion even though Superset
+ already told the user it stopped.
+ """
+ from superset.common.db_query_status import QueryStatus
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ ClickHouseQueryStoppedBeforeDispatchError,
+ )
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ query.id = 1
+ query.status = QueryStatus.STOPPED
+ cursor_mock = Mock()
+
+ with patch("superset.db_engine_specs.clickhouse.db.session.refresh"):
+ with pytest.raises(ClickHouseQueryStoppedBeforeDispatchError):
+ spec.execute_with_cursor(cursor_mock, "SELECT 1", query)
+
+ # The real assertion: the statement must never reach the driver.
+ cursor_mock.execute.assert_not_called()
+
+
+def test_execute_with_cursor_dispatches_normally_when_not_stopped() -> None:
+ from superset.common.db_query_status import QueryStatus
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ query.id = 1
+ query.status = QueryStatus.RUNNING
+ cursor_mock = Mock()
+
+ with patch("superset.db_engine_specs.clickhouse.db.session.refresh"):
+ spec.execute_with_cursor(cursor_mock, "SELECT 1", query)
+
+ cursor_mock.execute.assert_called_once_with("SELECT 1", settings=None)
+
+
+def test_cancel_query_issues_kill_query_sync_for_the_recorded_id() -> None:
+ """
+ SYNC is required: the default ASYNC mode returns immediately with
+ kill_status="waiting", before ClickHouse has confirmed the query is
+ actually gone -- verified empirically against a real server (see
+ RCA.md). Without SYNC, cancel_query's return value would be meaningless.
+ """
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+ cancel_query_id = "d14b9e5b-47ec-478e-bc7c-f8b56be98576"
+ cursor_mock.fetchall.return_value = [
+ ["finished", cancel_query_id, "default", "SELECT 1"]
+ ]
+
+ assert spec.cancel_query(cursor_mock, Query(), cancel_query_id) is True
+ cursor_mock.execute.assert_called_once_with(
+ f"KILL QUERY WHERE query_id = '{cancel_query_id}' SYNC"
+ )
+
+
+def test_cancel_query_returns_false_when_kill_status_not_confirmed() -> None:
+ """
+ A real process matched our query_id, but ClickHouse itself reports it
+ couldn't confirm termination (e.g. `cant_cancel`) -- there's positive
+ evidence the query may still be running, so this must not report
+ success just because `cursor.execute()` didn't raise.
+ """
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ cancel_query_id = "d14b9e5b-47ec-478e-bc7c-f8b56be98576"
+
+ for kill_status in ("waiting", "cant_cancel", "pending", "unknown_status"):
+ cursor_mock = Mock()
+ cursor_mock.fetchall.return_value = [
+ [kill_status, cancel_query_id, "default", "SELECT 1"]
+ ]
+ assert spec.cancel_query(cursor_mock, Query(), cancel_query_id) is
False, (
+ f"kill_status={kill_status!r} must not report success"
+ )
+
+
+def test_cancel_query_returns_true_when_nothing_matched() -> None:
+ """
+ No row with a recognizable kill_status -- either a genuinely empty
+ result (no process matched, e.g. the query already finished on its
+ own) or clickhouse-connect==1.9.0's own synthesized "stats" row (its
+ DB-API cursor returns a 1-row summary derived from the
+ X-ClickHouse-Summary response header instead of an empty list when the
+ real HTTP body is empty -- verified empirically against a real server,
+ see RCA.md). Both mean "nothing to confirm", which is treated as
+ success so SQL Lab doesn't get stuck believing an already-finished
+ query is still running.
+
+ `Query()` here is transient (never added to a session), so the
+ best-effort `db.session.refresh(query)` follow-up check (see the next
+ two tests) raises and is caught, falling back to this same True --
+ exercising that fallback path, not bypassing it.
+ """
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ cancel_query_id = "d14b9e5b-47ec-478e-bc7c-f8b56be98576"
+
+ # Case 1: a genuinely empty result.
+ cursor_mock = Mock()
+ cursor_mock.fetchall.return_value = []
+ assert spec.cancel_query(cursor_mock, Query(), cancel_query_id) is True
+
+ # Case 2: clickhouse-connect's synthesized X-ClickHouse-Summary stats
+ # row -- same 9-column shape observed against a real server, first
+ # column is an int (a row count), never a kill_status string.
+ cursor_mock = Mock()
+ cursor_mock.fetchall.return_value = [[1, 834, 0, 0, 0, 0, 0, 1686330,
"some-id"]]
+ assert spec.cancel_query(cursor_mock, Query(), cancel_query_id) is True
+
+
+def
test_cancel_query_returns_false_when_query_finished_during_kill_round_trip() ->
(
+ None
+):
+ """
+ The narrow race this guards: KILL QUERY SYNC finds no match because the
+ query finished -- with a real result already committed -- *during* the
+ round-trip itself, after QueryDAO.stop_query()'s own already-complete
+ guard ran (it only checks once, before calling cancel_query(), and never
+ re-checks after). Returning True here would make that caller
+ unconditionally overwrite the just-committed status with STOPPED.
+ """
+ from superset.common.db_query_status import QueryStatus
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ cancel_query_id = "d14b9e5b-47ec-478e-bc7c-f8b56be98576"
+
+ for terminal_status in (
+ QueryStatus.SUCCESS,
+ QueryStatus.FAILED,
+ QueryStatus.TIMED_OUT,
+ ):
+ query = Query()
+ query.status = terminal_status
+ cursor_mock = Mock()
+ cursor_mock.fetchall.return_value = []
+
+ with patch("superset.db_engine_specs.clickhouse.db.session.refresh"):
+ assert spec.cancel_query(cursor_mock, query, cancel_query_id) is
False, (
+ f"status={terminal_status!r} must not be overwritten"
+ )
+
+
+def test_cancel_query_returns_true_when_refresh_shows_still_running() -> None:
+ """
+ The common case: the refresh succeeds and shows the query is still
+ genuinely RUNNING (or STOPPED already) -- nothing matched because it
+ hadn't registered yet, or some other benign reason, not because a
+ result was raced. This must stay True, or routine stops would start
+ spuriously failing.
+ """
+ from superset.common.db_query_status import QueryStatus
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ cancel_query_id = "d14b9e5b-47ec-478e-bc7c-f8b56be98576"
+
+ query = Query()
+ query.status = QueryStatus.RUNNING
+ cursor_mock = Mock()
+ cursor_mock.fetchall.return_value = []
+
+ with patch("superset.db_engine_specs.clickhouse.db.session.refresh"):
+ assert spec.cancel_query(cursor_mock, query, cancel_query_id) is True
+
+
+def test_cancel_query_rejects_non_uuid_ids() -> None:
+ """
+ Defense-in-depth against injection via a tampered/forged cancel id --
+ mirrors every other engine's own id-shaped validation (e.g. Postgres/
+ MySQL's `^\\d+$` for their integer ids).
+ """
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert (
+ spec.cancel_query(cursor_mock, Query(), "1; DROP TABLE
system.processes")
+ is False
+ )
+ cursor_mock.execute.assert_not_called()
+
+
+def test_cancel_query_returns_false_on_driver_error() -> None:
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseConnectEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+ cursor_mock.execute.side_effect = Exception("boom")
+
+ assert (
+ spec.cancel_query(cursor_mock, Query(),
"d14b9e5b-47ec-478e-bc7c-f8b56be98576")
+ is False
+ )
+
+
+def test_legacy_clickhouse_engine_spec_has_no_cancel_support() -> None:
+ """
+ The legacy `clickhouse-sqlalchemy` driver's cursor (`ClickHouseEngineSpec`,
+ the non-recommended "clickhouse" engine) has a different `execute()`
+ signature with no `settings` kwarg, and only generates its own query id
+ internally inside `execute()` -- it isn't available beforehand the way
+ `has_query_id_before_execute` (the base default both engines use)
+ requires. Passing a `settings` kwarg to that driver's cursor would raise
+ a TypeError, so the cancel overrides below live on
+ `ClickHouseConnectEngineSpec` only; this locks in that the legacy engine
+ still falls through to the inert base defaults rather than inheriting
+ something that would break it.
+ """
+ from superset.db_engine_specs.clickhouse import ( # noqa: N813
+ ClickHouseEngineSpec as spec,
+ )
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+ assert spec.get_cancel_query_id(cursor_mock, Query()) is None
+ assert spec.cancel_query(cursor_mock, Query(), "123") is False
diff --git a/tests/unit_tests/sql/execution/test_celery_task.py
b/tests/unit_tests/sql/execution/test_celery_task.py
index f48b6d4c789..30d7b696c41 100644
--- a/tests/unit_tests/sql/execution/test_celery_task.py
+++ b/tests/unit_tests/sql/execution/test_celery_task.py
@@ -1496,6 +1496,53 @@ def test_execute_sql_task_unhandled_exception(
assert result["status"] == QueryStatusEnum.FAILED
+def test_execute_sql_task_preserves_stopped_status_despite_dispatch_exception(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_query: MagicMock,
+) -> None:
+ """
+ Regression guard, not a characterization of a gap: an exception that
+ reaches `execute_sql_task`'s outer handler while `query.status` is
+ already STOPPED must not get overwritten to FAILED -- regardless of
+ engine, and regardless of why the exception was raised (a real
+ cancel-induced driver error for any engine whose cancellation mechanism
+ can make a subsequent/concurrent `cursor.execute()` fail, not something
+ specific to ClickHouse).
+
+ This used to be a real, disclosed gap in `_handle_query_error`, which
+ only special-cased TIMED_OUT and had no equivalent STOPPED check. It is
+ now fixed upstream on master (independently of this PR, picked up via
+ merge): `_handle_query_error` checks
+ `_make_check_stopped_fn(query, preserve_pending=True)()` first and
+ returns the recorded STOPPED outcome via `_stopped_query_payload()`
+ instead of overwriting it. Kept as a regression guard for this
+ property, since it's exactly the invariant this PR's own
+ `ClickHouseQueryStoppedBeforeDispatchError` (see
+ `db_engine_specs/clickhouse.py::execute_with_cursor`) relies on holding
+ on the Celery path too.
+ """
+ from superset.sql.execution.celery_task import execute_sql_task
+
+ mock_query.status = QueryStatusEnum.STOPPED
+
+ 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("driver raised during dispatch"),
+ )
+ 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.STOPPED
+ assert mock_query.status == QueryStatusEnum.STOPPED
+
+
def test_execute_sql_task_success_final_commit(
mocker: MockerFixture,
app_context: None,
diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py
index aa9b04400a9..fe91caa1cc7 100644
--- a/tests/unit_tests/sql_lab_test.py
+++ b/tests/unit_tests/sql_lab_test.py
@@ -126,6 +126,44 @@ def test_execute_query(mocker: MockerFixture, app: None)
-> None:
SupersetResultSet.assert_called_with([(42,)], cursor.description,
db_engine_spec)
+def test_execute_query_converts_stopped_exception_to_sqllab_query_stopped(
+ mocker: MockerFixture, app: None
+) -> None:
+ """
+ This is the generic mechanism `ClickHouseConnectEngineSpec.
+ execute_with_cursor` relies on to refuse dispatching an already-STOPPED
+ query without needing to import or construct `SqlLabQueryStoppedException`
+ itself (which would be circular -- `sql_lab.py` already imports from
+ `superset.db_engine_specs`): any exception a `db_engine_spec`'s
+ `execute_with_cursor` raises is caught here, `query` is refreshed, and if
+ its status is STOPPED it's converted into `SqlLabQueryStoppedException`
+ -- regardless of the exception's type or message. Verified directly
+ against `execute_query` (not assumed) so a future change to this
+ mechanism would be caught here first.
+ """
+ from superset.sql_lab import SqlLabQueryStoppedException
+
+ query = mocker.MagicMock()
+ query.executed_sql = "SELECT 42 AS answer"
+ query.status = QueryStatus.STOPPED
+ query.limit = 1
+ database = query.database
+ database.allow_dml = False
+ db_engine_spec = database.db_engine_spec
+ db_engine_spec.execute_with_cursor.side_effect = Exception(
+ "refused to dispatch an already-stopped query"
+ )
+
+ cursor = mocker.MagicMock()
+ # Refreshing a MagicMock is a no-op; `query.status` is already set above
+ # to simulate QueryDAO.stop_query() having committed STOPPED during the
+ # window between the per-block loop's own earlier check and this call.
+ mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
+
+ with pytest.raises(SqlLabQueryStoppedException):
+ execute_query(query, cursor=cursor, log_params={})
+
+
def test_get_query_rolls_back_session_before_retrying(
mocker: MockerFixture, app: SupersetApp
) -> None:
