Create dag graph with nested groups and join_ids

airflow/ui/package.json

@@ -24,10 +24,14 @@
     "@tanstack/react-table": "^8.20.1",
     "@uiw/codemirror-themes-all": "^4.23.5",
     "@uiw/react-codemirror": "^4.23.5",
+    "@visx/group": "^3.12.0",
+    "@visx/shape": "^3.12.0",
+    "@xyflow/react": "^12.3.5",
     "axios": "^1.7.7",
     "chakra-react-select": "6.0.0-next.2",
     "chart.js": "^4.4.6",
     "dayjs": "^1.11.13",
+    "elkjs": "^0.9.3",
     "next-themes": "^0.3.0",
     "react": "^18.3.1",
     "react-chartjs-2": "^5.2.0",
@@ -69,6 +73,8 @@
     "typescript": "~5.5.4",
     "typescript-eslint": "^8.5.0",
     "vite": "^5.4.6",
-    "vitest": "^2.1.1"
+    "vite-plugin-css-injected-by-js": "^3.5.2",
+    "vitest": "^2.1.1",
+    "web-worker": "^1.3.0"
   }
 }

airflow/ui/src/components/ui/Dialog/CloseTrigger.tsx

@@ -19,21 +19,24 @@
 import { Dialog as ChakraDialog } from "@chakra-ui/react";
 import { forwardRef } from "react";
 
-import { CloseButton } from "../CloseButton";
+import { CloseButton, type CloseButtonProps } from "../CloseButton";
 
-export const CloseTrigger = forwardRef<
-  HTMLButtonElement,
-  ChakraDialog.CloseTriggerProps
->((props, ref) => (
-  <ChakraDialog.CloseTrigger
-    insetEnd="2"
-    position="absolute"
-    top="2"
-    {...props}
-    asChild
-  >
-    <CloseButton ref={ref} size="sm">
-      {props.children}
-    </CloseButton>
-  </ChakraDialog.CloseTrigger>
-));
+type Props = {
+  closeButtonProps?: CloseButtonProps;
+} & ChakraDialog.CloseTriggerProps;
+
+export const CloseTrigger = forwardRef<HTMLButtonElement, Props>(
+  ({ children, closeButtonProps, ...rest }, ref) => (
+    <ChakraDialog.CloseTrigger
+      insetEnd="2"
+      position="absolute"
+      top="2"
+      {...rest}
+      asChild
+    >
+      <CloseButton ref={ref} size="sm" {...closeButtonProps}>
+        {children}
+      </CloseButton>
+    </ChakraDialog.CloseTrigger>
+  ),
+);

airflow/ui/src/context/colorMode/useColorMode.tsx

@@ -25,7 +25,7 @@ export const useColorMode = () => {
   };
 
   return {
-    colorMode: resolvedTheme,
+    colorMode: resolvedTheme as "dark" | "light" | undefined,
     setColorMode: setTheme,
     toggleColorMode,
   };

airflow/ui/src/context/openGroups/OpenGroupsProvider.tsx

@@ -0,0 +1,69 @@
+/*!
+ * 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 {
+  createContext,
+  useCallback,
+  useMemo,
+  type PropsWithChildren,
+} from "react";
+import { useLocalStorage } from "usehooks-ts";
+
+export type OpenGroupsContextType = {
+  openGroupIds: Array<string>;
+  setOpenGroupIds: (groupIds: Array<string>) => void;
+  toggleGroupId: (groupId: string) => void;
+};
+
+export const OpenGroupsContext = createContext<
+  OpenGroupsContextType | undefined
+>(undefined);
+
+type Props = {
+  readonly dagId: string;
+} & PropsWithChildren;
+
+export const OpenGroupsProvider = ({ children, dagId }: Props) => {
+  const openGroupsKey = `${dagId}/open-groups`;
+  const [openGroupIds, setOpenGroupIds] = useLocalStorage<Array<string>>(
+    openGroupsKey,
+    [],
+  );
+
+  const toggleGroupId = useCallback(
+    (groupId: string) => {
+      if (openGroupIds.includes(groupId)) {
+        setOpenGroupIds(openGroupIds.filter((id) => id !== groupId));
+      } else {
+        setOpenGroupIds([...openGroupIds, groupId]);
+      }
+    },
+    [openGroupIds, setOpenGroupIds],
+  );
+
+  const value = useMemo<OpenGroupsContextType>(
+    () => ({ openGroupIds, setOpenGroupIds, toggleGroupId }),
+    [openGroupIds, setOpenGroupIds, toggleGroupId],
+  );
+
+  return (
+    <OpenGroupsContext.Provider value={value}>
+      {children}
+    </OpenGroupsContext.Provider>
+  );
+};

airflow/ui/src/context/openGroups/index.ts

@@ -0,0 +1,21 @@
+/*!
+ * 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.
+ */
+
+export * from "./OpenGroupsProvider";
+export * from "./useOpenGroups";

airflow/ui/src/context/openGroups/useOpenGroups.ts

@@ -0,0 +1,34 @@
+/*!
+ * 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 { useContext } from "react";
+
+import {
+  OpenGroupsContext,
+  type OpenGroupsContextType,
+} from "./OpenGroupsProvider";
+
+export const useOpenGroups = (): OpenGroupsContextType => {
+  const context = useContext(OpenGroupsContext);
+
+  if (context === undefined) {
+    throw new Error("useOpenGroup must be used within a OpenGroupsProvider");
+  }
+
+  return context;
+};

airflow/ui/src/pages/DagsList/Dag/Dag.tsx

@@ -16,26 +16,20 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Box, Button, Tabs } from "@chakra-ui/react";
+import { Box, Button } from "@chakra-ui/react";
 import { FiChevronsLeft } from "react-icons/fi";
-import {
-  Outlet,
-  Link as RouterLink,
-  useLocation,
-  useParams,
-} from "react-router-dom";
+import { Outlet, Link as RouterLink, useParams } from "react-router-dom";
 
 import {
   useDagServiceGetDagDetails,
   useDagsServiceRecentDagRuns,
 } from "openapi/queries";
 import { ErrorAlert } from "src/components/ErrorAlert";
 import { ProgressBar } from "src/components/ui";
-import { capitalize } from "src/utils";
+import { OpenGroupsProvider } from "src/context/openGroups";
 
 import { Header } from "./Header";
-
-const tabs = ["runs", "tasks", "events", "code"];
+import { DagTabs } from "./Tabs";
 
 export const Dag = () => {
   const { dagId } = useParams();
@@ -57,17 +51,12 @@ export const Dag = () => {
     enabled: Boolean(dagId),
   });
 
-  const { pathname } = useLocation();
-
   const runs =
     runsData?.dags.find((dagWithRuns) => dagWithRuns.dag_id === dagId)
       ?.latest_dag_runs ?? [];
 
-  const activeTab =
-    tabs.find((tab) => pathname.endsWith(`/${tab}`)) ?? "overview";
-
   return (
-    <>
+    <OpenGroupsProvider dagId={dagId ?? ""}>
       <Box>
         <Button asChild colorPalette="blue" variant="ghost">
           <RouterLink to="/dags">
@@ -81,24 +70,11 @@ export const Dag = () => {
           size="xs"
           visibility={isLoading || isLoadingRuns ? "visible" : "hidden"}
         />
-        <Tabs.Root value={activeTab}>
-          <Tabs.List>
-            <Tabs.Trigger asChild value="overview">
-              <RouterLink to={`/dags/${dagId}`}>Overview</RouterLink>
-            </Tabs.Trigger>
-            {tabs.map((tab) => (
-              <Tabs.Trigger asChild key={tab} value={tab}>
-                <RouterLink to={`/dags/${dagId}/${tab}`}>
-                  {capitalize(tab)}
-                </RouterLink>
-              </Tabs.Trigger>
-            ))}
-          </Tabs.List>
-        </Tabs.Root>
+        <DagTabs dag={dag} />
       </Box>
       <Box overflow="auto">
         <Outlet />
       </Box>
-    </>
+    </OpenGroupsProvider>
   );
 };

airflow/ui/src/pages/DagsList/Dag/DagVizModal.tsx

@@ -0,0 +1,48 @@
+/*!
+ * 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 { Heading } from "@chakra-ui/react";
+
+import type { DAGResponse } from "openapi/requests/types.gen";
+import { Dialog } from "src/components/ui";
+
+import { Graph } from "./Graph";
+
+type TriggerDAGModalProps = {
+  dagDisplayName: DAGResponse["dag_display_name"];
+  onClose: () => void;
+  open: boolean;
+};
+
+export const DagVizModal: React.FC<TriggerDAGModalProps> = ({
+  dagDisplayName,
+  onClose,
+  open,
+}) => (
+  <Dialog.Root onOpenChange={onClose} open={open} size="full">
+    <Dialog.Content backdrop>
+      <Dialog.Header bg="blue.muted">
+        <Heading size="xl">{dagDisplayName}</Heading>
+        <Dialog.CloseTrigger closeButtonProps={{ size: "xl" }} />
+      </Dialog.Header>
+      <Dialog.Body display="flex">
+        <Graph />
+      </Dialog.Body>
+    </Dialog.Content>
+  </Dialog.Root>
+);