Link graph to assets

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<NodeType<CustomNodeProps, "asset">>) => {
   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 (
     <NodeWrapper>
       <Flex
@@ -67,7 +69,9 @@ export const AssetNode = ({
           <Heading ml={-2} size="sm">
             <FiDatabase />
           </Heading>
-          {label}
+          <Link asChild color="fg.info" mb={2}>
+            <RouterLink to={`/assets/${assetId}`}>{label}</RouterLink>
+          </Link>
         </HStack>
         {datasetEvent === undefined ? undefined : (
           <>

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;

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<string>;
   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],
   });

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

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: [],
   });

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<UseQueryOptions<BaseGraphResponse, unknown>, "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;
+};