experimenting with await defer

PrefectHQ · Dec 29, 2024 · 230935e · 230935e
1 parent 3809ab2
commit 230935e
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 76 deletions.
diff --git a/...ponents/concurrency/task-run-concurrency-limits/task-run-concurrency-limit-page/index.tsx b/...ponents/concurrency/task-run-concurrency-limits/task-run-concurrency-limit-page/index.tsx
@@ -5,6 +5,7 @@ import { useState } from "react";
 import { TaskRunConcurrencyLimitActiveTaskRuns } from "@/components/concurrency/task-run-concurrency-limits/task-run-concurrency-limit-active-task-runs";
 import { TaskRunConcurrencyLimitDetails } from "@/components/concurrency/task-run-concurrency-limits/task-run-concurrency-limit-details";
 import { useSuspenseQuery } from "@tanstack/react-query";
+import { Await } from "@tanstack/react-router";
 import {
 	type Dialogs,
 	TaskRunConcurrencyLimitDialog,
@@ -43,11 +44,13 @@ export const TaskRunConcurrencyLimitPage = ({ id }: Props) => {
 					onReset={handleOpenResetDialog}
 				/>
 				<div className="grid gap-4" style={{ gridTemplateColumns: "3fr 1fr" }}>
-					<TaskRunConcurrencyLimitTabNavigation
-						activeTaskRuns={
-							<TaskRunConcurrencyLimitActiveTaskRuns data={activeTaskRuns} />
-						}
-					/>
+					<TaskRunConcurrencyLimitTabNavigation>
+						<Await promise={activeTaskRuns} fallback={<div>Loading...</div>}>
+							{(promiseData) => (
+								<TaskRunConcurrencyLimitActiveTaskRuns data={promiseData} />
+							)}
+						</Await>
+					</TaskRunConcurrencyLimitTabNavigation>
 					<TaskRunConcurrencyLimitDetails data={taskRunConcurrencyLimit} />
 				</div>
 			</div>

diff --git a/...ency-limits/task-run-concurrency-limit-page/task-run-concurrency-limit-tab-navigation.tsx b/...ency-limits/task-run-concurrency-limit-page/task-run-concurrency-limit-tab-navigation.tsx
@@ -20,13 +20,12 @@ const TAB_OPTIONS: Record<TabOptions, TabOptionValues> = {
 } as const;
 
 type Props = {
-	activeTaskRuns: React.ReactNode;
+	/** Should add ActiveTaskRun component */
+	children: React.ReactNode;
 };
 
 // TODO: Move Tabs for navigation to a generic styled component
-export const TaskRunConcurrencyLimitTabNavigation = ({
-	activeTaskRuns,
-}: Props) => {
+export const TaskRunConcurrencyLimitTabNavigation = ({ children }: Props) => {
 	const { tab } = routeApi.useSearch();
 	const navigate = routeApi.useNavigate();
 
@@ -48,7 +47,7 @@ export const TaskRunConcurrencyLimitTabNavigation = ({
 				</TabsTrigger>
 			</TabsList>
 			<TabsContent value={TAB_OPTIONS["active-task-runs"].tabSearchValue}>
-				{activeTaskRuns}
+				{children}
 			</TabsContent>
 		</Tabs>
 	);

diff --git a/ui-v2/src/hooks/task-run-concurrency-limits.test.tsx b/ui-v2/src/hooks/task-run-concurrency-limits.test.tsx
@@ -299,30 +299,6 @@ describe("buildConcurrenyLimitDetailsActiveRunsQuery()", () => {
 		);
 	};
 
-	/**
-	 * Data Management:
-	 * - Asserts fetch API to get details on the task run concurrency limit
-	 * - Other APIs will not be fired because there are no active task runs
-	 */
-	it("will fetch only necessary data when there is no active task runs", async () => {
-		const MOCK_DATA = { ...seedData(), active_slots: [] };
-		const MOCK_ID = MOCK_DATA.id;
-		mockFetchDetailsAPI(MOCK_DATA);
-
-		const { result } = renderHook(
-			() =>
-				useSuspenseQuery(buildConcurrenyLimitDetailsActiveRunsQuery(MOCK_ID)),
-			{ wrapper: createWrapper() },
-		);
-
-		// ------------ Assert
-		await waitFor(() => expect(result.current.isSuccess).toBe(true));
-		expect(result.current.data).toMatchObject({
-			taskRunConcurrencyLimit: MOCK_DATA,
-			activeTaskRuns: [],
-		});
-	});
-
 	/**
 	 * Data Management:
 	 * - Asserts waterfall of APIs will fire to get details on the task run concurrency limit
@@ -344,15 +320,16 @@ describe("buildConcurrenyLimitDetailsActiveRunsQuery()", () => {
 
 		// ------------ Assert
 		await waitFor(() => expect(result.current.isSuccess).toBe(true));
-		expect(result.current.data).toMatchObject({
-			taskRunConcurrencyLimit: MOCK_DATA,
-			activeTaskRuns: [
-				{
-					flow: MOCK_FLOW,
-					flowRun: MOCK_FLOW_RUNS[0],
-					taskRun: MOCK_TASK_RUNS[0],
-				},
-			],
-		});
+		expect(result.current.data.taskRunConcurrencyLimit).toMatchObject(
+			MOCK_DATA,
+		);
+		const activeTaskRunsResult = await result.current.data.activeTaskRuns;
+		expect(activeTaskRunsResult).toEqual([
+			{
+				flow: MOCK_FLOW,
+				flowRun: MOCK_FLOW_RUNS[0],
+				taskRun: MOCK_TASK_RUNS[0],
+			},
+		]);
 	});
 });
diff --git a/ui-v2/src/hooks/task-run-concurrency-limits.ts b/ui-v2/src/hooks/task-run-concurrency-limits.ts
@@ -55,15 +55,21 @@ export const buildListTaskRunConcurrencyLimitsQuery = (
 		refetchInterval: 30_000,
 	});
 
+const fetchTaskRunConcurrencyLimit = async (id: string) => {
+	// GET task-run-concurrency-limit by id
+	const res = await getQueryService().GET("/concurrency_limits/{id}", {
+		params: { path: { id } },
+	});
+	if (!res.data) {
+		throw new Error("'data' expected");
+	}
+	return res.data;
+};
+
 export const buildDetailTaskRunConcurrencyLimitsQuery = (id: string) =>
 	queryOptions({
 		queryKey: queryKeyFactory.detail(id),
-		queryFn: async () => {
-			const res = await getQueryService().GET("/concurrency_limits/{id}", {
-				params: { path: { id } },
-			});
-			return res.data as TaskRunConcurrencyLimit; // Expecting data to be truthy;
-		},
+		queryFn: () => fetchTaskRunConcurrencyLimit(id),
 	});
 
 /**
@@ -212,17 +218,6 @@ export const useResetTaskRunConcurrencyLimitTag = () => {
 	};
 };
 
-const fetchTaskRunConcurrencyLimit = async (id: string) => {
-	// GET task-run-concurrency-limit by id
-	const res = await getQueryService().GET("/concurrency_limits/{id}", {
-		params: { path: { id } },
-	});
-	if (!res.data) {
-		throw new Error("'data' expected");
-	}
-	return res.data;
-};
-
 const fetchActiveTaskRunDetails = async (activeSlots: Array<string>) => {
 	const taskRuns = await getQueryService().POST("/task_runs/filter", {
 		body: {
@@ -248,6 +243,18 @@ const fetchActiveTaskRunDetails = async (activeSlots: Array<string>) => {
 		}
 	});
 
+	const activeTaskRunsWithoutFlows = taskRunsOnly.map((taskRun) => ({
+		taskRun,
+		flowRun: null,
+		flow: null,
+	}));
+
+	// Early exit if there's no task with parent flows
+	if (taskRunsWithFlows.length === 0) {
+		return activeTaskRunsWithoutFlows;
+	}
+
+	// Now get parent flow information for tasks with parent flows
 	const flowRunsIds = taskRunsWithFlows.map(
 		(taskRun) => taskRun.flow_run_id as string,
 	);
@@ -307,12 +314,6 @@ const fetchActiveTaskRunDetails = async (activeSlots: Array<string>) => {
 		};
 	});
 
-	const activeTaskRunsWithoutFlows = taskRunsOnly.map((taskRun) => ({
-		taskRun,
-		flowRun: null,
-		flow: null,
-	}));
-
 	return [...activeTaskRunsWithFlows, ...activeTaskRunsWithoutFlows];
 };
 
@@ -329,20 +330,14 @@ export const buildConcurrenyLimitDetailsActiveRunsQuery = (id: string) =>
 			if (!taskRunConcurrencyLimit.active_slots) {
 				throw new Error("'active_slots' expected");
 			}
-			// Early exit because there are no active task runs
-			if (taskRunConcurrencyLimit.active_slots.length === 0) {
-				return {
-					taskRunConcurrencyLimit: taskRunConcurrencyLimit,
-					activeTaskRuns: [],
-				};
-			}
-			const activeTaskRuns = await fetchActiveTaskRunDetails(
+
+			const activeTaskRuns = fetchActiveTaskRunDetails(
 				taskRunConcurrencyLimit.active_slots,
 			);
 
 			return {
 				taskRunConcurrencyLimit,
-				activeTaskRuns,
+				activeTaskRuns, // defer to return promise
 			};
 		},
 	});