Add filter task upstream/downstream to grid view (#29885)

* Add filter task upstream/downstream to grid view

* Update button text, move reset root to dropdown

GitOrigin-RevId: 0a91a0bcc26f402fcd3b09c83554d31908b2d847

Author: bbovenzi

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

@@ -22,12 +22,16 @@ import axios, { AxiosResponse } from "axios";
 
 import { getMetaValue } from "src/utils";
 import type { DepNode } from "src/types";
+import useFilters, {
+  FILTER_DOWNSTREAM_PARAM,
+  FILTER_UPSTREAM_PARAM,
+  ROOT_PARAM,
+} from "src/dag/useFilters";
 
 const DAG_ID_PARAM = "dag_id";
 
 const dagId = getMetaValue(DAG_ID_PARAM);
 const graphDataUrl = getMetaValue("graph_data_url");
-const urlRoot = getMetaValue("root");
 
 interface GraphData {
   edges: WebserverEdge[];
@@ -40,15 +44,23 @@ export interface WebserverEdge {
   targetId: string;
 }
 
-const useGraphData = () =>
-  useQuery("graphData", async () => {
-    const params = {
-      [DAG_ID_PARAM]: dagId,
-      root: urlRoot || undefined,
-      filter_upstream: true,
-      filter_downstream: true,
-    };
-    return axios.get<AxiosResponse, GraphData>(graphDataUrl, { params });
-  });
+const useGraphData = () => {
+  const {
+    filters: { root, filterDownstream, filterUpstream },
+  } = useFilters();
+
+  return useQuery(
+    ["graphData", root, filterUpstream, filterDownstream],
+    async () => {
+      const params = {
+        [DAG_ID_PARAM]: dagId,
+        [ROOT_PARAM]: root,
+        [FILTER_UPSTREAM_PARAM]: filterUpstream,
+        [FILTER_DOWNSTREAM_PARAM]: filterDownstream,
+      };
+      return axios.get<AxiosResponse, GraphData>(graphDataUrl, { params });
+    }
+  );
+};
 
 export default useGraphData;

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

@@ -29,6 +29,9 @@ import useFilters, {
   RUN_STATE_PARAM,
   RUN_TYPE_PARAM,
   now,
+  FILTER_DOWNSTREAM_PARAM,
+  FILTER_UPSTREAM_PARAM,
+  ROOT_PARAM,
 } from "src/dag/useFilters";
 import type { Task, DagRun, RunOrdering } from "src/types";
 import { camelCase } from "lodash";
@@ -38,7 +41,6 @@ const DAG_ID_PARAM = "dag_id";
 // dagId comes from dag.html
 const dagId = getMetaValue(DAG_ID_PARAM);
 const gridDataUrl = getMetaValue("grid_data_url");
-const urlRoot = getMetaValue("root");
 
 export interface GridData {
   dagRuns: DagRun[];
@@ -68,14 +70,33 @@ const useGridData = () => {
   const { isRefreshOn, stopRefresh } = useAutoRefresh();
   const errorToast = useErrorToast();
   const {
-    filters: { baseDate, numRuns, runType, runState },
+    filters: {
+      baseDate,
+      numRuns,
+      runType,
+      runState,
+      root,
+      filterDownstream,
+      filterUpstream,
+    },
   } = useFilters();
 
   const query = useQuery(
-    ["gridData", baseDate, numRuns, runType, runState],
+    [
+      "gridData",
+      baseDate,
+      numRuns,
+      runType,
+      runState,
+      root,
+      filterUpstream,
+      filterDownstream,
+    ],
     async () => {
       const params = {
-        root: urlRoot || undefined,
+        [ROOT_PARAM]: root,
+        [FILTER_UPSTREAM_PARAM]: filterUpstream,
+        [FILTER_DOWNSTREAM_PARAM]: filterDownstream,
         [DAG_ID_PARAM]: dagId,
         [BASE_DATE_PARAM]: baseDate === now ? undefined : baseDate,
         [NUM_RUNS_PARAM]: numRuns,

airflow/www/static/js/dag/details/FilterTasks.tsx

@@ -0,0 +1,91 @@
+/*!
+ * 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 React from "react";
+import {
+  Flex,
+  Button,
+  Menu,
+  MenuButton,
+  MenuItem,
+  MenuList,
+} from "@chakra-ui/react";
+import useFilters from "src/dag/useFilters";
+import { MdArrowDropDown } from "react-icons/md";
+
+interface Props {
+  taskId: string;
+}
+
+const FilterTasks = ({ taskId }: Props) => {
+  const {
+    filters: { root },
+    onFilterTasksChange,
+    resetRoot,
+  } = useFilters();
+
+  const onFilterUpstream = () =>
+    onFilterTasksChange({
+      root: taskId,
+      filterUpstream: true,
+      filterDownstream: false,
+    });
+
+  const onFilterDownstream = () =>
+    onFilterTasksChange({
+      root: taskId,
+      filterUpstream: false,
+      filterDownstream: true,
+    });
+
+  const onFilterAll = () =>
+    onFilterTasksChange({
+      root: taskId,
+      filterUpstream: true,
+      filterDownstream: true,
+    });
+
+  const label = "Filter upstream and/or downstream of a task";
+
+  return (
+    <Menu>
+      <MenuButton
+        as={Button}
+        variant="outline"
+        colorScheme="blue"
+        transition="all 0.2s"
+        title={label}
+        aria-label={label}
+      >
+        <Flex>
+          {!root ? "Filter Tasks " : "Clear Task Filter "}
+          <MdArrowDropDown size="16px" />
+        </Flex>
+      </MenuButton>
+      <MenuList>
+        <MenuItem onClick={onFilterUpstream}>Filter Upstream</MenuItem>
+        <MenuItem onClick={onFilterDownstream}>Filter Downstream</MenuItem>
+        <MenuItem onClick={onFilterAll}>Filter Upstream & Downstream</MenuItem>
+        {!!root && <MenuItem onClick={resetRoot}>Reset Root</MenuItem>}
+      </MenuList>
+    </Menu>
+  );
+};
+
+export default FilterTasks;

airflow/www/static/js/dag/details/graph/index.tsx

@@ -37,6 +37,7 @@ import { useOffsetTop } from "src/utils";
 import { useGraphLayout } from "src/utils/graph";
 import Tooltip from "src/components/Tooltip";
 import { useContainerRef } from "src/context/containerRef";
+import useFilters from "src/dag/useFilters";
 
 import Edge from "./Edge";
 import Node, { CustomNodeProps } from "./Node";
@@ -56,6 +57,10 @@ const Graph = ({ openGroupIds, onToggleGroups }: Props) => {
   const { data } = useGraphData();
   const [arrange, setArrange] = useState(data?.arrange || "LR");
 
+  const {
+    filters: { root, filterDownstream, filterUpstream },
+  } = useFilters();
+
   useEffect(() => {
     setArrange(data?.arrange || "LR");
   }, [data?.arrange]);
@@ -71,9 +76,14 @@ const Graph = ({ openGroupIds, onToggleGroups }: Props) => {
     data: { dagRuns, groups },
   } = useGridData();
   const { colors } = useTheme();
-  const { setCenter } = useReactFlow();
+  const { setCenter, setViewport } = useReactFlow();
   const latestDagRunId = dagRuns[dagRuns.length - 1]?.runId;
 
+  // Reset viewport when tasks are filtered
+  useEffect(() => {
+    setViewport({ x: 0, y: 0, zoom: 1 });
+  }, [root, filterDownstream, filterUpstream, setViewport]);
+
   const offsetTop = useOffsetTop(graphRef);
 
   let nodes: ReactFlowNode<CustomNodeProps>[] = [];

airflow/www/static/js/dag/details/index.tsx

@@ -45,6 +45,7 @@ import Graph from "./graph";
 import MappedInstances from "./taskInstance/MappedInstances";
 import Logs from "./taskInstance/Logs";
 import BackToTaskSummary from "./taskInstance/BackToTaskSummary";
+import FilterTasks from "./FilterTasks";
 
 const dagId = getMetaValue("dag_id")!;
 
@@ -145,7 +146,10 @@ const Details = ({ openGroupIds, onToggleGroups }: Props) => {
 
   return (
     <Flex flexDirection="column" pl={3} height="100%">
-      <Header />
+      <Flex alignItems="center" justifyContent="space-between">
+        <Header />
+        <Flex>{taskId && runId && <FilterTasks taskId={taskId} />}</Flex>
+      </Flex>
       <Divider my={2} />
       <Tabs
         size="lg"

airflow/www/static/js/dag/details/taskInstance/Nav.tsx

@@ -22,23 +22,19 @@ import { Flex } from "@chakra-ui/react";
 
 import { getMetaValue, appendSearchParams } from "src/utils";
 import LinkButton from "src/components/LinkButton";
-import type { Task, DagRun } from "src/types";
+import type { Task } from "src/types";
 import URLSearchParamsWrapper from "src/utils/URLSearchParamWrapper";
 
 const dagId = getMetaValue("dag_id");
 const isK8sExecutor = getMetaValue("k8s_or_k8scelery_executor") === "True";
-const numRuns = getMetaValue("num_runs");
-const baseDate = getMetaValue("base_date");
 const taskInstancesUrl = getMetaValue("task_instances_list_url");
 const renderedK8sUrl = getMetaValue("rendered_k8s_url");
 const renderedTemplatesUrl = getMetaValue("rendered_templates_url");
 const xcomUrl = getMetaValue("xcom_url");
 const taskUrl = getMetaValue("task_url");
 const gridUrl = getMetaValue("grid_url");
-const gridUrlNoRoot = getMetaValue("grid_url_no_root");
 
 interface Props {
-  runId: DagRun["runId"];
   taskId: Task["id"];
   executionDate: string;
   operator?: string;
@@ -47,10 +43,7 @@ interface Props {
 }
 
 const Nav = forwardRef<HTMLDivElement, Props>(
-  (
-    { runId, taskId, executionDate, operator, isMapped = false, mapIndex },
-    ref
-  ) => {
+  ({ taskId, executionDate, operator, isMapped = false, mapIndex }, ref) => {
     if (!taskId) return null;
     const params = new URLSearchParamsWrapper({
       task_id: taskId,
@@ -70,23 +63,11 @@ const Nav = forwardRef<HTMLDivElement, Props>(
       execution_date: executionDate,
     }).toString();
 
-    const filterParams = new URLSearchParamsWrapper({
-      task_id: taskId,
-      dag_run_id: runId,
-      root: taskId,
-    });
-
     if (mapIndex !== undefined && mapIndex >= 0)
       listParams.append("_flt_0_map_index", mapIndex.toString());
-    if (baseDate) filterParams.append("base_date", baseDate);
-    if (numRuns) filterParams.append("num_runs", numRuns);
 
     const allInstancesLink = `${taskInstancesUrl}?${listParams.toString()}`;
 
-    const filterUpstreamLink = appendSearchParams(
-      gridUrlNoRoot,
-      filterParams.toString()
-    );
     const subDagLink = appendSearchParams(
       gridUrl.replace(dagId, `${dagId}.${taskId}`),
       subDagParams
@@ -116,7 +97,6 @@ const Nav = forwardRef<HTMLDivElement, Props>(
         >
           List Instances, all runs
         </LinkButton>
-        <LinkButton href={filterUpstreamLink}>Filter Upstream</LinkButton>
       </Flex>
     );
   }

airflow/www/static/js/dag/details/taskInstance/index.tsx

@@ -92,7 +92,6 @@ const TaskInstance = ({ taskId, runId, mapIndex }: Props) => {
       {!isGroup && (
         <TaskNav
           taskId={taskId}
-          runId={runId}
           isMapped={isMapped}
           mapIndex={mapIndex}
           executionDate={executionDate}

airflow/www/static/js/dag/grid/ResetRoot.tsx

@@ -26,8 +26,7 @@ import AutoRefresh from "src/components/AutoRefresh";
 
 import { useTimezone } from "src/context/timezone";
 import { isoFormatWithoutTZ } from "src/datetime_utils";
-import useFilters from "../useFilters";
-import ResetRoot from "../grid/ResetRoot";
+import useFilters from "src/dag/useFilters";
 
 declare const filtersOptions: {
   dagStates: RunState[];
@@ -123,10 +122,7 @@ const FilterBar = () => {
           </Button>
         </Box>
       </Flex>
-      <Flex>
-        <AutoRefresh />
-        <ResetRoot />
-      </Flex>
+      <AutoRefresh />
     </Flex>
   );
 };