From 4fa4d9b159b3e1461fd390c5f115ca797193c697 Mon Sep 17 00:00:00 2001 From: Pierre Jeambrun Date: Thu, 19 Dec 2024 22:15:27 +0800 Subject: [PATCH] Add clear dag run to UI (#45039) * Add clear dag run to UI More code * Small adjustments * Update following code reviews --- .../core_api/datamodels/dag_run.py | 2 +- .../core_api/openapi/v1-generated.yaml | 4 +- .../ui/openapi-gen/requests/schemas.gen.ts | 9 +- airflow/ui/openapi-gen/requests/types.gen.ts | 2 +- .../components/ClearRun/ClearRunButton.tsx | 84 +++++++++++ .../components/ClearRun/ClearRunDialog.tsx | 136 ++++++++++++++++++ .../ClearRun/ClearRunTaskAccordion.tsx | 107 ++++++++++++++ airflow/ui/src/components/ClearRun/index.tsx | 20 +++ .../ui/src/components/ui/SegmentedControl.tsx | 61 ++++++++ airflow/ui/src/pages/Run/Header.tsx | 4 + airflow/ui/src/queries/useClearRun.ts | 85 +++++++++++ 11 files changed, 501 insertions(+), 13 deletions(-) create mode 100644 airflow/ui/src/components/ClearRun/ClearRunButton.tsx create mode 100644 airflow/ui/src/components/ClearRun/ClearRunDialog.tsx create mode 100644 airflow/ui/src/components/ClearRun/ClearRunTaskAccordion.tsx create mode 100644 airflow/ui/src/components/ClearRun/index.tsx create mode 100644 airflow/ui/src/components/ui/SegmentedControl.tsx create mode 100644 airflow/ui/src/queries/useClearRun.ts diff --git a/airflow/api_fastapi/core_api/datamodels/dag_run.py b/airflow/api_fastapi/core_api/datamodels/dag_run.py index c24b42b109d8ba..48c92d2a83cb00 100644 --- a/airflow/api_fastapi/core_api/datamodels/dag_run.py +++ b/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 diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 980b045e34d06b..cb1b88de54fc2d 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/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 diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index defc76ec7345f7..501f76d0deaed0 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/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: { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 9aeda695c3554f..f7f0c32a2bc1da 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/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; diff --git a/airflow/ui/src/components/ClearRun/ClearRunButton.tsx b/airflow/ui/src/components/ClearRun/ClearRunButton.tsx new file mode 100644 index 00000000000000..7c3060765ccf84 --- /dev/null +++ b/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({ + task_instances: [], + total_entries: 0, + }); + + const { isPending, mutate } = useClearDagRun({ + dagId, + dagRunId, + onSuccessConfirm: onClose, + onSuccessDryRun: setAffectedTasks, + }); + + return ( + + + + + + ); +}; + +export default ClearRunButton; diff --git a/airflow/ui/src/components/ClearRun/ClearRunDialog.tsx b/airflow/ui/src/components/ClearRun/ClearRunDialog.tsx new file mode 100644 index 00000000000000..1bab63727499b7 --- /dev/null +++ b/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 ( + + + + + Clear DagRun - {dagRunId} + + + + + + + + + + + + + + + + + ); +}; + +export default ClearRunDialog; diff --git a/airflow/ui/src/components/ClearRun/ClearRunTaskAccordion.tsx b/airflow/ui/src/components/ClearRun/ClearRunTaskAccordion.tsx new file mode 100644 index 00000000000000..c579fa4d6fa78b --- /dev/null +++ b/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> = [ + { + accessorKey: "task_display_name", + cell: ({ row: { original } }) => ( + + -1 ? `?map_index=${original.map_index}` : ""}`} + > + {original.task_display_name} + + + ), + enableSorting: false, + header: "Task ID", + }, + { + accessorKey: "state", + cell: ({ + row: { + original: { state }, + }, + }) => {state}, + 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) => ( + + + + + Affected Tasks: {affectedTasks?.total_entries ?? 0} + + + + + + + + + +); + +export default ClearRunTasksAccordion; diff --git a/airflow/ui/src/components/ClearRun/index.tsx b/airflow/ui/src/components/ClearRun/index.tsx new file mode 100644 index 00000000000000..1e9362a626484d --- /dev/null +++ b/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"; diff --git a/airflow/ui/src/components/ui/SegmentedControl.tsx b/airflow/ui/src/components/ui/SegmentedControl.tsx new file mode 100644 index 00000000000000..8c972eb35edca5 --- /dev/null +++ b/airflow/ui/src/components/ui/SegmentedControl.tsx @@ -0,0 +1,61 @@ +/*! + * 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 { Tabs, For, type TabsRootProps } from "@chakra-ui/react"; + +type Option = { + disabled?: boolean; + label: string; + value: string; +}; + +type SegmentedControlProps = { + readonly onValueChange: (value: string) => void; + readonly options: Array