diff --git a/docs/developer_docs/extensions/tasks.md b/docs/developer_docs/extensions/tasks.md index 58c00cb2408..7e3e8c42fd8 100644 --- a/docs/developer_docs/extensions/tasks.md +++ b/docs/developer_docs/extensions/tasks.md @@ -333,6 +333,33 @@ assert task.uuid == task2.uuid # True print(task2.status) # "success" (terminal status) ``` +## Task Dependencies + +Tasks can declare prerequisite tasks, forming a directed acyclic graph (DAG). Pass the prerequisite `Task` objects (returned by `.schedule()`) via `depends_on`: + +```python +from superset_core.tasks.types import TaskOptions + +totals = totals_task.schedule(options=TaskOptions(task_key="totals_123")) + +# `dependent` only runs once `totals` has finished successfully. +dependent = dependent_task.schedule( + options=TaskOptions(depends_on=[totals]) +) +``` + +Passing the `Task` object is the canonical pattern. For convenience, a prerequisite's `UUID` (or UUID string) is also accepted where you don't hold the `Task` itself. + +**Semantics (`all_success`).** A task runs only once **every** direct prerequisite has reached a terminal `SUCCESS`. If **any** prerequisite ends in a non-`SUCCESS` terminal state (`FAILURE`, `ABORTED`, or `TIMED_OUT`), the dependent does **not** run and is transitioned to `FAILURE`. This propagates transitively: because a failed dependent is itself non-`SUCCESS`, its own dependents fail in turn, so a failure anywhere short-circuits everything downstream. + +**Scheduling model (block-and-wait).** All tasks in a DAG are enqueued immediately. Each dependent's worker blocks — holding its worker slot — until its prerequisites finish; while waiting, the task remains `PENDING` (shown as "waiting on N prerequisites" in the Task List). Tasks are enqueued in dependency order, so a dependent is rarely dequeued before its prerequisites. + +:::warning Worker fleet sizing +Because dependents hold a worker slot while awaiting their prerequisites, a deep or wide DAG can occupy many workers simultaneously. Deployments that use chained tasks heavily must size their Celery worker fleet large enough to absorb the idle waiting, or a large DAG can exhaust the pool and deadlock. +::: + +Cycles (including self-dependencies) are rejected at schedule time. Dependency edges are removed automatically when either endpoint task is pruned. + ## Task Scopes ```python @@ -409,13 +436,15 @@ By default, abort detection and sync join-and-wait use database polling. Configu TaskOptions( task_key: str | None = None, task_name: str | None = None, - timeout: int | None = None + timeout: int | None = None, + depends_on: list[Task | UUID | str] | None = None ) ``` - `task_key`: Deduplication key (also used as display name if `task_name` is not set) - `task_name`: Human-readable display name for the Task List UI - `timeout`: Timeout in seconds (overrides decorator default) +- `depends_on`: Prerequisite tasks to wait for before running. Pass the scheduled `Task` objects (canonical); a `UUID` or UUID string is also accepted (see [Task Dependencies](#task-dependencies)) :::tip Provide a descriptive `task_name` for better readability in the Task List UI. While `task_key` is used for deduplication and may be technical (e.g., `chart_export_123`), `task_name` can be user-friendly (e.g., `"Export Sales Chart 123"`). diff --git a/superset-core/src/superset_core/tasks/models.py b/superset-core/src/superset_core/tasks/models.py index 1dcf70ad1a8..075f2e0c899 100644 --- a/superset-core/src/superset_core/tasks/models.py +++ b/superset-core/src/superset_core/tasks/models.py @@ -163,7 +163,28 @@ class TaskSubscriber(CoreModel): changed_by_fk: int | None -__all__ = [ - "Task", - "TaskSubscriber", -] +class TaskDependency(CoreModel): + """ + Abstract TaskDependency model interface. + + Host implementations will replace this class during initialization + with concrete implementation providing actual functionality. + + This model represents a directed edge in the task dependency graph (DAG): + the task identified by ``task_id`` depends on the prerequisite task + identified by ``depends_on_task_id``. A task only runs once all of its + prerequisites have reached a terminal SUCCESS. + """ + + __abstract__ = True + + # Type hints for expected attributes (no actual field definitions) + id: int + task_id: int # The dependent task + depends_on_task_id: int # The prerequisite task + + # Audit fields from AuditMixinNullable + created_on: datetime | None + changed_on: datetime | None + created_by_fk: int | None + changed_by_fk: int | None diff --git a/superset-core/src/superset_core/tasks/types.py b/superset-core/src/superset_core/tasks/types.py index 7e4f886f10b..dcffbca027c 100644 --- a/superset-core/src/superset_core/tasks/types.py +++ b/superset-core/src/superset_core/tasks/types.py @@ -20,7 +20,11 @@ from __future__ import annotations from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum -from typing import Any, Callable, Literal, TypedDict +from typing import Any, Callable, Literal, TYPE_CHECKING, TypedDict, Union +from uuid import UUID + +if TYPE_CHECKING: + from superset_core.tasks.models import Task class TaskStatus(str, Enum): @@ -122,11 +126,24 @@ class TaskOptions: task = long_task.schedule( options=TaskOptions(timeout=600) # 10 minute timeout ) + + # Task that waits for prerequisite tasks to succeed before running. + # Pass the scheduled Task objects (canonical); UUIDs are also accepted. + parent = parent_task.schedule() + task = dependent_task.schedule( + options=TaskOptions(depends_on=[parent]) + ) """ task_key: str | None = None task_name: str | None = None timeout: int | None = None # Timeout in seconds + # Prerequisite tasks this task depends on. Each entry may be a scheduled + # Task, its UUID, or a UUID string. The task only runs once every + # prerequisite has reached a terminal SUCCESS; if any prerequisite ends in a + # non-SUCCESS terminal state the task fails without running (all_success + # semantics). + depends_on: list[Union["Task", UUID, str]] | None = None class TaskContext(ABC): @@ -226,12 +243,3 @@ class TaskContext(ABC): cleanup_partial_work() """ ... - - -__all__ = [ - "TaskStatus", - "TaskScope", - "TaskProperties", - "TaskContext", - "TaskOptions", -] diff --git a/superset-frontend/src/features/tasks/TaskDependenciesPopover.test.tsx b/superset-frontend/src/features/tasks/TaskDependenciesPopover.test.tsx new file mode 100644 index 00000000000..be4e3f14bc6 --- /dev/null +++ b/superset-frontend/src/features/tasks/TaskDependenciesPopover.test.tsx @@ -0,0 +1,45 @@ +/** + * 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. + */ +import { render, screen, fireEvent } from 'spec/helpers/testing-library'; +import { TaskStatus } from './types'; +import TaskDependenciesPopover from './TaskDependenciesPopover'; + +const dependencies = [ + { uuid: 'prereq-1', task_name: 'Totals Query', status: TaskStatus.Success }, + { uuid: 'prereq-2', task_name: null, status: TaskStatus.InProgress }, +]; + +test('renders the chain-link trigger icon', () => { + render(, { + useRedux: true, + }); + expect(screen.getByRole('img', { name: 'link' })).toBeInTheDocument(); +}); + +test('lists prerequisite tasks in the popover on hover', async () => { + render(, { + useRedux: true, + }); + + fireEvent.mouseEnter(screen.getByRole('img', { name: 'link' })); + + // Named prerequisite shows its name; the unnamed one falls back to its uuid + expect(await screen.findByText('Totals Query')).toBeInTheDocument(); + expect(screen.getByText('prereq-2')).toBeInTheDocument(); +}); diff --git a/superset-frontend/src/features/tasks/TaskDependenciesPopover.tsx b/superset-frontend/src/features/tasks/TaskDependenciesPopover.tsx new file mode 100644 index 00000000000..ba098c04d9f --- /dev/null +++ b/superset-frontend/src/features/tasks/TaskDependenciesPopover.tsx @@ -0,0 +1,94 @@ +/** + * 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. + */ + +import { useState } from 'react'; +import { t } from '@apache-superset/core/translation'; +import { styled } from '@apache-superset/core/theme'; +import { Popover } from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import TaskStatusIcon from './TaskStatusIcon'; +import { TaskDependency, TaskStatus } from './types'; + +const DependenciesContainer = styled.div` + max-width: 400px; + max-height: 300px; + overflow: auto; + padding: ${({ theme }) => theme.sizeUnit}px 0; +`; + +const DependencyRow = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + padding: ${({ theme }) => theme.sizeUnit / 2}px + ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const DependencyName = styled.span` + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +`; + +const LinkIconWrapper = styled.span` + cursor: pointer; + color: ${({ theme }) => theme.colorIcon}; + + &:hover { + color: ${({ theme }) => theme.colorPrimary}; + } +`; + +interface TaskDependenciesPopoverProps { + dependencies: TaskDependency[]; +} + +export default function TaskDependenciesPopover({ + dependencies, +}: TaskDependenciesPopoverProps) { + const [visible, setVisible] = useState(false); + + const content = ( + + {dependencies.map(dependency => ( + + + + {dependency.task_name || dependency.uuid} + + + ))} + + ); + + return ( + + + + + + ); +} diff --git a/superset-frontend/src/features/tasks/types.ts b/superset-frontend/src/features/tasks/types.ts index 27d129cb856..1dcd5b62ac9 100644 --- a/superset-frontend/src/features/tasks/types.ts +++ b/superset-frontend/src/features/tasks/types.ts @@ -24,6 +24,16 @@ export interface TaskSubscriber { subscribed_at: string; } +/** + * A prerequisite task in the dependency graph (DAG). The dependent task only + * runs once every prerequisite reaches a terminal SUCCESS. + */ +export interface TaskDependency { + uuid: string; + task_name: string | null; + status: TaskStatus; +} + export enum TaskScope { Private = 'private', Shared = 'shared', @@ -85,6 +95,8 @@ export interface Task { duration_seconds: number | null; subscriber_count: number; subscribers: TaskSubscriber[]; + // Prerequisite tasks this task depends on (all_success DAG semantics). + depends_on?: TaskDependency[]; } // Derived status helpers (frontend computes these from status and properties) diff --git a/superset-frontend/src/pages/TaskList/TaskList.test.tsx b/superset-frontend/src/pages/TaskList/TaskList.test.tsx index 6fef614d92c..3f8ee9e3c11 100644 --- a/superset-frontend/src/pages/TaskList/TaskList.test.tsx +++ b/superset-frontend/src/pages/TaskList/TaskList.test.tsx @@ -147,6 +147,13 @@ const mockTasks = [ subscribed_at: '2024-01-15T12:00:01Z', }, ], + depends_on: [ + { + uuid: 'prereq-uuid-1', + task_name: 'Totals Query', + status: TaskStatus.InProgress, + }, + ], properties: { is_abortable: null, progress_percent: null, @@ -324,3 +331,14 @@ test('displays empty state when no tasks', async () => { response: { result: mockTasks, count: 3 }, }); }); + +test('shows "waiting on" indicator and chain icon for pending tasks with unmet prerequisites', async () => { + renderTaskList(); + await screen.findByText('Shared Bulk Task'); + + // The pending task depends on an in-progress (unmet) prerequisite + expect(await screen.findByText('Waiting on 1')).toBeInTheDocument(); + // The chain-link dependency icon is rendered (antd LinkOutlined -> name "link"). + // Popover contents on hover are covered by TaskDependenciesPopover.test.tsx. + expect(screen.getByRole('img', { name: 'link' })).toBeInTheDocument(); +}); diff --git a/superset-frontend/src/pages/TaskList/index.tsx b/superset-frontend/src/pages/TaskList/index.tsx index 97d29daa685..22ae2f99f54 100644 --- a/superset-frontend/src/pages/TaskList/index.tsx +++ b/superset-frontend/src/pages/TaskList/index.tsx @@ -48,6 +48,7 @@ import { createErrorHandler, createFetchRelated } from 'src/views/CRUD/utils'; import TaskStatusIcon from 'src/features/tasks/TaskStatusIcon'; import TaskPayloadPopover from 'src/features/tasks/TaskPayloadPopover'; import TaskStackTracePopover from 'src/features/tasks/TaskStackTracePopover'; +import TaskDependenciesPopover from 'src/features/tasks/TaskDependenciesPopover'; import { formatDuration } from 'src/features/tasks/timeUtils'; import { Task, @@ -396,6 +397,50 @@ function TaskList({ addDangerToast, addSuccessToast, user }: TaskListProps) { id: 'duration_seconds', disableSortBy: true, }, + { + Cell: ({ + row: { + original: { depends_on, status }, + }, + }: TaskCellProps) => { + if (!depends_on || depends_on.length === 0) { + return null; + } + // "waiting on N" surfaces the block-and-wait gate: a PENDING task + // parked until its unmet (non-SUCCESS) prerequisites finish. + const unmet = depends_on.filter( + dep => dep.status !== TaskStatus.Success, + ).length; + const showWaiting = status === TaskStatus.Pending && unmet > 0; + return ( +
+ + {showWaiting && ( + + + + )} +
+ ); + }, + accessor: 'depends_on', + Header: t('Dependencies'), + size: 'sm', + id: 'depends_on', + disableSortBy: true, + }, { Cell: ({ row: { diff --git a/superset/commands/tasks/exceptions.py b/superset/commands/tasks/exceptions.py index a54030dd5c4..98b94fdcc06 100644 --- a/superset/commands/tasks/exceptions.py +++ b/superset/commands/tasks/exceptions.py @@ -38,6 +38,13 @@ class TaskInvalidError(CommandInvalidError): message = _("Task parameters are invalid.") +class TaskCyclicDependencyError(TaskInvalidError): + """A task's declared dependencies would introduce a cycle in the DAG.""" + + status = 400 + message = _("Task dependencies would introduce a cycle.") + + class TaskCreateFailedError(CreateFailedError): """Task creation failed.""" diff --git a/superset/commands/tasks/submit.py b/superset/commands/tasks/submit.py index 69f4388b5f0..f3e627ce325 100644 --- a/superset/commands/tasks/submit.py +++ b/superset/commands/tasks/submit.py @@ -20,6 +20,7 @@ import logging import uuid from functools import partial from typing import Any, TYPE_CHECKING +from uuid import UUID from flask import current_app from marshmallow import ValidationError @@ -28,6 +29,7 @@ from superset_core.tasks.types import TaskScope from superset.commands.base import BaseCommand from superset.commands.tasks.exceptions import ( TaskCreateFailedError, + TaskCyclicDependencyError, TaskInvalidError, ) from superset.daos.exceptions import DAOCreateFailedError @@ -38,6 +40,7 @@ from superset.utils.core import get_user_id from superset.utils.decorators import on_error, transaction if TYPE_CHECKING: + from superset.daos.tasks import TaskDAO from superset.models.tasks import Task logger = logging.getLogger(__name__) @@ -137,11 +140,79 @@ class SubmitTaskCommand(BaseCommand): payload=self._properties.get("payload", {}), properties=self._properties.get("properties", {}), ) + # Persist dependency edges (with cycle guard) for the new task. + # Joined/deduplicated tasks keep their original dependencies. + self._persist_dependencies(task, TaskDAO) stats_logger.incr("gtf.task.create") return task, True # is_new=True: created new task except DAOCreateFailedError as ex: raise TaskCreateFailedError() from ex + def _persist_dependencies(self, task: "Task", dao: type["TaskDAO"]) -> None: + """ + Resolve the declared ``depends_on`` references and write dependency edges. + + Runs inside the submit transaction and lock, after the task row is + flushed (so ``task.id``/``task.uuid`` are available), and only for a + freshly *created* task (never on a dedup join). Rejects self-dependencies + and unknown prerequisites. Prerequisite references are de-duplicated and + order-preserved. + + No transitive cycle check is needed here: a brand-new task has no + incoming edges, so its new ``task -> prerequisite`` edges cannot close a + cycle (nothing points back to it). The only cycle a create can express is + a direct self-dependency, rejected in-memory below. A future API that + adds edges to *existing* tasks would need a transitive check. + + This resolves all prerequisites in one query and inserts all edges in one + flush — 2 round-trips regardless of the number of dependencies. + + :param task: The newly created dependent task + :param dao: TaskDAO (passed to avoid re-importing) + :raises TaskInvalidError: if a prerequisite reference is malformed/unknown + :raises TaskCyclicDependencyError: on a direct self-dependency + """ + raw = self._properties.get("depends_on") or [] + if not raw: + return + + uuids: list[UUID] = [] + seen: set[UUID] = set() + for item in raw: + # Accept a scheduled Task entity, a UUID, or a UUID string, and + # normalize to a UUID. Task entities are the natural output of + # .schedule(), so passing them straight through is the common case. + if isinstance(item, UUID): + dep_uuid = item + elif hasattr(item, "uuid"): + raw_uuid = item.uuid + dep_uuid = ( + raw_uuid if isinstance(raw_uuid, UUID) else UUID(str(raw_uuid)) + ) + else: + try: + dep_uuid = UUID(str(item)) + except (ValueError, AttributeError, TypeError) as ex: + raise TaskInvalidError( + f"Invalid prerequisite task reference: {item!r}" + ) from ex + if dep_uuid == task.uuid: + raise TaskCyclicDependencyError( + f"A task cannot depend on itself ({dep_uuid})." + ) + if dep_uuid not in seen: + seen.add(dep_uuid) + uuids.append(dep_uuid) + + prerequisites = {p.uuid: p for p in dao.find_by_uuids(uuids)} + missing = [str(u) for u in uuids if u not in prerequisites] + if missing: + raise TaskInvalidError( + f"Unknown prerequisite task(s): {', '.join(missing)}" + ) + + dao.add_dependencies(task.id, [prerequisites[u].id for u in uuids]) + def validate(self) -> None: """Validate command parameters.""" exceptions: list[ValidationError] = [] diff --git a/superset/daos/tasks.py b/superset/daos/tasks.py index 8e2e74370f4..41001ac94ea 100644 --- a/superset/daos/tasks.py +++ b/superset/daos/tasks.py @@ -26,6 +26,7 @@ from superset_core.tasks.types import TaskProperties, TaskScope, TaskStatus from superset.daos.base import BaseDAO from superset.daos.exceptions import DAODeleteFailedError 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 @@ -332,6 +333,53 @@ class TaskDAO(BaseDAO[Task]): f"Failed to remove subscription for task {task_id}, user {user_id}" ) from ex + # Dependency (DAG) management methods + + @classmethod + def find_by_uuids(cls, uuids: list[UUID]) -> list[Task]: + """ + Resolve a list of task UUIDs to Task instances. + + Used when persisting dependency edges, which are declared with public + UUIDs but stored against the internal integer ``id``. The base filter is + intentionally skipped: dependency resolution is a structural operation on + tasks the caller is wiring together (typically its own), not a + user-facing listing. + + :param uuids: Task UUIDs to resolve + :returns: Matching Task instances (order not guaranteed; missing UUIDs + are simply absent from the result) + """ + if not uuids: + return [] + return db.session.query(Task).filter(Task.uuid.in_(uuids)).all() + + @classmethod + def add_dependencies(cls, task_id: int, depends_on_task_ids: list[int]) -> None: + """ + Bulk-insert prerequisite edges: ``task_id`` depends on each id given. + + Only ever called for a freshly created task (see ``SubmitTaskCommand``) + with already-deduplicated prerequisite ids, so no edge can pre-exist — + the per-row existence check that ``add_subscriber`` needs is unnecessary + here, and the edges are inserted in a single flush (one INSERT). + + :param task_id: ID of the dependent task + :param depends_on_task_ids: IDs of the prerequisite tasks + """ + if not depends_on_task_ids: + return + db.session.add_all( + TaskDependency(task_id=task_id, depends_on_task_id=prerequisite_id) + for prerequisite_id in depends_on_task_ids + ) + db.session.flush() + logger.info( + "Added %d dependencies to task %s", + len(depends_on_task_ids), + task_id, + ) + @classmethod def set_properties_and_payload( cls, 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 new file mode 100644 index 00000000000..6e9b08b0fca --- /dev/null +++ b/superset/migrations/versions/2026-08-21_12-00_7e2c9a4f1b83_create_task_dependencies_table.py @@ -0,0 +1,136 @@ +# 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. +"""Create task_dependencies table for Global Task Framework (GTF) task DAG + +Revision ID: 7e2c9a4f1b83 +Revises: 1072de5ed955 +Create Date: 2026-08-21 12:00:00.000000 + +""" + +from sqlalchemy import ( + Column, + DateTime, + Integer, + UniqueConstraint, +) + +from superset.migrations.shared.utils import ( + create_fks_for_table, + create_index, + create_table, + drop_fks_for_table, + drop_index, + drop_table, +) + +# revision identifiers, used by Alembic. +revision = "7e2c9a4f1b83" +down_revision = "1072de5ed955" + +TASKS_TABLE = "tasks" +TASK_DEPENDENCIES_TABLE = "task_dependencies" + + +def upgrade(): + """ + Create the task_dependencies junction table for the task dependency graph. + + Each row is a directed edge: ``task_id`` (the dependent) depends on + ``depends_on_task_id`` (the prerequisite). Both foreign keys reference + ``tasks.id`` with ``ON DELETE CASCADE`` so edges are removed when either + endpoint task is pruned (task pruning uses a bulk core DELETE that bypasses + the ORM cascade, so the database-level cascade is required for cleanup). + """ + create_table( + TASK_DEPENDENCIES_TABLE, + Column("id", Integer, primary_key=True), + Column("task_id", Integer, nullable=False), + Column("depends_on_task_id", Integer, nullable=False), + # AuditMixinNullable columns + Column("created_on", DateTime, nullable=True), + Column("changed_on", DateTime, nullable=True), + Column("created_by_fk", Integer, nullable=True), + Column("changed_by_fk", Integer, nullable=True), + # Unique constraint defined as part of table creation (SQLite compatible). + # The leading task_id column also serves forward (prerequisite) lookups. + UniqueConstraint( + "task_id", + "depends_on_task_id", + name="uq_task_dependencies_task_depends_on", + ), + ) + + # Index for reverse (dependents) lookups by prerequisite. + create_index( + TASK_DEPENDENCIES_TABLE, + "idx_task_dependencies_depends_on", + ["depends_on_task_id"], + ) + + # Both edge endpoints cascade-delete with their task. + create_fks_for_table( + foreign_key_name="fk_task_dependencies_task_id_tasks", + table_name=TASK_DEPENDENCIES_TABLE, + referenced_table=TASKS_TABLE, + local_cols=["task_id"], + remote_cols=["id"], + ondelete="CASCADE", + ) + + create_fks_for_table( + foreign_key_name="fk_task_dependencies_depends_on_task_id_tasks", + table_name=TASK_DEPENDENCIES_TABLE, + referenced_table=TASKS_TABLE, + local_cols=["depends_on_task_id"], + remote_cols=["id"], + ondelete="CASCADE", + ) + + create_fks_for_table( + foreign_key_name="fk_task_dependencies_created_by_fk_ab_user", + table_name=TASK_DEPENDENCIES_TABLE, + referenced_table="ab_user", + local_cols=["created_by_fk"], + remote_cols=["id"], + ondelete="SET NULL", + ) + + create_fks_for_table( + foreign_key_name="fk_task_dependencies_changed_by_fk_ab_user", + table_name=TASK_DEPENDENCIES_TABLE, + referenced_table="ab_user", + local_cols=["changed_by_fk"], + remote_cols=["id"], + ondelete="SET NULL", + ) + + +def downgrade(): + """Drop the task_dependencies table and its indexes and foreign keys.""" + drop_fks_for_table( + TASK_DEPENDENCIES_TABLE, + [ + "fk_task_dependencies_task_id_tasks", + "fk_task_dependencies_depends_on_task_id_tasks", + "fk_task_dependencies_created_by_fk_ab_user", + "fk_task_dependencies_changed_by_fk_ab_user", + ], + ) + + drop_index(TASK_DEPENDENCIES_TABLE, "idx_task_dependencies_depends_on") + drop_table(TASK_DEPENDENCIES_TABLE) diff --git a/superset/models/task_dependencies.py b/superset/models/task_dependencies.py new file mode 100644 index 00000000000..cc696785000 --- /dev/null +++ b/superset/models/task_dependencies.py @@ -0,0 +1,69 @@ +# 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. +"""TaskDependency model for the Global Task Framework (GTF) task DAG""" + +from __future__ import annotations + +from flask_appbuilder import Model +from sqlalchemy import Column, ForeignKey, Integer, UniqueConstraint +from superset_core.tasks.models import TaskDependency as CoreTaskDependency + +from superset.models.helpers import AuditMixinNullable + + +class TaskDependency(CoreTaskDependency, AuditMixinNullable, Model): + """ + A directed edge in the task dependency graph (DAG). + + The dependent task (``task_id``) waits for the prerequisite task + (``depends_on_task_id``) to reach a terminal state before it runs. A task + only executes once *every* prerequisite has reached a terminal SUCCESS; if + any prerequisite ends in a non-SUCCESS terminal state the dependent fails + without running (``all_success`` semantics), which cascades transitively to + its own dependents. + + This is a pure edge table: prerequisite ``Task`` entities are read through + ``Task.dependencies`` (a self-referential many-to-many over this table). + Both foreign keys use ``ON DELETE CASCADE`` so edges are cleaned up when + either endpoint task is pruned. This is required because + ``TaskPruneCommand`` deletes tasks via a bulk core ``DELETE`` that bypasses + the ORM cascade. + """ + + __tablename__ = "task_dependencies" + + id = Column(Integer, primary_key=True) + task_id = Column( + Integer, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False + ) + depends_on_task_id = Column( + Integer, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False + ) + + __table_args__ = ( + UniqueConstraint( + "task_id", + "depends_on_task_id", + name="uq_task_dependencies_task_depends_on", + ), + ) + + def __repr__(self) -> str: + return ( + f"" + ) diff --git a/superset/models/tasks.py b/superset/models/tasks.py index 679f950aa2b..89b5da3c58a 100644 --- a/superset/models/tasks.py +++ b/superset/models/tasks.py @@ -36,6 +36,7 @@ from superset_core.tasks.models import Task as CoreTask from superset_core.tasks.types import TaskProperties, TaskStatus from superset.models.helpers import AuditMixinNullable +from superset.models.task_dependencies import TaskDependency from superset.models.task_subscribers import TaskSubscriber from superset.tasks.constants import TERMINAL_STATES from superset.tasks.utils import ( @@ -111,6 +112,23 @@ class Task(CoreTask, AuditMixinNullable, Model): lazy="selectin", ) + # Prerequisite tasks (self-referential many-to-many over task_dependencies). + # `dependencies` is the list of Task entities this task depends on, so the + # full prerequisite tasks load in a single selectin fetch alongside the task + # (no per-edge round trips). It is viewonly: edges are written through + # TaskDAO.add_dependency, not by mutating this collection. Cleanup on task + # deletion relies on the DB-level FK ON DELETE CASCADE (see the migration), + # since TaskPruneCommand bulk-deletes via core DELETE, not the ORM. + dependencies = relationship( + "Task", + secondary=TaskDependency.__table__, + primaryjoin=id == TaskDependency.task_id, + secondaryjoin=id == TaskDependency.depends_on_task_id, + order_by=TaskDependency.id, + lazy="selectin", + viewonly=True, + ) + def __repr__(self) -> str: return f"" diff --git a/superset/tasks/decorators.py b/superset/tasks/decorators.py index bc43353145d..d8a0e5806ea 100644 --- a/superset/tasks/decorators.py +++ b/superset/tasks/decorators.py @@ -227,6 +227,7 @@ class TaskWrapper(Generic[P]): task_key=self.default_options.task_key, task_name=self.default_options.task_name, timeout=self.default_timeout, # Use decorator default + depends_on=self.default_options.depends_on, ) # Merge: use override if provided, otherwise use default @@ -238,6 +239,7 @@ class TaskWrapper(Generic[P]): timeout=override_options.timeout if override_options.timeout is not None else self.default_timeout, + depends_on=override_options.depends_on or self.default_options.depends_on, ) def _validate_task(self, options: TaskOptions) -> None: @@ -625,4 +627,5 @@ class TaskWrapper(Generic[P]): timeout=options.timeout, args=args, kwargs=kwargs, + depends_on=options.depends_on, ) diff --git a/superset/tasks/manager.py b/superset/tasks/manager.py index 564df0ac271..5c88435bdd9 100644 --- a/superset/tasks/manager.py +++ b/superset/tasks/manager.py @@ -271,6 +271,7 @@ class TaskManager: timeout: int | None, args: tuple[Any, ...], kwargs: dict[str, Any], + depends_on: list[Task | UUID | str] | None = None, ) -> "Task": """ Create task entry and schedule for async execution. @@ -291,6 +292,10 @@ class TaskManager: :param timeout: Optional timeout in seconds :param args: Positional arguments for the task function :param kwargs: Keyword arguments for the task function + :param depends_on: Optional prerequisite tasks (as Task entities, UUIDs, + or UUID strings). The task is still enqueued immediately + (block-and-wait model); ordering is enforced in the scheduler, which + waits for prerequisites before running the body. :returns: Task model representing the scheduled task """ from superset.commands.tasks.submit import SubmitTaskCommand @@ -316,6 +321,7 @@ class TaskManager: "task_name": task_name, "scope": scope.value, "properties": properties, + "depends_on": depends_on, } ).run_with_info() diff --git a/superset/tasks/scheduler.py b/superset/tasks/scheduler.py index 5a7485497f3..6a0db9b2fdc 100644 --- a/superset/tasks/scheduler.py +++ b/superset/tasks/scheduler.py @@ -18,7 +18,7 @@ from __future__ import annotations import logging from datetime import datetime, timezone -from typing import Any +from typing import Any, TYPE_CHECKING from uuid import UUID from celery import Task @@ -51,6 +51,9 @@ from superset.utils.core import LoggerLevel from superset.utils.log import get_logger_from_status from superset.utils.report_execution import get_report_task_timeout_options +if TYPE_CHECKING: + from superset.models.tasks import Task as TaskModel + logger = logging.getLogger(__name__) @@ -277,6 +280,48 @@ def prune_key_value( logger.exception("An error occurred while pruning the key-value store: %s", ex) +def _resolve_failed_prerequisite(task: "TaskModel") -> "TaskModel | None": + """ + Block until every direct prerequisite of ``task`` reaches a terminal state. + + Implements the ``all_success`` trigger rule for the task DAG: the task may + run only if *all* of its direct prerequisites ended in ``SUCCESS``. + + ``task.dependencies`` is already ``selectin``-loaded (in one query) with the + task, so a prerequisite that is *already* terminal in that snapshot is + evaluated with **no extra database reads** — a terminal status never changes, + so the snapshot is authoritative for it (the common case under Model A, where + FIFO enqueue order means parents usually finish before the dependent runs). + Only prerequisites that are not yet terminal fall through to + ``TaskManager.wait_for_completion`` (wake-on-completion else poll), and they + are awaited one at a time (≈1 read/poll-interval total, not per-prerequisite). + Transitive failure propagation is emergent — a dependent that fails here is + itself non-SUCCESS, so its own dependents fail in turn. + + :param task: The dependent task about to run (with ``dependencies`` loaded) + :returns: The first prerequisite that did not end in ``SUCCESS``, or ``None`` + if the task has no prerequisites or all of them succeeded + """ + prerequisites = list(task.dependencies) + if not prerequisites: + return None + + for prerequisite in prerequisites: + # Trust an already-terminal status from the loaded snapshot (no extra + # read); otherwise block on a fresh wait until it becomes terminal. + if prerequisite.status not in TERMINAL_STATES: + try: + prerequisite = TaskManager.wait_for_completion(prerequisite.uuid) + except ValueError: + # Prerequisite no longer exists (e.g. pruned mid-wait) — treat as + # a failed prerequisite rather than blocking or crashing. + return prerequisite + if prerequisite.status != TaskStatus.SUCCESS.value: + return prerequisite + + return None + + @celery_app.task(name="tasks.execute", bind=True) def execute_task( # noqa: C901 self: Any, # Celery task instance @@ -336,6 +381,35 @@ def execute_task( # noqa: C901 ).run() return {"status": TaskStatus.ABORTED.value, "task_uuid": task_uuid} + # DAG gate: wait for prerequisites before claiming the task. The task stays + # PENDING while waiting, so the "waiting on prerequisites" indicator applies + # and an abort mid-wait is caught by the PENDING → IN_PROGRESS transition + # below. If any prerequisite did not succeed, fail without running the body + # (all_success semantics); the failure then cascades to this task's own + # dependents. + if (failed_prerequisite := _resolve_failed_prerequisite(task)) is not None: + logger.info( + "Task %s (uuid=%s) failing: prerequisite %s did not succeed (status=%s)", + task_type, + task_uuid, + failed_prerequisite.uuid, + failed_prerequisite.status, + ) + InternalStatusTransitionCommand( + task_uuid=native_uuid, + new_status=TaskStatus.FAILURE, + expected_status=[TaskStatus.PENDING, TaskStatus.ABORTING], + set_ended_at=True, + properties={ + "error_message": ( + f"Prerequisite task {failed_prerequisite.uuid} did not " + f"succeed (status={failed_prerequisite.status})" + ) + }, + ).run() + TaskManager.publish_completion(native_uuid, TaskStatus.FAILURE.value) + return {"status": TaskStatus.FAILURE.value, "task_uuid": task_uuid} + # Atomic transition: PENDING → IN_PROGRESS (set started_at for duration tracking) if not InternalStatusTransitionCommand( task_uuid=native_uuid, diff --git a/superset/tasks/schemas.py b/superset/tasks/schemas.py index 31537cb83b2..6bd011cdb90 100644 --- a/superset/tasks/schemas.py +++ b/superset/tasks/schemas.py @@ -61,6 +61,10 @@ subscriber_count_description = ( "Number of users subscribed to this task (for shared tasks)" ) subscribers_description = "List of users subscribed to this task (for shared tasks)" +depends_on_description = ( + "Prerequisite tasks this task depends on. The task only runs once all of " + "them reach a terminal SUCCESS (all_success semantics)." +) class UserSchema(Schema): @@ -117,6 +121,9 @@ class TaskResponseSchema(Schema): subscribers = Method( "get_subscribers", metadata={"description": subscribers_description} ) + depends_on = Method( + "get_depends_on", metadata={"description": depends_on_description} + ) def get_payload_dict(self, obj: object) -> dict[str, object] | None: """Get payload as dictionary""" @@ -168,6 +175,17 @@ class TaskResponseSchema(Schema): ) return subscribers + def get_depends_on(self, obj: object) -> list[dict[str, object]]: + """Get prerequisite tasks (uuid, name, status) for DAG display.""" + return [ + { + "uuid": str(prerequisite.uuid), + "task_name": prerequisite.task_name, + "status": prerequisite.status, + } + for prerequisite in obj.dependencies # type: ignore[attr-defined] + ] + class TaskStatusResponseSchema(Schema): """Schema for task status response (lightweight for polling)""" diff --git a/tests/integration_tests/tasks/api_tests.py b/tests/integration_tests/tasks/api_tests.py index 45462ec3900..fdf6b42f513 100644 --- a/tests/integration_tests/tasks/api_tests.py +++ b/tests/integration_tests/tasks/api_tests.py @@ -463,6 +463,7 @@ class TestTaskApi(SupersetTestCase): "scope", "subscriber_count", "subscribers", + "depends_on", ] for field in expected_fields: @@ -472,6 +473,49 @@ class TestTaskApi(SupersetTestCase): properties = result["properties"] assert isinstance(properties, dict) + # A task with no declared prerequisites serializes an empty depends_on + assert result["depends_on"] == [] + + def test_task_depends_on_serialization(self): + """ + Task API: Test depends_on serializes prerequisite tasks (uuid/name/status) + """ + from superset.commands.tasks import SubmitTaskCommand + + self.login(ADMIN_USERNAME) + + prerequisite = SubmitTaskCommand( + data={ + "task_type": "test_type", + "task_key": "api_dep_prereq", + "task_name": "Prerequisite Task", + } + ).run() + dependent = None + try: + dependent = SubmitTaskCommand( + data={ + "task_type": "test_type", + "task_key": "api_dep_child", + "depends_on": [prerequisite.uuid], + } + ).run() + + rv = self.client.get(f"{self.TASK_API_BASE}/{dependent.uuid}") + assert rv.status_code == 200 + result = json.loads(rv.data.decode("utf-8"))["result"] + + assert len(result["depends_on"]) == 1 + dep = result["depends_on"][0] + assert dep["uuid"] == str(prerequisite.uuid) + assert dep["task_name"] == "Prerequisite Task" + assert dep["status"] == prerequisite.status + finally: + if dependent is not None: + db.session.delete(dependent) + db.session.delete(prerequisite) + db.session.commit() + def test_task_payload_serialization(self): """ Task API: Test payload is properly serialized as dict diff --git a/tests/integration_tests/tasks/commands/test_submit.py b/tests/integration_tests/tasks/commands/test_submit.py index 286701e9c28..f528c10e876 100644 --- a/tests/integration_tests/tasks/commands/test_submit.py +++ b/tests/integration_tests/tasks/commands/test_submit.py @@ -54,6 +54,8 @@ def test_submit_task_success(app_context, login_as, get_user) -> None: db.session.refresh(result) assert result.id is not None assert result.uuid is not None + # A task submitted without depends_on has no prerequisites + assert result.dependencies == [] finally: # Cleanup db.session.delete(result) @@ -233,3 +235,51 @@ def test_submit_task_run_with_info_returns_is_new_false( # Cleanup db.session.delete(task1) db.session.commit() + + +def test_submit_task_with_depends_on_persists_edges( + app_context, login_as, get_user +) -> None: + """Test depends_on (passed as a Task entity) persists dependency edges.""" + login_as("admin") + + prerequisite = SubmitTaskCommand( + data={"task_type": "test-type", "task_key": "dep-prereq"} + ).run() + + dependent = None + try: + dependent = SubmitTaskCommand( + data={ + "task_type": "test-type", + "task_key": "dep-child", + # Pass the Task entity directly (the natural flow) + "depends_on": [prerequisite], + } + ).run() + db.session.refresh(dependent) + + # dependencies resolves to the prerequisite Task entities + assert [dep.uuid for dep in dependent.dependencies] == [prerequisite.uuid] + finally: + # Deleting the dependent removes its edge rows via FK ON DELETE CASCADE + if dependent is not None: + db.session.delete(dependent) + db.session.delete(prerequisite) + db.session.commit() + + +def test_submit_task_unknown_dependency_rejected(app_context, login_as) -> None: + """Test depends_on referencing an unknown task is rejected (and rolls back).""" + from uuid import uuid4 + + login_as("admin") + + with pytest.raises(TaskInvalidError): + SubmitTaskCommand( + data={ + "task_type": "test-type", + "task_key": "dep-unknown", + "depends_on": [uuid4()], + } + ).run() diff --git a/tests/unit_tests/daos/test_tasks.py b/tests/unit_tests/daos/test_tasks.py index a24ad870fd5..f35597ed174 100644 --- a/tests/unit_tests/daos/test_tasks.py +++ b/tests/unit_tests/daos/test_tasks.py @@ -109,6 +109,8 @@ def test_find_by_task_key_active(session_with_task: Session) -> None: assert result.task_key == TEST_TASK_KEY assert result.task_type == TEST_TASK_TYPE assert result.status == TaskStatus.PENDING.value + # A task with no declared prerequisites has empty dependencies + assert result.dependencies == [] def test_find_by_task_key_not_found(session_with_task: Session) -> None: @@ -510,3 +512,30 @@ def test_conditional_status_update_terminal_state_updates_dedup_key( assert task.dedup_key != original_dedup_key, ( f"dedup_key should have changed for {terminal_state.value}" ) + + +def test_add_dependencies_bulk_inserts_edges(session_with_task: Session) -> None: + """add_dependencies bulk-inserts edges; Task.dependencies resolves them.""" + from superset.daos.tasks import TaskDAO + + parent1 = create_task(session_with_task, task_key="parent1") + parent2 = create_task(session_with_task, task_key="parent2") + child = create_task(session_with_task, task_key="child") + + TaskDAO.add_dependencies(child.id, [parent1.id, parent2.id]) + + # Task.dependencies resolves to the prerequisite Task entities + session_with_task.refresh(child) + assert {t.id for t in child.dependencies} == {parent1.id, parent2.id} + # Prerequisites themselves have no dependencies + assert parent1.dependencies == [] + + +def test_add_dependencies_empty_is_noop(session_with_task: Session) -> None: + """add_dependencies with no ids does nothing.""" + from superset.daos.tasks import TaskDAO + + task = create_task(session_with_task, task_key="lonely") + TaskDAO.add_dependencies(task.id, []) + session_with_task.refresh(task) + assert task.dependencies == [] diff --git a/tests/unit_tests/tasks/test_decorators.py b/tests/unit_tests/tasks/test_decorators.py index 9e0818b6576..988c9891fb3 100644 --- a/tests/unit_tests/tasks/test_decorators.py +++ b/tests/unit_tests/tasks/test_decorators.py @@ -170,6 +170,7 @@ class TestTaskWrapperMergeOptions: merged = merge_task_1._merge_options(None) assert merged.task_key == "default_key" assert merged.task_name == "Default Name" + assert merged.depends_on is None def test_merge_options_override_task_key(self): """Test overriding task_key at call time""" @@ -215,10 +216,32 @@ class TestTaskWrapperMergeOptions: override = TaskOptions( task_key="override_key", task_name="Override Name", + depends_on=[TEST_UUID], ) merged = merge_task_4._merge_options(override) assert merged.task_key == "override_key" assert merged.task_name == "Override Name" + assert merged.depends_on == [TEST_UUID] + + def test_merge_options_depends_on_from_override(self): + """Test depends_on is carried through from call-time options""" + + @task(name="test_merge_depends_on_unique") + def merge_task_deps() -> None: + pass + + override = TaskOptions(depends_on=[TEST_UUID]) + merged = merge_task_deps._merge_options(override) + assert merged.depends_on == [TEST_UUID] + + def test_merge_options_depends_on_default_none(self): + """Test depends_on defaults to None when not provided""" + + @task(name="test_merge_depends_on_default_unique") + def merge_task_no_deps() -> None: + pass + + assert merge_task_no_deps._merge_options(None).depends_on is None class TestTaskWrapperSchedule: @@ -256,10 +279,11 @@ class TestTaskWrapperSchedule: schedule_task_2.schedule(123) - # Verify PRIVATE scope was used (default) + # Verify PRIVATE scope was used (default) and no dependencies forwarded mock_submit.assert_called_once() call_args = mock_submit.call_args assert call_args[1]["scope"] == TaskScope.PRIVATE + assert call_args[1]["depends_on"] is None @patch("superset.tasks.decorators.TaskManager.submit_task") def test_schedule_with_custom_options(self, mock_submit): @@ -270,18 +294,23 @@ class TestTaskWrapperSchedule: def schedule_task_3(arg1: int) -> None: pass - # Use custom task key and name + # Use custom task key, name, and a prerequisite dependency schedule_task_3.schedule( 123, - options=TaskOptions(task_key="custom_key", task_name="Custom Task Name"), + options=TaskOptions( + task_key="custom_key", + task_name="Custom Task Name", + depends_on=[TEST_UUID], + ), ) - # Verify scope from decorator and options from call time + # Verify scope from decorator and options from call time are forwarded mock_submit.assert_called_once() call_args = mock_submit.call_args assert call_args[1]["scope"] == TaskScope.SYSTEM assert call_args[1]["task_key"] == "custom_key" assert call_args[1]["task_name"] == "Custom Task Name" + assert call_args[1]["depends_on"] == [TEST_UUID] @patch("superset.tasks.decorators.TaskManager.submit_task") def test_schedule_with_no_decorator_options(self, mock_submit): diff --git a/tests/unit_tests/tasks/test_dependencies.py b/tests/unit_tests/tasks/test_dependencies.py new file mode 100644 index 00000000000..d817ee1f688 --- /dev/null +++ b/tests/unit_tests/tasks/test_dependencies.py @@ -0,0 +1,165 @@ +# 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 task dependencies (DAG): scheduler gate + cycle guard.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest +from superset_core.tasks.types import TaskStatus + +from superset.commands.tasks.exceptions import ( + TaskCyclicDependencyError, + TaskInvalidError, +) +from superset.commands.tasks.submit import SubmitTaskCommand + + +def _task(status: str = TaskStatus.PENDING.value, **kwargs): + return SimpleNamespace(uuid=uuid4(), status=status, task_name=None, **kwargs) + + +class TestResolveFailedPrerequisite: + """Tests for the scheduler's all_success prerequisite gate.""" + + def test_no_dependencies_returns_none(self): + from superset.tasks.scheduler import _resolve_failed_prerequisite + + assert _resolve_failed_prerequisite(SimpleNamespace(dependencies=[])) is None + + @patch("superset.tasks.scheduler.TaskManager.wait_for_completion") + def test_all_success_returns_none(self, mock_wait): + from superset.tasks.scheduler import _resolve_failed_prerequisite + + task = SimpleNamespace(dependencies=[_task(), _task()]) + mock_wait.return_value = SimpleNamespace(status=TaskStatus.SUCCESS.value) + + assert _resolve_failed_prerequisite(task) is None + assert mock_wait.call_count == 2 + + @patch("superset.tasks.scheduler.TaskManager.wait_for_completion") + def test_already_terminal_prerequisites_skip_wait(self, mock_wait): + """Prerequisites already terminal in the loaded snapshot need no DB wait.""" + from superset.tasks.scheduler import _resolve_failed_prerequisite + + # One already succeeded, one already failed → decided from the snapshot + succeeded = _task(status=TaskStatus.SUCCESS.value) + failed = _task(status=TaskStatus.FAILURE.value) + assert ( + _resolve_failed_prerequisite( + SimpleNamespace(dependencies=[succeeded, failed]) + ) + is failed + ) + # All-terminal snapshot → wait_for_completion is never called + assert ( + _resolve_failed_prerequisite(SimpleNamespace(dependencies=[succeeded])) + is None + ) + mock_wait.assert_not_called() + + @patch("superset.tasks.scheduler.TaskManager.wait_for_completion") + def test_first_non_success_is_returned(self, mock_wait): + from superset.tasks.scheduler import _resolve_failed_prerequisite + + task = SimpleNamespace(dependencies=[_task(), _task()]) + failed = SimpleNamespace(uuid=uuid4(), status=TaskStatus.FAILURE.value) + mock_wait.side_effect = [ + SimpleNamespace(status=TaskStatus.SUCCESS.value), + failed, + ] + + assert _resolve_failed_prerequisite(task) is failed + + @patch("superset.tasks.scheduler.TaskManager.wait_for_completion") + def test_missing_prerequisite_treated_as_failed(self, mock_wait): + from superset.tasks.scheduler import _resolve_failed_prerequisite + + prerequisite = _task() + task = SimpleNamespace(dependencies=[prerequisite]) + mock_wait.side_effect = ValueError("gone") + + # The (stale) prerequisite is returned rather than blocking/raising. + assert _resolve_failed_prerequisite(task) is prerequisite + + +class TestPersistDependencies: + """Tests for SubmitTaskCommand._persist_dependencies.""" + + def test_no_depends_on_is_noop(self): + dao = MagicMock() + SubmitTaskCommand({})._persist_dependencies(_task(id=1), dao) + dao.add_dependencies.assert_not_called() + + def test_self_dependency_by_uuid_raises(self): + dao = MagicMock() + task = _task(id=1) + cmd = SubmitTaskCommand({"depends_on": [task.uuid]}) + with pytest.raises(TaskCyclicDependencyError): + cmd._persist_dependencies(task, dao) + + def test_unknown_prerequisite_raises(self): + dao = MagicMock() + dao.find_by_uuids.return_value = [] # nothing resolves + cmd = SubmitTaskCommand({"depends_on": [uuid4()]}) + with pytest.raises(TaskInvalidError): + cmd._persist_dependencies(_task(id=1), dao) + + def test_happy_path_bulk_inserts_edges(self): + u1, u2 = uuid4(), uuid4() + dao = MagicMock() + dao.find_by_uuids.return_value = [ + SimpleNamespace(uuid=u1, id=101), + SimpleNamespace(uuid=u2, id=102), + ] + + cmd = SubmitTaskCommand({"depends_on": [u1, u2]}) + cmd._persist_dependencies(_task(id=1), dao) + + # One bulk insert of all prerequisite ids (no per-edge round trips) + dao.add_dependencies.assert_called_once_with(1, [101, 102]) + + def test_duplicate_uuids_deduped(self): + u1 = uuid4() + dao = MagicMock() + dao.find_by_uuids.return_value = [SimpleNamespace(uuid=u1, id=101)] + + cmd = SubmitTaskCommand({"depends_on": [u1, u1]}) + cmd._persist_dependencies(_task(id=1), dao) + + dao.add_dependencies.assert_called_once_with(1, [101]) + + def test_accepts_task_entities_and_uuids(self): + """depends_on accepts Task entities, UUIDs, and UUID strings.""" + u1, u2, u3 = uuid4(), uuid4(), uuid4() + dao = MagicMock() + dao.find_by_uuids.return_value = [ + SimpleNamespace(uuid=u1, id=101), # referenced by entity + SimpleNamespace(uuid=u2, id=102), # referenced by UUID + SimpleNamespace(uuid=u3, id=103), # referenced by str + ] + + # A Task-like entity (has .uuid), a raw UUID, and a UUID string. + entity = SimpleNamespace(uuid=u1, id=101) + cmd = SubmitTaskCommand({"depends_on": [entity, u2, str(u3)]}) + cmd._persist_dependencies(_task(id=1), dao) + + dao.add_dependencies.assert_called_once_with(1, [101, 102, 103]) + # find_by_uuids received normalized UUIDs, not entities/strings + (resolved,), _ = dao.find_by_uuids.call_args + assert set(resolved) == {u1, u2, u3}