AIP-38 Plug Graph on backend endpoint (apache#44572)

* WIP functionnaly OK * Fix CI
VladaZakharova · Dec 4, 2024 · 99130ab · 99130ab
1 parent 2e8f456
commit 99130ab
Show file tree

Hide file tree

Showing 13 changed files with 69 additions and 284 deletions.
diff --git a/airflow/api_fastapi/core_api/datamodels/ui/structure.py b/airflow/api_fastapi/core_api/datamodels/ui/structure.py
@@ -34,12 +34,12 @@ class NodeResponse(BaseModel):
     """Node serializer for responses."""
 
     children: list[NodeResponse] | None = None
-    id: str | None
+    id: str
     is_mapped: bool | None = None
-    label: str | None = None
+    label: str
     tooltip: str | None = None
     setup_teardown_type: Literal["setup", "teardown"] | None = None
-    type: Literal["join", "sensor", "task", "task_group"]
+    type: Literal["join", "sensor", "task", "asset_condition"]
 
 
 class StructureDataResponse(BaseModel):

diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
@@ -7782,19 +7782,15 @@ components:
           - type: 'null'
           title: Children
         id:
-          anyOf:
-          - type: string
-          - type: 'null'
+          type: string
           title: Id
         is_mapped:
           anyOf:
           - type: boolean
           - type: 'null'
           title: Is Mapped
         label:
-          anyOf:
-          - type: string
-          - type: 'null'
+          type: string
           title: Label
         tooltip:
           anyOf:
@@ -7815,11 +7811,12 @@ components:
           - join
           - sensor
           - task
-          - task_group
+          - asset_condition
           title: Type
       type: object
       required:
       - id
+      - label
       - type
       title: NodeResponse
       description: Node serializer for responses.

diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -3128,14 +3128,7 @@ export const $NodeResponse = {
       title: "Children",
     },
     id: {
-      anyOf: [
-        {
-          type: "string",
-        },
-        {
-          type: "null",
-        },
-      ],
+      type: "string",
       title: "Id",
     },
     is_mapped: {
@@ -3150,14 +3143,7 @@ export const $NodeResponse = {
       title: "Is Mapped",
     },
     label: {
-      anyOf: [
-        {
-          type: "string",
-        },
-        {
-          type: "null",
-        },
-      ],
+      type: "string",
       title: "Label",
     },
     tooltip: {
@@ -3185,12 +3171,12 @@ export const $NodeResponse = {
     },
     type: {
       type: "string",
-      enum: ["join", "sensor", "task", "task_group"],
+      enum: ["join", "sensor", "task", "asset_condition"],
       title: "Type",
     },
   },
   type: "object",
-  required: ["id", "type"],
+  required: ["id", "label", "type"],
   title: "NodeResponse",
   description: "Node serializer for responses.",
 } as const;

diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -761,15 +761,15 @@ export type JobResponse = {
  */
 export type NodeResponse = {
   children?: Array<NodeResponse> | null;
-  id: string | null;
+  id: string;
   is_mapped?: boolean | null;
-  label?: string | null;
+  label: string;
   tooltip?: string | null;
   setup_teardown_type?: "setup" | "teardown" | null;
-  type: "join" | "sensor" | "task" | "task_group";
+  type: "join" | "sensor" | "task" | "asset_condition";
 };
 
-export type type = "join" | "sensor" | "task" | "task_group";
+export type type = "join" | "sensor" | "task" | "asset_condition";
 
 /**
  * Request body for Clear Task Instances endpoint.

diff --git a/airflow/ui/src/pages/DagsList/Dag/DagVizModal.tsx b/airflow/ui/src/pages/DagsList/Dag/DagVizModal.tsx
@@ -24,24 +24,28 @@ import { Dialog } from "src/components/ui";
 import { Graph } from "./Graph";
 
 type TriggerDAGModalProps = {
-  dagDisplayName: DAGResponse["dag_display_name"];
+  dagDisplayName?: DAGResponse["dag_display_name"];
+  dagId?: DAGResponse["dag_id"];
   onClose: () => void;
   open: boolean;
 };
 
 export const DagVizModal: React.FC<TriggerDAGModalProps> = ({
   dagDisplayName,
+  dagId,
   onClose,
   open,
 }) => (
   <Dialog.Root onOpenChange={onClose} open={open} size="full">
     <Dialog.Content backdrop>
       <Dialog.Header bg="blue.muted">
-        <Heading size="xl">{dagDisplayName}</Heading>
+        <Heading size="xl">
+          {Boolean(dagDisplayName) ? dagDisplayName : "Dag Undefined"}
+        </Heading>
         <Dialog.CloseTrigger closeButtonProps={{ size: "xl" }} />
       </Dialog.Header>
       <Dialog.Body display="flex">
-        <Graph />
+        {dagId === undefined ? undefined : <Graph dagId={dagId} />}
       </Dialog.Body>
     </Dialog.Content>
   </Dialog.Root>

diff --git a/airflow/ui/src/pages/DagsList/Dag/Graph/Graph.tsx b/airflow/ui/src/pages/DagsList/Dag/Graph/Graph.tsx
@@ -20,13 +20,14 @@ import { Flex } from "@chakra-ui/react";
 import { ReactFlow, Controls, Background, MiniMap } from "@xyflow/react";
 import "@xyflow/react/dist/style.css";
 
+import { useStructureServiceStructureData } from "openapi/queries";
+import type { DAGResponse } from "openapi/requests/types.gen";
 import { useColorMode } from "src/context/colorMode";
 import { useOpenGroups } from "src/context/openGroups";
 
 import Edge from "./Edge";
 import { JoinNode } from "./JoinNode";
 import { TaskNode } from "./TaskNode";
-import { graphData } from "./data";
 import { useGraphLayout } from "./useGraphLayout";
 
 const nodeTypes = {
@@ -35,10 +36,16 @@ const nodeTypes = {
 };
 const edgeTypes = { custom: Edge };
 
-export const Graph = () => {
+export const Graph = ({ dagId }: { readonly dagId: DAGResponse["dag_id"] }) => {
   const { colorMode } = useColorMode();
 
   const { openGroupIds } = useOpenGroups();
+
+  const { data: graphData = { arrange: "LR", edges: [], nodes: [] } } =
+    useStructureServiceStructureData({
+      dagId,
+    });
+
   const { data } = useGraphLayout({
     ...graphData,
     openGroupIds,

diff --git a/airflow/ui/src/pages/DagsList/Dag/Graph/TaskName.tsx b/airflow/ui/src/pages/DagsList/Dag/Graph/TaskName.tsx
@@ -20,7 +20,7 @@ import { Text, type TextProps } from "@chakra-ui/react";
 import type { CSSProperties } from "react";
 import { FiArrowUpRight, FiArrowDownRight } from "react-icons/fi";
 
-import type { Node } from "./data";
+import type { NodeResponse } from "openapi/requests/types.gen";
 
 type Props = {
   readonly id: string;
@@ -29,7 +29,7 @@ type Props = {
   readonly isOpen?: boolean;
   readonly isZoomedOut?: boolean;
   readonly label: string;
-  readonly setupTeardownType?: Node["setup_teardown_type"];
+  readonly setupTeardownType?: NodeResponse["setup_teardown_type"];
 } & TextProps;
 
 export const TaskName = ({