Skip to content

Commit

Permalink
Create dag graph with nested groups and join_ids (#44199)
Browse files Browse the repository at this point in the history
* Create dag graph with nested groups and join_ids

* Add edge labels, setup/teardown tasks, and mapped tasks

* move all graph formatting utils inside of useGraphLayout hook

* move opengroup logic to a local context provider

* Add vite plugin to fix graph css

* Remove commented out code
  • Loading branch information
bbovenzi authored Nov 21, 2024
1 parent 970cb27 commit 5a0272c
Show file tree
Hide file tree
Showing 21 changed files with 1,823 additions and 50 deletions.
8 changes: 7 additions & 1 deletion airflow/ui/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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"
}
}
455 changes: 455 additions & 0 deletions airflow/ui/pnpm-lock.yaml

Large diffs are not rendered by default.

37 changes: 20 additions & 17 deletions airflow/ui/src/components/ui/Dialog/CloseTrigger.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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>
),
);
2 changes: 1 addition & 1 deletion airflow/ui/src/context/colorMode/useColorMode.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ export const useColorMode = () => {
};

return {
colorMode: resolvedTheme,
colorMode: resolvedTheme as "dark" | "light" | undefined,
setColorMode: setTheme,
toggleColorMode,
};
Expand Down
69 changes: 69 additions & 0 deletions airflow/ui/src/context/openGroups/OpenGroupsProvider.tsx
Original file line number Diff line number Diff line change
@@ -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>
);
};
21 changes: 21 additions & 0 deletions airflow/ui/src/context/openGroups/index.ts
Original file line number Diff line number Diff line change
@@ -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";
34 changes: 34 additions & 0 deletions airflow/ui/src/context/openGroups/useOpenGroups.ts
Original file line number Diff line number Diff line change
@@ -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;
};
38 changes: 7 additions & 31 deletions airflow/ui/src/pages/DagsList/Dag/Dag.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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">
Expand All @@ -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>
);
};
48 changes: 48 additions & 0 deletions airflow/ui/src/pages/DagsList/Dag/DagVizModal.tsx
Original file line number Diff line number Diff line change
@@ -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>
);
Loading

0 comments on commit 5a0272c

Please sign in to comment.