Add clear dag run to UI (#45039)

* Add clear dag run to UI

More code

* Small adjustments

* Update following code reviews

airflow/api_fastapi/core_api/datamodels/dag_run.py

@@ -54,7 +54,7 @@ class DAGRunClearBody(BaseModel):
 class DAGRunResponse(BaseModel):
     """DAG Run serializer for responses."""
 
-    dag_run_id: str | None = Field(validation_alias="run_id")
+    dag_run_id: str = Field(validation_alias="run_id")
     dag_id: str
     logical_date: datetime | None
     queued_at: datetime | None

airflow/api_fastapi/core_api/openapi/v1-generated.yaml

@@ -7097,9 +7097,7 @@ components:
     DAGRunResponse:
       properties:
         dag_run_id:
-          anyOf:
-          - type: string
-          - type: 'null'
+          type: string
           title: Dag Run Id
         dag_id:
           type: string

airflow/ui/openapi-gen/requests/schemas.gen.ts

@@ -1723,14 +1723,7 @@ export const $DAGRunPatchStates = {
 export const $DAGRunResponse = {
   properties: {
     dag_run_id: {
-      anyOf: [
-        {
-          type: "string",
-        },
-        {
-          type: "null",
-        },
-      ],
+      type: "string",
       title: "Dag Run Id",
     },
     dag_id: {

airflow/ui/openapi-gen/requests/types.gen.ts

@@ -402,7 +402,7 @@ export type DAGRunPatchStates = "queued" | "success" | "failed";
  * DAG Run serializer for responses.
  */
 export type DAGRunResponse = {
-  dag_run_id: string | null;
+  dag_run_id: string;
   dag_id: string;
   logical_date: string | null;
   queued_at: string | null;

airflow/ui/src/components/ClearRun/ClearRunButton.tsx

@@ -0,0 +1,84 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { Box, useDisclosure } from "@chakra-ui/react";
+import { useState } from "react";
+import { FiRefreshCw } from "react-icons/fi";
+
+import type { TaskInstanceCollectionResponse } from "openapi/requests/types.gen";
+import { Button } from "src/components/ui";
+import { useClearDagRun } from "src/queries/useClearRun";
+
+import ClearRunDialog from "./ClearRunDialog";
+
+type Props = {
+  readonly dagId: string;
+  readonly dagRunId: string;
+};
+
+const ClearRunButton = ({ dagId, dagRunId }: Props) => {
+  const { onClose, onOpen, open } = useDisclosure();
+
+  const [onlyFailed, setOnlyFailed] = useState(false);
+
+  const [affectedTasks, setAffectedTasks] =
+    useState<TaskInstanceCollectionResponse>({
+      task_instances: [],
+      total_entries: 0,
+    });
+
+  const { isPending, mutate } = useClearDagRun({
+    dagId,
+    dagRunId,
+    onSuccessConfirm: onClose,
+    onSuccessDryRun: setAffectedTasks,
+  });
+
+  return (
+    <Box>
+      <Button
+        onClick={() => {
+          onOpen();
+          mutate({
+            dagId,
+            dagRunId,
+            requestBody: { dry_run: true, only_failed: onlyFailed },
+          });
+        }}
+        variant="outline"
+      >
+        <FiRefreshCw height={5} width={5} />
+        Clear Run
+      </Button>
+
+      <ClearRunDialog
+        affectedTasks={affectedTasks}
+        dagId={dagId}
+        dagRunId={dagRunId}
+        isPending={isPending}
+        mutate={mutate}
+        onClose={onClose}
+        onlyFailed={onlyFailed}
+        open={open}
+        setOnlyFailed={setOnlyFailed}
+      />
+    </Box>
+  );
+};
+
+export default ClearRunButton;

airflow/ui/src/components/ClearRun/ClearRunDialog.tsx

@@ -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.
+ */
+import { Flex, Heading, VStack } from "@chakra-ui/react";
+import { FiRefreshCw } from "react-icons/fi";
+
+import type {
+  DAGRunClearBody,
+  TaskInstanceCollectionResponse,
+} from "openapi/requests/types.gen";
+import { Button, Dialog } from "src/components/ui";
+
+import SegmentedControl from "../ui/SegmentedControl";
+import ClearRunTasksAccordion from "./ClearRunTaskAccordion";
+
+type Props = {
+  readonly affectedTasks: TaskInstanceCollectionResponse;
+  readonly dagId: string;
+  readonly dagRunId: string;
+  readonly isPending: boolean;
+  readonly mutate: ({
+    dagId,
+    dagRunId,
+    requestBody,
+  }: {
+    dagId: string;
+    dagRunId: string;
+    requestBody: DAGRunClearBody;
+  }) => void;
+  readonly onClose: () => void;
+  readonly onlyFailed: boolean;
+  readonly open: boolean;
+  readonly setOnlyFailed: (value: boolean) => void;
+};
+
+const ClearRunDialog = ({
+  affectedTasks,
+  dagId,
+  dagRunId,
+  isPending,
+  mutate,
+  onClose,
+  onlyFailed,
+  open,
+  setOnlyFailed,
+}: Props) => {
+  const onChange = (value: string) => {
+    switch (value) {
+      case "existing_tasks":
+        setOnlyFailed(false);
+        mutate({
+          dagId,
+          dagRunId,
+          requestBody: { dry_run: true, only_failed: false },
+        });
+        break;
+      case "only_failed":
+        setOnlyFailed(true);
+        mutate({
+          dagId,
+          dagRunId,
+          requestBody: { dry_run: true, only_failed: true },
+        });
+        break;
+      default:
+        // TODO: Handle this `new_tasks` case
+        break;
+    }
+  };
+
+  return (
+    <Dialog.Root onOpenChange={onClose} open={open} size="xl">
+      <Dialog.Content backdrop>
+        <Dialog.Header>
+          <VStack align="start" gap={4}>
+            <Heading size="xl">Clear DagRun - {dagRunId} </Heading>
+          </VStack>
+        </Dialog.Header>
+
+        <Dialog.CloseTrigger />
+
+        <Dialog.Body width="full">
+          <Flex justifyContent="center">
+            <SegmentedControl
+              mb={3}
+              onValueChange={onChange}
+              options={[
+                { label: "Clear existing tasks", value: "existing_tasks" },
+                { label: "Clear only failed tasks", value: "only_failed" },
+                {
+                  disabled: true,
+                  label: "Queued up new tasks",
+                  value: "new_tasks",
+                },
+              ]}
+              value={onlyFailed ? "only_failed" : "existing_tasks"}
+            />
+          </Flex>
+          <ClearRunTasksAccordion affectedTasks={affectedTasks} />
+          <Flex justifyContent="end" mt={3}>
+            <Button
+              colorPalette="blue"
+              loading={isPending}
+              onClick={() => {
+                mutate({
+                  dagId,
+                  dagRunId,
+                  requestBody: { dry_run: false, only_failed: onlyFailed },
+                });
+              }}
+            >
+              <FiRefreshCw /> Confirm
+            </Button>
+          </Flex>
+        </Dialog.Body>
+      </Dialog.Content>
+    </Dialog.Root>
+  );
+};
+
+export default ClearRunDialog;

airflow/ui/src/components/ClearRun/ClearRunTaskAccordion.tsx

@@ -0,0 +1,107 @@
+/*!
+ * 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 { Box, Text } from "@chakra-ui/react";
+import { Link } from "@chakra-ui/react";
+import type { ColumnDef } from "@tanstack/react-table";
+import { Link as RouterLink } from "react-router-dom";
+
+import type {
+  TaskInstanceCollectionResponse,
+  TaskInstanceResponse,
+} from "openapi/requests/types.gen";
+import { DataTable } from "src/components/DataTable";
+import { Status } from "src/components/ui";
+
+import { Accordion } from "../ui";
+
+const columns: Array<ColumnDef<TaskInstanceResponse>> = [
+  {
+    accessorKey: "task_display_name",
+    cell: ({ row: { original } }) => (
+      <Link asChild color="fg.info" fontWeight="bold">
+        <RouterLink
+          to={`/dags/${original.dag_id}/runs/${original.dag_run_id}/tasks/${original.task_id}${original.map_index > -1 ? `?map_index=${original.map_index}` : ""}`}
+        >
+          {original.task_display_name}
+        </RouterLink>
+      </Link>
+    ),
+    enableSorting: false,
+    header: "Task ID",
+  },
+  {
+    accessorKey: "state",
+    cell: ({
+      row: {
+        original: { state },
+      },
+    }) => <Status state={state}>{state}</Status>,
+    enableSorting: false,
+    header: () => "State",
+  },
+  {
+    accessorKey: "map_index",
+    enableSorting: false,
+    header: "Map Index",
+  },
+
+  {
+    accessorKey: "dag_run_id",
+    enableSorting: false,
+    header: "Run Id",
+  },
+];
+
+type Props = {
+  readonly affectedTasks?: TaskInstanceCollectionResponse;
+};
+
+// Table is in memory, pagination and sorting are disabled.
+// TODO: Make a front-end only unconnected table component with client side ordering and pagination
+const ClearRunTasksAccordion = ({ affectedTasks }: Props) => (
+  <Accordion.Root collapsible variant="enclosed">
+    <Accordion.Item key="tasks" value="tasks">
+      <Accordion.ItemTrigger>
+        <Text fontWeight="bold">
+          Affected Tasks: {affectedTasks?.total_entries ?? 0}
+        </Text>
+      </Accordion.ItemTrigger>
+      <Accordion.ItemContent>
+        <Box maxH="400px" overflowY="scroll">
+          <DataTable
+            columns={columns}
+            data={affectedTasks?.task_instances ?? []}
+            displayMode="table"
+            initialState={{
+              pagination: {
+                pageIndex: 0,
+                pageSize: affectedTasks?.total_entries ?? 0,
+              },
+              sorting: [],
+            }}
+            modelName="Task Instance"
+            total={affectedTasks?.total_entries}
+          />
+        </Box>
+      </Accordion.ItemContent>
+    </Accordion.Item>
+  </Accordion.Root>
+);
+
+export default ClearRunTasksAccordion;

airflow/ui/src/components/ClearRun/index.tsx

@@ -0,0 +1,20 @@
+/*!
+ * 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.
+ */
+
+export { default } from "./ClearRunButton";