You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by vi...@apache.org on 2026/08/28 00:40:19 UTC
(superset) branch villebro/gtf-heartbeat-cancellation created (now 48a97c9da1b)
This is an automated email from the ASF dual-hosted git repository.
villebro pushed a change to branch villebro/gtf-heartbeat-cancellation
in repository https://gitbox.apache.org/repos/asf/superset.git
at 48a97c9da1b feat(gtf): heartbeat orphan reaping + cancellable chart-data queries
This branch includes the following new commits:
new 48a97c9da1b feat(gtf): heartbeat orphan reaping + cancellable chart-data queries
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(superset) 01/01: feat(gtf): heartbeat orphan reaping + cancellable chart-data queries
Posted by vi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
villebro pushed a commit to branch villebro/gtf-heartbeat-cancellation
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 48a97c9da1b24fa34e085f68a740b41b59dea0f4
Author: Ville Brofeldt <vi...@gmail.com>
AuthorDate: Thu Aug 27 17:39:36 2026 -0700
feat(gtf): heartbeat orphan reaping + cancellable chart-data queries
Add a worker-liveness heartbeat and fold orphan reaping into the existing
prune_tasks cron so tasks abandoned by a dead worker are revoked and marked
FAILURE, and make async chart-data query tasks cancellable at the engine level
with a restored per-query timeout.
---
UPDATING.md | 18 +++
superset/commands/tasks/reap.py | 135 ++++++++++++++++++
superset/config.py | 28 +++-
superset/daos/tasks.py | 109 ++++++++++++++-
..._7e2c9a4f1b83_create_task_dependencies_table.py | 13 ++
superset/models/core.py | 7 +
superset/models/tasks.py | 7 +
superset/tasks/async_queries.py | 60 +++++++-
superset/tasks/heartbeat.py | 84 +++++++++++
superset/tasks/query_cancel.py | 131 ++++++++++++++++++
superset/tasks/scheduler.py | 130 ++++++++++++-----
.../integration_tests/tasks/commands/test_reap.py | 144 +++++++++++++++++++
tests/unit_tests/tasks/test_dependencies.py | 10 +-
tests/unit_tests/tasks/test_heartbeat.py | 76 ++++++++++
tests/unit_tests/tasks/test_query_cancel.py | 154 +++++++++++++++++++++
tests/unit_tests/tasks/test_reap.py | 124 +++++++++++++++++
tests/unit_tests/tasks/test_scheduler_executor.py | 137 ++++++++++++++++++
17 files changed, 1316 insertions(+), 51 deletions(-)
diff --git a/UPDATING.md b/UPDATING.md
index 0b210481742..e4b7d25ebc4 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -108,6 +108,24 @@ channels (`entity-changes:*` broadcast nudges, `realtime:<channel_id>`
per-principal messages) instead of the removed async-events streams; see
`superset-websocket/README.md`.
+Orphaned GTF tasks (a worker killed mid-execution) are now detected and cleaned
+up server-side. While a worker holds a task it writes a liveness heartbeat
+(`tasks.last_heartbeat`, every `GTF_TASK_HEARTBEAT_INTERVAL` seconds, default
+`15`); the existing `prune_tasks` Celery job now reaps any active task whose
+heartbeat is older than `GTF_ORPHAN_TASK_TIMEOUT` (default `60`) — revoking its
+Celery job and marking it `FAILURE` so waiters unblock — before its usual
+deletion of old rows. Because reaping rides on `prune_tasks`, enable that beat
+schedule (and run it on a short interval, e.g. every minute) to bound how long an
+orphaned task and its warehouse query can linger. The heartbeat write is issued
+out-of-band and deliberately does not advance `changed_on`.
+
+Async chart-data query tasks are now cancellable: a per-query timeout
+(`GLOBAL_ASYNC_QUERIES_QUERY_TIMEOUT`, default `None` = unbounded) or a user
+cancel aborts the task, and on database engines that support query cancellation
+(e.g. PostgreSQL, MySQL, Snowflake, Redshift) the abort also cancels the running
+warehouse query over a fresh connection. Engines without cancel support are
+unaffected — the task is still freed, but the query runs to completion.
+
- `SAMPLES_ROW_LIMIT` is now the default for `/datasource/samples` requests without a valid explicit `per_page`, rather than a hard per-request ceiling; explicit limits are honored up to the existing global row-limit ceiling, matching `/chart/data` SAMPLES requests.
### MCP tool results preserve stored string values
diff --git a/superset/commands/tasks/reap.py b/superset/commands/tasks/reap.py
new file mode 100644
index 00000000000..30e19aa09d0
--- /dev/null
+++ b/superset/commands/tasks/reap.py
@@ -0,0 +1,135 @@
+# 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.
+"""Reap GTF tasks abandoned by a dead worker or wedged in ABORTING."""
+
+import logging
+from typing import cast
+
+from flask import current_app
+from superset_core.tasks.types import TaskProperties, TaskStatus
+
+from superset import db
+from superset.commands.base import BaseCommand
+from superset.daos.tasks import OrphanCandidate, TaskDAO
+from superset.extensions import celery_app
+from superset.stats_logger import BaseStatsLogger
+from superset.tasks.constants import ACTIVE_STATES
+from superset.tasks.utils import parse_properties
+
+logger = logging.getLogger(__name__)
+
+ORPHAN_ERROR_MESSAGE = "Task orphaned: worker heartbeat timed out"
+
+
+class ReapOrphanedTasksCommand(BaseCommand):
+ """Detect and clean up tasks abandoned by a dead or wedged worker.
+
+ Run from the prune cron before row deletion. Candidates come from
+ ``TaskDAO.find_orphaned``:
+
+ - **Orphans** (dead worker, stale heartbeat): revoke any lingering Celery
+ job, then force the row to FAILURE and publish completion so waiters
+ (sync joiners, DAG dependents, chart-data pollers) unblock — the worker is
+ gone and cannot finalize itself.
+ - **Wedged aborts** (live worker still ABORTING past the grace window): only
+ escalate with a forced revoke, letting the worker's own ``finally`` block
+ finalize the status once ``SoftTimeLimitExceeded`` is raised.
+
+ All Celery/DB operations are best-effort accelerators; the atomic
+ compare-and-swap transition is authoritative, so a worker that revives and
+ commits its own terminal status simply makes the reaper's CAS a no-op.
+ """
+
+ def run(self) -> int:
+ stats_logger: BaseStatsLogger = current_app.config["STATS_LOGGER"]
+ # A single threshold governs both "no heartbeat for this long → orphaned"
+ # and "ABORTING for this long despite a live heartbeat → escalate".
+ timeout = current_app.config["GTF_ORPHAN_TASK_TIMEOUT"]
+
+ candidates = TaskDAO.find_orphaned(timeout, timeout)
+ reaped = sum(self._process(candidate, stats_logger) for candidate in candidates)
+ if candidates:
+ logger.info(
+ "Orphan reaper processed %d task(s), reaped %d", len(candidates), reaped
+ )
+ return reaped
+
+ def _process(
+ self, candidate: OrphanCandidate, stats_logger: BaseStatsLogger
+ ) -> bool:
+ """Handle one candidate; return True if it was reaped to FAILURE."""
+ # Preserve existing runtime state (is_abortable, timeout, celery_task_id,
+ # ...); conditional_status_update replaces the whole properties column.
+ properties = cast(TaskProperties, dict(parse_properties(candidate.properties)))
+ self._revoke(properties.get("celery_task_id"), stats_logger)
+
+ if not candidate.is_orphan:
+ # Wedged abort: the live worker finalizes its own status once the
+ # forced revoke raises SoftTimeLimitExceeded in it.
+ stats_logger.incr("gtf.task.abort_escalated")
+ logger.warning(
+ "Escalated wedged abort for task %s via forced revoke", candidate.uuid
+ )
+ return False
+
+ properties["error_message"] = ORPHAN_ERROR_MESSAGE
+ properties["exception_type"] = "OrphanedTaskError"
+
+ if not TaskDAO.conditional_status_update(
+ candidate.uuid,
+ TaskStatus.FAILURE,
+ expected_status=list(ACTIVE_STATES),
+ properties=properties,
+ set_ended_at=True,
+ ):
+ # The worker revived and committed a terminal status first.
+ return False
+
+ db.session.commit()
+
+ from superset.tasks.manager import TaskManager
+
+ TaskManager.publish_completion(candidate.uuid, TaskStatus.FAILURE.value)
+ TaskManager.publish_entity_change(candidate.uuid)
+ stats_logger.incr("gtf.task.orphan_reaped")
+ logger.warning(
+ "Reaped orphaned task %s (worker heartbeat stale)", candidate.uuid
+ )
+ return True
+
+ def _revoke(
+ self, celery_task_id: str | None, stats_logger: BaseStatsLogger
+ ) -> None:
+ """Forcibly revoke a Celery job (SIGUSR1 → SoftTimeLimitExceeded).
+
+ ``terminate=True`` signals the worker child running the task so it raises
+ and reclaims the slot; for a not-yet-started job it prevents execution.
+ Best-effort — the DB transition is the source of truth.
+ """
+ if not celery_task_id:
+ return
+ try:
+ celery_app.control.revoke(celery_task_id, terminate=True, signal="SIGUSR1")
+ stats_logger.incr("gtf.task.revoke")
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ stats_logger.incr("gtf.task.revoke_failure")
+ logger.warning(
+ "Failed to revoke Celery task %s", celery_task_id, exc_info=True
+ )
+
+ def validate(self) -> None:
+ pass
diff --git a/superset/config.py b/superset/config.py
index 179f448ebf2..0f9d7a2a6e8 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -1844,10 +1844,13 @@ class CeleryConfig: # pylint: disable=too-few-public-methods
# "schedule": crontab(minute="*", hour="*"),
# "kwargs": {"retention_period_days": 180, "max_rows_per_run": 10000},
# },
- # Uncomment to enable pruning of the tasks table
+ # Uncomment to enable pruning of the tasks table. This job also reaps
+ # orphaned GTF tasks (abandoned by a dead worker) before deleting old
+ # rows, so run it on a short interval — reaping latency tracks this
+ # cadence, and an abandoned warehouse query keeps running until reaped.
# "prune_tasks": {
# "task": "prune_tasks",
- # "schedule": crontab(minute=0, hour=0),
+ # "schedule": crontab(minute="*", hour="*"),
# "kwargs": {"retention_period_days": 90, "max_rows_per_run": 10000},
# },
# Uncomment to enable pruning of expired entries from the key-value store
@@ -2944,6 +2947,15 @@ GLOBAL_ASYNC_QUERIES_POLLING_STALE_TIMEOUT = int(
# ``/chart/data`` requests are unaffected even when GLOBAL_ASYNC_QUERIES is on.
GLOBAL_ASYNC_QUERIES_MIN_CACHE_TTL = int(timedelta(minutes=5).total_seconds())
+# Timeout (seconds) for an async chart-data query task. When reached, the task is
+# aborted; on engines that support query cancellation the abort handler also
+# cancels the underlying warehouse query (over a fresh connection), so the task
+# ends promptly as TIMED_OUT. On engines without cancel support the query is not
+# interrupted (the task is freed once the query returns on its own). Default None
+# leaves async chart-data queries unbounded (matching prior behavior); set an int
+# to enforce a ceiling.
+GLOBAL_ASYNC_QUERIES_QUERY_TIMEOUT: int | None = None
+
# Deployment default for whether the UI runs chart-data queries asynchronously when
# GLOBAL_ASYNC_QUERIES is enabled. This is a FRONTEND-ONLY policy input: async is
# opt-in per request via an ``async_mode`` flag on ``/chart/data`` (an absent flag is
@@ -3234,6 +3246,18 @@ TASK_ABORT_POLLING_DEFAULT_INTERVAL = 10
# Set to 0 to disable throttling (write every update to DB).
TASK_PROGRESS_UPDATE_THROTTLE_INTERVAL = 2 # seconds
+# GTF worker-liveness heartbeat. While a worker holds a task it bumps
+# tasks.last_heartbeat every GTF_TASK_HEARTBEAT_INTERVAL seconds. The prune cron
+# (prune_tasks) reaps any ACTIVE task whose heartbeat is older than
+# GTF_ORPHAN_TASK_TIMEOUT — a worker that died mid-execution — by revoking its
+# Celery job and marking it FAILURE. The same threshold escalates a task stuck
+# in ABORTING (live worker ignoring a cooperative abort) with a forced revoke.
+# Keep the timeout comfortably larger than the interval (>= ~3x) so a brief GC
+# pause or stall does not look like a dead worker. Reaping only runs when the
+# prune_tasks beat schedule is enabled (see CELERY_CONFIG.beat_schedule).
+GTF_TASK_HEARTBEAT_INTERVAL = 15 # seconds
+GTF_ORPHAN_TASK_TIMEOUT = 60 # seconds
+
# ---------------------------------------------------
# Distributed Coordination Configuration
# ---------------------------------------------------
diff --git a/superset/daos/tasks.py b/superset/daos/tasks.py
index 496166e1968..8a0030d0ca0 100644
--- a/superset/daos/tasks.py
+++ b/superset/daos/tasks.py
@@ -17,10 +17,11 @@
"""Task DAO for Global Task Framework (GTF)"""
import logging
-from datetime import datetime
-from typing import Any, Literal, TypedDict
+from datetime import datetime, timedelta
+from typing import Any, Literal, NamedTuple, TypedDict
from uuid import UUID
+import sqlalchemy as sa
from sqlalchemy.sql.elements import ColumnElement
from superset_core.tasks.types import TaskProperties, TaskScope, TaskStatus
@@ -30,7 +31,7 @@ from superset.extensions import db
from superset.models.task_dependencies import TaskDependency
from superset.models.task_subscribers import TaskSubscriber
from superset.models.tasks import Task
-from superset.tasks.constants import ABORTABLE_STATES, TERMINAL_STATES
+from superset.tasks.constants import ABORTABLE_STATES, ACTIVE_STATES, TERMINAL_STATES
from superset.tasks.filters import TaskFilter
from superset.tasks.utils import (
get_active_dedup_key,
@@ -52,6 +53,23 @@ class TaskSubscriberPrincipal(TypedDict):
sub: str
+class OrphanCandidate(NamedTuple):
+ """A task the prune cron should reap or escalate.
+
+ ``is_orphan`` distinguishes the two cases the reaper handles differently:
+ ``True`` = a dead/orphaned worker (stale heartbeat) that must be fully reaped
+ (cancel its query, revoke, force the row terminal); ``False`` = a still-live
+ worker wedged in ABORTING past the grace window, which is only escalated with
+ a forced revoke so the worker's own ``finally`` block finalizes the status.
+ """
+
+ id: int
+ uuid: UUID
+ status: str
+ properties: str | None
+ is_orphan: bool
+
+
class TaskDAO(BaseDAO[Task]):
"""
Concrete TaskDAO for the Global Task Framework (GTF).
@@ -99,6 +117,91 @@ class TaskDAO(BaseDAO[Task]):
"""
return db.session.query(Task.id).filter(Task.uuid == task_uuid).scalar()
+ @classmethod
+ def touch_heartbeat(cls, task_id: int) -> None:
+ """Bump a task's ``last_heartbeat`` to now without touching ``changed_on``.
+
+ Called on an interval by the executing worker's heartbeat thread so the
+ prune cron can tell a live task from an orphaned one. Uses a raw textual
+ UPDATE on the single column: ``changed_on`` carries a client-side
+ ``onupdate`` default that ORM/Core updates would fire, and any bump to
+ ``changed_on`` would resurface the task in ``get_statuses_changed_since``
+ every heartbeat (resetting client polling backoff). A textual statement
+ bypasses that default processing entirely. Bound by integer ``id`` to
+ sidestep ``UUIDType`` dialect handling.
+ """
+ db.session.execute(
+ sa.text("UPDATE tasks SET last_heartbeat = :ts WHERE id = :id"),
+ {"ts": naive_utcnow(), "id": task_id},
+ )
+ # Deliberate standalone commit: this is a single out-of-band column write,
+ # not a unit of work, and must not be wrapped in the ORM transaction flow.
+ db.session.commit() # pylint: disable=consider-using-transaction
+
+ @classmethod
+ def find_orphaned(
+ cls,
+ orphan_timeout_seconds: int,
+ abort_grace_seconds: int,
+ ) -> list[OrphanCandidate]:
+ """Return active tasks the prune cron should reap or escalate.
+
+ Two disjoint sets, keyed off the worker liveness heartbeat:
+
+ - **Orphans** (``is_orphan=True``): any ACTIVE task
+ (PENDING/IN_PROGRESS/ABORTING) whose heartbeat has gone stale
+ (``last_heartbeat < now - orphan_timeout_seconds``) — no live worker is
+ holding it. A NULL heartbeat means no worker has picked the task up yet
+ (still legitimately queued), so it is excluded.
+ - **Wedged aborts** (``is_orphan=False``): a task still ABORTING with a
+ *fresh* heartbeat that entered ABORTING (``changed_on``) more than
+ ``abort_grace_seconds`` ago — a live worker not honoring the cooperative
+ abort, to be escalated with a forced Celery revoke. Requiring a fresh
+ heartbeat keeps this set disjoint from the orphan set.
+
+ Skips the base filter: internal maintenance over all tasks, not a
+ user-facing listing.
+
+ :returns: candidates to reap (dead worker) or escalate (wedged abort)
+ """
+ now = naive_utcnow()
+ stale_before = now - timedelta(seconds=orphan_timeout_seconds)
+ # last_heartbeat / started_at / ended_at are stored as naive UTC, but
+ # FAB's changed_on is naive *local* (its onupdate is datetime.now). Use a
+ # matching local threshold for the changed_on comparison below.
+ abort_before = datetime.now() - timedelta(seconds=abort_grace_seconds)
+
+ columns = (Task.id, Task.uuid, Task.status, Task.properties)
+
+ orphan_rows = (
+ db.session.query(*columns)
+ .filter(
+ Task.status.in_(ACTIVE_STATES),
+ Task.last_heartbeat.isnot(None),
+ Task.last_heartbeat < stale_before,
+ )
+ .all()
+ )
+ wedged_rows = (
+ db.session.query(*columns)
+ .filter(
+ Task.status == TaskStatus.ABORTING.value,
+ Task.last_heartbeat.isnot(None),
+ Task.last_heartbeat >= stale_before,
+ # changed_on is shadowed by CoreTask's bare annotation; reference
+ # the real column so the comparison targets the timestamp.
+ Task.__table__.c.changed_on < abort_before,
+ )
+ .all()
+ )
+ return [
+ OrphanCandidate(id_, uuid, status, properties, is_orphan=True)
+ for id_, uuid, status, properties in orphan_rows
+ ] + [
+ OrphanCandidate(id_, uuid, status, properties, is_orphan=False)
+ for id_, uuid, status, properties in wedged_rows
+ ]
+
@classmethod
def get_statuses_changed_since(
cls, cursor: datetime | None, task_type: str | None = None
diff --git a/superset/migrations/versions/2026-08-21_12-00_7e2c9a4f1b83_create_task_dependencies_table.py b/superset/migrations/versions/2026-08-21_12-00_7e2c9a4f1b83_create_task_dependencies_table.py
index 09cfa973e04..656fbdf2e24 100644
--- a/superset/migrations/versions/2026-08-21_12-00_7e2c9a4f1b83_create_task_dependencies_table.py
+++ b/superset/migrations/versions/2026-08-21_12-00_7e2c9a4f1b83_create_task_dependencies_table.py
@@ -148,9 +148,22 @@ def upgrade():
"uq_task_subscribers_task_guest", ["task_id", "guest_key"]
)
+ # Liveness marker for orphan detection: the executing worker bumps
+ # ``tasks.last_heartbeat`` on a background thread, and the prune cron reaps
+ # ACTIVE tasks whose heartbeat has gone stale (a dead/orphaned worker). The
+ # index backs the reaper's ``last_heartbeat < now - timeout`` scan.
+ add_columns(
+ TASKS_TABLE,
+ Column("last_heartbeat", DateTime, nullable=True),
+ )
+ create_index(TASKS_TABLE, "ix_tasks_last_heartbeat", ["last_heartbeat"])
+
def downgrade():
"""Drop task_dependencies and revert the task_subscribers.guest_key change."""
+ drop_index(TASKS_TABLE, "ix_tasks_last_heartbeat")
+ drop_columns(TASKS_TABLE, "last_heartbeat")
+
# Guest subscriptions cannot be represented without the column; drop those
# rows first so restoring user_id NOT NULL does not fail on NULL user_id.
op.execute(sa.text("DELETE FROM task_subscribers WHERE user_id IS NULL"))
diff --git a/superset/models/core.py b/superset/models/core.py
index a941fd52acd..2be6654f93a 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -927,6 +927,13 @@ class Database(CoreDatabase, AuditMixinNullable, ImportExportMixin): # pylint:
rows = None
description = None
+ # Give an active GTF chart-data task a chance to capture an engine
+ # cancel id off the live cursor before the (blocking) execute below,
+ # so a concurrent abort/timeout can kill the query. No-op otherwise.
+ from superset.tasks.query_cancel import notify_cursor
+
+ notify_cursor(cursor)
+
for i, statement in enumerate(script.statements):
# For a single statement, execute the original SQL as-is. Re-rendering
# via statement.format() would round-trip through sqlglot
diff --git a/superset/models/tasks.py b/superset/models/tasks.py
index 11d111f8296..9c41b4e2b44 100644
--- a/superset/models/tasks.py
+++ b/superset/models/tasks.py
@@ -87,6 +87,13 @@ class Task(CoreTask, AuditMixinNullable, Model):
started_at = Column(DateTime, nullable=True)
ended_at = Column(DateTime, nullable=True)
+ # Liveness marker bumped by the executing worker's heartbeat thread while it
+ # holds the task. The prune cron reaps ACTIVE tasks whose heartbeat has gone
+ # stale (a dead/orphaned worker). Written out-of-band via a raw UPDATE (see
+ # TaskDAO.touch_heartbeat) so a heartbeat never advances changed_on and thus
+ # never resurfaces the task in the status-change poll.
+ last_heartbeat = Column(DateTime, nullable=True, index=True)
+
# User context for execution
user_id = Column(Integer, nullable=True)
diff --git a/superset/tasks/async_queries.py b/superset/tasks/async_queries.py
index 97b7e0f8cf3..f4d609af3d1 100644
--- a/superset/tasks/async_queries.py
+++ b/superset/tasks/async_queries.py
@@ -17,10 +17,12 @@
from __future__ import annotations
import logging
+from contextlib import contextmanager
from datetime import datetime
-from typing import Any, TYPE_CHECKING
+from typing import Any, Iterator, TYPE_CHECKING
from uuid import UUID
+from flask import current_app
from flask_appbuilder.security.sqla.models import User
from superset_core.tasks.types import TaskOptions, TaskScope
@@ -36,6 +38,7 @@ from superset.extensions import (
)
from superset.tasks.ambient_context import get_context
from superset.tasks.decorators import task
+from superset.tasks.query_cancel import cancel_chart_query, capture_cancel_id
from superset.utils.core import override_user
if TYPE_CHECKING:
@@ -120,10 +123,47 @@ def _get_dependency_cache_key() -> str:
raise SupersetException("Prerequisite task did not publish a cache key")
-# No timeout is set on these tasks yet: GTF enforces a timeout by aborting the
-# task, but chart-data queries have no abort handler to cancel the underlying
-# warehouse query, so a timeout would only mark the task failed while the query
-# kept running. Restore a timeout once query cancellation is implemented.
+@contextmanager
+def _capture_query_cancellation(query_context: "QueryContext") -> Iterator[None]:
+ """Enable engine-level cancellation of this task's warehouse query.
+
+ Captures the engine cancel id off the live cursor (for engines that expose
+ one before execution) and registers an abort handler that kills the backend
+ session over a fresh connection, unblocking the task's ``get_df``. Engines
+ without cancel support capture nothing and the task stays non-abortable, so
+ an abort/timeout simply frees the task without killing the (uncancellable)
+ query — matching the pre-cancellation behavior for those engines.
+ """
+ database = getattr(query_context.datasource, "database", None)
+ if database is None:
+ yield
+ return
+
+ ctx = get_context()
+ app = current_app._get_current_object() # noqa: SLF001
+ captured = False
+
+ def _sink(cursor: Any) -> None:
+ nonlocal captured
+ if captured:
+ return
+ # query is unused by the explicit-id specs; the chart path has no Query.
+ cancel_id = database.db_engine_spec.get_cancel_query_id(cursor, None)
+ if cancel_id is None:
+ return
+ captured = True
+
+ # Registering the first abort handler marks the task abortable and starts
+ # the abort listener; on abort it cancels the query on a fresh connection.
+ def _cancel() -> None:
+ cancel_chart_query(database, cancel_id, app)
+
+ ctx.on_abort(_cancel)
+
+ with capture_cancel_id(_sink):
+ yield
+
+
@task(name=CHART_QUERY_TASK, scope=TaskScope.SHARED)
def execute_chart_query(
serialized_query: SerializedQuery,
@@ -136,7 +176,8 @@ def execute_chart_query(
The atomic async unit: reconstruct the one query (canonical serialization),
optionally inject contribution totals from a prerequisite totals task, then run
the existing per-query execution/caching path so a re-request reads the same
- DATA-cache entry.
+ DATA-cache entry. The query runs under ``_capture_query_cancellation`` so an
+ abort/timeout can cancel it on engines that support query cancellation.
"""
with override_user(_resolve_user(user_id, guest_token), force=False):
query_context = load_serialized_query(serialized_query)
@@ -147,7 +188,8 @@ def execute_chart_query(
if requires_totals:
_inject_contribution_totals(query_obj, _get_dependency_cache_key())
# Executes on cache miss and writes CacheRegion.DATA under query_cache_key.
- result = query_context.get_df_payload_result(query_obj)
+ with _capture_query_cancellation(query_context):
+ result = query_context.get_df_payload_result(query_obj)
if cache_key := result.payload.get(CACHE_KEY_PAYLOAD_KEY):
# Write synchronously: a dependent contribution query reads this
# cache key via get_dependency_payloads once the DAG gate releases,
@@ -253,6 +295,10 @@ def submit_chart_data_query_tasks(
task_key=query_cache_keys[index],
task_name=_task_name(index),
depends_on=depends_on,
+ # Abort the query after this many seconds; for cancellable
+ # engines the abort handler kills the warehouse query too.
+ # None (default) leaves it unbounded.
+ timeout=current_app.config.get("GLOBAL_ASYNC_QUERIES_QUERY_TIMEOUT"),
),
)
diff --git a/superset/tasks/heartbeat.py b/superset/tasks/heartbeat.py
new file mode 100644
index 00000000000..770a01ef127
--- /dev/null
+++ b/superset/tasks/heartbeat.py
@@ -0,0 +1,84 @@
+# 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.
+"""Worker liveness heartbeat for the Global Task Framework (GTF)."""
+
+from __future__ import annotations
+
+import logging
+import threading
+from contextlib import contextmanager
+from typing import Iterator, TYPE_CHECKING
+
+from superset.stats_logger import BaseStatsLogger
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def task_heartbeat(task_id: int, app: "Flask") -> Iterator[None]:
+ """Keep ``tasks.last_heartbeat`` fresh while a worker holds the task.
+
+ Stamps an initial heartbeat synchronously (so the task reads as "picked up
+ by a worker" immediately — a NULL heartbeat means still queued), then runs a
+ daemon thread that bumps it every ``GTF_TASK_HEARTBEAT_INTERVAL`` seconds
+ until the context exits. The prune cron uses this to tell a live task (fresh
+ heartbeat) from one abandoned by a dead worker (stale heartbeat).
+
+ Best-effort: a failed write is counted and logged, and the loop continues —
+ the reaper only acts after several missed intervals. The thread is a daemon
+ so it never blocks worker shutdown (matching the existing timeout/abort
+ threads) and relies on DB drivers releasing the GIL during query I/O, so a
+ task blocked in a long warehouse query still heartbeats and is not reaped.
+
+ :param task_id: integer primary key of the running task
+ :param app: Flask app for DB access from the background thread
+ """
+ interval: float = app.config["GTF_TASK_HEARTBEAT_INTERVAL"]
+ stats_logger: BaseStatsLogger = app.config.get("STATS_LOGGER", BaseStatsLogger())
+ stop = threading.Event()
+
+ def _write() -> None:
+ from superset.daos.tasks import TaskDAO
+
+ try:
+ TaskDAO.touch_heartbeat(task_id)
+ stats_logger.incr("gtf.task.heartbeat")
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ stats_logger.incr("gtf.task.heartbeat_failure")
+ logger.warning(
+ "Heartbeat write failed for task id=%s", task_id, exc_info=True
+ )
+
+ # Initial stamp runs in the worker's existing app context.
+ _write()
+
+ def _beat() -> None:
+ while not stop.wait(interval):
+ with app.app_context():
+ _write()
+
+ thread = threading.Thread(
+ target=_beat, name=f"gtf-heartbeat-{task_id}", daemon=True
+ )
+ thread.start()
+ try:
+ yield
+ finally:
+ stop.set()
diff --git a/superset/tasks/query_cancel.py b/superset/tasks/query_cancel.py
new file mode 100644
index 00000000000..9e001675bbe
--- /dev/null
+++ b/superset/tasks/query_cancel.py
@@ -0,0 +1,131 @@
+# 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.
+"""Engine-level cancellation for GTF chart-data query tasks.
+
+The chart-data execution path (``Database.get_df``) has none of SQL Lab's
+cancel plumbing, so a GTF abort/timeout could only mark the task terminal while
+the warehouse query kept running. This module adds the missing seam:
+
+- ``capture_cancel_id`` registers a cursor sink for the duration of a task's
+ query; ``notify_cursor`` (called from ``get_df`` before the blocking execute)
+ hands the live cursor to that sink so the task can capture an engine cancel id
+ via ``db_engine_spec.get_cancel_query_id`` — the same contract SQL Lab uses.
+- ``cancel_chart_query`` kills the backend session over a *fresh* connection
+ (``db_engine_spec.cancel_query``), which unblocks the task's blocked ``get_df``.
+
+Only engines that return a cancel id before execution participate; others
+capture nothing and are simply not cancellable (the abort still frees the task).
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import closing, contextmanager
+from contextvars import ContextVar
+from typing import Any, Callable, Iterator, TYPE_CHECKING
+
+from superset.stats_logger import BaseStatsLogger
+
+if TYPE_CHECKING:
+ from flask import Flask
+
+ from superset.models.core import Database
+
+logger = logging.getLogger(__name__)
+
+# Set by a chart-data task for the span of its query. When present,
+# Database._execute_sql_with_mutation_and_logging hands the sink the live cursor
+# before the blocking execute so the task can capture an engine cancel id.
+# Absent (None) for every other get_df caller, so this is a no-op elsewhere.
+_cancel_id_sink: ContextVar["Callable[[Any], None] | None"] = ContextVar(
+ "gtf_cancel_id_sink", default=None
+)
+
+
+@contextmanager
+def capture_cancel_id(sink: "Callable[[Any], None]") -> Iterator[None]:
+ """Register a cursor sink for the duration of a chart-data query execution."""
+ token = _cancel_id_sink.set(sink)
+ try:
+ yield
+ finally:
+ _cancel_id_sink.reset(token)
+
+
+def notify_cursor(cursor: Any) -> None:
+ """Hand the live cursor to the active sink, if any.
+
+ Called from ``get_df`` before the query executes. Best-effort: a capture
+ failure must never break query execution — it only forfeits cancellability.
+ """
+ sink = _cancel_id_sink.get()
+ if sink is None:
+ return
+ try:
+ sink(cursor)
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ logger.warning("Cancel-id capture failed", exc_info=True)
+
+
+def cancel_chart_query(
+ database: "Database", cancel_query_id: str, app: "Flask | None" = None
+) -> bool:
+ """Cancel a running chart-data warehouse query over a fresh connection.
+
+ Runs ``db_engine_spec.cancel_query`` against a new connection to the same
+ database, terminating the backend session that the task's blocked ``get_df``
+ is waiting on. Invoked from the task's abort handler (on the abort-listener
+ thread), so it opens its own connection rather than touching the busy one.
+ Best-effort and fully logged; the task's terminal transition is authoritative.
+
+ :param database: the database the query is running against
+ :param cancel_query_id: engine cancel handle captured at query start
+ :param app: Flask app for config/DB access from the background thread
+ :returns: True if the engine reported the query cancelled
+ """
+ from flask import current_app
+
+ stats_logger: BaseStatsLogger = (app or current_app).config.get(
+ "STATS_LOGGER", BaseStatsLogger()
+ )
+ spec = database.db_engine_spec
+ try:
+ with database.get_sqla_engine() as engine:
+ with closing(engine.raw_connection()) as conn:
+ with closing(conn.cursor()) as cursor:
+ # query is unused by the explicit-id specs (they cancel by
+ # the captured id alone); the chart path has no Query model.
+ cancelled = spec.cancel_query(cursor, None, cancel_query_id) # type: ignore[arg-type]
+ if cancelled:
+ stats_logger.incr("gtf.query.cancel")
+ logger.info(
+ "Cancelled warehouse query on database %s (id=%s)",
+ database.id,
+ cancel_query_id,
+ )
+ else:
+ stats_logger.incr("gtf.query.cancel_failed")
+ return cancelled
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ stats_logger.incr("gtf.query.cancel_failed")
+ logger.warning(
+ "Failed to cancel warehouse query on database %s (id=%s)",
+ database.id,
+ cancel_query_id,
+ exc_info=True,
+ )
+ return False
diff --git a/superset/tasks/scheduler.py b/superset/tasks/scheduler.py
index 5d016068a91..dea01adbfb3 100644
--- a/superset/tasks/scheduler.py
+++ b/superset/tasks/scheduler.py
@@ -35,9 +35,10 @@ from superset.commands.report.execute import AsyncExecuteReportScheduleCommand
from superset.commands.report.log_prune import AsyncPruneReportScheduleLogCommand
from superset.commands.sql_lab.query import QueryPruneCommand
from superset.commands.tasks.prune import TaskPruneCommand
+from superset.commands.tasks.reap import ReapOrphanedTasksCommand
from superset.daos.report import ReportScheduleDAO
from superset.daos.tasks import TaskDAO
-from superset.extensions import celery_app
+from superset.extensions import celery_app, db
from superset.key_value.commands.prune import KeyValuePruneCommand
from superset.reports.models import ReportScheduleType
from superset.stats_logger import BaseStatsLogger
@@ -45,6 +46,7 @@ from superset.tasks.ambient_context import use_context
from superset.tasks.constants import ABORT_STATES, TERMINAL_STATES
from superset.tasks.context import TaskContext
from superset.tasks.cron_util import cron_schedule_window
+from superset.tasks.heartbeat import task_heartbeat
from superset.tasks.manager import TaskManager
from superset.tasks.registry import TaskRegistry
from superset.utils.core import LoggerLevel
@@ -259,6 +261,13 @@ def prune_tasks(
"retention period, please use `kwargs` instead."
)
+ # Reap first: revoke + fail tasks abandoned by a dead/wedged worker so they
+ # reach a terminal state, then let the retention pass below delete old rows.
+ try:
+ ReapOrphanedTasksCommand().run()
+ except CommandException as ex:
+ logger.exception("An error occurred while reaping orphaned tasks: %s", ex)
+
try:
TaskPruneCommand(retention_period_days, max_rows_per_run).run()
except CommandException as ex:
@@ -324,8 +333,25 @@ def _resolve_failed_prerequisite(task: "TaskModel") -> "TaskModel | None":
return None
+def _persist_celery_task_id(task: "TaskModel", celery_task_id: str | None) -> None:
+ """Record the Celery job id on the task so the reaper can revoke it.
+
+ Mutates the in-memory task (so the ``TaskContext`` built later carries the id
+ through its property writes) and commits. Written once at pickup, before the
+ DAG gate and the status transition, so the reaper can revoke even a task
+ still waiting on prerequisites. The id is dropped on the terminal transition
+ along with the rest of ``properties``, which is fine — only ACTIVE tasks are
+ ever reaped.
+ """
+ if not celery_task_id:
+ return
+ task.update_properties({"celery_task_id": celery_task_id})
+ # One-shot write at pickup, outside the lifecycle transaction below.
+ db.session.commit() # pylint: disable=consider-using-transaction
+
+
@celery_app.task(name="tasks.execute", bind=True)
-def execute_task( # noqa: C901
+def execute_task(
self: Any, # Celery task instance
task_uuid: str,
task_type: str,
@@ -335,17 +361,10 @@ def execute_task( # noqa: C901
"""
Generic task executor for GTF tasks.
- This executor:
- 1. Checks if task was aborted before execution starts
- 2. Fetches task from metastore
- 3. Builds context (task + user) and sets ambient context via contextvars
- 4. Executes the task function (which accesses context via get_context())
- 5. Updates task status throughout lifecycle using atomic conditional updates
- 6. Runs cleanup handlers on task end (success/failure/abortion)
- 7. Resets context after execution
-
- Uses atomic conditional status updates to prevent race conditions with
- concurrent abort operations.
+ Loads the task, records the Celery job id, and runs the lifecycle body under
+ a liveness heartbeat that spans the whole time this worker holds the task
+ (DAG wait, IN_PROGRESS, ABORTING) so the prune cron can revoke and reap it if
+ this worker dies. See ``_execute_task_body`` for the lifecycle itself.
:param task_uuid: UUID of the task to execute
:param task_type: Type of the task (for registry lookup)
@@ -353,8 +372,6 @@ def execute_task( # noqa: C901
:param kwargs: Keyword arguments for the task function
:returns: Dict with status and task_uuid
"""
- from superset.commands.tasks.internal_update import InternalStatusTransitionCommand
-
# Convert string UUID to native UUID (Celery deserializes as string)
native_uuid = UUID(task_uuid)
@@ -367,6 +384,37 @@ def execute_task( # noqa: C901
logger.error("Task %s not found in metastore", task_uuid)
return {"status": "error", "message": "Task not found"}
+ _persist_celery_task_id(task, self.request.id)
+ app = current_app._get_current_object() # noqa: SLF001
+ with task_heartbeat(task.id, app):
+ return _execute_task_body(task, native_uuid, task_type, args, kwargs)
+
+
+def _execute_task_body( # noqa: C901
+ task: "TaskModel",
+ native_uuid: UUID,
+ task_type: str,
+ args: tuple[Any, ...],
+ kwargs: dict[str, Any],
+) -> dict[str, Any]:
+ """
+ Run a claimed GTF task through its lifecycle.
+
+ This body:
+ 1. Checks if task was aborted before execution starts
+ 2. Builds context (task + user) and sets ambient context via contextvars
+ 3. Executes the task function (which accesses context via get_context())
+ 4. Updates task status throughout lifecycle using atomic conditional updates
+ 5. Runs cleanup handlers on task end (success/failure/abortion)
+ 6. Resets context after execution
+
+ Uses atomic conditional status updates to prevent race conditions with
+ concurrent abort operations.
+ """
+ from superset.commands.tasks.internal_update import InternalStatusTransitionCommand
+
+ task_uuid = str(native_uuid)
+
# AUTOMATIC PRE-EXECUTION CHECK: Don't execute if already aborted/aborting
if task.status in ABORT_STATES:
logger.info(
@@ -509,26 +557,38 @@ def execute_task( # noqa: C901
# Mark execution as completed to prevent late abort handlers
ctx.mark_execution_completed()
- # Atomic transition to FAILURE (only if still IN_PROGRESS or ABORTING)
- InternalStatusTransitionCommand(
- task_uuid=native_uuid,
- new_status=TaskStatus.FAILURE,
- expected_status=[TaskStatus.IN_PROGRESS, TaskStatus.ABORTING],
- properties={"error_message": str(ex)},
- set_ended_at=True,
- ).run()
-
- logger.error(
- "Task %s (uuid=%s) failed with error: %s",
- task_type,
- task_uuid,
- str(ex),
- exc_info=True,
- )
-
- # Emit stats metric for failure
- stats_logger = current_app.config["STATS_LOGGER"]
- stats_logger.incr("gtf.task.failure")
+ # An abort/timeout that actually cancels the work (e.g. an abort handler
+ # killing the underlying warehouse query) surfaces here as an exception.
+ # That is a successful abort, not a failure: leave the task in ABORTING
+ # so the finally block finalizes it as ABORTED/TIMED_OUT. Only a genuine
+ # error (no abort in flight) transitions to FAILURE.
+ if ctx._abort_detected or ctx.timeout_triggered: # noqa: SLF001
+ logger.info(
+ "Task %s (uuid=%s) raised while aborting; finalizing as aborted",
+ task_type,
+ task_uuid,
+ )
+ else:
+ # Atomic transition to FAILURE (only if still IN_PROGRESS or ABORTING)
+ InternalStatusTransitionCommand(
+ task_uuid=native_uuid,
+ new_status=TaskStatus.FAILURE,
+ expected_status=[TaskStatus.IN_PROGRESS, TaskStatus.ABORTING],
+ properties={"error_message": str(ex)},
+ set_ended_at=True,
+ ).run()
+
+ logger.error(
+ "Task %s (uuid=%s) failed with error: %s",
+ task_type,
+ task_uuid,
+ str(ex),
+ exc_info=True,
+ )
+
+ # Emit stats metric for failure
+ stats_logger = current_app.config["STATS_LOGGER"]
+ stats_logger.incr("gtf.task.failure")
finally:
# ALWAYS run cleanup handlers (also stops timeout timer)
diff --git a/tests/integration_tests/tasks/commands/test_reap.py b/tests/integration_tests/tasks/commands/test_reap.py
new file mode 100644
index 00000000000..db57f176161
--- /dev/null
+++ b/tests/integration_tests/tasks/commands/test_reap.py
@@ -0,0 +1,144 @@
+# 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.
+
+from datetime import datetime, timedelta
+from unittest.mock import patch
+
+import sqlalchemy as sa
+from superset_core.tasks.types import TaskScope, TaskStatus
+
+from superset import db
+from superset.commands.tasks.reap import ORPHAN_ERROR_MESSAGE, ReapOrphanedTasksCommand
+from superset.daos.tasks import TaskDAO
+from superset.models.tasks import Task
+from superset.tasks.utils import naive_utcnow
+
+
+def _make_task(admin, key, status, *, heartbeat_offset=None, celery_task_id=None):
+ """Create a task in a given state with an optional heartbeat age (seconds)."""
+ task = TaskDAO.create_task(
+ task_type="test_type",
+ task_key=key,
+ scope=TaskScope.PRIVATE,
+ user_id=admin.id,
+ )
+ task.created_by = admin
+ task.set_status(status)
+ if heartbeat_offset is not None:
+ task.last_heartbeat = naive_utcnow() - timedelta(seconds=heartbeat_offset)
+ if celery_task_id is not None:
+ task.update_properties({"celery_task_id": celery_task_id})
+ db.session.commit()
+ return task
+
+
+def _cleanup(*tasks):
+ for task in tasks:
+ existing = db.session.get(Task, task.id)
+ if existing:
+ db.session.delete(existing)
+ db.session.commit()
+
+
+def test_touch_heartbeat_does_not_advance_changed_on(
+ app_context, get_user, login_as
+) -> None:
+ """A heartbeat write must not bump changed_on (would reset polling backoff)."""
+ login_as("admin")
+ admin = get_user("admin")
+ task = _make_task(admin, "hb_task", TaskStatus.IN_PROGRESS)
+ original_changed_on = task.changed_on
+ try:
+ TaskDAO.touch_heartbeat(task.id)
+ db.session.refresh(task)
+ assert task.last_heartbeat is not None
+ assert task.changed_on == original_changed_on
+ finally:
+ _cleanup(task)
+
+
+def test_find_orphaned_selects_only_stale_active_tasks(
+ app_context, get_user, login_as
+) -> None:
+ """Orphans = ACTIVE tasks with a stale, non-null heartbeat; nothing else."""
+ login_as("admin")
+ admin = get_user("admin")
+
+ stale = _make_task(admin, "stale", TaskStatus.IN_PROGRESS, heartbeat_offset=600)
+ fresh = _make_task(admin, "fresh", TaskStatus.IN_PROGRESS, heartbeat_offset=0)
+ never_started = _make_task(admin, "queued", TaskStatus.PENDING) # null heartbeat
+ finished = _make_task(admin, "done", TaskStatus.SUCCESS, heartbeat_offset=600)
+ try:
+ found = {c.uuid for c in TaskDAO.find_orphaned(60, 60)}
+ assert stale.uuid in found
+ assert fresh.uuid not in found
+ assert never_started.uuid not in found
+ assert finished.uuid not in found
+ finally:
+ _cleanup(stale, fresh, never_started, finished)
+
+
+def test_find_orphaned_selects_wedged_aborting_tasks(
+ app_context, get_user, login_as
+) -> None:
+ """Wedged aborts = ABORTING + fresh heartbeat + old changed_on (live worker)."""
+ login_as("admin")
+ admin = get_user("admin")
+
+ # ABORTING with a fresh heartbeat; backdate changed_on past the grace window.
+ # changed_on is naive *local* (FAB onupdate), so use datetime.now() here.
+ wedged = _make_task(admin, "wedged", TaskStatus.ABORTING, heartbeat_offset=0)
+ db.session.execute(
+ sa.text("UPDATE tasks SET changed_on = :c WHERE id = :id"),
+ {"c": datetime.now() - timedelta(seconds=600), "id": wedged.id},
+ )
+ db.session.commit()
+ # A task that just entered ABORTING (recent changed_on) must NOT be selected.
+ recent = _make_task(admin, "recent_abort", TaskStatus.ABORTING, heartbeat_offset=0)
+ try:
+ found = {c.uuid: c.is_orphan for c in TaskDAO.find_orphaned(60, 60)}
+ assert found.get(wedged.uuid) is False # selected, flagged wedged (not orphan)
+ assert recent.uuid not in found
+ finally:
+ _cleanup(wedged, recent)
+
+
+def test_reap_marks_orphan_failed_and_revokes(app_context, get_user, login_as) -> None:
+ """End-to-end: the reaper forces a stale in-progress task to FAILURE."""
+ login_as("admin")
+ admin = get_user("admin")
+ orphan = _make_task(
+ admin,
+ "orphan",
+ TaskStatus.IN_PROGRESS,
+ heartbeat_offset=600,
+ celery_task_id="celery-xyz",
+ )
+ try:
+ with patch("superset.commands.tasks.reap.celery_app") as celery:
+ reaped = ReapOrphanedTasksCommand().run()
+
+ assert reaped == 1
+ celery.control.revoke.assert_called_once_with(
+ "celery-xyz", terminate=True, signal="SIGUSR1"
+ )
+ db.session.refresh(orphan)
+ assert orphan.status == TaskStatus.FAILURE.value
+ assert orphan.ended_at is not None
+ assert orphan.properties_dict["error_message"] == ORPHAN_ERROR_MESSAGE
+ finally:
+ _cleanup(orphan)
diff --git a/tests/unit_tests/tasks/test_dependencies.py b/tests/unit_tests/tasks/test_dependencies.py
index 70d9cfb5998..a556084668a 100644
--- a/tests/unit_tests/tasks/test_dependencies.py
+++ b/tests/unit_tests/tasks/test_dependencies.py
@@ -103,7 +103,7 @@ class TestExecuteTaskFailedPrerequisite:
@staticmethod
def _run_with_transition(*, transitioned: bool, refreshed_status: str):
- """Drive execute_task through the failed-prerequisite branch.
+ """Drive the task body through the failed-prerequisite branch.
Returns ``(result, publish_completion_mock)`` with the status transition
forced to ``transitioned`` and the post-transition re-read reporting
@@ -120,7 +120,7 @@ class TestExecuteTaskFailedPrerequisite:
with (
patch(
"superset.tasks.scheduler.TaskDAO.find_one_or_none",
- side_effect=[pending, refreshed],
+ side_effect=[refreshed],
),
patch(
"superset.tasks.scheduler._resolve_failed_prerequisite",
@@ -133,9 +133,11 @@ class TestExecuteTaskFailedPrerequisite:
transition,
),
):
- from superset.tasks.scheduler import execute_task
+ from superset.tasks.scheduler import _execute_task_body
- result = execute_task.run(str(native), "some.task", (), {})
+ # Call the lifecycle body directly with the already-loaded task; the
+ # heartbeat/celery-id wrapper (execute_task) is covered separately.
+ result = _execute_task_body(pending, native, "some.task", (), {}) # type: ignore[arg-type]
return result, task_manager.publish_completion
def test_failure_published_when_transition_commits(self):
diff --git a/tests/unit_tests/tasks/test_heartbeat.py b/tests/unit_tests/tasks/test_heartbeat.py
new file mode 100644
index 00000000000..a1b39d56702
--- /dev/null
+++ b/tests/unit_tests/tasks/test_heartbeat.py
@@ -0,0 +1,76 @@
+# 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.
+"""Unit tests for the GTF worker-liveness heartbeat."""
+
+import time
+from unittest.mock import MagicMock, patch
+
+from superset.tasks.heartbeat import task_heartbeat
+
+
+def _mock_app(interval: float, stats: MagicMock) -> MagicMock:
+ app = MagicMock()
+ app.config = {"GTF_TASK_HEARTBEAT_INTERVAL": interval, "STATS_LOGGER": stats}
+ app.app_context.return_value.__enter__ = MagicMock(return_value=None)
+ app.app_context.return_value.__exit__ = MagicMock(return_value=None)
+ return app
+
+
+def test_heartbeat_stamps_immediately_and_stops_on_exit() -> None:
+ """An initial heartbeat is written synchronously; none after the context ends."""
+ stats = MagicMock()
+ # Large interval so the background thread never beats within the test window;
+ # only the synchronous initial write should occur.
+ app = _mock_app(interval=100, stats=stats)
+
+ with patch("superset.daos.tasks.TaskDAO.touch_heartbeat") as touch:
+ with task_heartbeat(42, app):
+ assert touch.call_count == 1
+ touch.assert_called_once_with(42)
+ stats.incr.assert_any_call("gtf.task.heartbeat")
+
+ # No further writes after the context exits.
+ assert touch.call_count == 1
+
+
+def test_heartbeat_beats_periodically() -> None:
+ """The background thread keeps bumping the heartbeat on the interval."""
+ stats = MagicMock()
+ app = _mock_app(interval=0.02, stats=stats)
+
+ with patch("superset.daos.tasks.TaskDAO.touch_heartbeat") as touch:
+ with task_heartbeat(7, app):
+ deadline = time.time() + 2.0
+ while touch.call_count < 3 and time.time() < deadline:
+ time.sleep(0.02)
+ assert touch.call_count >= 3
+
+
+def test_heartbeat_failure_is_counted_and_swallowed() -> None:
+ """A failed write emits the failure metric and does not raise."""
+ stats = MagicMock()
+ app = _mock_app(interval=100, stats=stats)
+
+ with patch(
+ "superset.daos.tasks.TaskDAO.touch_heartbeat",
+ side_effect=RuntimeError("db down"),
+ ):
+ # Must not raise even though the initial write fails.
+ with task_heartbeat(1, app):
+ pass
+
+ stats.incr.assert_any_call("gtf.task.heartbeat_failure")
diff --git a/tests/unit_tests/tasks/test_query_cancel.py b/tests/unit_tests/tasks/test_query_cancel.py
new file mode 100644
index 00000000000..370fbc903b0
--- /dev/null
+++ b/tests/unit_tests/tasks/test_query_cancel.py
@@ -0,0 +1,154 @@
+# 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.
+"""Unit tests for GTF chart-data query cancellation (Phase 2)."""
+
+from unittest.mock import MagicMock, patch
+
+from superset.tasks.query_cancel import (
+ _cancel_id_sink,
+ cancel_chart_query,
+ capture_cancel_id,
+ notify_cursor,
+)
+
+
+def _mock_database(cancel_result: bool = True):
+ database = MagicMock()
+ database.id = 5
+ cursor = MagicMock()
+ conn = MagicMock()
+ conn.cursor.return_value = cursor
+ engine = MagicMock()
+ engine.raw_connection.return_value = conn
+ database.get_sqla_engine.return_value.__enter__ = MagicMock(return_value=engine)
+ database.get_sqla_engine.return_value.__exit__ = MagicMock(return_value=None)
+ database.db_engine_spec.cancel_query.return_value = cancel_result
+ return database, cursor
+
+
+def _app_with_stats(stats: MagicMock) -> MagicMock:
+ app = MagicMock()
+ app.config = {"STATS_LOGGER": stats}
+ return app
+
+
+# --- capture seam ---------------------------------------------------------
+
+
+def test_notify_cursor_is_noop_without_a_sink() -> None:
+ # No sink registered: must not raise and must not touch the cursor.
+ assert _cancel_id_sink.get() is None
+ notify_cursor(MagicMock())
+
+
+def test_capture_cancel_id_sets_and_resets_the_sink() -> None:
+ sink = MagicMock()
+ with capture_cancel_id(sink):
+ cursor = MagicMock()
+ notify_cursor(cursor)
+ sink.assert_called_once_with(cursor)
+ assert _cancel_id_sink.get() is None
+
+
+def test_notify_cursor_swallows_sink_errors() -> None:
+ with capture_cancel_id(MagicMock(side_effect=RuntimeError("boom"))):
+ # A capture failure only forfeits cancellability; it must not raise.
+ notify_cursor(MagicMock())
+
+
+# --- cancel helper --------------------------------------------------------
+
+
+def test_cancel_chart_query_success() -> None:
+ stats = MagicMock()
+ database, cursor = _mock_database(cancel_result=True)
+
+ assert cancel_chart_query(database, "123", _app_with_stats(stats)) is True
+ database.db_engine_spec.cancel_query.assert_called_once_with(cursor, None, "123")
+ stats.incr.assert_any_call("gtf.query.cancel")
+
+
+def test_cancel_chart_query_reports_failure_when_engine_declines() -> None:
+ stats = MagicMock()
+ database, _ = _mock_database(cancel_result=False)
+
+ assert cancel_chart_query(database, "123", _app_with_stats(stats)) is False
+ stats.incr.assert_any_call("gtf.query.cancel_failed")
+
+
+def test_cancel_chart_query_swallows_exceptions() -> None:
+ stats = MagicMock()
+ database, _ = _mock_database()
+ database.get_sqla_engine.side_effect = RuntimeError("no connection")
+
+ assert cancel_chart_query(database, "123", _app_with_stats(stats)) is False
+ stats.incr.assert_any_call("gtf.query.cancel_failed")
+
+
+# --- capture wiring in execute_chart_query --------------------------------
+
+
+def _query_context(cancel_id):
+ qc = MagicMock()
+ qc.datasource.database.db_engine_spec.get_cancel_query_id.return_value = cancel_id
+ return qc
+
+
+def test_capture_registers_abort_handler_once_when_id_available() -> None:
+ from superset.tasks.async_queries import _capture_query_cancellation
+
+ ctx = MagicMock()
+ qc = _query_context("42")
+ with (
+ patch("superset.tasks.async_queries.get_context", return_value=ctx),
+ patch("superset.tasks.async_queries.current_app"),
+ ):
+ with _capture_query_cancellation(qc):
+ sink = _cancel_id_sink.get()
+ assert sink is not None
+ sink(MagicMock()) # first cursor -> registers handler
+ sink(MagicMock()) # second cursor -> no double registration
+
+ ctx.on_abort.assert_called_once()
+
+
+def test_capture_skips_abort_handler_when_engine_has_no_cancel_id() -> None:
+ from superset.tasks.async_queries import _capture_query_cancellation
+
+ ctx = MagicMock()
+ qc = _query_context(None)
+ with (
+ patch("superset.tasks.async_queries.get_context", return_value=ctx),
+ patch("superset.tasks.async_queries.current_app"),
+ ):
+ with _capture_query_cancellation(qc):
+ sink = _cancel_id_sink.get()
+ assert sink is not None
+ sink(MagicMock())
+
+ ctx.on_abort.assert_not_called()
+
+
+def test_capture_is_noop_without_a_database() -> None:
+ from superset.tasks.async_queries import _capture_query_cancellation
+
+ qc = MagicMock()
+ qc.datasource.database = None
+ with patch("superset.tasks.async_queries.get_context") as get_context:
+ with _capture_query_cancellation(qc):
+ assert _cancel_id_sink.get() is None
+ get_context.assert_not_called()
diff --git a/tests/unit_tests/tasks/test_reap.py b/tests/unit_tests/tasks/test_reap.py
new file mode 100644
index 00000000000..5c3609f7592
--- /dev/null
+++ b/tests/unit_tests/tasks/test_reap.py
@@ -0,0 +1,124 @@
+# 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.
+"""Unit tests for ReapOrphanedTasksCommand."""
+
+from contextlib import contextmanager
+from unittest.mock import MagicMock, patch
+from uuid import UUID
+
+from superset_core.tasks.types import TaskStatus
+
+from superset.commands.tasks.reap import ORPHAN_ERROR_MESSAGE, ReapOrphanedTasksCommand
+from superset.daos.tasks import OrphanCandidate
+from superset.utils import json
+
+TEST_UUID = UUID("b8b61b7b-1cd3-4a31-a74a-0a95341afc06")
+
+
+def _candidate(
+ *, is_orphan: bool, celery_task_id: str | None = "celery-1"
+) -> OrphanCandidate:
+ properties = json.dumps(
+ {"celery_task_id": celery_task_id} if celery_task_id else {}
+ )
+ return OrphanCandidate(
+ id=1,
+ uuid=TEST_UUID,
+ status=TaskStatus.IN_PROGRESS.value,
+ properties=properties,
+ is_orphan=is_orphan,
+ )
+
+
+@contextmanager
+def _patched(candidates, cas_result=True):
+ """Patch the reaper's collaborators and yield the mocks it interacts with."""
+ stats = MagicMock()
+ app = MagicMock()
+ app.config = {"STATS_LOGGER": stats, "GTF_ORPHAN_TASK_TIMEOUT": 60}
+ with (
+ patch("superset.commands.tasks.reap.current_app", app),
+ patch("superset.commands.tasks.reap.db"),
+ patch("superset.daos.tasks.TaskDAO.find_orphaned", return_value=candidates),
+ patch(
+ "superset.daos.tasks.TaskDAO.conditional_status_update",
+ return_value=cas_result,
+ ) as cas,
+ patch("superset.commands.tasks.reap.celery_app") as celery,
+ patch("superset.tasks.manager.TaskManager.publish_completion") as publish,
+ patch("superset.tasks.manager.TaskManager.publish_entity_change"),
+ ):
+ yield stats, cas, celery, publish
+
+
+def test_orphan_is_revoked_failed_and_published() -> None:
+ """A dead-worker orphan is revoked, forced to FAILURE, and completion published."""
+ with _patched([_candidate(is_orphan=True)]) as (stats, cas, celery, publish):
+ reaped = ReapOrphanedTasksCommand().run()
+
+ assert reaped == 1
+ celery.control.revoke.assert_called_once_with(
+ "celery-1", terminate=True, signal="SIGUSR1"
+ )
+ # Forced to FAILURE, preserving the orphan error in properties.
+ _args, kwargs = cas.call_args
+ assert kwargs["properties"]["error_message"] == ORPHAN_ERROR_MESSAGE
+ publish.assert_called_once_with(TEST_UUID, TaskStatus.FAILURE.value)
+ stats.incr.assert_any_call("gtf.task.orphan_reaped")
+ stats.incr.assert_any_call("gtf.task.revoke")
+
+
+def test_wedged_abort_is_only_escalated() -> None:
+ """A live wedged-ABORTING worker is force-revoked but not marked itself."""
+ with _patched([_candidate(is_orphan=False)]) as (stats, cas, celery, publish):
+ reaped = ReapOrphanedTasksCommand().run()
+
+ assert reaped == 0
+ celery.control.revoke.assert_called_once()
+ cas.assert_not_called() # the worker finalizes its own status
+ publish.assert_not_called()
+ stats.incr.assert_any_call("gtf.task.abort_escalated")
+
+
+def test_cas_noop_does_not_publish() -> None:
+ """A revived worker that already committed a terminal status wins the CAS race."""
+ with _patched([_candidate(is_orphan=True)], cas_result=False) as (
+ _stats,
+ _cas,
+ celery,
+ publish,
+ ):
+ reaped = ReapOrphanedTasksCommand().run()
+
+ assert reaped == 0
+ celery.control.revoke.assert_called_once()
+ publish.assert_not_called()
+
+
+def test_missing_celery_id_skips_revoke_but_still_reaps() -> None:
+ """A task with no stored Celery id is still reaped; revoke is simply skipped."""
+ with _patched([_candidate(is_orphan=True, celery_task_id=None)]) as (
+ _stats,
+ _cas,
+ celery,
+ publish,
+ ):
+ reaped = ReapOrphanedTasksCommand().run()
+
+ assert reaped == 1
+ celery.control.revoke.assert_not_called()
+ publish.assert_called_once_with(TEST_UUID, TaskStatus.FAILURE.value)
diff --git a/tests/unit_tests/tasks/test_scheduler_executor.py b/tests/unit_tests/tasks/test_scheduler_executor.py
new file mode 100644
index 00000000000..747e156a738
--- /dev/null
+++ b/tests/unit_tests/tasks/test_scheduler_executor.py
@@ -0,0 +1,137 @@
+# 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.
+"""Unit tests for the GTF executor body (scheduler._execute_task_body)."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+from uuid import uuid4
+
+from superset_core.tasks.types import TaskStatus
+
+
+def _run_body_with_raising_executor(*, timeout_triggered: bool, abort_detected: bool):
+ """Drive _execute_task_body where the task body raises mid-abort/timeout.
+
+ Returns ``(result, transition_mock, task_manager_mock, stats_logger_mock)`` for
+ the caller to assert on the terminal handling.
+ """
+ native = uuid4()
+ task = SimpleNamespace(
+ uuid=native,
+ status=TaskStatus.PENDING.value,
+ properties_dict={}, # no "timeout" key -> no timeout timer started
+ )
+ # What the finally block's terminal transition commits, read back at the end.
+ refreshed = SimpleNamespace(
+ status=(
+ TaskStatus.TIMED_OUT.value
+ if timeout_triggered
+ else TaskStatus.ABORTED.value
+ )
+ )
+
+ ctx = MagicMock()
+ ctx._abort_detected = abort_detected
+ ctx.timeout_triggered = timeout_triggered
+ ctx.abort_handlers_completed = True
+
+ transition = MagicMock()
+ transition.return_value.run.return_value = True # PENDING->IN_PROGRESS succeeds
+
+ def _raise(*_args, **_kwargs):
+ # Simulates the abort handler cancelling the warehouse query, which makes
+ # the blocked get_df raise inside the task body.
+ raise RuntimeError("query cancelled")
+
+ stats_logger = MagicMock()
+ app = MagicMock()
+ app.config = {"STATS_LOGGER": stats_logger}
+
+ with (
+ patch(
+ "superset.tasks.scheduler.TaskDAO.find_one_or_none", return_value=refreshed
+ ),
+ patch(
+ "superset.tasks.scheduler._resolve_failed_prerequisite", return_value=None
+ ),
+ patch("superset.tasks.scheduler.TaskContext", return_value=ctx),
+ patch(
+ "superset.tasks.scheduler.TaskRegistry.get_executor", return_value=_raise
+ ),
+ patch("superset.tasks.scheduler.TaskManager") as task_manager,
+ patch("superset.tasks.scheduler.current_app", app),
+ patch(
+ "superset.commands.tasks.internal_update.InternalStatusTransitionCommand",
+ transition,
+ ),
+ ):
+ from superset.tasks.scheduler import _execute_task_body
+
+ # SimpleNamespace stands in for the Task ORM row the body only reads from.
+ result = _execute_task_body(task, native, "some.task", (), {}) # type: ignore[arg-type]
+ return result, transition, task_manager, stats_logger
+
+
+def _requested_statuses(transition: MagicMock):
+ return [call.kwargs.get("new_status") for call in transition.call_args_list]
+
+
+def test_exception_during_timeout_finalizes_timed_out_not_failure() -> None:
+ """A body that raises while timing out must end TIMED_OUT, never FAILURE."""
+ result, transition, task_manager, stats_logger = _run_body_with_raising_executor(
+ timeout_triggered=True, abort_detected=False
+ )
+
+ assert TaskStatus.FAILURE not in _requested_statuses(transition)
+ assert result["status"] == TaskStatus.TIMED_OUT.value
+ task_manager.publish_completion.assert_called_once()
+ # The failure metric must not fire for a successful cancellation.
+ assert ("gtf.task.failure",) not in [
+ c.args for c in stats_logger.incr.call_args_list
+ ]
+
+
+def test_exception_during_abort_finalizes_aborted_not_failure() -> None:
+ """A body that raises while aborting must end ABORTED, never FAILURE."""
+ result, transition, _task_manager, _stats = _run_body_with_raising_executor(
+ timeout_triggered=False, abort_detected=True
+ )
+
+ assert TaskStatus.FAILURE not in _requested_statuses(transition)
+ assert result["status"] == TaskStatus.ABORTED.value
+
+
+def test_persist_celery_task_id_writes_and_commits() -> None:
+ from superset.tasks.scheduler import _persist_celery_task_id
+
+ task = MagicMock()
+ with patch("superset.tasks.scheduler.db") as db:
+ _persist_celery_task_id(task, "celery-9")
+
+ task.update_properties.assert_called_once_with({"celery_task_id": "celery-9"})
+ db.session.commit.assert_called_once()
+
+
+def test_persist_celery_task_id_is_noop_without_an_id() -> None:
+ from superset.tasks.scheduler import _persist_celery_task_id
+
+ task = MagicMock()
+ with patch("superset.tasks.scheduler.db") as db:
+ _persist_celery_task_id(task, None)
+
+ task.update_properties.assert_not_called()
+ db.session.commit.assert_not_called()