Add dataset event info to dag graph

airflow/www/static/js/api/useDatasetEvents.ts

@@ -17,14 +17,21 @@
  * under the License.
  */
 
-import axios, { AxiosResponse } from "axios";
-import { useQuery } from "react-query";
+import axios from "axios";
+import { useQuery, UseQueryOptions } from "react-query";
 
 import { getMetaValue } from "src/utils";
-import type { API } from "src/types";
 import URLSearchParamsWrapper from "src/utils/URLSearchParamWrapper";
+import type {
+  DatasetEventCollection,
+  GetDatasetEventsVariables,
+} from "src/types/api-generated";
 
-export default function useDatasetEvents({
+interface Props extends GetDatasetEventsVariables {
+  options?: UseQueryOptions<DatasetEventCollection>;
+}
+
+const useDatasetEvents = ({
   datasetId,
   sourceDagId,
   sourceRunId,
@@ -33,8 +40,9 @@ export default function useDatasetEvents({
   limit,
   offset,
   orderBy,
-}: API.GetDatasetEventsVariables) {
-  const query = useQuery(
+  options,
+}: Props) => {
+  const query = useQuery<DatasetEventCollection>(
     [
       "datasets-events",
       datasetId,
@@ -61,16 +69,19 @@ export default function useDatasetEvents({
       if (sourceMapIndex)
         params.set("source_map_index", sourceMapIndex.toString());
 
-      return axios.get<AxiosResponse, API.DatasetEventCollection>(datasetsUrl, {
+      return axios.get(datasetsUrl, {
         params,
       });
     },
     {
       keepPreviousData: true,
+      ...options,
     }
   );
   return {
     ...query,
     data: query.data ?? { datasetEvents: [], totalEntries: 0 },
   };
-}
+};
+
+export default useDatasetEvents;

airflow/www/static/js/api/useTaskInstance.ts

@@ -17,7 +17,7 @@
  * under the License.
  */
 
-import axios, { AxiosResponse } from "axios";
+import axios from "axios";
 import type { API } from "src/types";
 import { useQuery, UseQueryOptions } from "react-query";
 import { useAutoRefresh } from "src/context/autorefresh";
@@ -55,7 +55,7 @@ const useTaskInstance = ({
 
   return useQuery<API.TaskInstance>(
     ["taskInstance", dagId, dagRunId, taskId, mapIndex],
-    () => axios.get<AxiosResponse, API.TaskInstance>(url),
+    () => axios.get(url),
     {
       refetchInterval: isRefreshOn && (autoRefreshInterval || 1) * 1000,
       ...options,

airflow/www/static/js/api/useUpstreamDatasetEvents.ts

@@ -17,29 +17,35 @@
  * under the License.
  */
 
-import axios, { AxiosResponse } from "axios";
-import { useQuery } from "react-query";
+import axios from "axios";
+import { useQuery, UseQueryOptions } from "react-query";
 
 import { getMetaValue } from "src/utils";
-import type { API } from "src/types";
+import type {
+  DatasetEventCollection,
+  GetUpstreamDatasetEventsVariables,
+} from "src/types/api-generated";
 
-interface Props {
-  runId: string;
+interface Props extends GetUpstreamDatasetEventsVariables {
+  options?: UseQueryOptions<DatasetEventCollection>;
 }
 
-export default function useUpstreamDatasetEvents({ runId }: Props) {
-  const query = useQuery(["upstreamDatasetEvents", runId], () => {
-    const dagId = getMetaValue("dag_id");
-    const upstreamEventsUrl = (
-      getMetaValue("upstream_dataset_events_api") ||
-      `api/v1/dags/${dagId}/dagRuns/_DAG_RUN_ID_/upstreamDatasetEvents`
-    ).replace("_DAG_RUN_ID_", encodeURIComponent(runId));
-    return axios.get<AxiosResponse, API.DatasetEventCollection>(
-      upstreamEventsUrl
-    );
-  });
+const useUpstreamDatasetEvents = ({ dagId, dagRunId, options }: Props) => {
+  const upstreamEventsUrl = (
+    getMetaValue("upstream_dataset_events_api") ||
+    `api/v1/dags/${dagId}/dagRuns/_DAG_RUN_ID_/upstreamDatasetEvents`
+  ).replace("_DAG_RUN_ID_", encodeURIComponent(dagRunId));
+
+  const query = useQuery<DatasetEventCollection>(
+    ["upstreamDatasetEvents", dagRunId],
+    () => axios.get(upstreamEventsUrl),
+    options
+  );
+
   return {
     ...query,
-    data: query.data || { datasetEvents: [], totalEntries: 0 },
+    data: query.data ?? { datasetEvents: [], totalEntries: 0 },
   };
-}
+};
+
+export default useUpstreamDatasetEvents;

airflow/www/static/js/components/DatasetEventCard.tsx

@@ -32,23 +32,20 @@ import {
   Link,
 } from "@chakra-ui/react";
 import { HiDatabase } from "react-icons/hi";
-import { FiLink } from "react-icons/fi";
 import { useSearchParams } from "react-router-dom";
 
 import { getMetaValue } from "src/utils";
 import Time from "src/components/Time";
 import { useContainerRef } from "src/context/containerRef";
-import { SimpleStatus } from "src/dag/StatusBox";
-import { formatDuration, getDuration } from "src/datetime_utils";
 import RenderedJsonField from "src/components/RenderedJsonField";
 
 import SourceTaskInstance from "./SourceTaskInstance";
+import TriggeredDagRuns from "./TriggeredDagRuns";
 
 type CardProps = {
   datasetEvent: DatasetEvent;
 };
 
-const gridUrl = getMetaValue("grid_url");
 const datasetsUrl = getMetaValue("datasets_url");
 
 const DatasetEventCard = ({ datasetEvent }: CardProps) => {
@@ -116,51 +113,7 @@ const DatasetEventCard = ({ datasetEvent }: CardProps) => {
           {!!datasetEvent?.createdDagruns?.length && (
             <>
               Triggered Dag Runs:
-              <Flex alignItems="center">
-                {datasetEvent?.createdDagruns.map((run) => {
-                  const runId = (run as any).dagRunId; // For some reason the type is wrong here
-                  const url = `${gridUrl?.replace(
-                    "__DAG_ID__",
-                    run.dagId || ""
-                  )}?dag_run_id=${encodeURIComponent(runId)}`;
-
-                  return (
-                    <Tooltip
-                      key={runId}
-                      label={
-                        <Box>
-                          <Text>DAG Id: {run.dagId}</Text>
-                          <Text>Status: {run.state || "no status"}</Text>
-                          <Text>
-                            Duration:{" "}
-                            {formatDuration(
-                              getDuration(run.startDate, run.endDate)
-                            )}
-                          </Text>
-                          <Text>
-                            Start Date: <Time dateTime={run.startDate} />
-                          </Text>
-                          {run.endDate && (
-                            <Text>
-                              End Date: <Time dateTime={run.endDate} />
-                            </Text>
-                          )}
-                        </Box>
-                      }
-                      portalProps={{ containerRef }}
-                      hasArrow
-                      placement="top"
-                    >
-                      <Flex width="30px">
-                        <SimpleStatus state={run.state} mx={1} />
-                        <Link color="blue.600" href={url}>
-                          <FiLink size="12px" />
-                        </Link>
-                      </Flex>
-                    </Tooltip>
-                  );
-                })}
-              </Flex>
+              <TriggeredDagRuns createdDagRuns={datasetEvent?.createdDagruns} />
             </>
           )}
         </GridItem>

airflow/www/static/js/components/SourceTaskInstance.tsx

@@ -31,11 +31,15 @@ import { getMetaValue } from "src/utils";
 
 type SourceTIProps = {
   datasetEvent: DatasetEvent;
+  showLink?: boolean;
 };
 
 const gridUrl = getMetaValue("grid_url");
 
-const SourceTaskInstance = ({ datasetEvent }: SourceTIProps) => {
+const SourceTaskInstance = ({
+  datasetEvent,
+  showLink = true,
+}: SourceTIProps) => {
   const containerRef = useContainerRef();
   const { sourceDagId, sourceRunId, sourceTaskId, sourceMapIndex } =
     datasetEvent;
@@ -80,11 +84,13 @@ const SourceTaskInstance = ({ datasetEvent }: SourceTIProps) => {
           hasArrow
           placement="top"
         >
-          <Flex width="30px">
+          <Flex width={showLink ? "30px" : "16px"}>
             <SimpleStatus state={taskInstance.state} mx={1} />
-            <Link color="blue.600" href={url}>
-              <FiLink size="12px" />
-            </Link>
+            {showLink && (
+              <Link color="blue.600" href={url}>
+                <FiLink size="12px" />
+              </Link>
+            )}
           </Flex>
         </Tooltip>
       )}

airflow/www/static/js/components/TriggeredDagRuns.tsx

@@ -0,0 +1,91 @@
+/*!
+ * 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 React from "react";
+
+import type { DAGRun } from "src/types/api-generated";
+import { Box, Flex, Tooltip, Text, Link } from "@chakra-ui/react";
+import { FiLink } from "react-icons/fi";
+
+import { getMetaValue } from "src/utils";
+import Time from "src/components/Time";
+import { useContainerRef } from "src/context/containerRef";
+import { SimpleStatus } from "src/dag/StatusBox";
+import { formatDuration, getDuration } from "src/datetime_utils";
+
+type CardProps = {
+  createdDagRuns: DAGRun[];
+  showLink?: boolean;
+};
+
+const gridUrl = getMetaValue("grid_url");
+
+const TriggeredDagRuns = ({ createdDagRuns, showLink = true }: CardProps) => {
+  const containerRef = useContainerRef();
+
+  return (
+    <Flex alignItems="center">
+      {createdDagRuns.map((run) => {
+        const runId = (run as any).dagRunId; // For some reason the type is wrong here
+        const url = `${gridUrl?.replace(
+          "__DAG_ID__",
+          run.dagId || ""
+        )}?dag_run_id=${encodeURIComponent(runId)}`;
+
+        return (
+          <Tooltip
+            key={runId}
+            label={
+              <Box>
+                <Text>DAG Id: {run.dagId}</Text>
+                <Text>Status: {run.state || "no status"}</Text>
+                <Text>
+                  Duration:{" "}
+                  {formatDuration(getDuration(run.startDate, run.endDate))}
+                </Text>
+                <Text>
+                  Start Date: <Time dateTime={run.startDate} />
+                </Text>
+                {run.endDate && (
+                  <Text>
+                    End Date: <Time dateTime={run.endDate} />
+                  </Text>
+                )}
+              </Box>
+            }
+            portalProps={{ containerRef }}
+            hasArrow
+            placement="top"
+          >
+            <Flex width={showLink ? "30px" : "16px"}>
+              <SimpleStatus state={run.state} mx={1} />
+              {showLink && (
+                <Link color="blue.600" href={url}>
+                  <FiLink size="12px" />
+                </Link>
+              )}
+            </Flex>
+          </Tooltip>
+        );
+      })}
+    </Flex>
+  );
+};
+
+export default TriggeredDagRuns;

airflow/www/static/js/dag/details/dagRun/DatasetTriggerEvents.tsx

@@ -24,11 +24,14 @@ import type { DagRun as DagRunType } from "src/types";
 import { CardDef, CardList } from "src/components/Table";
 import type { DatasetEvent } from "src/types/api-generated";
 import DatasetEventCard from "src/components/DatasetEventCard";
+import { getMetaValue } from "src/utils";
 
 interface Props {
   runId: DagRunType["runId"];
 }
 
+const dagId = getMetaValue("dag_id");
+
 const cardDef: CardDef<DatasetEvent> = {
   card: ({ row }) => <DatasetEventCard datasetEvent={row} />,
 };
@@ -37,7 +40,7 @@ const DatasetTriggerEvents = ({ runId }: Props) => {
   const {
     data: { datasetEvents = [] },
     isLoading,
-  } = useUpstreamDatasetEvents({ runId });
+  } = useUpstreamDatasetEvents({ dagRunId: runId, dagId });
 
   const columns = useMemo(
     () => [