[v3-0-test] Update TaskLogContent to support virtualized rendering (#50746) (#51202)

* Fix OpenAPI schema for `get_log` API (#50547)

* Fix openapi schema for get_log API

* Fix test_log

(cherry picked from commit 08cc57d)

* [v3-0-test] Update `TaskLogContent` to support virtualized rendering (#50746)

* Update TaskLogContent to support virtualized rendering

* Update TaskLogPreview and Logs to handle undefined parsedLogs
(cherry picked from commit 813f3e3)

Co-authored-by: Guan Ming(Wesley) Chiu <[email protected]>

---------

Co-authored-by: LIU ZHE YOU <[email protected]>
Co-authored-by: Guan Ming(Wesley) Chiu <[email protected]>

Author: 3 people

airflow-core/src/airflow/ui/package.json

@@ -22,6 +22,7 @@
     "@emotion/react": "^11.14.0",
     "@tanstack/react-query": "^5.75.1",
     "@tanstack/react-table": "^8.21.3",
+    "@tanstack/react-virtual": "^3.13.8",
     "@types/debounce-promise": "^3.1.9",
     "@uiw/codemirror-themes-all": "^4.23.12",
     "@uiw/react-codemirror": "^4.23.12",

airflow-core/src/airflow/ui/pnpm-lock.yaml

@@ -67,7 +67,7 @@ export const TaskLogPreview = ({
           error={error}
           isLoading={isLoading}
           logError={error}
-          parsedLogs={data.parsedLogs}
+          parsedLogs={data.parsedLogs ?? []}
           wrap={wrap}
         />
       </Box>

airflow-core/src/airflow/ui/src/pages/Dag/Overview/TaskLogPreview.tsx

@@ -17,18 +17,25 @@
  * under the License.
  */
 import "@testing-library/jest-dom";
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { setupServer, type SetupServerApi } from "msw/node";
 import { afterEach, describe, it, expect, beforeAll, afterAll } from "vitest";
 
 import { handlers } from "src/mocks/handlers";
 import { AppWrapper } from "src/utils/AppWrapper";
 
 let server: SetupServerApi;
+const ITEM_HEIGHT = 20;
 
 beforeAll(() => {
   server = setupServer(...handlers);
   server.listen({ onUnhandledRequest: "bypass" });
+  Object.defineProperty(HTMLElement.prototype, "offsetHeight", {
+    value: ITEM_HEIGHT,
+  });
+  Object.defineProperty(HTMLElement.prototype, "offsetWidth", {
+    value: 800,
+  });
 });
 
 afterEach(() => server.resetHandlers());
@@ -39,14 +46,18 @@ describe("Task log grouping", () => {
     render(
       <AppWrapper initialEntries={["/dags/log_grouping/runs/manual__2025-02-18T12:19/tasks/generate"]} />,
     );
+    await waitFor(() => expect(screen.queryByTestId("virtualized-list")).toBeInTheDocument());
+    await waitFor(() => expect(screen.queryByTestId("virtualized-item-0")).toBeInTheDocument());
+    await waitFor(() => expect(screen.queryByTestId("virtualized-item-10")).toBeInTheDocument());
 
-    await waitFor(() => expect(screen.queryByTestId("summary-Pre task execution logs")).toBeInTheDocument(), {
-      timeout: 10_000,
-    });
+    fireEvent.scroll(screen.getByTestId("virtualized-list"), { target: { scrollTop: ITEM_HEIGHT * 6 } });
+    await waitFor(() => expect(screen.queryByTestId("virtualized-item-16")).toBeInTheDocument());
+
+    await waitFor(() => expect(screen.queryByTestId("summary-Pre task execution logs")).toBeInTheDocument());
     await waitFor(() => expect(screen.getByTestId("summary-Pre task execution logs")).toBeVisible());
     await waitFor(() => expect(screen.queryByText(/Task instance is in running state/iu)).not.toBeVisible());
 
     await waitFor(() => screen.getByTestId("summary-Pre task execution logs").click());
     await waitFor(() => expect(screen.queryByText(/Task instance is in running state/iu)).toBeVisible());
-  });
+  }, 10_000);
 });

airflow-core/src/airflow/ui/src/pages/TaskInstance/Logs/Logs.test.tsx

@@ -88,7 +88,7 @@ export const Logs = () => {
   });
 
   return (
-    <Box p={2}>
+    <Box display="flex" flexDirection="column" h="100%" p={2}>
       <TaskLogHeader
         onSelectTryNumber={onSelectTryNumber}
         sourceOptions={data.sources}
@@ -102,7 +102,7 @@ export const Logs = () => {
         error={error}
         isLoading={isLoading || isLoadingLogs}
         logError={logError}
-        parsedLogs={data.parsedLogs}
+        parsedLogs={data.parsedLogs ?? []}
         wrap={wrap}
       />
       <Dialog.Root onOpenChange={onOpenChange} open={fullscreen} scrollBehavior="inside" size="full">
@@ -124,12 +124,12 @@ export const Logs = () => {
 
           <Dialog.CloseTrigger />
 
-          <Dialog.Body>
+          <Dialog.Body display="flex" flexDirection="column">
             <TaskLogContent
               error={error}
               isLoading={isLoading || isLoadingLogs}
               logError={logError}
-              parsedLogs={data.parsedLogs}
+              parsedLogs={data.parsedLogs ?? []}
               wrap={wrap}
             />
           </Dialog.Body>

airflow-core/src/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx

@@ -17,8 +17,8 @@
  * under the License.
  */
 import { Box, Code, VStack, useToken } from "@chakra-ui/react";
-import type { ReactNode } from "react";
-import { useLayoutEffect } from "react";
+import { useVirtualizer } from "@tanstack/react-virtual";
+import { useLayoutEffect, useRef } from "react";
 
 import { ErrorAlert } from "src/components/ErrorAlert";
 import { ProgressBar } from "src/components/ui";
@@ -27,12 +27,19 @@ type Props = {
   readonly error: unknown;
   readonly isLoading: boolean;
   readonly logError: unknown;
-  readonly parsedLogs: ReactNode;
+  readonly parsedLogs: Array<JSX.Element | string | undefined>;
   readonly wrap: boolean;
 };
 
 export const TaskLogContent = ({ error, isLoading, logError, parsedLogs, wrap }: Props) => {
   const [bgLine] = useToken("colors", ["blue.emphasized"]);
+  const parentRef = useRef(null);
+  const rowVirtualizer = useVirtualizer({
+    count: parsedLogs.length,
+    estimateSize: () => 20,
+    getScrollElement: () => parentRef.current,
+    overscan: 10,
+  });
 
   useLayoutEffect(() => {
     if (location.hash) {
@@ -53,7 +60,7 @@ export const TaskLogContent = ({ error, isLoading, logError, parsedLogs, wrap }:
   }, [isLoading, bgLine]);
 
   return (
-    <Box>
+    <Box display="flex" flexDirection="column" flexGrow={1} h="100%" minHeight={0}>
       <ErrorAlert error={error ?? logError} />
       <ProgressBar size="xs" visibility={isLoading ? "visible" : "hidden"} />
       <Code
@@ -62,13 +69,32 @@ export const TaskLogContent = ({ error, isLoading, logError, parsedLogs, wrap }:
             bg: "blue.subtle",
           },
         }}
+        data-testid="virtualized-list"
+        flexGrow={1}
+        h="auto"
         overflow="auto"
+        position="relative"
         py={3}
+        ref={parentRef}
         textWrap={wrap ? "pre" : "nowrap"}
         width="100%"
       >
-        <VStack alignItems="flex-start" gap={0}>
-          {parsedLogs}
+        <VStack alignItems="flex-start" gap={0} h={`${rowVirtualizer.getTotalSize()}px`}>
+          {rowVirtualizer.getVirtualItems().map((virtualRow) => (
+            <Box
+              data-index={virtualRow.index}
+              data-testid={`virtualized-item-${virtualRow.index}`}
+              key={virtualRow.key}
+              left={0}
+              position="absolute"
+              ref={rowVirtualizer.measureElement}
+              top={0}
+              transform={`translateY(${virtualRow.start}px)`}
+              width={wrap ? "100%" : "max-content"}
+            >
+              {parsedLogs[virtualRow.index] ?? undefined}
+            </Box>
+          ))}
         </VStack>
       </Code>
     </Box>

airflow-core/src/airflow/ui/src/pages/TaskInstance/Logs/TaskLogContent.tsx

@@ -28,6 +28,7 @@ import { isStatePending, useAutoRefresh } from "src/utils";
 import { getTaskInstanceLink } from "src/utils/links";
 
 type Props = {
+  accept?: "*/*" | "application/json" | "application/x-ndjson";
   dagId: string;
   logLevelFilters?: Array<string>;
   sourceFilters?: Array<string>;
@@ -120,13 +121,14 @@ const parseLogs = ({ data, logLevelFilters, sourceFilters, taskInstance, tryNumb
 };
 
 export const useLogs = (
-  { dagId, logLevelFilters, sourceFilters, taskInstance, tryNumber = 1 }: Props,
+  { accept = "application/json", dagId, logLevelFilters, sourceFilters, taskInstance, tryNumber = 1 }: Props,
   options?: Omit<UseQueryOptions<TaskInstancesLogResponse>, "queryFn" | "queryKey">,
 ) => {
   const refetchInterval = useAutoRefresh({ dagId });
 
   const { data, ...rest } = useTaskInstanceServiceGetLog(
     {
+      accept,
       dagId,
       dagRunId: taskInstance?.dag_run_id ?? "",
       mapIndex: taskInstance?.map_index ?? -1,