Add icons to task states (apache#46028)

* Add icons to task states

* Move stateColor to theme

* Fix graph

airflow/ui/src/components/DagRunInfo.tsx

@@ -74,7 +74,7 @@ const DagRunInfo = ({
     >
       <Box>
         <Time datetime={dataIntervalStart} mr={2} showTooltip={false} />
-        {state === undefined ? undefined : <Status state={state}>{state}</Status>}
+        {state === undefined ? undefined : <Status state={state} />}
       </Box>
     </Tooltip>
   ) : undefined;

airflow/ui/src/components/MetricsBadge.tsx

@@ -17,15 +17,16 @@
  * under the License.
  */
 import { Badge, type BadgeProps } from "@chakra-ui/react";
+import type { ReactNode } from "react";
 
 type MetricBadgeProps = {
-  readonly backgroundColor: string;
-  readonly color?: string;
+  readonly icon?: ReactNode;
   readonly runs?: number;
 } & BadgeProps;
 
-export const MetricsBadge = ({ backgroundColor, color = "fg.inverted", runs }: MetricBadgeProps) => (
-  <Badge bg={backgroundColor} borderRadius={15} color={color} minWidth={10} mr={1} px={4} py={1} size="md">
+export const MetricsBadge = ({ icon, runs, ...rest }: MetricBadgeProps) => (
+  <Badge {...rest} borderRadius={15} minWidth={10} mr={1} px={4} py={1} size="md" variant="solid">
+    {icon}
     {runs}
   </Badge>
 );

airflow/ui/src/components/StateIcon.tsx

@@ -0,0 +1,70 @@
+/*!
+ * 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 type { IconBaseProps } from "react-icons";
+import {
+  FiActivity,
+  FiAlertCircle,
+  FiAlertOctagon,
+  FiCalendar,
+  FiCheckCircle,
+  FiCircle,
+  FiRepeat,
+  FiSkipForward,
+  FiSlash,
+  FiWatch,
+} from "react-icons/fi";
+import { LuCalendarSync, LuRedo2 } from "react-icons/lu";
+import { PiQueue } from "react-icons/pi";
+
+import type { TaskInstanceState } from "openapi/requests/types.gen";
+
+type Props = {
+  readonly state?: TaskInstanceState | null;
+} & IconBaseProps;
+
+export const StateIcon = ({ state, ...rest }: Props) => {
+  switch (state) {
+    case "deferred":
+      return <FiWatch {...rest} />;
+    case "failed":
+      return <FiAlertOctagon {...rest} />;
+    case "queued":
+      return <PiQueue {...rest} />;
+    case "removed":
+      return <FiSlash {...rest} />;
+    case "restarting":
+      return <FiRepeat {...rest} />;
+    case "running":
+      return <FiActivity {...rest} />;
+    case "scheduled":
+      return <FiCalendar {...rest} />;
+    case "skipped":
+      return <FiSkipForward {...rest} />;
+    case "success":
+      return <FiCheckCircle {...rest} />;
+    case "up_for_reschedule":
+      return <LuCalendarSync {...rest} />;
+    case "up_for_retry":
+      return <LuRedo2 {...rest} />;
+    case "upstream_failed":
+      return <FiAlertCircle {...rest} />;
+    default:
+      return <FiCircle {...rest} />;
+  }
+};

airflow/ui/src/components/TaskInstanceTooltip.tsx

@@ -38,6 +38,7 @@ const TaskInstanceTooltip = ({ children, positioning, taskInstance, ...rest }: P
       {...rest}
       content={
         <Box>
+          <Text>State: {taskInstance.state}</Text>
           {"dag_run_id" in taskInstance ? <Text>Run ID: {taskInstance.dag_run_id}</Text> : undefined}
           <Text>
             Start Date: <Time datetime={taskInstance.start_date} />

airflow/ui/src/components/ui/Status.tsx

@@ -16,26 +16,24 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Status as ChakraStatus } from "@chakra-ui/react";
+import { Status as ChakraStatus, chakra } from "@chakra-ui/react";
 import * as React from "react";
 
 import type { DagRunState, TaskInstanceState } from "openapi/requests/types.gen";
-import { stateColor } from "src/utils/stateColor";
+
+import { StateIcon } from "../StateIcon";
 
 type StatusValue = DagRunState | TaskInstanceState | null;
 
 export type StatusProps = {
-  state: StatusValue;
+  state?: StatusValue;
 } & ChakraStatus.RootProps;
 
-export const Status = React.forwardRef<HTMLDivElement, StatusProps>(({ children, state, ...rest }, ref) => {
-  // "null" is actually a string on stateColor
-  const colorPalette = stateColor[state ?? "null"];
-
-  return (
-    <ChakraStatus.Root ref={ref} {...rest}>
-      <ChakraStatus.Indicator bg={colorPalette} />
-      {children}
-    </ChakraStatus.Root>
-  );
-});
+export const Status = React.forwardRef<HTMLDivElement, StatusProps>(({ children, state, ...rest }, ref) => (
+  <ChakraStatus.Root ref={ref} {...rest}>
+    <chakra.span color={`${state}.solid`}>
+      <StateIcon state={state} />
+    </chakra.span>
+    {children}
+  </ChakraStatus.Root>
+));

airflow/ui/src/layouts/Details/Graph/Graph.tsx

@@ -24,7 +24,6 @@ import { useParams } from "react-router-dom";
 import { useGridServiceGridData, useStructureServiceStructureData } from "openapi/queries";
 import { useColorMode } from "src/context/colorMode";
 import { useOpenGroups } from "src/context/openGroups";
-import { stateColor } from "src/utils/stateColor";
 
 import Edge from "./Edge";
 import { JoinNode } from "./JoinNode";
@@ -42,7 +41,7 @@ const nodeColor = (
   }
 
   if (taskInstance?.state !== undefined && !isOpen) {
-    return stateColor[taskInstance.state ?? "null"];
+    return `var(--chakra-colors-${taskInstance.state}-solid)`;
   }
 
   if (isOpen && depth !== undefined && depth % 2 === 0) {

airflow/ui/src/layouts/Details/Graph/TaskNode.tsx

@@ -24,7 +24,6 @@ import TaskInstanceTooltip from "src/components/TaskInstanceTooltip";
 import { Status } from "src/components/ui";
 import { useOpenGroups } from "src/context/openGroups";
 import { pluralize } from "src/utils";
-import { stateColor } from "src/utils/stateColor";
 
 import { NodeWrapper } from "./NodeWrapper";
 import { TaskName } from "./TaskName";
@@ -67,9 +66,7 @@ export const TaskNode = ({
           <Flex
             // Alternate background color for nested open groups
             bg={isOpen && depth !== undefined && depth % 2 === 0 ? "bg.muted" : "bg"}
-            borderColor={
-              taskInstance?.state === undefined ? "border" : stateColor[taskInstance.state ?? "null"]
-            }
+            borderColor={taskInstance?.state ? `${taskInstance.state}.solid` : "border"}
             borderRadius={5}
             borderWidth={isSelected ? 6 : 2}
             height={`${height}px`}
@@ -87,12 +84,12 @@ export const TaskNode = ({
                 label={label}
                 setupTeardownType={setupTeardownType}
               />
-              <Text color="fg.muted" fontSize="xs" mb={-1} mt={2} textTransform="capitalize">
+              <Text color="fg.muted" fontSize="sm" mt={1} textTransform="capitalize">
                 {isGroup ? "Task Group" : operator}
               </Text>
               {taskInstance === undefined ? undefined : (
                 <HStack>
-                  <Status fontSize="xs" state={taskInstance.state}>
+                  <Status fontSize="sm" state={taskInstance.state}>
                     {taskInstance.state}
                   </Status>
                   {taskInstance.try_number > 1 ? <MdRefresh /> : undefined}

airflow/ui/src/main.tsx

@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { ChakraProvider, defaultSystem } from "@chakra-ui/react";
+import { ChakraProvider } from "@chakra-ui/react";
 import { QueryClientProvider } from "@tanstack/react-query";
 import axios, { type AxiosError } from "axios";
 import { StrictMode } from "react";
@@ -28,6 +28,7 @@ import { TimezoneProvider } from "src/context/timezone";
 import { router } from "src/router";
 
 import { queryClient } from "./queryClient";
+import { system } from "./theme";
 
 // redirect to login page if the API responds with unauthorized or forbidden errors
 axios.interceptors.response.use(
@@ -46,7 +47,7 @@ axios.interceptors.response.use(
 
 createRoot(document.querySelector("#root") as HTMLDivElement).render(
   <StrictMode>
-    <ChakraProvider value={defaultSystem}>
+    <ChakraProvider value={system}>
       <ColorModeProvider>
         <QueryClientProvider client={queryClient}>
           <TimezoneProvider>

airflow/ui/src/pages/Dag/Overview/Overview.tsx

@@ -24,7 +24,6 @@ import { useParams } from "react-router-dom";
 import { useDagRunServiceGetDagRuns, useTaskInstanceServiceGetTaskInstances } from "openapi/queries";
 import TimeRangeSelector from "src/components/TimeRangeSelector";
 import { TrendCountButton } from "src/components/TrendCountButton";
-import { stateColor } from "src/utils/stateColor";
 
 const defaultHour = "168";
 
@@ -64,7 +63,7 @@ export const Overview = () => {
       </Box>
       <HStack>
         <TrendCountButton
-          colorPalette={stateColor.failed}
+          colorPalette="failed"
           count={failedTasks?.total_entries ?? 0}
           endDate={endDate}
           events={(failedTasks?.task_instances ?? []).map((ti) => ({
@@ -79,7 +78,7 @@ export const Overview = () => {
           startDate={startDate}
         />
         <TrendCountButton
-          colorPalette={stateColor.failed}
+          colorPalette="failed"
           count={failedRuns?.total_entries ?? 0}
           endDate={endDate}
           events={(failedRuns?.dag_runs ?? []).map((dr) => ({

airflow/ui/src/pages/Dag/Tasks/TaskCard.tsx

@@ -63,9 +63,7 @@ export const TaskCard = ({ dagId, task, taskInstances }: Props) => (
             <Link asChild color="fg.info" fontSize="sm">
               <RouterLink to={getTaskInstanceLink(taskInstances[0])}>
                 <Time datetime={taskInstances[0].start_date} />
-                {taskInstances[0].state === null ? undefined : (
-                  <Status state={taskInstances[0].state}>{taskInstances[0].state}</Status>
-                )}
+                {taskInstances[0].state === null ? undefined : <Status state={taskInstances[0].state} />}
               </RouterLink>
             </Link>
           </TaskInstanceTooltip>

airflow/ui/src/pages/Dag/Tasks/TaskRecentRuns.tsx

@@ -24,7 +24,6 @@ import { Link } from "react-router-dom";
 import type { TaskInstanceResponse } from "openapi/requests/types.gen";
 import TaskInstanceTooltip from "src/components/TaskInstanceTooltip";
 import { getTaskInstanceLink } from "src/utils/links";
-import { stateColor } from "src/utils/stateColor";
 
 dayjs.extend(duration);
 
@@ -58,7 +57,7 @@ export const TaskRecentRuns = ({
             <Link to={getTaskInstanceLink(taskInstance)}>
               <Box p={1}>
                 <Box
-                  bg={stateColor[taskInstance.state]}
+                  bg={`${taskInstance.state}.solid`}
                   borderRadius="4px"
                   height={`${(taskInstance.duration / max) * BAR_HEIGHT}px`}
                   minHeight={1}

airflow/ui/src/pages/DagsList/RecentRuns.tsx

@@ -24,7 +24,6 @@ import { Link } from "react-router-dom";
 import type { DAGWithLatestDagRunsResponse } from "openapi/requests/types.gen";
 import Time from "src/components/Time";
 import { Tooltip } from "src/components/ui";
-import { stateColor } from "src/utils/stateColor";
 
 dayjs.extend(duration);
 
@@ -74,7 +73,7 @@ export const RecentRuns = ({
           <Link to={`/dags/${run.dag_id}/runs/${run.dag_run_id}/`}>
             <Box px={1}>
               <Box
-                bg={stateColor[run.state]}
+                bg={`${run.state}.solid`}
                 borderRadius="4px"
                 height={`${(run.duration / max) * BAR_HEIGHT}px`}
                 minHeight={1}

airflow/ui/src/pages/Dashboard/Health/HealthTag.tsx

@@ -16,8 +16,9 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Skeleton, TagLabel, Text } from "@chakra-ui/react";
+import { Skeleton, TagLabel, Text, HStack } from "@chakra-ui/react";
 
+import { StateIcon } from "src/components/StateIcon";
 import Time from "src/components/Time";
 import { Tag, Tooltip } from "src/components/ui";
 import { capitalize } from "src/utils";
@@ -37,6 +38,8 @@ export const HealthTag = ({
     return <Skeleton borderRadius="full" height={8} width={24} />;
   }
 
+  const state = status === "healthy" ? "success" : "failed";
+
   return (
     <Tooltip
       content={
@@ -49,8 +52,11 @@ export const HealthTag = ({
       }
       disabled={!Boolean(latestHeartbeat)}
     >
-      <Tag borderRadius="full" colorPalette={status === "healthy" ? "green" : "red"} size="lg">
-        <TagLabel>{title}</TagLabel>
+      <Tag borderRadius="full" colorPalette={state} size="lg">
+        <HStack>
+          <StateIcon state={state} />
+          <TagLabel>{title}</TagLabel>
+        </HStack>
       </Tag>
     </Tooltip>
   );

airflow/ui/src/pages/Dashboard/HistoricalMetrics/AssetEvents.tsx

@@ -51,7 +51,7 @@ export const AssetEvents = ({ assetSortBy, endDate, setAssetSortBy, startDate }:
     <Box borderRadius={5} borderWidth={1} ml={2} pb={2}>
       <Flex justify="space-between" mr={1} mt={0} pl={3} pt={1}>
         <HStack>
-          <MetricsBadge backgroundColor="blue.solid" runs={isLoading ? 0 : data?.total_entries} />
+          <MetricsBadge colorPalette="blue" runs={isLoading ? 0 : data?.total_entries} />
           <Heading marginEnd="auto" size="md">
             Asset Events
           </Heading>

airflow/ui/src/pages/Dashboard/HistoricalMetrics/DagRunMetrics.tsx

@@ -33,7 +33,7 @@ const DAGRUN_STATES: Array<keyof DAGRunStates> = ["queued", "running", "success"
 export const DagRunMetrics = ({ dagRunStates, total }: DagRunMetricsProps) => (
   <Box borderRadius={5} borderWidth={1} p={2}>
     <HStack mb={4}>
-      <MetricsBadge backgroundColor="blue.solid" runs={total} />
+      <MetricsBadge colorPalette="blue" runs={total} />
       <Heading size="md">Dag Runs</Heading>
     </HStack>
     {DAGRUN_STATES.map((state) => (

airflow/ui/src/pages/Dashboard/HistoricalMetrics/MetricSection.tsx

@@ -20,7 +20,6 @@ import { Box, Flex, HStack, VStack, Text } from "@chakra-ui/react";
 
 import { MetricsBadge } from "src/components/MetricsBadge";
 import { capitalize } from "src/utils";
-import { stateColor } from "src/utils/stateColor";
 
 const BAR_WIDTH = 100;
 const BAR_HEIGHT = 5;
@@ -42,14 +41,14 @@ export const MetricSection = ({ runs, state, total }: MetricSectionProps) => {
     <VStack align="left" gap={1} mb={4} ml={0} pl={0}>
       <Flex justify="space-between">
         <HStack>
-          <MetricsBadge backgroundColor={stateColor[state as keyof typeof stateColor]} runs={runs} />
+          <MetricsBadge colorPalette={state} runs={runs} />
           <Text> {capitalize(state)} </Text>
         </HStack>
         <Text color="fg.muted"> {statePercent}% </Text>
       </Flex>
       <HStack gap={0} mt={2}>
         <Box
-          bg={stateColor[state as keyof typeof stateColor]}
+          bg={`${state}.solid`}
           borderLeftRadius={5}
           height={`${BAR_HEIGHT}px`}
           minHeight={2}

airflow/ui/src/pages/Dashboard/HistoricalMetrics/TaskInstanceMetrics.tsx

@@ -39,7 +39,7 @@ const TASK_STATES: Array<keyof TaskInstanceStateCount> = [
 export const TaskInstanceMetrics = ({ taskInstanceStates, total }: TaskInstanceMetricsProps) => (
   <Box borderRadius={5} borderWidth={1} mt={2} p={2}>
     <HStack mb={4}>
-      <MetricsBadge backgroundColor="blue.solid" runs={total} />
+      <MetricsBadge colorPalette="blue" runs={total} />
       <Heading size="md">Task Instances</Heading>
     </HStack>
 

airflow/ui/src/pages/Dashboard/Stats/DAGImportErrors.tsx

@@ -22,7 +22,6 @@ import { FiChevronRight } from "react-icons/fi";
 import { useImportErrorServiceGetImportErrors } from "openapi/queries";
 import { ErrorAlert } from "src/components/ErrorAlert";
 import { MetricsBadge } from "src/components/MetricsBadge";
-import { stateColor } from "src/utils/stateColor";
 
 import { DAGImportErrorsModal } from "./DAGImportErrorsModal";
 
@@ -49,7 +48,7 @@ export const DAGImportErrors = () => {
           onClick={onOpen}
           variant="outline"
         >
-          <MetricsBadge backgroundColor={stateColor.failed} runs={importErrorsCount} />
+          <MetricsBadge colorPalette="failed" runs={importErrorsCount} />
           <Box alignItems="center" display="flex" gap={1}>
             <Text fontWeight="bold">Dag Import Errors</Text>
             <FiChevronRight />