Add multiselect to run state in grid view (#35403)

* Add multiselect to run state in grid view

* Fix tests

* Multiselect for run types, UI nits, refactor

* Fix tests, refactor

* Simplify multiselect value

* Nits and refactor

* Use arrays instead of serializing to csv

* Fix tests and global axios paramsSerializer to null

airflow/www/jest-setup.js

@@ -58,6 +58,12 @@ global.stateColors = {
 
 global.defaultDagRunDisplayNumber = 245;
 
+global.filtersOptions = {
+  // Must stay in sync with airflow/www/static/js/types/index.ts
+  dagStates: ["success", "running", "queued", "failed"],
+  runTypes: ["manual", "backfill", "scheduled", "dataset_triggered"],
+};
+
 global.moment = moment;
 
 global.standaloneDagProcessor = true;

airflow/www/static/js/api/index.ts

@@ -49,6 +49,13 @@ import useDags from "./useDags";
 import useDagRuns from "./useDagRuns";
 import useHistoricalMetricsData from "./useHistoricalMetricsData";
 
+axios.interceptors.request.use((config) => {
+  config.paramsSerializer = {
+    indexes: null,
+  };
+  return config;
+});
+
 axios.interceptors.response.use((res: AxiosResponse) =>
   res.data ? camelcaseKeys(res.data, { deep: true }) : res
 );

airflow/www/static/js/dag/nav/FilterBar.tsx

@@ -20,9 +20,12 @@
 /* global moment */
 
 import { Box, Button, Flex, Input, Select } from "@chakra-ui/react";
+import MultiSelect from "src/components/MultiSelect";
 import React from "react";
 import type { DagRun, RunState, TaskState } from "src/types";
 import AutoRefresh from "src/components/AutoRefresh";
+import type { Size } from "chakra-react-select";
+import { useChakraSelectProps } from "chakra-react-select";
 
 import { useTimezone } from "src/context/timezone";
 import { isoFormatWithoutTZ } from "src/datetime_utils";
@@ -43,6 +46,7 @@ const FilterBar = () => {
     onRunTypeChange,
     onRunStateChange,
     clearFilters,
+    transformArrayToMultiSelectOptions,
   } = useFilters();
 
   const { timezone } = useTimezone();
@@ -51,7 +55,26 @@ const FilterBar = () => {
   // @ts-ignore
   const formattedTime = time.tz(timezone).format(isoFormatWithoutTZ);
 
-  const inputStyles = { backgroundColor: "white", size: "lg" };
+  const inputStyles: { backgroundColor: string; size: Size } = {
+    backgroundColor: "white",
+    size: "lg",
+  };
+
+  const multiSelectBoxStyle = { minWidth: "160px", zIndex: 3 };
+  const multiSelectStyles = useChakraSelectProps({
+    ...inputStyles,
+    isMulti: true,
+    tagVariant: "solid",
+    hideSelectedOptions: false,
+    isClearable: false,
+    selectedOptionStyle: "check",
+    chakraStyles: {
+      container: (provided) => ({
+        ...provided,
+        bg: "white",
+      }),
+    },
+  });
 
   return (
     <Flex
@@ -83,38 +106,44 @@ const FilterBar = () => {
             ))}
           </Select>
         </Box>
-        <Box px={2}>
-          <Select
-            {...inputStyles}
-            value={filters.runType || ""}
-            onChange={(e) => onRunTypeChange(e.target.value)}
-          >
-            <option value="" key="all">
-              All Run Types
-            </option>
-            {filtersOptions.runTypes.map((value) => (
-              <option value={value.toString()} key={value}>
-                {value}
-              </option>
-            ))}
-          </Select>
+        <Box px={2} style={multiSelectBoxStyle}>
+          <MultiSelect
+            {...multiSelectStyles}
+            value={transformArrayToMultiSelectOptions(filters.runType)}
+            onChange={(typeOptions) => {
+              if (
+                Array.isArray(typeOptions) &&
+                typeOptions.every((typeOption) => "value" in typeOption)
+              ) {
+                onRunTypeChange(
+                  typeOptions.map((typeOption) => typeOption.value)
+                );
+              }
+            }}
+            options={transformArrayToMultiSelectOptions(filters.runTypeOptions)}
+            placeholder="All Run Types"
+          />
         </Box>
         <Box />
-        <Box px={2}>
-          <Select
-            {...inputStyles}
-            value={filters.runState || ""}
-            onChange={(e) => onRunStateChange(e.target.value)}
-          >
-            <option value="" key="all">
-              All Run States
-            </option>
-            {filtersOptions.dagStates.map((value) => (
-              <option value={value} key={value}>
-                {value}
-              </option>
-            ))}
-          </Select>
+        <Box px={2} style={multiSelectBoxStyle}>
+          <MultiSelect
+            {...multiSelectStyles}
+            value={transformArrayToMultiSelectOptions(filters.runState)}
+            onChange={(stateOptions) => {
+              if (
+                Array.isArray(stateOptions) &&
+                stateOptions.every((stateOption) => "value" in stateOption)
+              ) {
+                onRunStateChange(
+                  stateOptions.map((stateOption) => stateOption.value)
+                );
+              }
+            }}
+            options={transformArrayToMultiSelectOptions(
+              filters.runStateOptions
+            )}
+            placeholder="All Run States"
+          />
         </Box>
         <Box px={2}>
           <Button

airflow/www/static/js/dag/useFilters.test.tsx

@@ -21,11 +21,16 @@
 import { act, renderHook } from "@testing-library/react";
 
 import { RouterWrapper } from "src/utils/testUtils";
+import type { DagRun, RunState } from "src/types";
 
 declare global {
   namespace NodeJS {
     interface Global {
       defaultDagRunDisplayNumber: number;
+      filtersOptions: {
+        dagStates: RunState[];
+        runTypes: DagRun["runType"][];
+      };
     }
   }
 }
@@ -62,8 +67,8 @@ describe("Test useFilters hook", () => {
 
     expect(baseDate).toBe(date.toISOString());
     expect(numRuns).toBe(global.defaultDagRunDisplayNumber.toString());
-    expect(runType).toBeNull();
-    expect(runState).toBeNull();
+    expect(runType).toEqual([]);
+    expect(runState).toEqual([]);
     expect(root).toBeUndefined();
     expect(filterUpstream).toBeUndefined();
     expect(filterDownstream).toBeUndefined();
@@ -84,12 +89,22 @@ describe("Test useFilters hook", () => {
     {
       fnName: "onRunTypeChange" as keyof UtilFunctions,
       paramName: "runType" as keyof Filters,
-      paramValue: "manual",
+      paramValue: ["manual"],
+    },
+    {
+      fnName: "onRunTypeChange" as keyof UtilFunctions,
+      paramName: "runType" as keyof Filters,
+      paramValue: ["manual", "backfill"],
     },
     {
       fnName: "onRunStateChange" as keyof UtilFunctions,
       paramName: "runState" as keyof Filters,
-      paramValue: "success",
+      paramValue: ["success"],
+    },
+    {
+      fnName: "onRunStateChange" as keyof UtilFunctions,
+      paramName: "runState" as keyof Filters,
+      paramValue: ["success", "failed", "queued"],
     },
   ])("Test $fnName functions", async ({ fnName, paramName, paramValue }) => {
     const { result } = renderHook<FilterHookReturn, undefined>(
@@ -98,10 +113,12 @@ describe("Test useFilters hook", () => {
     );
 
     await act(async () => {
-      result.current[fnName](paramValue as "string" & FilterTasksProps);
+      result.current[fnName](
+        paramValue as "string" & string[] & FilterTasksProps
+      );
     });
 
-    expect(result.current.filters[paramName]).toBe(paramValue);
+    expect(result.current.filters[paramName]).toEqual(paramValue);
 
     // clearFilters
     await act(async () => {
@@ -115,7 +132,7 @@ describe("Test useFilters hook", () => {
         global.defaultDagRunDisplayNumber.toString()
       );
     } else {
-      expect(result.current.filters[paramName]).toBeNull();
+      expect(result.current.filters[paramName]).toEqual([]);
     }
   });
 

airflow/www/static/js/dag/useFilters.tsx

@@ -21,17 +21,27 @@
 
 import { useSearchParams } from "react-router-dom";
 import URLSearchParamsWrapper from "src/utils/URLSearchParamWrapper";
+import type { DagRun, RunState, TaskState } from "src/types";
 
 declare const defaultDagRunDisplayNumber: number;
 
+declare const filtersOptions: {
+  dagStates: RunState[];
+  numRuns: number[];
+  runTypes: DagRun["runType"][];
+  taskStates: TaskState[];
+};
+
 export interface Filters {
   root: string | undefined;
   filterUpstream: boolean | undefined;
   filterDownstream: boolean | undefined;
   baseDate: string | null;
   numRuns: string | null;
-  runType: string | null;
-  runState: string | null;
+  runType: string[] | null;
+  runTypeOptions: string[] | null;
+  runState: string[] | null;
+  runStateOptions: string[] | null;
 }
 
 export interface FilterTasksProps {
@@ -43,9 +53,12 @@ export interface FilterTasksProps {
 export interface UtilFunctions {
   onBaseDateChange: (value: string) => void;
   onNumRunsChange: (value: string) => void;
-  onRunTypeChange: (value: string) => void;
-  onRunStateChange: (value: string) => void;
+  onRunTypeChange: (values: string[]) => void;
+  onRunStateChange: (values: string[]) => void;
   onFilterTasksChange: (args: FilterTasksProps) => void;
+  transformArrayToMultiSelectOptions: (
+    options: string[] | null
+  ) => { label: string; value: string }[];
   clearFilters: () => void;
   resetRoot: () => void;
 }
@@ -83,8 +96,12 @@ const useFilters = (): FilterHookReturn => {
   const baseDate = searchParams.get(BASE_DATE_PARAM) || now;
   const numRuns =
     searchParams.get(NUM_RUNS_PARAM) || defaultDagRunDisplayNumber.toString();
-  const runType = searchParams.get(RUN_TYPE_PARAM);
-  const runState = searchParams.get(RUN_STATE_PARAM);
+
+  const runTypeOptions = filtersOptions.runTypes;
+  const runType = searchParams.getAll(RUN_TYPE_PARAM);
+
+  const runStateOptions = filtersOptions.dagStates;
+  const runState = searchParams.getAll(RUN_STATE_PARAM);
 
   const makeOnChangeFn =
     (paramName: string, formatFn?: (arg: string) => string) =>
@@ -98,14 +115,40 @@ const useFilters = (): FilterHookReturn => {
       setSearchParams(params);
     };
 
+  const makeMultiSelectOnChangeFn =
+    (paramName: string, options: string[]) => (values: string[]) => {
+      const params = new URLSearchParamsWrapper(searchParams);
+      if (values.length === options.length || values.length === 0) {
+        params.delete(paramName);
+      } else {
+        // Delete and reinsert anew each time; otherwise, there will be duplicates
+        params.delete(paramName);
+        values.forEach((value) => params.append(paramName, value));
+      }
+      setSearchParams(params);
+    };
+
+  const transformArrayToMultiSelectOptions = (
+    options: string[] | null
+  ): { label: string; value: string }[] =>
+    options === null
+      ? []
+      : options.map((option) => ({ label: option, value: option }));
+
   const onBaseDateChange = makeOnChangeFn(
     BASE_DATE_PARAM,
     // @ts-ignore
     (localDate: string) => moment(localDate).utc().format()
   );
   const onNumRunsChange = makeOnChangeFn(NUM_RUNS_PARAM);
-  const onRunTypeChange = makeOnChangeFn(RUN_TYPE_PARAM);
-  const onRunStateChange = makeOnChangeFn(RUN_STATE_PARAM);
+  const onRunTypeChange = makeMultiSelectOnChangeFn(
+    RUN_TYPE_PARAM,
+    filtersOptions.runTypes
+  );
+  const onRunStateChange = makeMultiSelectOnChangeFn(
+    RUN_STATE_PARAM,
+    filtersOptions.dagStates
+  );
 
   const onFilterTasksChange = ({
     root: newRoot,
@@ -154,7 +197,9 @@ const useFilters = (): FilterHookReturn => {
       baseDate,
       numRuns,
       runType,
+      runTypeOptions,
       runState,
+      runStateOptions,
     },
     onBaseDateChange,
     onNumRunsChange,
@@ -163,6 +208,7 @@ const useFilters = (): FilterHookReturn => {
     onFilterTasksChange,
     clearFilters,
     resetRoot,
+    transformArrayToMultiSelectOptions,
   };
 };
 

airflow/www/views.py

@@ -3527,13 +3527,13 @@ def grid_data(self):
         with create_session() as session:
             query = select(DagRun).where(DagRun.dag_id == dag.dag_id, DagRun.execution_date <= base_date)
 
-        run_type = request.args.get("run_type")
-        if run_type:
-            query = query.where(DagRun.run_type == run_type)
+        run_types = request.args.getlist("run_type")
+        if run_types:
+            query = query.where(DagRun.run_type.in_(run_types))
 
-        run_state = request.args.get("run_state")
-        if run_state:
-            query = query.where(DagRun.state == run_state)
+        run_states = request.args.getlist("run_state")
+        if run_states:
+            query = query.where(DagRun.state.in_(run_states))
 
         dag_runs = wwwutils.sorted_dag_runs(
             query, ordering=dag.timetable.run_ordering, limit=num_runs, session=session

tests/www/views/test_views_grid.py

@@ -162,6 +162,27 @@ def test_no_runs(admin_client, dag_without_runs):
     }
 
 
+def test_grid_data_filtered_on_run_type_and_run_state(admin_client, dag_with_runs):
+    for uri_params, expected_run_types, expected_run_states in [
+        ("run_state=success&run_state=queued", ["scheduled"], ["success"]),
+        ("run_state=running&run_state=failed", ["scheduled"], ["running"]),
+        ("run_type=scheduled&run_type=manual", ["scheduled", "scheduled"], ["success", "running"]),
+        ("run_type=backfill&run_type=manual", [], []),
+        ("run_state=running&run_type=failed&run_type=backfill&run_type=manual", [], []),
+        (
+            "run_state=running&run_type=failed&run_type=scheduled&run_type=backfill&run_type=manual",
+            ["scheduled"],
+            ["running"],
+        ),
+    ]:
+        resp = admin_client.get(f"/object/grid_data?dag_id={DAG_ID}&{uri_params}", follow_redirects=True)
+        assert resp.status_code == 200, resp.json
+        actual_run_types = list(map(lambda x: x["run_type"], resp.json["dag_runs"]))
+        actual_run_states = list(map(lambda x: x["state"], resp.json["dag_runs"]))
+        assert actual_run_types == expected_run_types
+        assert actual_run_states == expected_run_states
+
+
 # Create this as a fixture so that it is applied before the `dag_with_runs` fixture is!
 @pytest.fixture
 def freeze_time_for_dagruns(time_machine):