Feature(UI) : Add upstream/downstream task filtering in Graph & Grid views (#53458)

airflow-core/src/airflow/ui/public/i18n/locales/en/components.json

@@ -94,6 +94,14 @@
   "sortedAscending": "sorted ascending",
   "sortedDescending": "sorted descending",
   "sortedUnsorted": "unsorted",
+  "taskFilter": {
+    "action_prefix": "Filter",
+    "all": "All",
+    "both": "Both upstream & downstream",
+    "button_all": "Filter DAG by task",
+    "downstream": "Only downstream",
+    "upstream": "Only upstream"
+  },
   "taskTries": "Task Tries",
   "toggleCardView": "Show card view",
   "toggleTableView": "Show table view",

airflow-core/src/airflow/ui/src/components/FilterTaskButton.tsx

@@ -0,0 +1,122 @@
+/*!
+ * 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 { useTranslation } from "react-i18next";
+import { MdFilterList } from "react-icons/md";
+import { useSearchParams, useParams } from "react-router-dom";
+
+import { Menu } from "src/components/ui";
+import { Button } from "src/components/ui";
+import ActionButton from "src/components/ui/ActionButton";
+
+const FILTER_PARAM = "task_filter";
+
+const OPTIONS = [
+  { buttonKey: "taskFilter.button_all", itemKey: "taskFilter.all", value: "all" },
+  { buttonKey: "taskFilter.upstream", itemKey: "taskFilter.upstream", value: "upstream" },
+  { buttonKey: "taskFilter.downstream", itemKey: "taskFilter.downstream", value: "downstream" },
+  { buttonKey: "taskFilter.both", itemKey: "taskFilter.both", value: "both" },
+] as const;
+
+type FilterValue = (typeof OPTIONS)[number]["value"];
+
+type Props = {
+  readonly withText?: boolean;
+};
+
+const FilterTaskButton = ({ withText = true }: Props) => {
+  const { t: translate } = useTranslation("components");
+  const [searchParams, setSearchParams] = useSearchParams();
+  const { taskId } = useParams<{ taskId?: string }>();
+
+  const rawFilter = searchParams.get(FILTER_PARAM) as FilterValue | null;
+  const currentFilter = rawFilter ?? "all";
+
+  const isActiveForThisTask = currentFilter !== "all";
+
+  const handleSelect = (value: FilterValue) => {
+    const next = new URLSearchParams(searchParams.toString());
+
+    if (value === "all") {
+      next.delete(FILTER_PARAM);
+    } else {
+      if (taskId === undefined) {
+        return; // optional: show tooltip instead
+      }
+      next.set(FILTER_PARAM, value);
+    }
+
+    setSearchParams(next, { replace: true });
+  };
+
+  const selectedOption = OPTIONS.find((opt) => opt.value === currentFilter) ?? OPTIONS[0];
+  const buttonLabel = translate(selectedOption.buttonKey);
+
+  return (
+    <Box display="inline-block" position="relative">
+      <Menu.Root positioning={{ gutter: 0, placement: "bottom" }}>
+        <Menu.Trigger asChild>
+          <Box position="relative">
+            <ActionButton
+              actionName={`${translate("taskFilter.action_prefix", { defaultValue: "Filter" })}: ${buttonLabel}`}
+              flexDirection="row-reverse"
+              icon={<MdFilterList />}
+              text={buttonLabel}
+              withText={withText}
+            />
+            {isActiveForThisTask ? (
+              <Box
+                bg="blue.500"
+                borderRadius="full"
+                height={2.5}
+                position="absolute"
+                right={1} // adjust to align correctly
+                top={1} // adjust to align correctly
+                width={2.5}
+              />
+            ) : undefined}
+          </Box>
+        </Menu.Trigger>
+
+        <Menu.Content>
+          {OPTIONS.map((option) => (
+            <Menu.Item
+              asChild
+              disabled={currentFilter === option.value}
+              key={option.value}
+              value={option.value}
+            >
+              <Button
+                justifyContent="start"
+                onClick={() => handleSelect(option.value)}
+                size="sm"
+                variant="ghost"
+                width="100%"
+              >
+                {translate(option.itemKey)}
+              </Button>
+            </Menu.Item>
+          ))}
+        </Menu.Content>
+      </Menu.Root>
+    </Box>
+  );
+};
+
+export default FilterTaskButton;

airflow-core/src/airflow/ui/src/layouts/Details/Gantt/Gantt.tsx

@@ -36,12 +36,15 @@ import annotationPlugin from "chartjs-plugin-annotation";
 import { useMemo, useRef, useDeferredValue } from "react";
 import { Bar } from "react-chartjs-2";
 import { useTranslation } from "react-i18next";
-import { useParams, useNavigate, useLocation } from "react-router-dom";
+import { useParams, useNavigate, useLocation, useSearchParams } from "react-router-dom";
+import { useLocalStorage } from "usehooks-ts";
 
 import { useTaskInstanceServiceGetTaskInstances } from "openapi/queries";
+import { useStructureServiceStructureData } from "openapi/queries";
 import { useColorMode } from "src/context/colorMode";
 import { useOpenGroups } from "src/context/openGroups";
 import { useTimezone } from "src/context/timezone";
+import useSelectedVersion from "src/hooks/useSelectedVersion";
 import { flattenNodes } from "src/layouts/Details/Grid/utils";
 import { useGridRuns } from "src/queries/useGridRuns";
 import { useGridStructure } from "src/queries/useGridStructure";
@@ -50,6 +53,7 @@ import { getComputedCSSVariableValue } from "src/theme";
 import { isStatePending, useAutoRefresh } from "src/utils";
 import { formatDate } from "src/utils/datetimeUtils";
 
+import { buildEdges, filterNodesByDirection } from "../Grid/utils";
 import { createHandleBarClick, createChartOptions } from "./utils";
 
 ChartJS.register(
@@ -120,13 +124,70 @@ export const Gantt = ({ limit }: Props) => {
     },
   );
 
+  const [searchParams] = useSearchParams();
+  const rawTaskFilter = (searchParams.get("task_filter") ?? undefined) as
+    | "all"
+    | "both"
+    | "downstream"
+    | "upstream"
+    | null;
+  const taskFilter = rawTaskFilter ?? "all";
+
+  const selectedVersion = useSelectedVersion();
+  const [dependencies] = useLocalStorage<"all" | "immediate" | "tasks">(`dependencies-${dagId}`, "tasks");
+
+  // derive server flags
+  const includeUpstream = taskFilter === "upstream" || taskFilter === "both";
+  const includeDownstream = taskFilter === "downstream" || taskFilter === "both";
+
+  type StructurePruneParams = {
+    includeDownstream?: boolean;
+    includeUpstream?: boolean;
+    root?: string;
+  };
+
+  const structureParams: StructurePruneParams =
+    taskFilter !== "all" && selectedTaskId !== undefined && selectedTaskId !== ""
+      ? {
+          includeDownstream,
+          includeUpstream,
+          root: selectedTaskId,
+        }
+      : {};
+
+  const { data: structureData = { edges: [], nodes: [] } } = useStructureServiceStructureData(
+    {
+      dagId,
+      externalDependencies: dependencies === "immediate",
+      versionNumber: selectedVersion,
+      ...structureParams,
+    } as unknown as Parameters<typeof useStructureServiceStructureData>[0],
+    undefined,
+    { enabled: selectedVersion !== undefined },
+  );
+
   const { flatNodes } = useMemo(
     () => flattenNodes(dagStructure, deferredOpenGroupIds),
     [dagStructure, deferredOpenGroupIds],
   );
 
   const isLoading = runsLoading || structureLoading || summariesLoading || tiLoading;
 
+  // Build edges (same shape used by Grid)
+  const edges = useMemo(() => buildEdges(structureData.edges), [structureData.edges]);
+
+  // Filter flatNodes using the same helper as Grid
+  const visibleNodes = useMemo(
+    () =>
+      filterNodesByDirection({
+        edges,
+        filter: taskFilter,
+        flatNodes,
+        taskId: selectedTaskId,
+      }),
+    [flatNodes, edges, selectedTaskId, taskFilter],
+  );
+
   const data = useMemo(() => {
     if (isLoading || runId === "") {
       return [];
@@ -135,7 +196,7 @@ export const Gantt = ({ limit }: Props) => {
     const gridSummaries = gridTiSummaries?.task_instances ?? [];
     const taskInstances = taskInstancesData?.task_instances ?? [];
 
-    return flatNodes
+    return visibleNodes
       .map((node) => {
         const gridSummary = gridSummaries.find((ti) => ti.task_id === node.id);
 
@@ -174,7 +235,7 @@ export const Gantt = ({ limit }: Props) => {
         return undefined;
       })
       .filter((item) => item !== undefined);
-  }, [flatNodes, gridTiSummaries, taskInstancesData, selectedTimezone, isLoading, runId]);
+  }, [visibleNodes, gridTiSummaries, taskInstancesData, selectedTimezone, isLoading, runId]);
 
   // Get all unique states and their colors
   const states = [...new Set(data.map((item) => item.state ?? "none"))];
@@ -204,7 +265,7 @@ export const Gantt = ({ limit }: Props) => {
     [data, flatNodes, stateColorMap],
   );
 
-  const fixedHeight = flatNodes.length * CHART_ROW_HEIGHT + CHART_PADDING;
+  const fixedHeight = visibleNodes.length * CHART_ROW_HEIGHT + CHART_PADDING;
   const selectedId = selectedTaskId ?? selectedGroupId;
 
   const handleBarClick = useMemo(
@@ -247,7 +308,7 @@ export const Gantt = ({ limit }: Props) => {
         options={chartOptions}
         ref={ref}
         style={{
-          paddingTop: flatNodes.length === 1 ? 15 : 1.5,
+          paddingTop: visibleNodes.length === 1 ? 15 : 1.5,
         }}
       />
     </Box>

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

@@ -20,7 +20,7 @@ import { useToken } from "@chakra-ui/react";
 import { ReactFlow, Controls, Background, MiniMap, type Node as ReactFlowNode } from "@xyflow/react";
 import "@xyflow/react/dist/style.css";
 import { useEffect } from "react";
-import { useParams } from "react-router-dom";
+import { useParams, useSearchParams } from "react-router-dom";
 import { useLocalStorage } from "usehooks-ts";
 
 import { useStructureServiceStructureData } from "openapi/queries";
@@ -58,9 +58,12 @@ const nodeColor = (
   return "";
 };
 
+type TaskFilter = "all" | "both" | "downstream" | "upstream";
+
 export const Graph = () => {
   const { colorMode = "light" } = useColorMode();
   const { dagId = "", runId = "", taskId } = useParams();
+  const [searchParams] = useSearchParams();
 
   const selectedVersion = useSelectedVersion();
 
@@ -79,13 +82,41 @@ export const Graph = () => {
   const [dependencies] = useLocalStorage<"all" | "immediate" | "tasks">(`dependencies-${dagId}`, "tasks");
   const [direction] = useLocalStorage<Direction>(`direction-${dagId}`, "RIGHT");
 
+  const rawFilter = (searchParams.get("task_filter") ?? undefined) as TaskFilter | null;
+
+  // default to "all" by default (no pruning). You can change to "both" if you want previous UX.
+  const filter: TaskFilter = rawFilter ?? "all";
+
+  // derive server flags
+  const includeUpstream = filter === "upstream" || filter === "both";
+  const includeDownstream = filter === "downstream" || filter === "both";
+
   const selectedColor = colorMode === "dark" ? selectedDarkColor : selectedLightColor;
+
+  type StructurePruneParams = {
+    includeDownstream?: boolean;
+    includeUpstream?: boolean;
+    root?: string;
+  };
+
+  // Only include server pruning params when filter !== 'all' and a root is provided
+  const structureParams: StructurePruneParams =
+    filter !== "all" && taskId !== undefined && taskId !== ""
+      ? {
+          includeDownstream,
+          includeUpstream,
+          root: taskId,
+        }
+      : {};
+
+  // server structure — pass pruning params conditionally
   const { data: graphData = { edges: [], nodes: [] } } = useStructureServiceStructureData(
     {
       dagId,
       externalDependencies: dependencies === "immediate",
       versionNumber: selectedVersion,
-    },
+      ...structureParams,
+    } as unknown as Parameters<typeof useStructureServiceStructureData>[0],
     undefined,
     { enabled: selectedVersion !== undefined },
   );