From 97cb24f0f57c4475862b855688af292f3a8b397a Mon Sep 17 00:00:00 2001
From: Karthikeyan Singaravelan <tir.karthi@gmail.com>
Date: Fri, 6 Dec 2024 20:50:49 +0530
Subject: [PATCH] Implement card view for tasks view of a dag (#44604)

* Initial commit for tasks tab.

* Remove unused columns for card view.

* Fetch latest dag run and then fetch corresponding task instances.

* Make font to be consistent.

* Pass dagRunsLimit to limit the number of dagruns.

* Add recent task instance plot to card.

* Fix merge conflicts and use Status component.

* Make task name bold
Change Last Run to Last Instance in header name
Use start date and wrap it with Time tag
Show try number only when greater than 1
Use pluralize for task

* Refactor tooltip to a separate component. Fix variable name casing.

* Fix duration when null.

* Use children prop to simplify TaskInstanceTooltip wrapping.

* Omit content to reuse tooltip type.
---
 .../ui/src/components/TaskInstanceTooltip.tsx |  62 ++++++++++
 .../src/pages/DagsList/Dag/Tasks/TaskCard.tsx |  88 ++++++++++++++
 .../DagsList/Dag/Tasks/TaskRecentRuns.tsx     |  76 ++++++++++++
 .../ui/src/pages/DagsList/Dag/Tasks/Tasks.tsx | 114 ++++++++++++++++++
 .../ui/src/pages/DagsList/Dag/Tasks/index.ts  |  20 +++
 airflow/ui/src/router.tsx                     |   3 +-
 6 files changed, 362 insertions(+), 1 deletion(-)
 create mode 100644 airflow/ui/src/components/TaskInstanceTooltip.tsx
 create mode 100644 airflow/ui/src/pages/DagsList/Dag/Tasks/TaskCard.tsx
 create mode 100644 airflow/ui/src/pages/DagsList/Dag/Tasks/TaskRecentRuns.tsx
 create mode 100644 airflow/ui/src/pages/DagsList/Dag/Tasks/Tasks.tsx
 create mode 100644 airflow/ui/src/pages/DagsList/Dag/Tasks/index.ts

diff --git a/airflow/ui/src/components/TaskInstanceTooltip.tsx b/airflow/ui/src/components/TaskInstanceTooltip.tsx
new file mode 100644
index 0000000000000..0b82567ae83e7
--- /dev/null
+++ b/airflow/ui/src/components/TaskInstanceTooltip.tsx
@@ -0,0 +1,62 @@
+/*!
+ * 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 type { TaskInstanceResponse } from "openapi/requests/types.gen";
+import Time from "src/components/Time";
+import { Tooltip, type TooltipProps } from "src/components/ui";
+
+type Props = {
+  readonly taskInstance: TaskInstanceResponse;
+} & Omit<TooltipProps, "content">;
+
+const TaskInstanceTooltip = ({ children, taskInstance }: Props) => (
+  <Tooltip
+    content={
+      <Box>
+        <Text>Run ID: {taskInstance.dag_run_id}</Text>
+        <Text>Logical Date: {taskInstance.logical_date}</Text>
+        <Text>
+          Start Date: <Time datetime={taskInstance.start_date} />
+        </Text>
+        <Text>
+          End Date: <Time datetime={taskInstance.end_date} />
+        </Text>
+        {taskInstance.try_number > 1 && (
+          <Text>Try Number: {taskInstance.try_number}</Text>
+        )}
+        <Text>Duration: {taskInstance.duration?.toFixed(2) ?? 0}s</Text>
+        <Text>State: {taskInstance.state}</Text>
+      </Box>
+    }
+    key={taskInstance.dag_run_id}
+    positioning={{
+      offset: {
+        crossAxis: 5,
+        mainAxis: 5,
+      },
+      placement: "bottom-start",
+    }}
+    showArrow
+  >
+    {children}
+  </Tooltip>
+);
+
+export default TaskInstanceTooltip;
diff --git a/airflow/ui/src/pages/DagsList/Dag/Tasks/TaskCard.tsx b/airflow/ui/src/pages/DagsList/Dag/Tasks/TaskCard.tsx
new file mode 100644
index 0000000000000..05726989a91d6
--- /dev/null
+++ b/airflow/ui/src/pages/DagsList/Dag/Tasks/TaskCard.tsx
@@ -0,0 +1,88 @@
+/*!
+ * 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 {
+  Heading,
+  VStack,
+  HStack,
+  Box,
+  SimpleGrid,
+  Text,
+} from "@chakra-ui/react";
+
+import type {
+  TaskResponse,
+  TaskInstanceResponse,
+} from "openapi/requests/types.gen";
+import TaskInstanceTooltip from "src/components/TaskInstanceTooltip";
+import Time from "src/components/Time";
+import { Status } from "src/components/ui";
+
+import { TaskRecentRuns } from "./TaskRecentRuns.tsx";
+
+type Props = {
+  readonly task: TaskResponse;
+  readonly taskInstances: Array<TaskInstanceResponse>;
+};
+
+export const TaskCard = ({ task, taskInstances }: Props) => (
+  <Box
+    borderColor="border.emphasized"
+    borderRadius={8}
+    borderWidth={1}
+    overflow="hidden"
+  >
+    <Text bg="bg.info" color="fg.info" fontWeight="bold" p={2}>
+      {task.task_display_name ?? task.task_id}
+      {task.is_mapped ? "[]" : undefined}
+    </Text>
+    <SimpleGrid columns={4} gap={4} height={20} px={3} py={2}>
+      <VStack align="flex-start" gap={1}>
+        <Heading color="fg.muted" fontSize="xs">
+          Operator
+        </Heading>
+        <Text fontSize="sm">{task.operator_name}</Text>
+      </VStack>
+      <VStack align="flex-start" gap={1}>
+        <Heading color="fg.muted" fontSize="xs">
+          Trigger Rule
+        </Heading>
+        <Text fontSize="sm">{task.trigger_rule}</Text>
+      </VStack>
+      <VStack align="flex-start" gap={1}>
+        <Heading color="fg.muted" fontSize="xs">
+          Last Instance
+        </Heading>
+        {taskInstances[0] ? (
+          <TaskInstanceTooltip taskInstance={taskInstances[0]}>
+            <HStack fontSize="sm">
+              <Time datetime={taskInstances[0].start_date} />
+              {taskInstances[0].state === null ? undefined : (
+                <Status state={taskInstances[0].state}>
+                  {taskInstances[0].state}
+                </Status>
+              )}
+            </HStack>
+          </TaskInstanceTooltip>
+        ) : undefined}
+      </VStack>
+      {/* TODO: Handled mapped tasks to not plot each map index as a task instance */}
+      {!task.is_mapped && <TaskRecentRuns taskInstances={taskInstances} />}
+    </SimpleGrid>
+  </Box>
+);
diff --git a/airflow/ui/src/pages/DagsList/Dag/Tasks/TaskRecentRuns.tsx b/airflow/ui/src/pages/DagsList/Dag/Tasks/TaskRecentRuns.tsx
new file mode 100644
index 0000000000000..a69d7d72790e2
--- /dev/null
+++ b/airflow/ui/src/pages/DagsList/Dag/Tasks/TaskRecentRuns.tsx
@@ -0,0 +1,76 @@
+/*!
+ * 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 } from "@chakra-ui/react";
+import { Flex } from "@chakra-ui/react";
+import dayjs from "dayjs";
+import duration from "dayjs/plugin/duration";
+
+import type { TaskInstanceResponse } from "openapi/requests/types.gen";
+import TaskInstanceTooltip from "src/components/TaskInstanceTooltip";
+import { stateColor } from "src/utils/stateColor";
+
+dayjs.extend(duration);
+
+const BAR_HEIGHT = 60;
+
+export const TaskRecentRuns = ({
+  taskInstances,
+}: {
+  readonly taskInstances: Array<TaskInstanceResponse>;
+}) => {
+  if (!taskInstances.length) {
+    return undefined;
+  }
+
+  const taskInstancesWithDuration = taskInstances.map((taskInstance) => ({
+    ...taskInstance,
+    duration:
+      dayjs
+        .duration(dayjs(taskInstance.end_date).diff(taskInstance.start_date))
+        .asSeconds() || 0,
+  }));
+
+  const max = Math.max.apply(
+    undefined,
+    taskInstancesWithDuration.map((taskInstance) => taskInstance.duration),
+  );
+
+  return (
+    <Flex alignItems="flex-end" flexDirection="row-reverse">
+      {taskInstancesWithDuration.map((taskInstance) =>
+        taskInstance.state === null ? undefined : (
+          <TaskInstanceTooltip
+            key={taskInstance.dag_run_id}
+            taskInstance={taskInstance}
+          >
+            <Box p={1}>
+              <Box
+                bg={stateColor[taskInstance.state]}
+                borderRadius="4px"
+                height={`${(taskInstance.duration / max) * BAR_HEIGHT}px`}
+                minHeight={1}
+                width="4px"
+              />
+            </Box>
+          </TaskInstanceTooltip>
+        ),
+      )}
+    </Flex>
+  );
+};
diff --git a/airflow/ui/src/pages/DagsList/Dag/Tasks/Tasks.tsx b/airflow/ui/src/pages/DagsList/Dag/Tasks/Tasks.tsx
new file mode 100644
index 0000000000000..4b844ad2b64f9
--- /dev/null
+++ b/airflow/ui/src/pages/DagsList/Dag/Tasks/Tasks.tsx
@@ -0,0 +1,114 @@
+/*!
+ * 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 { Heading, Skeleton, Box } from "@chakra-ui/react";
+import { useParams } from "react-router-dom";
+
+import {
+  useTaskServiceGetTasks,
+  useTaskInstanceServiceGetTaskInstances,
+  useDagsServiceRecentDagRuns,
+} from "openapi/queries";
+import type {
+  TaskResponse,
+  TaskInstanceResponse,
+} from "openapi/requests/types.gen";
+import { DataTable } from "src/components/DataTable";
+import type { CardDef } from "src/components/DataTable/types";
+import { ErrorAlert } from "src/components/ErrorAlert";
+import { pluralize } from "src/utils";
+
+import { TaskCard } from "./TaskCard";
+
+const cardDef = (
+  taskInstances?: Array<TaskInstanceResponse>,
+): CardDef<TaskResponse> => ({
+  card: ({ row }) => (
+    <TaskCard
+      task={row}
+      taskInstances={
+        taskInstances
+          ? taskInstances.filter(
+              (instance: TaskInstanceResponse) =>
+                instance.task_id === row.task_id,
+            )
+          : []
+      }
+    />
+  ),
+  meta: {
+    customSkeleton: <Skeleton height="120px" width="100%" />,
+  },
+});
+
+export const Tasks = () => {
+  const { dagId } = useParams();
+  const {
+    data,
+    error: tasksError,
+    isFetching,
+    isLoading,
+  } = useTaskServiceGetTasks({
+    dagId: dagId ?? "",
+  });
+
+  // TODO: Replace dagIdPattern with dagId once supported for better matching
+  const { data: runsData } = useDagsServiceRecentDagRuns(
+    { dagIdPattern: dagId ?? "", dagRunsLimit: 14 },
+    undefined,
+    {
+      enabled: Boolean(dagId),
+    },
+  );
+
+  const runs =
+    runsData?.dags.find((dagWithRuns) => dagWithRuns.dag_id === dagId)
+      ?.latest_dag_runs ?? [];
+
+  // TODO: Revisit this endpoint since only 100 task instances are returned and
+  // only duration is calculated with other attributes unused.
+  const { data: taskInstancesResponse } =
+    useTaskInstanceServiceGetTaskInstances(
+      {
+        dagId: dagId ?? "",
+        dagRunId: "~",
+        logicalDateGte: runs.at(-1)?.logical_date ?? "",
+      },
+      undefined,
+      { enabled: Boolean(runs[0]?.dag_run_id) },
+    );
+
+  return (
+    <Box>
+      <ErrorAlert error={tasksError} />
+      <Heading my={1} size="md">
+        {pluralize("Task", data ? data.total_entries : 0)}
+      </Heading>
+      <DataTable
+        cardDef={cardDef(taskInstancesResponse?.task_instances.reverse())}
+        columns={[]}
+        data={data ? data.tasks : []}
+        displayMode="card"
+        isFetching={isFetching}
+        isLoading={isLoading}
+        modelName="Task"
+        total={data ? data.total_entries : 0} // Todo : Disable pagination?
+      />
+    </Box>
+  );
+};
diff --git a/airflow/ui/src/pages/DagsList/Dag/Tasks/index.ts b/airflow/ui/src/pages/DagsList/Dag/Tasks/index.ts
new file mode 100644
index 0000000000000..c5e32f6ebc611
--- /dev/null
+++ b/airflow/ui/src/pages/DagsList/Dag/Tasks/index.ts
@@ -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 { Tasks } from "./Tasks";
diff --git a/airflow/ui/src/router.tsx b/airflow/ui/src/router.tsx
index 6c5081c35e2b4..95b881e019e75 100644
--- a/airflow/ui/src/router.tsx
+++ b/airflow/ui/src/router.tsx
@@ -24,6 +24,7 @@ import { Dag } from "src/pages/DagsList/Dag";
 import { Code } from "src/pages/DagsList/Dag/Code";
 import { Overview } from "src/pages/DagsList/Dag/Overview";
 import { Runs } from "src/pages/DagsList/Dag/Runs";
+import { Tasks } from "src/pages/DagsList/Dag/Tasks";
 import { Run } from "src/pages/DagsList/Run";
 import { Dashboard } from "src/pages/Dashboard";
 import { ErrorPage } from "src/pages/Error";
@@ -49,7 +50,7 @@ export const router = createBrowserRouter(
           children: [
             { element: <Overview />, index: true },
             { element: <Runs />, path: "runs" },
-            { element: <div>Tasks</div>, path: "tasks" },
+            { element: <Tasks />, path: "tasks" },
             { element: <Events />, path: "events" },
             { element: <Code />, path: "code" },
           ],