mirror of
https://github.com/apache/superset.git
synced 2026-04-14 05:34:38 +00:00
315 lines
10 KiB
Python
315 lines
10 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Unified cancel task command for GTF."""
|
|
|
|
import logging
|
|
from functools import partial
|
|
from typing import TYPE_CHECKING
|
|
from uuid import UUID
|
|
|
|
from flask import current_app
|
|
from superset_core.tasks.types import TaskScope, TaskStatus
|
|
|
|
from superset.commands.base import BaseCommand
|
|
from superset.commands.tasks.exceptions import (
|
|
TaskAbortFailedError,
|
|
TaskNotAbortableError,
|
|
TaskNotFoundError,
|
|
TaskPermissionDeniedError,
|
|
)
|
|
from superset.extensions import security_manager
|
|
from superset.stats_logger import BaseStatsLogger
|
|
from superset.tasks.locks import task_lock
|
|
from superset.tasks.utils import get_active_dedup_key
|
|
from superset.utils.core import get_user_id
|
|
from superset.utils.decorators import on_error, transaction
|
|
|
|
if TYPE_CHECKING:
|
|
from superset.models.tasks import Task
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CancelTaskCommand(BaseCommand):
|
|
"""
|
|
Unified command to cancel a task.
|
|
|
|
Behavior:
|
|
- For private tasks or single-subscriber tasks: aborts the task
|
|
- For shared tasks with multiple subscribers (non-admin): unsubscribes user
|
|
- For shared tasks with force=True (admin only): aborts for all subscribers
|
|
|
|
The term "cancel" is user-facing; internally this may abort or unsubscribe.
|
|
|
|
This command acquires a distributed lock before starting a transaction to
|
|
prevent race conditions with concurrent submit/cancel operations.
|
|
|
|
Permission checks are deferred to inside the lock to minimize SELECTs:
|
|
we only fetch the task once, then validate permissions on the fetched data.
|
|
"""
|
|
|
|
def __init__(self, task_uuid: UUID, force: bool = False):
|
|
"""
|
|
Initialize the cancel command.
|
|
|
|
:param task_uuid: UUID of the task to cancel
|
|
:param force: If True, force abort even with multiple subscribers (admin only)
|
|
"""
|
|
self._task_uuid = task_uuid
|
|
self._force = force
|
|
self._action_taken: str = (
|
|
"cancelled" # Will be set to 'aborted' or 'unsubscribed'
|
|
)
|
|
self._should_publish_abort: bool = False
|
|
|
|
def run(self) -> "Task":
|
|
"""
|
|
Execute the cancel command with distributed locking.
|
|
|
|
The lock is acquired BEFORE starting the transaction to avoid holding
|
|
a DB connection during lock acquisition. Uses dedup_key as lock key
|
|
to ensure Submit and Cancel operations use the same lock.
|
|
|
|
:returns: The updated task model
|
|
"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
# Lightweight fetch to compute dedup_key for locking
|
|
# This is needed to use the same lock key as SubmitTaskCommand
|
|
task = TaskDAO.find_one_or_none(
|
|
skip_base_filter=security_manager.is_admin(), uuid=self._task_uuid
|
|
)
|
|
|
|
if not task:
|
|
raise TaskNotFoundError()
|
|
|
|
# Compute dedup_key using the same logic as SubmitTaskCommand
|
|
dedup_key = get_active_dedup_key(
|
|
scope=task.scope,
|
|
task_type=task.task_type,
|
|
task_key=task.task_key,
|
|
user_id=task.user_id,
|
|
)
|
|
|
|
# Acquire lock BEFORE transaction starts
|
|
# Using dedup_key ensures Submit and Cancel use the same lock
|
|
with task_lock(dedup_key):
|
|
result = self._execute_with_transaction()
|
|
|
|
# Publish abort notification AFTER transaction commits
|
|
# This prevents race conditions where listeners check DB before commit
|
|
if self._should_publish_abort:
|
|
from superset.tasks.manager import TaskManager
|
|
|
|
TaskManager.publish_abort(self._task_uuid)
|
|
|
|
return result
|
|
|
|
@transaction(on_error=partial(on_error, reraise=TaskAbortFailedError))
|
|
def _execute_with_transaction(self) -> "Task":
|
|
"""
|
|
Execute the cancel operation inside a transaction.
|
|
|
|
Combines fetch + validation + execution in a single transaction,
|
|
reducing the number of SELECTs from 3 to 1 (plus DAO operations).
|
|
|
|
:returns: The updated task model
|
|
"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
# Check admin status (no DB access)
|
|
is_admin = security_manager.is_admin()
|
|
|
|
# Force flag requires admin
|
|
if self._force and not is_admin:
|
|
raise TaskPermissionDeniedError(
|
|
"Only administrators can force cancel a task"
|
|
)
|
|
|
|
# Single SELECT: fetch task and validate permissions on it
|
|
task = TaskDAO.find_one_or_none(skip_base_filter=is_admin, uuid=self._task_uuid)
|
|
|
|
if not task:
|
|
raise TaskNotFoundError()
|
|
|
|
# Validate permissions on the fetched task
|
|
self._validate_permissions(task, is_admin)
|
|
|
|
# Execute cancel and return updated task
|
|
return self._do_cancel(task, is_admin)
|
|
|
|
def _validate_permissions(self, task: "Task", is_admin: bool) -> None:
|
|
"""
|
|
Validate permissions on an already-fetched task.
|
|
|
|
Permission rules by scope:
|
|
- private: Only creator or admin (already filtered by base_filter)
|
|
- shared: Subscribers or admin
|
|
- system: Only admin
|
|
|
|
:param task: The task to validate permissions for
|
|
:param is_admin: Whether current user is admin
|
|
:raises TaskAbortFailedError: If task is not in cancellable state
|
|
:raises TaskPermissionDeniedError: If user lacks permission
|
|
"""
|
|
# Check if task is in a cancellable state
|
|
if task.status not in [
|
|
TaskStatus.PENDING.value,
|
|
TaskStatus.IN_PROGRESS.value,
|
|
TaskStatus.ABORTING.value, # Already aborting is OK (idempotent)
|
|
]:
|
|
raise TaskAbortFailedError()
|
|
|
|
# Admin can cancel anything
|
|
if is_admin:
|
|
return
|
|
|
|
# Non-admin permission checks by scope
|
|
user_id = get_user_id()
|
|
|
|
if task.scope == TaskScope.SYSTEM.value:
|
|
# System tasks are admin-only
|
|
raise TaskPermissionDeniedError(
|
|
"Only administrators can cancel system tasks"
|
|
)
|
|
|
|
if task.is_shared:
|
|
# Shared tasks: must be a subscriber
|
|
if not user_id or not task.has_subscriber(user_id):
|
|
raise TaskPermissionDeniedError(
|
|
"You must be subscribed to cancel this shared task"
|
|
)
|
|
|
|
# Private tasks: already filtered by base_filter (only creator can see)
|
|
# If we got here, user has permission
|
|
|
|
def _do_cancel(self, task: "Task", is_admin: bool) -> "Task":
|
|
"""
|
|
Execute the cancel operation (abort or unsubscribe).
|
|
|
|
:param task: The task to cancel
|
|
:param is_admin: Whether current user is admin
|
|
:returns: The updated task model
|
|
"""
|
|
user_id = get_user_id()
|
|
|
|
# Determine action based on task scope and force flag
|
|
should_abort = (
|
|
# Admin with force flag always aborts
|
|
(is_admin and self._force)
|
|
# Private tasks always abort (only one user)
|
|
or task.is_private
|
|
# System tasks always abort (admin only anyway)
|
|
or task.is_system
|
|
# Single or last subscriber - abort
|
|
or task.subscriber_count <= 1
|
|
)
|
|
|
|
if should_abort:
|
|
return self._do_abort(task, is_admin)
|
|
else:
|
|
return self._do_unsubscribe(task, user_id)
|
|
|
|
def _do_abort(self, task: "Task", is_admin: bool) -> "Task":
|
|
"""
|
|
Execute abort operation.
|
|
|
|
:param task: The task to abort
|
|
:param is_admin: Whether current user is admin
|
|
:returns: The updated task model
|
|
"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
try:
|
|
result: Task | None = TaskDAO.abort_task(
|
|
task.uuid, skip_base_filter=is_admin
|
|
)
|
|
except TaskNotAbortableError:
|
|
raise
|
|
|
|
if result is None:
|
|
# abort_task returned None - task wasn't aborted
|
|
# This can happen if task is already finished
|
|
raise TaskAbortFailedError()
|
|
|
|
self._action_taken = "aborted"
|
|
|
|
# Track if we need to publish abort after commit
|
|
if TaskStatus(result.status) == TaskStatus.ABORTING:
|
|
self._should_publish_abort = True
|
|
|
|
# Emit stats metric
|
|
stats_logger: BaseStatsLogger = current_app.config["STATS_LOGGER"]
|
|
stats_logger.incr("gtf.task.abort")
|
|
|
|
logger.info(
|
|
"Task aborted: %s (scope: %s, force: %s)",
|
|
task.uuid,
|
|
task.scope,
|
|
self._force,
|
|
)
|
|
|
|
return result
|
|
|
|
def _do_unsubscribe(self, task: "Task", user_id: int | None) -> "Task":
|
|
"""
|
|
Execute unsubscribe operation.
|
|
|
|
:param task: The task to unsubscribe from
|
|
:param user_id: ID of user to unsubscribe
|
|
:returns: The updated task model
|
|
"""
|
|
from superset.daos.tasks import TaskDAO
|
|
|
|
self._action_taken = "unsubscribed"
|
|
|
|
if not user_id or not task.has_subscriber(user_id):
|
|
# User not subscribed - they shouldn't be able to cancel
|
|
raise TaskPermissionDeniedError(
|
|
"You are not subscribed to this shared task"
|
|
)
|
|
|
|
result = TaskDAO.remove_subscriber(task.id, user_id)
|
|
if result is None:
|
|
raise TaskPermissionDeniedError(
|
|
"You are not subscribed to this shared task"
|
|
)
|
|
|
|
# Emit stats metric
|
|
stats_logger: BaseStatsLogger = current_app.config["STATS_LOGGER"]
|
|
stats_logger.incr("gtf.task.unsubscribe")
|
|
|
|
logger.info(
|
|
"User %s unsubscribed from shared task: %s",
|
|
user_id,
|
|
task.uuid,
|
|
)
|
|
|
|
return result
|
|
|
|
def validate(self) -> None:
|
|
pass
|
|
|
|
@property
|
|
def action_taken(self) -> str:
|
|
"""
|
|
Get the action that was taken.
|
|
|
|
:returns: 'aborted' or 'unsubscribed'
|
|
"""
|
|
return self._action_taken
|