Use reactflow for datasets graph (#31775)

apache · Jul 2, 2023 · 00bf2ab · 00bf2ab
1 parent 05bd90f
commit 00bf2ab
Show file tree

Hide file tree

Showing 15 changed files with 252 additions and 353 deletions.
diff --git a/airflow/www/package.json b/airflow/www/package.json
@@ -92,9 +92,7 @@
     "@emotion/react": "^11.9.3",
     "@emotion/styled": "^11",
     "@visx/group": "^2.10.0",
-    "@visx/marker": "^2.12.2",
     "@visx/shape": "^2.12.2",
-    "@visx/zoom": "^2.10.0",
     "axios": "^0.26.0",
     "bootstrap-3-typeahead": "^4.0.2",
     "camelcase-keys": "^7.0.0",

diff --git a/airflow/www/static/js/api/useDatasetDependencies.ts b/airflow/www/static/js/api/useDatasetDependencies.ts
@@ -22,11 +22,10 @@ import { useQuery } from "react-query";
 import ELK, { ElkShape, ElkExtendedEdge } from "elkjs";
 
 import { getMetaValue } from "src/utils";
-import type { DepEdge, DepNode } from "src/types";
-import type { NodeType } from "src/datasets/Graph/Node";
-
 import { getTextWidth } from "src/utils/graph";
 
+import type { NodeType, DepEdge, DepNode } from "src/types";
+
 interface DatasetDependencies {
   edges: DepEdge[];
   nodes: DepNode[];

diff --git a/airflow/www/static/js/api/useGraphData.ts b/airflow/www/static/js/api/useGraphData.ts
@@ -21,12 +21,12 @@ import { useQuery } from "react-query";
 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";
+import type { WebserverEdge, DepNode } from "src/types";
 
 const DAG_ID_PARAM = "dag_id";
 
@@ -38,11 +38,6 @@ interface GraphData {
   nodes: DepNode;
   arrange: string;
 }
-export interface WebserverEdge {
-  label?: string;
-  sourceId: string;
-  targetId: string;
-}
 
 const useGraphData = () => {
   const {

diff --git a/.../www/static/js/dag/details/graph/Edge.tsx → ...w/www/static/js/components/Graph/Edge.tsx b/.../www/static/js/dag/details/graph/Edge.tsx → ...w/www/static/js/components/Graph/Edge.tsx
@@ -18,11 +18,11 @@
  */
 
 import React from "react";
-import { Text } from "@chakra-ui/react";
+import { Text, useTheme } from "@chakra-ui/react";
 import type { ElkEdgeSection, ElkLabel, ElkPoint, LayoutOptions } from "elkjs";
 
-import Edge from "src/datasets/Graph/Edge";
 import { Group } from "@visx/group";
+import { LinePath } from "@visx/shape";
 
 interface EdgeProps {
   data?: {
@@ -40,6 +40,7 @@ interface EdgeProps {
 }
 
 const CustomEdge = ({ data }: EdgeProps) => {
+  const { colors } = useTheme();
   if (!data) return null;
   const { rest } = data;
   return (
@@ -54,11 +55,16 @@ const CustomEdge = ({ data }: EdgeProps) => {
           </Group>
         );
       })}
-      <Edge
-        edge={rest}
-        showMarker={false}
-        isSelected={rest.isSelected || undefined}
-      />
+      {(rest.sections || []).map((s) => (
+        <LinePath
+          key={s.id}
+          stroke={rest.isSelected ? colors.blue[400] : colors.gray[400]}
+          strokeWidth={rest.isSelected ? 3 : 2}
+          x={(d) => d.x || 0}
+          y={(d) => d.y || 0}
+          data={[s.startPoint, ...(s.bendPoints || []), s.endPoint]}
+        />
+      ))}
     </>
   );
 };

diff --git a/airflow/www/static/js/dag/details/graph/index.tsx b/airflow/www/static/js/dag/details/graph/index.tsx
@@ -38,8 +38,8 @@ 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 "src/components/Graph/Edge";
 
-import Edge from "./Edge";
 import Node, { CustomNodeProps } from "./Node";
 import { buildEdges, nodeStrokeColor, nodeColor, flattenNodes } from "./utils";
 

diff --git a/airflow/www/static/js/dag/details/graph/utils.ts b/airflow/www/static/js/dag/details/graph/utils.ts
@@ -23,8 +23,7 @@ import type { ElkExtendedEdge } from "elkjs";
 
 import type { SelectionProps } from "src/dag/useSelection";
 import { getTask } from "src/utils";
-import type { Task, TaskInstance } from "src/types";
-import type { NodeType } from "src/datasets/Graph/Node";
+import type { Task, TaskInstance, NodeType } from "src/types";
 
 import type { CustomNodeProps } from "./Node";
 

diff --git a/airflow/www/static/js/datasets/Graph/DagNode.tsx b/airflow/www/static/js/datasets/Graph/DagNode.tsx
@@ -33,6 +33,7 @@ import {
   useTheme,
 } from "@chakra-ui/react";
 import { MdOutlineAccountTree } from "react-icons/md";
+
 import { useContainerRef } from "src/context/containerRef";
 import { getMetaValue } from "src/utils";
 
@@ -41,7 +42,7 @@ const DagNode = ({
   isHighlighted,
 }: {
   dagId: string;
-  isHighlighted: boolean;
+  isHighlighted?: boolean;
 }) => {
   const { colors } = useTheme();
   const containerRef = useContainerRef();
@@ -63,7 +64,7 @@ const DagNode = ({
           alignItems="center"
         >
           <MdOutlineAccountTree size="16px" />
-          <Text>{dagId}</Text>
+          <Text ml={2}>{dagId}</Text>
         </Flex>
       </PopoverTrigger>
       <Portal containerRef={containerRef}>

diff --git a/airflow/www/static/js/datasets/Graph/Edge.tsx b/airflow/www/static/js/datasets/Graph/Edge.tsx
diff --git a/airflow/www/static/js/datasets/Graph/Legend.tsx b/airflow/www/static/js/datasets/Graph/Legend.tsx
@@ -18,59 +18,29 @@
  */
 
 import React from "react";
-import { Flex, Box, IconButton, Text } from "@chakra-ui/react";
-import {
-  MdOutlineZoomOutMap,
-  MdFilterCenterFocus,
-  MdOutlineAccountTree,
-} from "react-icons/md";
+import { Flex, Box, Text } from "@chakra-ui/react";
+import { MdOutlineAccountTree } from "react-icons/md";
 import { HiDatabase } from "react-icons/hi";
 
-interface Props {
-  zoom: any;
-  center: () => void;
-}
-
-const Legend = ({ zoom, center }: Props) => (
-  <Flex height="100%" flexDirection="column" justifyContent="space-between">
-    <Box>
-      <IconButton
-        onClick={zoom.reset}
-        fontSize="2xl"
-        m={2}
-        title="Reset zoom"
-        aria-label="Reset zoom"
-        icon={<MdOutlineZoomOutMap />}
-      />
-      <IconButton
-        onClick={center}
-        fontSize="2xl"
-        m={2}
-        title="Center"
-        aria-label="Center"
-        icon={<MdFilterCenterFocus />}
-      />
-    </Box>
-    <Box
-      backgroundColor="white"
-      p={2}
-      borderColor="gray.200"
-      borderRightWidth={1}
-      borderTopWidth={1}
-    >
-      <Text>Legend</Text>
-      <Flex>
-        <Flex mr={2} alignItems="center">
-          <MdOutlineAccountTree size="16px" />
-          <Text ml={1}>DAG</Text>
-        </Flex>
-        <Flex alignItems="center">
-          <HiDatabase size="16px" />
-          <Text ml={1}>Dataset</Text>
-        </Flex>
+const Legend = () => (
+  <Box
+    backgroundColor="white"
+    p={2}
+    borderColor="gray.200"
+    borderWidth={1}
+    fontSize={14}
+  >
+    <Flex>
+      <Flex mr={2} alignItems="center">
+        <MdOutlineAccountTree size="14px" />
+        <Text ml={1}>DAG</Text>
+      </Flex>
+      <Flex alignItems="center">
+        <HiDatabase size="14px" />
+        <Text ml={1}>Dataset</Text>
       </Flex>
-    </Box>
-  </Flex>
+    </Flex>
+  </Box>
 );
 
 export default Legend;