Add task instance state and name filter to task instances tab in a da…

…grun (apache#45215) * Add state filter to task instances tab in dagrun page. * Add tasks search in task instances page. * Make dropdown clearable. * Rebase and run format for line number rule. * Add support for multiple state selection. * Support no status for filtering. * Fix rebase errors.
PolideaInternal · Jan 20, 2025 · c68083e · c68083e
1 parent 0f8707e
commit c68083e
Show file tree

Hide file tree

Showing 3 changed files with 144 additions and 10 deletions.
diff --git a/airflow/ui/src/components/SearchBar.tsx b/airflow/ui/src/components/SearchBar.tsx
@@ -29,11 +29,19 @@ type Props = {
   readonly buttonProps?: ButtonProps;
   readonly defaultValue: string;
   readonly groupProps?: InputGroupProps;
+  readonly hideAdvanced?: boolean;
   readonly onChange: (value: string) => void;
   readonly placeHolder: string;
 };
 
-export const SearchBar = ({ buttonProps, defaultValue, groupProps, onChange, placeHolder }: Props) => {
+export const SearchBar = ({
+  buttonProps,
+  defaultValue,
+  groupProps,
+  hideAdvanced = false,
+  onChange,
+  placeHolder,
+}: Props) => {
   const handleSearchChange = useDebouncedCallback((val: string) => onChange(val), debounceDelay);
 
   const [value, setValue] = useState(defaultValue);
@@ -61,9 +69,11 @@ export const SearchBar = ({ buttonProps, defaultValue, groupProps, onChange, pla
               size="xs"
             />
           ) : undefined}
-          <Button fontWeight="normal" height="1.75rem" variant="ghost" width={140} {...buttonProps}>
-            Advanced Search
-          </Button>
+          {Boolean(hideAdvanced) ? undefined : (
+            <Button fontWeight="normal" height="1.75rem" variant="ghost" width={140} {...buttonProps}>
+              Advanced Search
+            </Button>
+          )}
         </>
       }
       startElement={<FiSearch />}

diff --git a/airflow/ui/src/pages/Run/TaskInstances.tsx b/airflow/ui/src/pages/Run/TaskInstances.tsx
@@ -16,19 +16,29 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Box, Flex, Link } from "@chakra-ui/react";
+import {
+  Box,
+  Flex,
+  Link,
+  createListCollection,
+  HStack,
+  type SelectValueChangeDetails,
+} from "@chakra-ui/react";
 import type { ColumnDef } from "@tanstack/react-table";
-import { Link as RouterLink, useParams } from "react-router-dom";
+import { useCallback, useState } from "react";
+import { Link as RouterLink, useParams, useSearchParams } from "react-router-dom";
 
 import { useTaskInstanceServiceGetTaskInstances } from "openapi/queries";
-import type { TaskInstanceResponse } from "openapi/requests/types.gen";
+import type { TaskInstanceResponse, TaskInstanceState } from "openapi/requests/types.gen";
 import { ClearTaskInstanceButton } from "src/components/Clear";
 import { DataTable } from "src/components/DataTable";
 import { useTableURLState } from "src/components/DataTable/useTableUrlState";
 import { ErrorAlert } from "src/components/ErrorAlert";
+import { SearchBar } from "src/components/SearchBar";
 import Time from "src/components/Time";
-import { Status } from "src/components/ui";
-import { getDuration } from "src/utils";
+import { Select, Status } from "src/components/ui";
+import { SearchParamsKeys, type SearchParamsKeysType } from "src/constants/searchParams";
+import { capitalize, getDuration } from "src/utils";
 import { getTaskInstanceLink } from "src/utils/links";
 
 const columns: Array<ColumnDef<TaskInstanceResponse>> = [
@@ -96,12 +106,74 @@ const columns: Array<ColumnDef<TaskInstanceResponse>> = [
   },
 ];
 
+const stateOptions = createListCollection<{ label: string; value: TaskInstanceState | "all" | "none" }>({
+  items: [
+    { label: "All States", value: "all" },
+    { label: "Scheduled", value: "scheduled" },
+    { label: "Queued", value: "queued" },
+    { label: "Running", value: "running" },
+    { label: "Success", value: "success" },
+    { label: "Restarting", value: "restarting" },
+    { label: "Failed", value: "failed" },
+    { label: "Up For Retry", value: "up_for_retry" },
+    { label: "Up For Reschedule", value: "up_for_reschedule" },
+    { label: "Upstream failed", value: "upstream_failed" },
+    { label: "Skipped", value: "skipped" },
+    { label: "Deferred", value: "deferred" },
+    { label: "Removed", value: "removed" },
+    { label: "No Status", value: "none" },
+  ],
+});
+
+const STATE_PARAM = "state";
+
 export const TaskInstances = () => {
   const { dagId = "", runId = "" } = useParams();
+  const [searchParams, setSearchParams] = useSearchParams();
   const { setTableURLState, tableURLState } = useTableURLState();
   const { pagination, sorting } = tableURLState;
   const [sort] = sorting;
   const orderBy = sort ? `${sort.desc ? "-" : ""}${sort.id}` : "-start_date";
+  const filteredState = searchParams.getAll(STATE_PARAM);
+  const hasFilteredState = filteredState.length > 0;
+  const { NAME_PATTERN: NAME_PATTERN_PARAM }: SearchParamsKeysType = SearchParamsKeys;
+
+  const [taskDisplayNamePattern, setTaskDisplayNamePattern] = useState(
+    searchParams.get(NAME_PATTERN_PARAM) ?? undefined,
+  );
+
+  const handleStateChange = useCallback(
+    ({ value }: SelectValueChangeDetails<string>) => {
+      const [val, ...rest] = value;
+
+      if ((val === undefined || val === "all") && rest.length === 0) {
+        searchParams.delete(STATE_PARAM);
+      } else {
+        searchParams.delete(STATE_PARAM);
+        value.filter((state) => state !== "all").map((state) => searchParams.append(STATE_PARAM, state));
+      }
+      setTableURLState({
+        pagination: { ...pagination, pageIndex: 0 },
+        sorting,
+      });
+      setSearchParams(searchParams);
+    },
+    [pagination, searchParams, setSearchParams, setTableURLState, sorting],
+  );
+
+  const handleSearchChange = (value: string) => {
+    if (value) {
+      searchParams.set(NAME_PATTERN_PARAM, value);
+    } else {
+      searchParams.delete(NAME_PATTERN_PARAM);
+    }
+    setTableURLState({
+      pagination: { ...pagination, pageIndex: 0 },
+      sorting,
+    });
+    setTaskDisplayNamePattern(value);
+    setSearchParams(searchParams);
+  };
 
   const { data, error, isFetching, isLoading } = useTaskInstanceServiceGetTaskInstances(
     {
@@ -110,13 +182,64 @@ export const TaskInstances = () => {
       limit: pagination.pageSize,
       offset: pagination.pageIndex * pagination.pageSize,
       orderBy,
+      state: hasFilteredState ? filteredState : undefined,
+      taskDisplayNamePattern: Boolean(taskDisplayNamePattern) ? taskDisplayNamePattern : undefined,
     },
     undefined,
     { enabled: !isNaN(pagination.pageSize) },
   );
 
   return (
-    <Box>
+    <Box pt={4}>
+      <HStack>
+        <Select.Root
+          collection={stateOptions}
+          maxW="250px"
+          multiple
+          onValueChange={handleStateChange}
+          value={hasFilteredState ? filteredState : ["all"]}
+        >
+          <Select.Trigger
+            {...(hasFilteredState ? { clearable: true } : {})}
+            colorPalette="blue"
+            isActive={Boolean(filteredState)}
+          >
+            <Select.ValueText>
+              {() =>
+                hasFilteredState ? (
+                  <HStack gap="10px">
+                    {filteredState.map((state) => (
+                      <Status key={state} state={state as TaskInstanceState}>
+                        {state === "none" ? "No Status" : capitalize(state)}
+                      </Status>
+                    ))}
+                  </HStack>
+                ) : (
+                  "All States"
+                )
+              }
+            </Select.ValueText>
+          </Select.Trigger>
+          <Select.Content>
+            {stateOptions.items.map((option) => (
+              <Select.Item item={option} key={option.label}>
+                {option.value === "all" ? (
+                  option.label
+                ) : (
+                  <Status state={option.value as TaskInstanceState}>{option.label}</Status>
+                )}
+              </Select.Item>
+            ))}
+          </Select.Content>
+        </Select.Root>
+        <SearchBar
+          buttonProps={{ disabled: true }}
+          defaultValue={taskDisplayNamePattern ?? ""}
+          hideAdvanced
+          onChange={handleSearchChange}
+          placeHolder="Search Tasks"
+        />
+      </HStack>
       <DataTable
         columns={columns}
         data={data?.task_instances ?? []}

diff --git a/airflow/ui/src/utils/stateColor.ts b/airflow/ui/src/utils/stateColor.ts
@@ -22,6 +22,7 @@
 export const stateColor = {
   deferred: "mediumpurple",
   failed: "red",
+  none: "lightblue",
   null: "lightblue",
   queued: "gray",
   removed: "lightgrey",