From ba1cb2c03e0b6b5404ddefb3fc5c96b552c26ac1 Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Sat, 3 Oct 2026 06:58:44 +0100 Subject: [PATCH] Show current task executions in the UI A loop pass or a mapped expansion reuses a task id and map index, so the UI cannot identify an execution by (task_id, map_index): the logs, XCom, state or HITL page opened for one pass could show another's. Pages and mutations therefore carry the region id and index, and row selection is keyed by attempt UUID so two executions that share public coordinates stay distinct even when they are selected together. The run-level Execution tab lists live task instances grouped by loop iteration and mapped expansion, because that is how a user thinks about what is running. Clearing a selection goes by UUID through a dialog that offers the later-iteration and whole-expansion choices the clear API takes, so a user can decide how far a clear reaches. Gantt segments link to the exact region and try, and a retired try opens only when both are given, since a try number alone is ambiguous. Only the current view is covered. Selecting history by clear boundary is deliberately absent, so no trajectory parameter is read or written. --- .../core_api/datamodels/task_instances.py | 1 + .../core_api/openapi/_private_ui.yaml | 4 + .../openapi/v2-rest-api-generated.yaml | 4 + .../services/public/task_coordinates.py | 2 + .../src/airflow/models/task_coordinates.py | 12 +- .../ui/openapi-gen/requests/schemas.gen.ts | 6 +- .../ui/openapi-gen/requests/types.gen.ts | 1 + .../ui/public/i18n/locales/en/common.json | 1 + .../ui/public/i18n/locales/en/dag.json | 11 + .../ui/public/i18n/locales/en/dags.json | 4 +- .../components/ActionAccordion/columns.tsx | 4 +- .../ClearGroupTaskInstanceDialog.tsx | 1 - .../ClearTaskInstanceConfirmationDialog.tsx | 33 +- .../ClearTaskInstanceDialog.test.tsx | 338 +++++++++++- .../TaskInstance/ClearTaskInstanceDialog.tsx | 383 ++++++++----- .../ui/src/components/DurationChart.test.tsx | 36 +- .../ui/src/components/DurationChart.tsx | 10 +- .../ui/src/components/Graph/TaskLink.test.tsx | 6 +- .../ui/src/components/Graph/TaskLink.tsx | 2 + .../MarkTaskInstanceAsDialog.test.tsx | 63 ++- .../TaskInstance/MarkTaskInstanceAsDialog.tsx | 14 +- .../components/TaskInstanceTooltip.test.tsx | 3 +- .../ui/src/components/TaskInstanceTooltip.tsx | 16 +- .../ui/src/components/TaskTrySelect.test.tsx | 18 +- .../ui/src/components/TaskTrySelect.tsx | 31 +- .../airflow/ui/src/constants/searchParams.ts | 3 + .../ui/src/hooks/useHITLReviewTabs.test.tsx | 58 ++ .../airflow/ui/src/hooks/useHITLReviewTabs.ts | 26 +- .../hooks/usePluginAppliesToContext.test.ts | 36 +- .../ui/src/hooks/usePluginAppliesToContext.ts | 4 +- .../ui/src/hooks/useRequiredActionTabs.ts | 55 +- .../ui/src/hooks/useSelectedVersion.test.tsx | 73 +++ .../ui/src/hooks/useSelectedVersion.ts | 3 + .../src/hooks/useTaskInstanceCoordinates.ts | 31 ++ .../ui/src/hooks/useTaskInstanceView.test.ts | 97 ++++ .../ui/src/hooks/useTaskInstanceView.ts | 68 +++ .../layouts/Details/DagBreadcrumb.test.tsx | 20 +- .../ui/src/layouts/Details/DagBreadcrumb.tsx | 6 +- .../ui/src/layouts/Details/Gantt/Gantt.tsx | 1 - .../Details/Gantt/GanttTimeline.test.tsx | 119 +++- .../layouts/Details/Gantt/GanttTimeline.tsx | 29 +- .../src/layouts/Details/Gantt/utils.test.ts | 122 +++-- .../ui/src/layouts/Details/Gantt/utils.ts | 45 +- .../ui/src/layouts/Details/Grid/Bar.tsx | 7 +- .../src/layouts/Details/Grid/GridButton.tsx | 7 +- .../src/layouts/Details/Grid/GridTI.test.tsx | 11 + .../ui/src/layouts/Details/Grid/GridTI.tsx | 11 +- .../ui/src/layouts/Details/Grid/TaskNames.tsx | 7 +- .../AssetStateStore/AssetStateStore.test.tsx | 123 +++++ .../Asset/AssetStateStore/AssetStateStore.tsx | 14 +- .../ui/src/pages/Events/Events.test.tsx | 61 +++ .../airflow/ui/src/pages/Events/Events.tsx | 8 +- .../HITLTaskInstances/HITLResponseForm.tsx | 2 + .../airflow/ui/src/pages/ReactPlugin.test.tsx | 28 +- .../src/airflow/ui/src/pages/ReactPlugin.tsx | 4 + .../ui/src/pages/Run/Execution.test.tsx | 512 ++++++++++++++++++ .../airflow/ui/src/pages/Run/Execution.tsx | 325 +++++++++++ .../src/airflow/ui/src/pages/Run/Run.tsx | 1 + .../pages/TaskInstance/AssetEvents.test.tsx | 61 +++ .../ui/src/pages/TaskInstance/AssetEvents.tsx | 8 +- .../src/pages/TaskInstance/BlockingDeps.tsx | 2 + .../ui/src/pages/TaskInstance/Details.tsx | 55 +- .../pages/TaskInstance/ExtraLinks.test.tsx | 20 +- .../ui/src/pages/TaskInstance/ExtraLinks.tsx | 3 + .../src/pages/TaskInstance/HITLResponse.tsx | 4 + .../ui/src/pages/TaskInstance/Header.test.tsx | 32 +- .../ui/src/pages/TaskInstance/Header.tsx | 24 +- .../TaskInstance/Logs/ExternalLogLink.tsx | 6 +- .../ui/src/pages/TaskInstance/Logs/Logs.tsx | 28 +- .../TaskInstance/Logs/LogsTryNumber.test.tsx | 25 + .../pages/TaskInstance/Logs/TaskLogHeader.tsx | 6 +- .../pages/TaskInstance/RenderedTemplates.tsx | 3 + .../pages/TaskInstance/TaskInstance.test.tsx | 253 ++++++++- .../src/pages/TaskInstance/TaskInstance.tsx | 116 ++-- .../BulkClearTaskInstancesButton.tsx | 126 +---- .../BulkDeleteTaskInstancesButton.tsx | 3 + .../BulkMarkTaskInstancesAsButton.tsx | 13 +- .../DeleteTaskInstanceButton.test.tsx | 58 ++ .../DeleteTaskInstanceButton.tsx | 3 + .../TaskInstances/TaskInstances.test.tsx | 6 +- .../src/pages/TaskInstances/TaskInstances.tsx | 2 +- .../ClearAllTaskStateStoreButton.tsx | 4 +- .../DeleteTaskStateStoreButton.tsx | 4 +- .../TaskStateStore/TaskStateStore.test.tsx | 114 ++++ .../pages/TaskStateStore/TaskStateStore.tsx | 4 + .../TaskStateStoreModal.test.tsx | 81 +++ .../TaskStateStore/TaskStateStoreModal.tsx | 5 +- .../ui/src/pages/XCom/AddXComButton.tsx | 4 + .../ui/src/pages/XCom/DeleteXComButton.tsx | 2 + .../ui/src/pages/XCom/EditXComButton.tsx | 2 + .../src/airflow/ui/src/pages/XCom/XCom.tsx | 24 +- .../ui/src/pages/XCom/XComEntry.test.tsx | 46 +- .../airflow/ui/src/pages/XCom/XComEntry.tsx | 15 +- .../ui/src/pages/XCom/XComModal.test.tsx | 82 +++ .../airflow/ui/src/pages/XCom/XComModal.tsx | 21 +- .../ui/src/queries/useBulkClearDryRun.ts | 114 ---- .../src/queries/useBulkClearTaskInstances.ts | 133 ----- .../src/queries/useBulkMarkAsDryRun.test.tsx | 76 +++ .../ui/src/queries/useBulkMarkAsDryRun.ts | 4 +- .../queries/useClearTaskInstances.test.tsx | 82 +++ .../ui/src/queries/useClearTaskInstances.ts | 38 +- .../useClearTaskInstancesDryRun.test.tsx | 88 +++ .../queries/useClearTaskInstancesDryRun.ts | 57 +- .../airflow/ui/src/queries/useLogs.test.ts | 44 +- .../src/airflow/ui/src/queries/useLogs.tsx | 22 +- .../src/queries/usePatchTaskInstance.test.tsx | 60 ++ .../ui/src/queries/usePatchTaskInstance.ts | 6 +- .../src/queries/usePatchTaskInstanceDryRun.ts | 2 + .../src/queries/useTaskInstanceNote.test.tsx | 54 ++ .../ui/src/queries/useTaskInstanceNote.ts | 2 +- .../ui/src/queries/useUpdateHITLDetail.ts | 6 + airflow-core/src/airflow/ui/src/router.tsx | 2 + .../src/airflow/ui/src/utils/links.test.ts | 15 + .../src/airflow/ui/src/utils/links.ts | 18 +- .../airflow/ui/src/utils/loopTaskInstance.ts | 22 + .../core_api/routes/public/test_hitl.py | 1 + .../routes/public/test_task_instances.py | 16 + .../core_api/routes/ui/test_gantt.py | 34 ++ .../airflowctl/api/datamodels/generated.py | 1 + .../tests/airflow_ctl/api/test_operations.py | 2 + .../ctl/commands/test_task_command.py | 3 + 121 files changed, 4216 insertions(+), 872 deletions(-) create mode 100644 airflow-core/src/airflow/ui/src/hooks/useHITLReviewTabs.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/hooks/useSelectedVersion.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/hooks/useTaskInstanceCoordinates.ts create mode 100644 airflow-core/src/airflow/ui/src/hooks/useTaskInstanceView.test.ts create mode 100644 airflow-core/src/airflow/ui/src/hooks/useTaskInstanceView.ts create mode 100644 airflow-core/src/airflow/ui/src/pages/Asset/AssetStateStore/AssetStateStore.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/pages/Events/Events.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/pages/Run/Execution.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/pages/Run/Execution.tsx create mode 100644 airflow-core/src/airflow/ui/src/pages/TaskInstance/AssetEvents.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/pages/TaskInstances/DeleteTaskInstanceButton.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/pages/TaskStateStore/TaskStateStore.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/pages/TaskStateStore/TaskStateStoreModal.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/pages/XCom/XComModal.test.tsx delete mode 100644 airflow-core/src/airflow/ui/src/queries/useBulkClearDryRun.ts delete mode 100644 airflow-core/src/airflow/ui/src/queries/useBulkClearTaskInstances.ts create mode 100644 airflow-core/src/airflow/ui/src/queries/useBulkMarkAsDryRun.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/queries/useClearTaskInstances.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/queries/useClearTaskInstancesDryRun.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/queries/useTaskInstanceNote.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/utils/loopTaskInstance.ts diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py index 7ad2fcbae5df2..baf4dc788df0e 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py @@ -60,6 +60,7 @@ class TaskInstanceResponse(OmitsMissingRegion, BaseModel): map_index: int region_id: RegionId = None region_index: RegionIndex = None + in_loop: bool logical_date: datetime | None run_after: datetime start_date: datetime | None diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml index 17c110656f7bc..d2ad5ae578771 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml @@ -5250,6 +5250,9 @@ components: region_index: type: integer title: Region Index + in_loop: + type: boolean + title: In Loop logical_date: anyOf: - type: string @@ -5401,6 +5404,7 @@ components: - dag_id - dag_run_id - map_index + - in_loop - logical_date - run_after - start_date diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 5a7520750fc64..3bacc8b2dc777 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -17920,6 +17920,9 @@ components: region_index: type: integer title: Region Index + in_loop: + type: boolean + title: In Loop logical_date: anyOf: - type: string @@ -18071,6 +18074,7 @@ components: - dag_id - dag_run_id - map_index + - in_loop - logical_date - run_after - start_date diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_coordinates.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_coordinates.py index 63d24c94f26e3..b9722e139d846 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_coordinates.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_coordinates.py @@ -53,6 +53,8 @@ def __getattr__(self, name: str) -> Any: if self.projected_map_index is not None: return self.projected_map_index return self.resolver.public_map_index(self.value) + if name == "in_loop": + return self.value.region_id != SENTINEL_REGION_ID and self.resolver._is_in_loop(self.value) return getattr(self.value, name) diff --git a/airflow-core/src/airflow/models/task_coordinates.py b/airflow-core/src/airflow/models/task_coordinates.py index fb4eef123a4fe..180153ca9c2d2 100644 --- a/airflow-core/src/airflow/models/task_coordinates.py +++ b/airflow-core/src/airflow/models/task_coordinates.py @@ -164,6 +164,7 @@ class TaskCoordinateResolver: _dags: dict[UUID, SerializedDAG] = attrs.field(factory=dict, init=False) _region_nodes: dict[UUID, str | None] = attrs.field(factory=dict, init=False) _regional_tasks: dict[tuple[str, str | None, str], bool] = attrs.field(factory=dict, init=False) + _looped_tasks: dict[tuple[UUID, str], bool] = attrs.field(factory=dict, init=False) @classmethod def for_dag(cls, dag: SerializedDAG | None, session: Session) -> TaskCoordinateResolver: @@ -201,7 +202,16 @@ def get_task( raise ValueError(f"Pinned Dag for {dag_id}/{run_id} not found") return dag.get_task(task_id) - def find_task(self, ti: TaskInstance) -> SerializedOperator | None: + def _is_in_loop(self, ti: TaskCoordinate) -> bool: + if ti.dag_version_id is None: + return False + key = (ti.dag_version_id, ti.task_id) + if key not in self._looped_tasks: + task = self.find_task(ti) + self._looped_tasks[key] = task is not None and enclosing_loop(task) is not None + return self._looped_tasks[key] + + def find_task(self, ti: TaskCoordinate) -> SerializedOperator | None: """Return the task of ``ti`` from its pinned Dag, or None when that Dag or the task is gone.""" if ti.dag_version_id is None or (dag := self.get_dag(ti.dag_version_id)) is None: return None diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 2ff0b3b399e85..74445fc1107b0 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -8067,6 +8067,10 @@ export const $TaskInstanceResponse = { type: 'integer', title: 'Region Index' }, + in_loop: { + type: 'boolean', + title: 'In Loop' + }, logical_date: { anyOf: [ { @@ -8351,7 +8355,7 @@ export const $TaskInstanceResponse = { } }, type: 'object', - required: ['id', 'task_id', 'dag_id', 'dag_run_id', 'map_index', 'logical_date', 'run_after', 'start_date', 'end_date', 'duration', 'state', 'try_number', 'max_tries', 'task_display_name', 'dag_display_name', 'hostname', 'unixname', 'pool', 'pool_slots', 'queue', 'priority_weight', 'operator', 'operator_name', 'queued_when', 'scheduled_when', 'pid', 'executor', 'executor_config', 'note', 'rendered_map_index', 'trigger', 'triggerer_job', 'dag_version'], + required: ['id', 'task_id', 'dag_id', 'dag_run_id', 'map_index', 'in_loop', 'logical_date', 'run_after', 'start_date', 'end_date', 'duration', 'state', 'try_number', 'max_tries', 'task_display_name', 'dag_display_name', 'hostname', 'unixname', 'pool', 'pool_slots', 'queue', 'priority_weight', 'operator', 'operator_name', 'queued_when', 'scheduled_when', 'pid', 'executor', 'executor_config', 'note', 'rendered_map_index', 'trigger', 'triggerer_job', 'dag_version'], title: 'TaskInstanceResponse', description: 'TaskInstance serializer for responses.' } as const; diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 0eb19aa0a36b6..0afe4e031adeb 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2199,6 +2199,7 @@ export type TaskInstanceResponse = { map_index: number; region_id?: string; region_index?: number; + in_loop: boolean; logical_date: string | null; run_after: string; start_date: string | null; diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json b/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json index 61141af84c6c8..ce77c25d18260 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json @@ -466,6 +466,7 @@ "executorConfig": "Executor Config", "hostname": "Hostname", "id": "ID", + "iteration": "Iteration", "maxTries": "Max Tries", "pid": "PID", "pool": "Pool", diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/en/dag.json b/airflow-core/src/airflow/ui/public/i18n/locales/en/dag.json index ce354726ad21e..5bc0277a88f0c 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/en/dag.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/dag.json @@ -68,6 +68,17 @@ "upcomingCount_one": "{{count}} Upcoming Deadline", "upcomingCount_other": "{{count}} Upcoming Deadlines" }, + "execution": { + "clearSelected": "Clear selected executions", + "empty": "No executions in this run.", + "iteration": "{{node}} · Iteration {{index}}", + "mapped_one": "{{task}} ({{count, number}} mapped task on this page)", + "mapped_other": "{{task}} ({{count, number}} mapped tasks on this page)", + "page": "Showing {{start, number}}–{{end, number}} of {{total, number}} executions", + "runTasks": "Run tasks", + "select": "Select {{task}}", + "title": "Execution" + }, "extraLinks": "Extra Links", "grid": { "buttons": { diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/en/dags.json b/airflow-core/src/airflow/ui/public/i18n/locales/en/dags.json index e39a81943325f..c0aef93f96452 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/en/dags.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/dags.json @@ -92,13 +92,15 @@ "existingTasks": "Clear existing tasks", "future": "Future", "keepTaskState": "Keep task state and resume", + "laterIterations": "Later iterations", "onlyFailed": "Clear only failed tasks", "past": "Past", "preventRunningTasks": "Prevent rerun if task is running", "queueNew": "Queue up new tasks", "runOnLatestVersion": "Run with latest bundle version", "runOnLatestVersionForced": "Always uses the latest — there's no earlier version to go back to", - "upstream": "Upstream" + "upstream": "Upstream", + "wholeExpansion": "Whole expansion" } }, "runStateCounts": { diff --git a/airflow-core/src/airflow/ui/src/components/ActionAccordion/columns.tsx b/airflow-core/src/airflow/ui/src/components/ActionAccordion/columns.tsx index a092a2b94c448..cdf4bb490acb6 100644 --- a/airflow-core/src/airflow/ui/src/components/ActionAccordion/columns.tsx +++ b/airflow-core/src/airflow/ui/src/components/ActionAccordion/columns.tsx @@ -27,10 +27,8 @@ import type { DataTableFeatures } from "src/components/DataTable/features"; import type { MetaColumn } from "src/components/DataTable/types"; import { StateBadge } from "src/components/StateBadge"; -// Stable per-row key; dag_run_id keeps the same task distinct across runs (past/future -// expansion), and map_index disambiguates the mapped instances of one task. export const taskInstanceKey = (ti: TaskInstanceResponse): string => - `${ti.dag_run_id}:${ti.task_id}:${ti.map_index}`; + `${ti.dag_run_id}:${ti.task_id}:${ti.map_index}:${ti.region_id}:${ti.region_index}`; // When provided, rows render a leading checkbox so the user can exclude individual // task instances from the action. Excluded keys are tracked by the caller. diff --git a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearGroupTaskInstanceDialog.tsx b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearGroupTaskInstanceDialog.tsx index 7f6c6fad8eedb..455bf93b527f2 100644 --- a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearGroupTaskInstanceDialog.tsx +++ b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearGroupTaskInstanceDialog.tsx @@ -75,7 +75,6 @@ export const ClearGroupTaskInstanceDialog = ({ onClose, open, taskInstance }: Pr }, [open, keepTaskStateDefault]); const { isPending, mutate } = useClearTaskInstances({ - dagId, dagRunId: runId, onSuccessConfirm: onCloseDialog, }); diff --git a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceConfirmationDialog.tsx b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceConfirmationDialog.tsx index c3b4a43e57f05..3d04624e1f602 100644 --- a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceConfirmationDialog.tsx +++ b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceConfirmationDialog.tsx @@ -30,17 +30,9 @@ import { useClearTaskInstancesDryRun } from "src/queries/useClearTaskInstancesDr import { useDurationFormat } from "src/utils"; type Props = { - readonly dagDetails?: { + readonly dryRun?: { dagId: string; - dagRunId: string; - downstream?: boolean; - future?: boolean; - mapIndex?: number; - onlyFailed?: boolean; - past?: boolean; - taskId: string; - taskIds?: ClearTaskInstancesBody["task_ids"]; - upstream?: boolean; + requestBody: ClearTaskInstancesBody; }; readonly onClose: () => void; readonly onConfirm?: () => void; @@ -49,7 +41,7 @@ type Props = { }; const ClearTaskInstanceConfirmationDialog = ({ - dagDetails, + dryRun, onClose, onConfirm, open, @@ -57,29 +49,16 @@ const ClearTaskInstanceConfirmationDialog = ({ }: Props) => { const { t: translate } = useTranslation(); const { formatRelative } = useDurationFormat(); - const useExplicitTaskIds = dagDetails?.taskIds !== undefined; const { data, isFetching } = useClearTaskInstancesDryRun({ - dagId: dagDetails?.dagId ?? "", + dagId: dryRun?.dagId ?? "", options: { - enabled: open && Boolean(dagDetails), + enabled: open && Boolean(dryRun), gcTime: 0, refetchOnMount: "always", refetchOnWindowFocus: false, staleTime: 0, }, - requestBody: { - dag_run_id: dagDetails?.dagRunId ?? "", - include_downstream: !useExplicitTaskIds && dagDetails?.downstream, - include_future: dagDetails?.future, - include_past: dagDetails?.past, - include_upstream: !useExplicitTaskIds && dagDetails?.upstream, - only_failed: dagDetails?.onlyFailed, - task_ids: useExplicitTaskIds - ? dagDetails.taskIds - : dagDetails?.mapIndex === undefined - ? [dagDetails?.taskId ?? ""] - : [[dagDetails.taskId, dagDetails.mapIndex]], - }, + requestBody: dryRun?.requestBody ?? {}, }); const [isReady, setIsReady] = useState(false); diff --git a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.test.tsx b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.test.tsx index 3418b8c2ef684..5b857456e976f 100644 --- a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.test.tsx +++ b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.test.tsx @@ -17,17 +17,23 @@ * under the License. */ import "@testing-library/jest-dom/vitest"; -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import { http, HttpResponse } from "msw"; import { setupServer, type SetupServer } from "msw/node"; import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import type { TaskInstanceResponse } from "openapi/requests/types.gen"; -import { CLEAR_KEEP_TASK_STATE_KEY } from "src/constants/localStorage"; +import BulkClearTaskInstancesButton from "src/pages/TaskInstances/BulkClearTaskInstancesButton"; + +import { + CLEAR_KEEP_TASK_STATE_KEY, + CLEAR_TASK_INSTANCE_DEFAULT_OPTIONS_KEY, +} from "src/constants/localStorage"; import { handlers } from "src/mocks/handlers"; import { Wrapper } from "src/utils/Wrapper"; +import { ClearTaskInstanceButton } from "./ClearTaskInstanceButton"; import ClearTaskInstanceDialog from "./ClearTaskInstanceDialog"; const DAG_ID = "test_dag"; @@ -45,6 +51,7 @@ const taskInstance: TaskInstanceResponse = { executor_config: "{}", hostname: null, id: "test_task_instance", + in_loop: false, logical_date: "2025-01-01T00:00:00Z", map_index: -1, max_tries: 0, @@ -57,8 +64,6 @@ const taskInstance: TaskInstanceResponse = { priority_weight: null, queue: null, queued_when: null, - region_id: "00000000-0000-0000-0000-000000000000", - region_index: -1, rendered_fields: undefined, rendered_map_index: null, run_after: "2025-01-01T00:00:00Z", @@ -75,7 +80,21 @@ const taskInstance: TaskInstanceResponse = { const affectedTasks = { task_instances: [taskInstance], total_entries: 1 }; +const loopPass = (id: string, regionIndex: number): TaskInstanceResponse => ({ + ...taskInstance, + id, + region_id: "11111111-1111-4111-8111-111111111111", + region_index: regionIndex, + task_id: "body.work", +}); + +const loopMember: TaskInstanceResponse = { + ...loopPass("loop-member", 1), + in_loop: true, +}; + const mutateMock = vi.fn(); +const dryRunsMock = vi.fn(); vi.mock("src/queries/useClearTaskInstances", () => ({ useClearTaskInstances: () => ({ isPending: false, mutate: mutateMock }), @@ -86,6 +105,11 @@ vi.mock("src/queries/useClearTaskInstances", () => ({ // same-tick response for both keeps the gate from blocking on a real fetch. vi.mock("src/queries/useClearTaskInstancesDryRun", () => ({ useClearTaskInstancesDryRun: () => ({ data: affectedTasks, isFetching: false, isPending: false }), + useClearTaskInstancesDryRuns: (args: unknown) => { + dryRunsMock(args); + + return { data: affectedTasks, error: null, isFetching: false }; + }, })); let server: SetupServer; @@ -99,6 +123,9 @@ beforeAll(() => { server.listen({ onUnhandledFrame: "bypass" }); }); afterEach(() => { + affectedTasks.task_instances = [taskInstance]; + mutateMock.mockClear(); + dryRunsMock.mockClear(); server.resetHandlers(); localStorage.clear(); }); @@ -126,4 +153,307 @@ describe("ClearTaskInstanceDialog", () => { expect(requestBody.keep_task_state).toBe(true); }); + + it("clears the kept loop passes by execution id when another pass is unticked", async () => { + affectedTasks.task_instances = [loopPass("pass-0", 0), loopPass("pass-2", 2)]; + + render(, { + wrapper: Wrapper, + }); + + const rows = await screen.findAllByRole("row"); + + fireEvent.click(within(rows[1] as HTMLElement).getByRole("checkbox")); + fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu })); + + await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1)); + const [{ requestBody }] = mutateMock.mock.calls[0] as [ + { requestBody: { task_ids?: unknown; task_instance_ids?: Array } }, + ]; + + expect(requestBody.task_instance_ids).toEqual(["pass-2"]); + expect(requestBody.task_ids).toBeUndefined(); + }); + + it("keeps an unticked loop pass excluded when its retry comes back with a new execution id", async () => { + affectedTasks.task_instances = [loopPass("pass-0", 0), loopPass("pass-2", 2)]; + + const dialog = ; + const { rerender } = render(dialog, { wrapper: Wrapper }); + + const rows = await screen.findAllByRole("row"); + + fireEvent.click(within(rows[1] as HTMLElement).getByRole("checkbox")); + + affectedTasks.task_instances = [loopPass("pass-0-retried", 0), loopPass("pass-2", 2)]; + rerender(dialog); + + fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu })); + + await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1)); + const [{ requestBody }] = mutateMock.mock.calls[0] as [ + { requestBody: { task_instance_ids?: Array } }, + ]; + + expect(requestBody.task_instance_ids).toEqual(["pass-2"]); + }); + + it.each([ + { inLoop: true, shown: true }, + { inLoop: false, shown: false }, + ])( + "shows the later-iterations option only for a loop member (in_loop=$inLoop)", + async ({ inLoop, shown }) => { + render( + , + { + wrapper: Wrapper, + }, + ); + + expect(await screen.findByRole("checkbox", { name: /keepTaskState/iu })).toBeVisible(); + expect(screen.queryByRole("button", { name: /options\.laterIterations/iu }) !== null).toBe(shown); + }, + ); + + it("includes later iterations in the first dry run and leaves the saved defaults untouched", async () => { + localStorage.setItem(CLEAR_TASK_INSTANCE_DEFAULT_OPTIONS_KEY, JSON.stringify(["downstream"])); + render(, { + wrapper: Wrapper, + }); + + fireEvent.click(await screen.findByRole("button", { name: /options\.laterIterations/iu })); + + expect(dryRunsMock.mock.calls[0]?.[0]).toMatchObject({ + requests: [{ requestBody: { include_later_loop_iterations: true } }], + }); + expect(localStorage.getItem(CLEAR_TASK_INSTANCE_DEFAULT_OPTIONS_KEY)).toBe( + JSON.stringify(["downstream"]), + ); + }); + + it("clears a loop member by execution id with the later-iterations choice and never sends past or future", async () => { + localStorage.setItem( + CLEAR_TASK_INSTANCE_DEFAULT_OPTIONS_KEY, + JSON.stringify(["past", "future", "downstream"]), + ); + render(, { + wrapper: Wrapper, + }); + + fireEvent.click(await screen.findByRole("button", { name: /options\.laterIterations/iu })); + fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu })); + + await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1)); + const [{ requestBody }] = mutateMock.mock.calls[0] as [{ requestBody: Record }]; + + expect(requestBody).toMatchObject({ + include_downstream: true, + include_later_loop_iterations: false, + task_instance_ids: ["loop-member"], + whole_expansion_ids: [], + }); + expect(requestBody).not.toHaveProperty("task_ids"); + expect(requestBody).not.toHaveProperty("include_past"); + expect(requestBody).not.toHaveProperty("include_future"); + }); + + it("offers whole-expansion clearing for a mapped loop member and sends the member's id", async () => { + const mappedMember = { ...loopMember, map_index: 2 }; + + render(, { + wrapper: Wrapper, + }); + + fireEvent.click(await screen.findByRole("button", { name: /options\.wholeExpansion/iu })); + fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu })); + + await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1)); + const [{ requestBody }] = mutateMock.mock.calls[0] as [{ requestBody: Record }]; + + expect(requestBody).toMatchObject({ + task_instance_ids: ["loop-member"], + whole_expansion_ids: ["loop-member"], + }); + }); + + it("offers no whole-expansion option when the loop member is not mapped", async () => { + render(, { + wrapper: Wrapper, + }); + + expect(await screen.findByRole("button", { name: /options\.laterIterations/iu })).toBeVisible(); + expect(screen.queryByRole("button", { name: /options\.wholeExpansion/iu })).not.toBeInTheDocument(); + }); + + it("does not offer the exclusion checkboxes for a loop clear", async () => { + affectedTasks.task_instances = [loopPass("pass-0", 0), loopPass("pass-2", 2)]; + + render(, { + wrapper: Wrapper, + }); + + const rows = await screen.findAllByRole("row"); + + expect(rows).toHaveLength(3); + expect(within(rows[1] as HTMLElement).queryByRole("checkbox")).not.toBeInTheDocument(); + }); + + it("clears a selection by task and map index when none of it is in a loop", async () => { + const other = { ...taskInstance, id: "other", map_index: 3, task_id: "task_2" }; + + render(, { + wrapper: Wrapper, + }); + + fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu })); + + await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1)); + const [{ requestBody }] = mutateMock.mock.calls[0] as [{ requestBody: Record }]; + + expect(requestBody).toMatchObject({ + dag_run_id: DAG_RUN_ID, + task_ids: [ + [TASK_ID, -1], + ["task_2", 3], + ], + }); + expect(requestBody).not.toHaveProperty("task_instance_ids"); + }); + + it("sends one run-scoped request per run when a selection spans runs", async () => { + const otherRun = { ...loopMember, dag_run_id: "run_2", id: "other-run-member" }; + + render(, { + wrapper: Wrapper, + }); + + fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu })); + + await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(2)); + expect( + mutateMock.mock.calls.map(([request]) => (request as { requestBody: unknown }).requestBody), + ).toMatchObject([ + { dag_run_id: DAG_RUN_ID, task_instance_ids: ["loop-member"] }, + { dag_run_id: "run_2", task_instance_ids: ["other-run-member"] }, + ]); + expect(dryRunsMock.mock.lastCall?.[0]).toMatchObject({ + requests: [ + { dagId: DAG_ID, requestBody: { dag_run_id: DAG_RUN_ID, task_instance_ids: ["loop-member"] } }, + { dagId: DAG_ID, requestBody: { dag_run_id: "run_2", task_instance_ids: ["other-run-member"] } }, + ], + }); + }); + + it("sends the note only when it was edited", async () => { + const withNote = { ...loopMember, note: "existing note" }; + + render(, { + wrapper: Wrapper, + }); + + const note = await screen.findByRole("textbox", { hidden: true }); + + expect(note).toHaveValue("existing note"); + fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu })); + await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1)); + expect((mutateMock.mock.calls[0] as [{ requestBody: { note?: string } }])[0].requestBody.note).toBe( + undefined, + ); + }); + + it("sends the note once the user edits it", async () => { + render(, { + wrapper: Wrapper, + }); + + fireEvent.change(await screen.findByRole("textbox", { hidden: true }), { + target: { value: "new reason" }, + }); + fireEvent.click(await screen.findByRole("button", { name: /modal\.confirm/iu })); + await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1)); + expect((mutateMock.mock.calls[0] as [{ requestBody: { note?: string } }])[0].requestBody.note).toBe( + "new reason", + ); + }); + + it("lists every affected loop pass, including later ones", async () => { + affectedTasks.task_instances = [loopPass("pass-1", 1), loopPass("pass-2", 2)]; + + render(, { + wrapper: Wrapper, + }); + + expect(await screen.findAllByText("body.work")).toHaveLength(2); + }); + + it.each([ + { expectedShown: true, version: { bundle_version: "bundle-1", id: "version-1", version_number: 1 } }, + { expectedShown: false, version: { bundle_version: "bundle-2", id: "version-2", version_number: 2 } }, + ])( + "shows the run-on-latest option for an execution row on an outdated version: $expectedShown", + async ({ expectedShown, version }) => { + const answered = new Set(); + + server.use( + http.get(`/api/v2/dags/${DAG_ID}/details`, () => { + answered.add("details"); + + return HttpResponse.json({ + bundle_version: "bundle-2", + latest_dag_version: { version_number: 2 }, + rerun_with_latest_version: null, + }); + }), + http.get(`/api/v2/dags/${DAG_ID}/dagRuns/${DAG_RUN_ID}`, () => { + answered.add("run"); + + return HttpResponse.json({ dag_versions: [version] }); + }), + ); + const execution = { + dag_id: DAG_ID, + dag_run_id: DAG_RUN_ID, + dag_version_id: version.id, + id: "execution", + in_loop: true, + map_index: -1, + region_index: 1, + task_id: "body.work", + }; + + render(, { + wrapper: Wrapper, + }); + + expect(await screen.findByRole("checkbox", { name: /keepTaskState/iu })).toBeVisible(); + await waitFor(() => expect(answered.size).toBe(2)); + await new Promise((resolve) => { + setTimeout(resolve, 50); + }); + expect(screen.queryByRole("checkbox", { name: /runOnLatestVersion/iu }) !== null).toBe(expectedShown); + }, + ); + + it.each([ + { + entry: "the task instance page", + ui: (member: TaskInstanceResponse) => , + }, + { + entry: "a bulk selection", + ui: (member: TaskInstanceResponse) => ( + + ), + }, + ])("offers the loop options for a loop member cleared from $entry", async ({ ui }) => { + render(ui(loopMember), { wrapper: Wrapper }); + + fireEvent.click(screen.getByRole("button", { name: /clear\.button/iu })); + + expect(await screen.findByRole("button", { name: /options\.laterIterations/iu })).toBeVisible(); + }); }); diff --git a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx index 8299596a423ec..c9c1de43558cf 100644 --- a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx +++ b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx @@ -16,20 +16,25 @@ * specific language governing permissions and limitations * under the License. */ -import { useEffect, useMemo, useState } from "react"; +import { useEffect, useState } from "react"; import { Button, Flex, useDisclosure } from "@chakra-ui/react"; import { useTranslation } from "react-i18next"; import { CgRedo } from "react-icons/cg"; import { useDagRunServiceGetDagRun, useDagServiceGetDagDetails } from "openapi/queries"; -import type { ClearTaskInstancesBody, TaskInstanceResponse } from "openapi/requests/types.gen"; +import type { + ClearTaskInstancesBody, + DagVersionResponse, + TaskInstanceResponse, +} from "openapi/requests/types.gen"; import { Checkbox, Modal, SegmentedControl } from "src/system-components"; import { ActionAccordion } from "src/components/ActionAccordion"; import { taskInstanceKey } from "src/components/ActionAccordion/columns"; import { useRerunWithLatestVersion } from "src/components/Clear/useRerunWithLatestVersion"; +import { ErrorAlert } from "src/components/ErrorAlert"; import Time from "src/components/Time"; import { @@ -38,16 +43,46 @@ import { useClearTaskInstanceDefaultOptions, } from "src/hooks/useUserSettings"; import { useClearTaskInstances } from "src/queries/useClearTaskInstances"; -import { useClearTaskInstancesDryRun } from "src/queries/useClearTaskInstancesDryRun"; +import { useClearTaskInstancesDryRuns } from "src/queries/useClearTaskInstancesDryRun"; import { isStatePending, useAutoRefresh } from "src/utils"; import ClearTaskInstanceConfirmationDialog from "./ClearTaskInstanceConfirmationDialog"; import { getRunOnLatestVersionState } from "./runOnLatestVersion"; +/** + * What the dialog needs to know about a task instance to clear it. A `TaskInstanceResponse` + * satisfies it, as does an Execution tab row once its loop membership is resolved. + */ +export type ClearTarget = { + readonly dag_id: string; + readonly dag_run_id: string; + readonly dag_version?: Pick | null; + readonly dag_version_id?: string | null; + readonly id: string; + readonly in_loop: boolean; + readonly logical_date?: string | null; + readonly map_index: number; + readonly note?: string | null; + readonly region_index?: number; + readonly task_id: string; +}; + +type RunGroup = { + readonly dagId: string; + readonly dagRunId: string; + readonly targets: Array; +}; + +type ClearRequest = { + dagId: string; + requestBody: ClearTaskInstancesBody; +}; + // Discriminated union: callers pass either `allMapped: true` together with -// `dagId`/`dagRunId`/`taskId` (clears every mapped TI of the task), or a full -// `taskInstance` (clears that single TI and reads its display fields). The -// two variants are mutually exclusive at the type level — no defensive +// `dagId`/`dagRunId`/`taskId` (clears every mapped TI of the task), a full +// `taskInstance` (clears that single TI and reads its display fields), or +// `taskInstances` (clears a selection, possibly spanning runs and Dags). The +// variants are mutually exclusive at the type level — no defensive // runtime fallback chains needed in the body. type Props = ( | { @@ -60,11 +95,36 @@ type Props = ( readonly allMapped?: false; readonly taskInstance: TaskInstanceResponse; } + | { + readonly allMapped?: false; + readonly taskInstances: Array; + } ) & { + readonly onCleared?: () => void; readonly onClose: () => void; readonly open: boolean; }; +const getRunGroups = (targets: Array): Array => { + const groups = new Map(); + + for (const target of targets) { + const key = JSON.stringify([target.dag_id, target.dag_run_id]); + const group = groups.get(key) ?? { dagId: target.dag_id, dagRunId: target.dag_run_id, targets: [] }; + + group.targets.push(target); + groups.set(key, group); + } + + return [...groups.values()]; +}; + +const LATER_ITERATIONS = "laterIterations"; +const WHOLE_EXPANSION = "wholeExpansion"; + +// An empty mapped expansion is a placeholder at region index -1. +const isMappedTarget = (target: ClearTarget) => target.map_index >= 0 || target.region_index === -1; + // react/destructuring-assignment expects every prop access via signature // destructure, but TypeScript's discriminated-union narrowing needs the union // kept whole on the parameter — otherwise the `props.allMapped` discriminator @@ -74,11 +134,18 @@ type Props = ( /* eslint-disable react/destructuring-assignment */ const ClearTaskInstanceDialog = (props: Props) => { const allMapped = props.allMapped === true; - const dagId = props.allMapped ? props.dagId : props.taskInstance.dag_id; - const dagRunId = props.allMapped ? props.dagRunId : props.taskInstance.dag_run_id; - const taskId = props.allMapped ? props.taskId : props.taskInstance.task_id; - const mapIndex: number | undefined = props.allMapped ? undefined : props.taskInstance.map_index; - const taskInstance: TaskInstanceResponse | undefined = props.allMapped ? undefined : props.taskInstance; + const allMappedTaskId = props.allMapped ? props.taskId : undefined; + const taskInstance: TaskInstanceResponse | undefined = + !props.allMapped && "taskInstance" in props ? props.taskInstance : undefined; + const targets: Array = props.allMapped + ? [] + : "taskInstances" in props + ? props.taskInstances + : [props.taskInstance]; + const groups: Array = props.allMapped + ? [{ dagId: props.dagId, dagRunId: props.dagRunId, targets: [] }] + : getRunGroups(targets); + const { onCleared } = props; const closeDialog = props.onClose; const openDialog = props.open; /* eslint-enable react/destructuring-assignment */ @@ -86,26 +153,38 @@ const ClearTaskInstanceDialog = (props: Props) => { const { onClose, onOpen, open } = useDisclosure(); const isDialogOpen = openDialog && !open; + // The Dag details and run versions behind the run-on-latest option are read for one run. + const singleRun = groups.length === 1; + const dagId = groups[0]?.dagId ?? ""; + const dagRunId = groups[0]?.dagRunId ?? ""; + const loopMode = targets.some((target) => target.in_loop); + const initialNote = targets.length === 1 ? (targets[0]?.note ?? null) : null; + const [clearTaskInstanceDefaultOptions] = useClearTaskInstanceDefaultOptions(); const [preventRunningTaskDefault] = useClearPreventRunningTaskDefault(); const [keepTaskStateDefault] = useClearKeepTaskStateDefault(); - const [selectedOptions, setSelectedOptions] = useState>(clearTaskInstanceDefaultOptions); + const initialOptions = loopMode + ? [...clearTaskInstanceDefaultOptions, LATER_ITERATIONS] + : clearTaskInstanceDefaultOptions; + const [selectedOptions, setSelectedOptions] = useState>(initialOptions); const onlyFailed = selectedOptions.includes("onlyFailed"); const past = selectedOptions.includes("past"); const future = selectedOptions.includes("future"); const upstream = selectedOptions.includes("upstream"); const downstream = selectedOptions.includes("downstream"); + const laterLoopIterations = selectedOptions.includes(LATER_ITERATIONS); + const wholeExpansions = selectedOptions.includes(WHOLE_EXPANSION); const [keepTaskState, setKeepTaskState] = useState(keepTaskStateDefault); const [preventRunningTask, setPreventRunningTask] = useState(preventRunningTaskDefault); - const [note, setNote] = useState(taskInstance?.note ?? null); + const [note, setNote] = useState(initialNote); useEffect(() => { if (openDialog) { - setNote(taskInstance?.note ?? null); + setNote(initialNote); } - }, [openDialog, taskInstance?.note]); + }, [openDialog, initialNote]); // Separate from the note effect above: this must only reset on open, not on every // note refetch, or a background note change while the dialog is open silently @@ -117,26 +196,39 @@ const ClearTaskInstanceDialog = (props: Props) => { }, [openDialog, keepTaskStateDefault]); const onCloseDialog = () => { - setNote(taskInstance?.note ?? null); + setNote(initialNote); setKeepTaskState(keepTaskStateDefault); closeDialog(); }; // Get current DAG's bundle version to compare with task instance's DAG version bundle version - const { data: dagDetails } = useDagServiceGetDagDetails({ - dagId, - }); + const { data: dagDetails } = useDagServiceGetDagDetails({ dagId }, undefined, { enabled: singleRun }); const { data: dagRun } = useDagRunServiceGetDagRun({ dagId, dagRunId }, undefined, { - enabled: openDialog, + enabled: openDialog && singleRun, }); + const runVersions = new Map((dagRun?.dag_versions ?? []).map((version) => [version.id, version])); + const selectedVersions = targets.flatMap((target) => { + const version = + target.dag_version ?? + (target.dag_version_id === null || target.dag_version_id === undefined + ? undefined + : runVersions.get(target.dag_version_id)); + + return version === undefined ? [] : [version]; + }); + const selectedVersion = + selectedVersions.find( + (version) => version.version_number !== dagDetails?.latest_dag_version?.version_number, + ) ?? selectedVersions[0]; + const { dagVersionsDiffer, runOnLatestVersionForced, shouldShowRunOnLatestOption } = getRunOnLatestVersionState({ latestBundleVersion: dagDetails?.bundle_version, latestDagVersionNumber: dagDetails?.latest_dag_version?.version_number, - selectedBundleVersion: taskInstance?.dag_version?.bundle_version, - selectedDagVersionNumber: taskInstance?.dag_version?.version_number, + selectedBundleVersion: selectedVersion?.bundle_version, + selectedDagVersionNumber: selectedVersion?.version_number, selectedVersionMissing: dagRun?.dag_versions.length === 0, }); @@ -146,17 +238,51 @@ const ClearTaskInstanceDialog = (props: Props) => { dagLevelConfig: dagDetails?.rerun_with_latest_version, fallback: dagVersionsDiffer, }); + const requestedRunOnLatestVersion = singleRun ? runOnLatestVersion : undefined; const { isPending, mutate } = useClearTaskInstances({ - dagId, dagRunId, - onSuccessConfirm: onCloseDialog, + onSuccessConfirm: () => { + onCleared?.(); + onCloseDialog(); + }, }); const refetchInterval = useAutoRefresh({ dagId }); - const { data } = useClearTaskInstancesDryRun({ - dagId, + // Loop members are addressed by exact execution id: their task id and map index alone + // do not say which pass is meant. Past and future cannot be combined with exact ids. + const previewRequests: Array = groups.map((group) => ({ + dagId: group.dagId, + requestBody: loopMode + ? { + dag_run_id: group.dagRunId, + include_downstream: downstream, + include_later_loop_iterations: laterLoopIterations, + include_upstream: upstream, + only_failed: onlyFailed, + run_on_latest_version: requestedRunOnLatestVersion, + task_instance_ids: group.targets.map((target) => target.id), + whole_expansion_ids: wholeExpansions + ? group.targets.filter(isMappedTarget).map((target) => target.id) + : [], + } + : { + dag_run_id: group.dagRunId, + include_downstream: downstream, + include_future: future, + include_past: past, + include_upstream: upstream, + only_failed: onlyFailed, + run_on_latest_version: requestedRunOnLatestVersion, + task_ids: + allMappedTaskId === undefined + ? group.targets.map((target): [string, number] => [target.task_id, target.map_index]) + : [allMappedTaskId], + }, + })); + + const { data, error: dryRunError } = useClearTaskInstancesDryRuns({ options: { enabled: openDialog, refetchInterval: (query) => @@ -165,23 +291,9 @@ const ClearTaskInstanceDialog = (props: Props) => { : false, refetchOnMount: "always", }, - requestBody: { - dag_run_id: dagRunId, - include_downstream: downstream, - include_future: future, - include_past: past, - include_upstream: upstream, - only_failed: onlyFailed, - run_on_latest_version: runOnLatestVersion, - task_ids: allMapped ? [taskId] : [[taskId, mapIndex as number]], - }, + requests: previewRequests, }); - const affectedTasks = data ?? { - task_instances: [], - total_entries: 0, - }; - // Tasks the user has unticked in the affected list; excluded from the clear. const [excludedKeys, setExcludedKeys] = useState>(new Set()); @@ -198,20 +310,65 @@ const ClearTaskInstanceDialog = (props: Props) => { return next; }); + // Loop clears expand on the server (later iterations, whole expansions), so the + // affected list cannot be turned back into an explicit selection there. + const hasExclusions = !loopMode && excludedKeys.size > 0; + const keptTaskInstances = data.task_instances.filter((ti) => !excludedKeys.has(taskInstanceKey(ti))); + // The dry run already resolved the full affected set, so on confirm we send those // task instances explicitly (minus the unticked ones) with the graph-expansion flags // off, instead of re-deriving them from the selected task + upstream/downstream. - const keptTaskInstances = useMemo( - () => affectedTasks.task_instances.filter((ti) => !excludedKeys.has(taskInstanceKey(ti))), - [affectedTasks.task_instances, excludedKeys], - ); + // The clear endpoint only targets one run per request, so group the kept instances by + // run and fire one run-scoped clear each. This honors per-run exclusions (e.g. keep task + // X in run 1 but drop it from run 2) that a single flat request cannot express. + const getKeptRequests = (): Array => { + const idsByRun = new Map }>(); + + for (const ti of keptTaskInstances) { + const key = JSON.stringify([ti.dag_id, ti.dag_run_id]); + const entry = idsByRun.get(key) ?? { dagId: ti.dag_id, dagRunId: ti.dag_run_id, ids: [] }; + + entry.ids.push(ti.id); + idsByRun.set(key, entry); + } - const checkedTaskIds = useMemo( - () => keptTaskInstances.map((ti) => (ti.map_index < 0 ? ti.task_id : [ti.task_id, ti.map_index])), - [keptTaskInstances], - ); + return [...idsByRun.values()].map(({ dagId: runDagId, dagRunId: runId, ids }) => ({ + dagId: runDagId, + requestBody: { + dag_run_id: runId, + include_downstream: false, + include_future: false, + include_past: false, + include_upstream: false, + only_failed: onlyFailed, + run_on_latest_version: requestedRunOnLatestVersion, + task_instance_ids: ids, + }, + })); + }; - const hasExclusions = excludedKeys.size > 0; + const confirmRequests = hasExclusions ? getKeptRequests() : previewRequests; + const gateRequest = confirmRequests.length === 1 ? confirmRequests[0] : undefined; + const confirmClear = () => { + const noteChanged = note !== initialNote; + + for (const { dagId: requestDagId, requestBody } of confirmRequests) { + mutate({ + dagId: requestDagId, + requestBody: { + ...requestBody, + dry_run: false, + note: noteChanged ? note : undefined, + ...(keepTaskState ? { keep_task_state: true } : {}), + ...(preventRunningTask ? { prevent_running_task: true } : {}), + }, + }); + } + onCloseDialog(); + }; + const hasNoLogicalDate = + targets.length > 0 && targets.every((target) => (target.logical_date ?? null) === null); + const hasMappedTargets = targets.some(isMappedTarget); return ( <> @@ -219,9 +376,9 @@ const ClearTaskInstanceDialog = (props: Props) => { footerActions={ <> @@ -263,33 +420,32 @@ const ClearTaskInstanceDialog = (props: Props) => { {allMapped ? translate("dags:runAndTaskActions.clearAllMapped.title") : translate("dags:runAndTaskActions.clear.title", { - type: translate("taskInstance_one"), + type: translate(taskInstance === undefined ? "taskInstance_other" : "taskInstance_one"), })} - : + {allMapped || taskInstance !== undefined ? ":" : undefined} {" "} - {allMapped ? ( - taskId - ) : ( - <> - {taskInstance?.task_display_name}