Add support of different graph directions in the Asset graph view.

airflow-core/src/airflow/ui/src/components/Graph/DirectionDropdown.tsx

@@ -0,0 +1,78 @@
+/*!
+ * 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 { createListCollection, Select, type SelectValueChangeDetails } from "@chakra-ui/react";
+import { useTranslation } from "react-i18next";
+import { useLocalStorage } from "usehooks-ts";
+
+import { directionKey } from "src/constants/localStorage";
+
+export type Direction = "DOWN" | "LEFT" | "RIGHT" | "UP";
+
+export const DirectionDropdown = ({ graphId }: { readonly graphId: string }) => {
+  const { t: translate } = useTranslation(["components", "dag"]);
+
+  const [direction, setDirection] = useLocalStorage<Direction>(directionKey(graphId), "RIGHT");
+
+  const directionOptions = () =>
+    createListCollection({
+      items: [
+        { label: translate("graph.directionRight"), value: "RIGHT" as Direction },
+        { label: translate("graph.directionLeft"), value: "LEFT" as Direction },
+        { label: translate("graph.directionUp"), value: "UP" as Direction },
+        { label: translate("graph.directionDown"), value: "DOWN" as Direction },
+      ],
+    });
+
+  const handleDirectionUpdate = (
+    event: SelectValueChangeDetails<{ label: string; value: Array<string> }>,
+  ) => {
+    if (event.value[0] !== undefined) {
+      setDirection(event.value[0] as Direction);
+    }
+  };
+
+  return (
+    <Select.Root
+      // @ts-expect-error The expected option type is incorrect
+      collection={directionOptions()}
+      onValueChange={handleDirectionUpdate}
+      size="sm"
+      value={[direction]}
+    >
+      <Select.Label fontSize="xs">{translate("dag:panel.graphDirection.label")}</Select.Label>
+      <Select.Control>
+        <Select.Trigger>
+          <Select.ValueText />
+        </Select.Trigger>
+        <Select.IndicatorGroup>
+          <Select.Indicator />
+        </Select.IndicatorGroup>
+      </Select.Control>
+      <Select.Positioner>
+        <Select.Content>
+          {directionOptions().items.map((option) => (
+            <Select.Item item={option} key={option.value}>
+              {option.label}
+            </Select.Item>
+          ))}
+        </Select.Content>
+      </Select.Positioner>
+    </Select.Root>
+  );
+};

airflow-core/src/airflow/ui/src/components/Graph/elkGraphUtils.ts

@@ -20,7 +20,7 @@ import type { ElkNode, ElkExtendedEdge, ElkShape } from "elkjs";
 
 import type { EdgeResponse, NodeResponse } from "openapi/requests/types.gen";
 
-import type { Direction } from "./useGraphLayout";
+import type { Direction } from "./DirectionDropdown";
 
 // ---------------------------------------------------------------------------
 // Types

airflow-core/src/airflow/ui/src/components/Graph/useGraphLayout.ts

@@ -16,7 +16,6 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { createListCollection } from "@chakra-ui/react";
 import { useQuery } from "@tanstack/react-query";
 import ELK, { type ElkNode } from "elkjs";
 // ?raw imports the file content as a plain string without any transformation.
@@ -26,10 +25,10 @@ import ELK, { type ElkNode } from "elkjs";
 // URL-based worker approaches (?worker, ?worker&inline, new URL()) resolve to
 // the Vite origin which browsers reject for Workers.
 import ElkWorkerSource from "elkjs/lib/elk-worker.min.js?raw";
-import type { TFunction } from "i18next";
 
 import type { NodeResponse, StructureDataResponse } from "openapi/requests/types.gen";
 
+import type { Direction } from "./DirectionDropdown";
 import { generateElkGraph } from "./elkGraphUtils";
 import { flattenGraph, formatFlowEdges } from "./reactflowUtils";
 
@@ -43,17 +42,6 @@ const elk = new ELK({
   workerFactory: () => new Worker(elkWorkerBlobUrl, { type: "classic" }),
 });
 
-export type Direction = "DOWN" | "LEFT" | "RIGHT" | "UP";
-export const directionOptions = (translate: TFunction) =>
-  createListCollection({
-    items: [
-      { label: translate("graph.directionRight"), value: "RIGHT" as Direction },
-      { label: translate("graph.directionLeft"), value: "LEFT" as Direction },
-      { label: translate("graph.directionUp"), value: "UP" as Direction },
-      { label: translate("graph.directionDown"), value: "DOWN" as Direction },
-    ],
-  });
-
 export type LayoutNode = ElkNode & NodeResponse;
 
 type LayoutProps = {

airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx

@@ -24,10 +24,11 @@ import { useParams } from "react-router-dom";
 import { useLocalStorage } from "usehooks-ts";
 
 import { useDagRunServiceGetDagRun, useStructureServiceStructureData } from "openapi/queries";
+import type { Direction } from "src/components/Graph/DirectionDropdown";
 import { DownloadButton } from "src/components/Graph/DownloadButton";
 import { edgeTypes, nodeTypes } from "src/components/Graph/graphTypes";
 import type { CustomNodeProps } from "src/components/Graph/reactflowUtils";
-import { type Direction, useGraphLayout } from "src/components/Graph/useGraphLayout";
+import { useGraphLayout } from "src/components/Graph/useGraphLayout";
 import { dependenciesKey, directionKey } from "src/constants/localStorage";
 import { useColorMode } from "src/context/colorMode";
 import { useGroups } from "src/context/groups";

airflow-core/src/airflow/ui/src/layouts/Details/PanelButtons.tsx

@@ -42,12 +42,12 @@ import { useParams } from "react-router-dom";
 import { useLocalStorage } from "usehooks-ts";
 
 import { DagVersionSelect } from "src/components/DagVersionSelect";
-import { directionOptions, type Direction } from "src/components/Graph/useGraphLayout";
+import { DirectionDropdown } from "src/components/Graph/DirectionDropdown";
 import { GraphTaskFilters } from "src/components/GraphTaskFilters";
 import { Tooltip } from "src/components/ui";
 import { type ButtonGroupOption, ButtonGroupToggle } from "src/components/ui/ButtonGroupToggle";
 import type { DagView } from "src/constants/dagView";
-import { dependenciesKey, directionKey } from "src/constants/localStorage";
+import { dependenciesKey } from "src/constants/localStorage";
 import type { VersionIndicatorOptions } from "src/constants/showVersionIndicatorOptions";
 import { useContainerWidth } from "src/utils/useContainerWidth";
 
@@ -118,7 +118,6 @@ export const PanelButtons = ({
     dependenciesKey(dagId),
     "tasks",
   );
-  const [direction, setDirection] = useLocalStorage<Direction>(directionKey(dagId), "RIGHT");
   const containerRef = useRef<HTMLDivElement>(null);
   const containerWidth = useContainerWidth(containerRef);
   const handleLimitChange = (event: SelectValueChangeDetails<{ label: string; value: Array<string> }>) => {
@@ -148,14 +147,6 @@ export const PanelButtons = ({
     }
   };
 
-  const handleDirectionUpdate = (
-    event: SelectValueChangeDetails<{ label: string; value: Array<string> }>,
-  ) => {
-    if (event.value[0] !== undefined) {
-      setDirection(event.value[0] as Direction);
-    }
-  };
-
   const handleFocus = (view: string) => {
     if (panelGroupRef.current) {
       const newLayout = view === "graph" ? [70, 30] : [30, 70];
@@ -279,34 +270,7 @@ export const PanelButtons = ({
                           </Select.Positioner>
                         </Select.Root>
 
-                        <Select.Root
-                          // @ts-expect-error The expected option type is incorrect
-                          collection={directionOptions(translate)}
-                          onValueChange={handleDirectionUpdate}
-                          size="sm"
-                          value={[direction]}
-                        >
-                          <Select.Label fontSize="xs">
-                            {translate("dag:panel.graphDirection.label")}
-                          </Select.Label>
-                          <Select.Control>
-                            <Select.Trigger>
-                              <Select.ValueText />
-                            </Select.Trigger>
-                            <Select.IndicatorGroup>
-                              <Select.Indicator />
-                            </Select.IndicatorGroup>
-                          </Select.Control>
-                          <Select.Positioner>
-                            <Select.Content>
-                              {directionOptions(translate).items.map((option) => (
-                                <Select.Item item={option} key={option.value}>
-                                  {option.label}
-                                </Select.Item>
-                              ))}
-                            </Select.Content>
-                          </Select.Positioner>
-                        </Select.Root>
+                        <DirectionDropdown graphId={dagId} />
                       </>
                     ) : (
                       <>

airflow-core/src/airflow/ui/src/pages/Asset/AssetGraph.tsx

@@ -20,12 +20,15 @@ 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 { useLocalStorage } from "usehooks-ts";
 
 import type { AssetResponse } from "openapi/requests/types.gen";
+import type { Direction } from "src/components/Graph/DirectionDropdown";
 import { DownloadButton } from "src/components/Graph/DownloadButton";
 import { edgeTypes, nodeTypes } from "src/components/Graph/graphTypes";
 import type { CustomNodeProps } from "src/components/Graph/reactflowUtils";
 import { useGraphLayout } from "src/components/Graph/useGraphLayout";
+import { directionKey } from "src/constants/localStorage";
 import { useColorMode } from "src/context/colorMode";
 import { useDependencyGraph } from "src/queries/useDependencyGraph";
 import { getReactFlowThemeStyle } from "src/theme";
@@ -45,9 +48,11 @@ export const AssetGraph = ({
     dependencyType,
   });
 
+  const [direction] = useLocalStorage<Direction>(directionKey(assetId ?? ""), "RIGHT");
+
   const { data: layoutData } = useGraphLayout({
     ...graphData,
-    direction: "RIGHT",
+    direction,
     openGroupIds: [],
   });
 

airflow-core/src/airflow/ui/src/pages/Asset/AssetPanelButtons.tsx

@@ -16,8 +16,12 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Box, Button, ButtonGroup } from "@chakra-ui/react";
+import { Box, Button, ButtonGroup, Flex, IconButton, Popover, Portal } from "@chakra-ui/react";
 import { useTranslation } from "react-i18next";
+import { MdSettings } from "react-icons/md";
+import { useParams } from "react-router-dom";
+
+import { DirectionDropdown } from "src/components/Graph/DirectionDropdown";
 
 type Props = {
   readonly dependencyType: "data" | "scheduling";
@@ -26,27 +30,60 @@ type Props = {
 
 export const AssetPanelButtons = ({ dependencyType, setDependencyType }: Props) => {
   const { t: translate } = useTranslation(["assets"]);
+  const { assetId } = useParams();
 
   return (
     <Box borderRadius="md" position="absolute" px={2} py={1} right={2} top={1} zIndex={1}>
-      <ButtonGroup attached size="sm" variant="outline">
-        <Button
-          bg={dependencyType === "scheduling" ? "brand.500" : "bg.subtle"}
-          color={dependencyType === "scheduling" ? "white" : "fg.default"}
-          colorPalette="brand"
-          onClick={() => setDependencyType("scheduling")}
-        >
-          {translate("assets:scheduling")}
-        </Button>
-        <Button
-          bg={dependencyType === "data" ? "brand.500" : "bg.subtle"}
-          color={dependencyType === "data" ? "white" : "fg.default"}
-          colorPalette="brand"
-          onClick={() => setDependencyType("data")}
-        >
-          {translate("assets:taskDependencies")}
-        </Button>
-      </ButtonGroup>
+      <Flex justifyContent="space-between">
+        <ButtonGroup attached size="sm" variant="outline">
+          <Button
+            bg={dependencyType === "scheduling" ? "brand.500" : "bg.subtle"}
+            color={dependencyType === "scheduling" ? "white" : "fg.default"}
+            colorPalette="brand"
+            onClick={() => setDependencyType("scheduling")}
+          >
+            {translate("assets:scheduling")}
+          </Button>
+          <Button
+            bg={dependencyType === "data" ? "brand.500" : "bg.subtle"}
+            color={dependencyType === "data" ? "white" : "fg.default"}
+            colorPalette="brand"
+            onClick={() => setDependencyType("data")}
+          >
+            {translate("assets:taskDependencies")}
+          </Button>
+        </ButtonGroup>
+        <Popover.Root positioning={{ placement: "bottom-end" }}>
+          <Popover.Trigger asChild>
+            <IconButton
+              aria-label={translate("dag:panel.buttons.options")}
+              colorPalette="brand"
+              size="md"
+              title={translate("dag:panel.buttons.options")}
+              variant="ghost"
+            >
+              <MdSettings />
+            </IconButton>
+          </Popover.Trigger>
+          <Portal>
+            <Popover.Positioner>
+              <Popover.Content>
+                <Popover.Arrow />
+                <Popover.Body
+                  display="flex"
+                  flexDirection="column"
+                  gap={4}
+                  maxH="70vh"
+                  overflowY="auto"
+                  p={2}
+                >
+                  <DirectionDropdown graphId={assetId ?? ""} />
+                </Popover.Body>
+              </Popover.Content>
+            </Popover.Positioner>
+          </Portal>
+        </Popover.Root>
+      </Flex>
     </Box>
   );
 };

airflow-core/src/airflow/ui/testsSetup.ts

@@ -28,7 +28,6 @@ import { handlers } from "src/mocks/handlers";
 // happy-dom. Mock useGraphLayout so the Worker is never constructed. Any
 // test that specifically exercises graph layout should override this mock.
 vi.mock("src/components/Graph/useGraphLayout", () => ({
-  directionOptions: () => ({ items: [] }),
   useGraphLayout: vi.fn().mockReturnValue({ data: undefined, isPending: false }),
 }));