get things working so we don't need to re-query the data when clicking onto a group or mapped ti. this requires passing state through the link. i have no clue how it works but it does (thanks chat gpt)

dstandish · dstandish · commit 6713e44784bb · 2025-06-26T15:41:22.000-07:00
diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/Grid/GridTI.tsx b/airflow-core/src/airflow/ui/src/layouts/Details/Grid/GridTI.tsx
@@ -21,7 +21,7 @@ import type { MouseEvent } from "react";
 import React from "react";
 import { Link, useParams } from "react-router-dom";
 
-import type { TaskInstanceState } from "openapi/requests/types.gen";
+import type { LightGridTaskInstanceSummary, TaskInstanceState } from "openapi/requests/types.gen";
 import { StateIcon } from "src/components/StateIcon";
 
 type Props = {
@@ -33,6 +33,7 @@ type Props = {
   readonly search: string;
   readonly state?: TaskInstanceState | null;
   readonly taskId: string;
+  readonly taskInstance: LightGridTaskInstanceSummary;
 };
 
 const onMouseEnter = (event: MouseEvent<HTMLDivElement>) => {
@@ -51,7 +52,7 @@ const onMouseLeave = (event: MouseEvent<HTMLDivElement>) => {
   });
 };
 
-const Instance = ({ dagId, isGroup, isMapped, runId, search, state, taskId }: Props) => {
+const Instance = ({ dagId, isGroup, isMapped, runId, search, state, taskId, taskInstance }: Props) => {
   const { groupId: selectedGroupId, taskId: selectedTaskId } = useParams();
 
   return (
@@ -71,6 +72,7 @@ const Instance = ({ dagId, isGroup, isMapped, runId, search, state, taskId }: Pr
     >
       <Link
         replace
+        state={{ taskInstance }}
         to={{
           pathname: `/dags/${dagId}/runs/${runId}/tasks/${isGroup ? "group/" : ""}${taskId}${isMapped ? "/mapped" : ""}`,
           search,
diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/Grid/TaskInstancesColumn.tsx b/airflow-core/src/airflow/ui/src/layouts/Details/Grid/TaskInstancesColumn.tsx
@@ -55,6 +55,7 @@ export const TaskInstancesColumn = ({ nodes, runId, taskInstances }: Props) => {
         search={search}
         state={taskInstance.state}
         taskId={node.id}
+        taskInstance={taskInstance}
       />
     );
   });
diff --git a/airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/GroupTaskInstance.tsx b/airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/GroupTaskInstance.tsx
@@ -18,42 +18,35 @@
  */
 import { ReactFlowProvider } from "@xyflow/react";
 import { MdOutlineTask } from "react-icons/md";
-import { useParams } from "react-router-dom";
+import { useLocation, useParams } from "react-router-dom";
 
-import { useDagServiceGetDagDetails } from "openapi/queries";
+import type { LightGridTaskInstanceSummary } from "openapi/requests";
 import { DetailsLayout } from "src/layouts/Details/DetailsLayout";
-import { useGridTiSummaries } from "src/queries/useGridTISummaries.ts";
 import { isStatePending, useAutoRefresh } from "src/utils";
 
 import { Header } from "./Header";
 
-export const GroupTaskInstance = () => {
-  const { dagId = "", groupId = "", runId = "" } = useParams();
-  const refetchInterval = useAutoRefresh({ dagId });
-
-  const {
-    data: dag,
-    error: dagError,
-    isLoading: isDagLoading,
-  } = useDagServiceGetDagDetails({
-    dagId,
-  });
+type LocationState = {
+  taskInstance: LightGridTaskInstanceSummary;
+};
 
-  const { data: gridTISummaries, error, isLoading } = useGridTiSummaries({ dagId, runId, state: undefined });
+export const GroupTaskInstance = () => {
+  const { dagId = "" } = useParams();
 
-  const taskInstance = gridTISummaries?.task_instances.find((ti) => ti.task_id === groupId);
+  const location = useLocation();
+  const state = location.state as LocationState;
+  const taskInstance: LightGridTaskInstanceSummary = state.taskInstance;
+  const refetchInterval = useAutoRefresh({ dagId });
 
   const tabs = [{ icon: <MdOutlineTask />, label: "Task Instances", value: "" }];
 
   return (
     <ReactFlowProvider>
-      <DetailsLayout dag={dag} error={error ?? dagError} isLoading={isLoading || isDagLoading} tabs={tabs}>
-        {taskInstance === undefined ? undefined : (
-          <Header
-            isRefreshing={Boolean(isStatePending(taskInstance.state) && Boolean(refetchInterval))}
-            taskInstance={taskInstance}
-          />
-        )}
+      <DetailsLayout tabs={tabs}>
+        <Header
+          isRefreshing={Boolean(isStatePending(taskInstance.state) && Boolean(refetchInterval))}
+          taskInstance={taskInstance}
+        />
       </DetailsLayout>
     </ReactFlowProvider>
   );
diff --git a/airflow-core/src/airflow/ui/src/pages/MappedTaskInstance/Header.tsx b/airflow-core/src/airflow/ui/src/pages/MappedTaskInstance/Header.tsx
@@ -20,7 +20,7 @@ import { Box } from "@chakra-ui/react";
 import type { ReactNode } from "react";
 import { MdOutlineTask } from "react-icons/md";
 
-import type { GridTaskInstanceSummary } from "openapi/requests/types.gen";
+import type { LightGridTaskInstanceSummary } from "openapi/requests/types.gen";
 import { HeaderCard } from "src/components/HeaderCard";
 import Time from "src/components/Time";
 import { getDuration } from "src/utils";
@@ -30,23 +30,21 @@ export const Header = ({
   taskInstance,
 }: {
   readonly isRefreshing?: boolean;
-  readonly taskInstance: GridTaskInstanceSummary;
+  readonly taskInstance: LightGridTaskInstanceSummary;
 }) => {
   const entries: Array<{ label: string; value: number | ReactNode | string }> = [];
+  let taskCount: number = 0;
 
-  if (taskInstance.child_states !== null) {
-    Object.entries(taskInstance.child_states).forEach(([state, count]) => {
-      if (count > 0) {
-        entries.push({ label: `Total ${state}`, value: count });
-      }
-    });
-  }
+  Object.entries(taskInstance.child_states ?? {}).forEach(([state, count]) => {
+    entries.push({ label: `Total ${state}`, value: count });
+    taskCount += count;
+  });
   const stats = [
     ...entries,
-    { label: "Start", value: <Time datetime={taskInstance.start_date} /> },
-    { label: "End", value: <Time datetime={taskInstance.end_date} /> },
-    ...(Boolean(taskInstance.start_date)
-      ? [{ label: "Duration", value: getDuration(taskInstance.start_date, taskInstance.end_date) }]
+    { label: "Start", value: <Time datetime={taskInstance.min_start_date} /> },
+    { label: "End", value: <Time datetime={taskInstance.max_end_date} /> },
+    ...(Boolean(taskInstance.max_end_date)
+      ? [{ label: "Duration", value: getDuration(taskInstance.min_start_date, taskInstance.max_end_date) }]
       : []),
   ];
 
@@ -57,8 +55,8 @@ export const Header = ({
         isRefreshing={isRefreshing}
         state={taskInstance.state}
         stats={stats}
-        subTitle={<Time datetime={taskInstance.start_date} />}
-        title={`${taskInstance.task_id} [${taskInstance.task_count}]`}
+        subTitle={<Time datetime={taskInstance.min_start_date} />}
+        title={`${taskInstance.task_id} [${taskCount}]`}
       />
     </Box>
   );
diff --git a/airflow-core/src/airflow/ui/src/pages/MappedTaskInstance/MappedTaskInstance.tsx b/airflow-core/src/airflow/ui/src/pages/MappedTaskInstance/MappedTaskInstance.tsx
@@ -18,73 +18,41 @@
  */
 import { ReactFlowProvider } from "@xyflow/react";
 import { MdOutlineTask } from "react-icons/md";
-import { useParams } from "react-router-dom";
+import { useLocation, useParams } from "react-router-dom";
 
-import {
-  useDagRunServiceGetDagRun,
-  useDagServiceGetDagDetails,
-  useGridServiceGridData,
-} from "openapi/queries";
+import type { LightGridTaskInstanceSummary } from "openapi/requests";
 import { DetailsLayout } from "src/layouts/Details/DetailsLayout";
 import { isStatePending, useAutoRefresh } from "src/utils";
 
 import { Header } from "./Header";
 
+type LocationState = {
+  taskInstance: LightGridTaskInstanceSummary;
+};
+
 export const MappedTaskInstance = () => {
-  const { dagId = "", runId = "", taskId = "" } = useParams();
+  const { dagId = "" } = useParams();
   const refetchInterval = useAutoRefresh({ dagId });
 
-  const {
-    data: dag,
-    error: dagError,
-    isLoading: isDagLoading,
-  } = useDagServiceGetDagDetails({
-    dagId,
-  });
+  const location = useLocation();
+  const state = location.state as LocationState;
+  const taskInstance: LightGridTaskInstanceSummary = state.taskInstance;
 
-  const { data: dagRun } = useDagRunServiceGetDagRun(
-    {
-      dagId,
-      dagRunId: runId,
-    },
-    undefined,
-    { enabled: runId !== "" },
-  );
+  let taskCount: number = 0;
 
-  // Filter grid data to get only a single dag run
-  const { data, error, isLoading } = useGridServiceGridData(
-    {
-      dagId,
-      limit: 1,
-      offset: 0,
-      runAfterGte: dagRun?.run_after,
-      runAfterLte: dagRun?.run_after,
-    },
-    undefined,
-    {
-      enabled: dagRun !== undefined,
-      refetchInterval: (query) =>
-        query.state.data?.dag_runs.some((dr) => isStatePending(dr.state)) && refetchInterval,
-    },
-  );
-
-  const taskInstance = data?.dag_runs
-    .find((dr) => dr.dag_run_id === runId)
-    ?.task_instances.find((ti) => ti.task_id === taskId);
+  Object.entries(taskInstance.child_states ?? {}).forEach(([, count]) => {
+    taskCount += count;
+  });
 
-  const tabs = [
-    { icon: <MdOutlineTask />, label: `Task Instances [${taskInstance?.task_count}]`, value: "" },
-  ];
+  const tabs = [{ icon: <MdOutlineTask />, label: `Task Instances [${taskCount}]`, value: "" }];
 
   return (
     <ReactFlowProvider>
-      <DetailsLayout dag={dag} error={error ?? dagError} isLoading={isLoading || isDagLoading} tabs={tabs}>
-        {taskInstance === undefined ? undefined : (
-          <Header
-            isRefreshing={Boolean(isStatePending(taskInstance.state) && Boolean(refetchInterval))}
-            taskInstance={taskInstance}
-          />
-        )}
+      <DetailsLayout tabs={tabs}>
+        <Header
+          isRefreshing={Boolean(isStatePending(taskInstance.state) && Boolean(refetchInterval))}
+          taskInstance={taskInstance}
+        />
       </DetailsLayout>
     </ReactFlowProvider>
   );
