This is an automated email from the ASF dual-hosted git repository.
bbovenzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0138a31bc91 Don't hit API for "all group ids" (#53116)
0138a31bc91 is described below
commit 0138a31bc91e29cb279fb44f4f371578fc33175d
Author: Daniel Standish <[email protected]>
AuthorDate: Wed Jul 16 16:23:03 2025 -0700
Don't hit API for "all group ids" (#53116)
* Add endpoint for group ids
We should not need to query the full grid structure here. All we need is
the group ids for the dag.
* fixup! Add endpoint for group ids
* use context provider to calculate group ids from existing data
* remove recently added graph group ids endpoint
* fix test
* revert debug stuff
* compare with value not just check length
---
.../airflow/ui/src/context/openGroups/Context.ts | 2 ++
.../src/context/openGroups/OpenGroupsProvider.tsx | 6 +++--
.../airflow/ui/src/layouts/Details/Graph/Graph.tsx | 28 +++++++++++++++++-----
.../airflow/ui/src/layouts/Details/Grid/utils.ts | 21 +++++++++++++++-
.../ui/src/layouts/Details/ToggleGroups.tsx | 15 +-----------
5 files changed, 49 insertions(+), 23 deletions(-)
diff --git a/airflow-core/src/airflow/ui/src/context/openGroups/Context.ts
b/airflow-core/src/airflow/ui/src/context/openGroups/Context.ts
index a344ddddd84..b33c5aba550 100644
--- a/airflow-core/src/airflow/ui/src/context/openGroups/Context.ts
+++ b/airflow-core/src/airflow/ui/src/context/openGroups/Context.ts
@@ -19,7 +19,9 @@
import { createContext } from "react";
export type OpenGroupsContextType = {
+ allGroupIds: Array<string>;
openGroupIds: Array<string>;
+ setAllGroupIds: (groupIds: Array<string>) => void;
setOpenGroupIds: (groupIds: Array<string>) => void;
toggleGroupId: (groupId: string) => void;
};
diff --git
a/airflow-core/src/airflow/ui/src/context/openGroups/OpenGroupsProvider.tsx
b/airflow-core/src/airflow/ui/src/context/openGroups/OpenGroupsProvider.tsx
index 3892bd8aea4..84c63a1236b 100644
--- a/airflow-core/src/airflow/ui/src/context/openGroups/OpenGroupsProvider.tsx
+++ b/airflow-core/src/airflow/ui/src/context/openGroups/OpenGroupsProvider.tsx
@@ -27,7 +27,9 @@ type Props = {
export const OpenGroupsProvider = ({ children, dagId }: Props) => {
const openGroupsKey = `${dagId}/open-groups`;
+ const allGroupsKey = `${dagId}/all-groups`;
const [openGroupIds, setOpenGroupIds] =
useLocalStorage<Array<string>>(openGroupsKey, []);
+ const [allGroupIds, setAllGroupIds] =
useLocalStorage<Array<string>>(allGroupsKey, []);
const toggleGroupId = useCallback(
(groupId: string) => {
@@ -41,8 +43,8 @@ export const OpenGroupsProvider = ({ children, dagId }:
Props) => {
);
const value = useMemo<OpenGroupsContextType>(
- () => ({ openGroupIds, setOpenGroupIds, toggleGroupId }),
- [openGroupIds, setOpenGroupIds, toggleGroupId],
+ () => ({ allGroupIds, openGroupIds, setAllGroupIds, setOpenGroupIds,
toggleGroupId }),
+ [allGroupIds, openGroupIds, setAllGroupIds, setOpenGroupIds,
toggleGroupId],
);
return <OpenGroupsContext.Provider
value={value}>{children}</OpenGroupsContext.Provider>;
diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx
b/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx
index 2d24b95a408..c61a63673b6 100644
--- a/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx
+++ b/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx
@@ -19,6 +19,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, useMemo } from "react";
import { useParams } from "react-router-dom";
import { useLocalStorage } from "usehooks-ts";
@@ -30,6 +31,7 @@ import { type Direction, useGraphLayout } from
"src/components/Graph/useGraphLay
import { useColorMode } from "src/context/colorMode";
import { useOpenGroups } from "src/context/openGroups";
import useSelectedVersion from "src/hooks/useSelectedVersion";
+import { flattenGraphNodes } from "src/layouts/Details/Grid/utils.ts";
import { useDependencyGraph } from "src/queries/useDependencyGraph";
import { useGridTiSummaries } from "src/queries/useGridTISummaries.ts";
@@ -71,18 +73,32 @@ export const Graph = () => {
"gray.800",
]);
- const { openGroupIds } = useOpenGroups();
+ const { allGroupIds, openGroupIds, setAllGroupIds } = useOpenGroups();
const [dependencies] = useLocalStorage<"all" | "immediate" |
"tasks">(`dependencies-${dagId}`, "tasks");
const [direction] = useLocalStorage<Direction>(`direction-${dagId}`,
"RIGHT");
const selectedColor = colorMode === "dark" ? selectedDarkColor :
selectedLightColor;
+ const { data: graphData = { edges: [], nodes: [] } } =
useStructureServiceStructureData(
+ {
+ dagId,
+ externalDependencies: dependencies === "immediate",
+ versionNumber: selectedVersion,
+ },
+ undefined,
+ { enabled: selectedVersion !== undefined },
+ );
- const { data: graphData = { edges: [], nodes: [] } } =
useStructureServiceStructureData({
- dagId,
- externalDependencies: dependencies === "immediate",
- versionNumber: selectedVersion,
- });
+ const { allGroupIds: observedGroupIds } = useMemo(
+ () => flattenGraphNodes(graphData.nodes),
+ [graphData.nodes],
+ );
+
+ useEffect(() => {
+ if (observedGroupIds !== allGroupIds) {
+ setAllGroupIds(observedGroupIds);
+ }
+ }, [allGroupIds, observedGroupIds, setAllGroupIds]);
const { data: dagDependencies = { edges: [], nodes: [] } } =
useDependencyGraph(`dag:${dagId}`, {
enabled: dependencies === "all",
diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/Grid/utils.ts
b/airflow-core/src/airflow/ui/src/layouts/Details/Grid/utils.ts
index 6901287d04c..019039a8401 100644
--- a/airflow-core/src/airflow/ui/src/layouts/Details/Grid/utils.ts
+++ b/airflow-core/src/airflow/ui/src/layouts/Details/Grid/utils.ts
@@ -16,7 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
-import type { GridNodeResponse } from "openapi/requests/types.gen";
+import type { GridNodeResponse, NodeResponse } from
"openapi/requests/types.gen";
export type GridTask = {
depth: number;
@@ -54,3 +54,22 @@ export const flattenNodes = (
return { allGroupIds, flatNodes };
};
+
+export const flattenGraphNodes = (
+ nodes: Array<NodeResponse>,
+ depth: number = 0,
+): { allGroupIds: Array<string> } => {
+ let allGroupIds: Array<string> = [];
+
+ nodes.forEach((node) => {
+ if (node.children) {
+ allGroupIds.push(node.id);
+
+ const { allGroupIds: childGroupIds } = flattenGraphNodes(node.children,
depth + 1);
+
+ allGroupIds = [...allGroupIds, ...childGroupIds];
+ }
+ });
+
+ return { allGroupIds };
+};
diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/ToggleGroups.tsx
b/airflow-core/src/airflow/ui/src/layouts/Details/ToggleGroups.tsx
index 69ffc13866a..756a3298a9b 100644
--- a/airflow-core/src/airflow/ui/src/layouts/Details/ToggleGroups.tsx
+++ b/airflow-core/src/airflow/ui/src/layouts/Details/ToggleGroups.tsx
@@ -17,27 +17,14 @@
* under the License.
*/
import { type ButtonGroupProps, IconButton, ButtonGroup } from
"@chakra-ui/react";
-import { useMemo } from "react";
import { useTranslation } from "react-i18next";
import { MdExpand, MdCompress } from "react-icons/md";
-import { useParams } from "react-router-dom";
-import { useLocalStorage } from "usehooks-ts";
import { useOpenGroups } from "src/context/openGroups";
-import { useGridStructure } from "src/queries/useGridStructure.ts";
-
-import { flattenNodes } from "./Grid/utils";
export const ToggleGroups = (props: ButtonGroupProps) => {
const { t: translate } = useTranslation();
- const { openGroupIds, setOpenGroupIds } = useOpenGroups();
- const { dagId = "" } = useParams();
- const [limit] = useLocalStorage<number>(`dag_runs_limit-${dagId}`, 10);
- const { data: dagStructure } = useGridStructure({ limit });
- const { allGroupIds } = useMemo(
- () => flattenNodes(dagStructure, openGroupIds),
- [dagStructure, openGroupIds],
- );
+ const { allGroupIds, openGroupIds, setOpenGroupIds } = useOpenGroups();
// Don't show button if the DAG has no task groups
if (!allGroupIds.length) {
