diff --git a/airflow-core/src/airflow/ui/src/components/Graph/AssetNode.tsx b/airflow-core/src/airflow/ui/src/components/Graph/AssetNode.tsx index 828cd382bbb86..025eafabfb8ee 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/AssetNode.tsx +++ b/airflow-core/src/airflow/ui/src/components/Graph/AssetNode.tsx @@ -16,10 +16,10 @@ * specific language governing permissions and limitations * under the License. */ -import { Flex, Heading, HStack, Text } from "@chakra-ui/react"; +import { Flex, Heading, HStack, Link, Text } from "@chakra-ui/react"; import type { NodeProps, Node as NodeType } from "@xyflow/react"; import { FiDatabase } from "react-icons/fi"; -import { useParams } from "react-router-dom"; +import { useParams, Link as RouterLink } from "react-router-dom"; import { useAssetServiceGetAssetEvents, useDagRunServiceGetUpstreamAssetEvents } from "openapi/queries"; import { pluralize } from "src/utils"; @@ -29,7 +29,7 @@ import { NodeWrapper } from "./NodeWrapper"; import type { CustomNodeProps } from "./reactflowUtils"; export const AssetNode = ({ - data: { height, isSelected, label, width }, + data: { height, id, isSelected, label, width }, }: NodeProps>) => { const { dagId = "", runId = "" } = useParams(); const { data: upstreamEventsData } = useDagRunServiceGetUpstreamAssetEvents( @@ -49,6 +49,8 @@ export const AssetNode = ({ ...(downstreamEventsData?.asset_events ?? []), ].find((event) => event.name === label); + const assetId = id.replace("asset:", ""); + return ( - {label} + + {label} + {datasetEvent === undefined ? undefined : ( <> diff --git a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts index 24df2f9541c08..cf704be094a5f 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts +++ b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts @@ -28,6 +28,7 @@ export type CustomNodeProps = { childCount?: number; depth?: number; height?: number; + id: string; isGroup?: boolean; isMapped?: boolean; isOpen?: boolean; diff --git a/airflow-core/src/airflow/ui/src/components/Graph/useGraphLayout.ts b/airflow-core/src/airflow/ui/src/components/Graph/useGraphLayout.ts index cef1d48925182..bb46eb7699652 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/useGraphLayout.ts +++ b/airflow-core/src/airflow/ui/src/components/Graph/useGraphLayout.ts @@ -19,12 +19,7 @@ import { useQuery } from "@tanstack/react-query"; import ELK, { type ElkNode, type ElkExtendedEdge, type ElkShape } from "elkjs"; -import type { - DAGResponse, - EdgeResponse, - NodeResponse, - StructureDataResponse, -} from "openapi/requests/types.gen"; +import type { EdgeResponse, NodeResponse, StructureDataResponse } from "openapi/requests/types.gen"; import { flattenGraph, formatFlowEdges } from "./reactflowUtils"; @@ -220,14 +215,12 @@ const generateElkGraph = ({ }; type LayoutProps = { - dagId: DAGResponse["dag_id"]; direction: Direction; openGroupIds: Array; versionNumber?: number; } & StructureDataResponse; export const useGraphLayout = ({ - dagId, direction = "RIGHT", edges, nodes, @@ -265,5 +258,5 @@ export const useGraphLayout = ({ return { edges: formattedEdges, nodes: flattenedData.nodes }; }, - queryKey: ["graphLayout", nodes, openGroupIds, dagId, versionNumber, edges], + queryKey: ["graphLayout", nodes, openGroupIds, versionNumber, edges], }); 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 f16e40671178f..8f4148c518926 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 @@ -24,7 +24,6 @@ import { useLocalStorage } from "usehooks-ts"; import { useDagRunServiceGetDagRun, - useDependenciesServiceGetDependencies, useGridServiceGridData, useStructureServiceStructureData, } from "openapi/queries"; @@ -40,6 +39,7 @@ import { useGraphLayout } from "src/components/Graph/useGraphLayout"; import { useColorMode } from "src/context/colorMode"; import { useOpenGroups } from "src/context/openGroups"; import useSelectedVersion from "src/hooks/useSelectedVersion"; +import { useDependencyGraph } from "src/queries/useDependencyGraph"; import { isStatePending, useAutoRefresh } from "src/utils"; const nodeColor = ( @@ -103,11 +103,9 @@ export const Graph = () => { versionNumber: selectedVersion, }); - const { data: dagDependencies = { edges: [], nodes: [] } } = useDependenciesServiceGetDependencies( - { nodeId: `dag:${dagId}` }, - undefined, - { enabled: dependencies === "all" }, - ); + const { data: dagDependencies = { edges: [], nodes: [] } } = useDependencyGraph(`dag:${dagId}`, { + enabled: dependencies === "all", + }); const { data: dagRun } = useDagRunServiceGetDagRun( { @@ -122,7 +120,6 @@ export const Graph = () => { const dagDepNodes = dependencies === "all" ? dagDependencies.nodes : []; const { data } = useGraphLayout({ - dagId, direction: "RIGHT", edges: [...graphData.edges, ...dagDepEdges], nodes: dagDepNodes.length diff --git a/airflow-core/src/airflow/ui/src/pages/Asset/AssetGraph.tsx b/airflow-core/src/airflow/ui/src/pages/Asset/AssetGraph.tsx index 72d52fbc10a42..7fe089ea3a597 100644 --- a/airflow-core/src/airflow/ui/src/pages/Asset/AssetGraph.tsx +++ b/airflow-core/src/airflow/ui/src/pages/Asset/AssetGraph.tsx @@ -19,8 +19,8 @@ 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 { useParams } from "react-router-dom"; -import { useDependenciesServiceGetDependencies } from "openapi/queries"; import type { AssetResponse } from "openapi/requests/types.gen"; import { AliasNode } from "src/components/Graph/AliasNode"; import { AssetNode } from "src/components/Graph/AssetNode"; @@ -29,6 +29,7 @@ import Edge from "src/components/Graph/Edge"; import type { CustomNodeProps } from "src/components/Graph/reactflowUtils"; import { useGraphLayout } from "src/components/Graph/useGraphLayout"; import { useColorMode } from "src/context/colorMode"; +import { useDependencyGraph } from "src/queries/useDependencyGraph"; const nodeTypes = { asset: AssetNode, @@ -38,17 +39,13 @@ const nodeTypes = { const edgeTypes = { custom: Edge }; export const AssetGraph = ({ asset }: { readonly asset?: AssetResponse }) => { + const { assetId } = useParams(); const { colorMode = "light" } = useColorMode(); - const { data = { edges: [], nodes: [] } } = useDependenciesServiceGetDependencies( - { nodeId: `asset:${asset?.id}` }, - undefined, - { enabled: Boolean(asset) && Boolean(asset?.name) }, - ); + const { data = { edges: [], nodes: [] } } = useDependencyGraph(`asset:${assetId}`); const { data: graphData } = useGraphLayout({ ...data, - dagId: asset?.name ?? "", direction: "RIGHT", openGroupIds: [], }); diff --git a/airflow-core/src/airflow/ui/src/queries/useDependencyGraph.ts b/airflow-core/src/airflow/ui/src/queries/useDependencyGraph.ts new file mode 100644 index 0000000000000..d781e0d169892 --- /dev/null +++ b/airflow-core/src/airflow/ui/src/queries/useDependencyGraph.ts @@ -0,0 +1,47 @@ +/*! + * 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 { useQueryClient, type UseQueryOptions } from "@tanstack/react-query"; + +import { + useDependenciesServiceGetDependencies, + UseDependenciesServiceGetDependenciesKeyFn, +} from "openapi/queries"; +import type { BaseGraphResponse } from "openapi/requests/types.gen"; + +export const useDependencyGraph = ( + nodeId: string, + options?: Omit, "queryFn" | "queryKey">, +) => { + const queryClient = useQueryClient(); + + const query = useDependenciesServiceGetDependencies( + { + nodeId, + }, + undefined, + options, + ); + + // Update the queries for all connected assets and dags so we save an API request + query.data?.nodes.forEach((node) => { + queryClient.setQueryData(UseDependenciesServiceGetDependenciesKeyFn({ nodeId: node.id }), query.data); + }); + + return query; +};