diff --git a/airflow/ui/src/pages/Dag/Code/Code.tsx b/airflow/ui/src/pages/Dag/Code/Code.tsx index f52215f6c61f3..2efd4e569ea52 100644 --- a/airflow/ui/src/pages/Dag/Code/Code.tsx +++ b/airflow/ui/src/pages/Dag/Code/Code.tsx @@ -16,14 +16,23 @@ * specific language governing permissions and limitations * under the License. */ -import { Box, Button, Heading, HStack } from "@chakra-ui/react"; -import { useState } from "react"; -import { useParams } from "react-router-dom"; +import { Box, Button, Heading, HStack, Field } from "@chakra-ui/react"; +import { useQueryClient } from "@tanstack/react-query"; +import type { OptionsOrGroups, GroupBase, SingleValue } from "chakra-react-select"; +import { AsyncSelect } from "chakra-react-select"; +import { useState, useCallback } from "react"; +import { useParams, useSearchParams } from "react-router-dom"; import { createElement, PrismLight as SyntaxHighlighter } from "react-syntax-highlighter"; import python from "react-syntax-highlighter/dist/esm/languages/prism/python"; import { oneLight, oneDark } from "react-syntax-highlighter/dist/esm/styles/prism"; -import { useDagServiceGetDagDetails, useDagSourceServiceGetDagSource } from "openapi/queries"; +import { + useDagServiceGetDagDetails, + useDagSourceServiceGetDagSource, + UseDagVersionServiceGetDagVersionsKeyFn, +} from "openapi/queries"; +import { DagVersionService } from "openapi/requests/services.gen"; +import type { DAGVersionCollectionResponse, DagVersionResponse } from "openapi/requests/types.gen"; import { ErrorAlert } from "src/components/ErrorAlert"; import Time from "src/components/Time"; import { ProgressBar } from "src/components/ui"; @@ -32,8 +41,18 @@ import { useConfig } from "src/queries/useConfig"; SyntaxHighlighter.registerLanguage("python", python); +const VERSION_PARAM = "version"; + +type Option = { + label: string; + value: string; +}; + export const Code = () => { const { dagId } = useParams(); + const [searchParams, setSearchParams] = useSearchParams(); + const selectedVersion = searchParams.get(VERSION_PARAM); + const queryClient = useQueryClient(); const { data: dag, @@ -49,6 +68,7 @@ export const Code = () => { isLoading: isCodeLoading, } = useDagSourceServiceGetDagSource({ dagId: dagId ?? "", + versionNumber: selectedVersion === null ? undefined : parseInt(selectedVersion, 10), }); const defaultWrap = Boolean(useConfig("default_wrap")); @@ -65,6 +85,42 @@ export const Code = () => { style['code[class*="language-"]'].whiteSpace = wrap ? "pre-wrap" : "pre"; } + const loadVersions = ( + _: string, + callback: (options: OptionsOrGroups>) => void, + ): Promise>> => + queryClient.fetchQuery({ + queryFn: () => + DagVersionService.getDagVersions({ + dagId: dagId ?? "", + }).then((data: DAGVersionCollectionResponse) => { + const options = data.dag_versions.map((version: DagVersionResponse) => { + const versionNumber = version.version_number.toString(); + + return { + label: versionNumber, + value: versionNumber, + }; + }); + + callback(options); + + return options; + }), + queryKey: UseDagVersionServiceGetDagVersionsKeyFn({ dagId: dagId ?? "" }), + staleTime: 0, + }); + + const handleStateChange = useCallback( + (version: SingleValue