diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactPreviewContent.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactPreviewContent.tsx index 89b4bbf7af..265ef64b62 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactPreviewContent.tsx +++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactPreviewContent.tsx @@ -63,6 +63,7 @@ interface PreviewContentProps { type: string; name: string; isFullscreen: boolean; + byteLength?: number; } export const PreviewContent = ({ @@ -70,6 +71,7 @@ export const PreviewContent = ({ type, name, isFullscreen, + byteLength, }: PreviewContentProps) => { const { backendUrl } = useBackend(); @@ -106,7 +108,11 @@ export const PreviewContent = ({ ); case "apacheparquet": return ( - + ); case "jsonobject": case "jsonarray": diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactVisualizer.test.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactVisualizer.test.tsx index a0f5ba180c..b0513ff37f 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactVisualizer.test.tsx +++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactVisualizer.test.tsx @@ -108,8 +108,18 @@ vi.mock("./UrlVisualizer", () => ({ })); vi.mock("./ParquetVisualizer", () => ({ - default: ({ signedUrl }: { signedUrl: string }) => ( -
+ default: ({ + signedUrl, + byteLength, + }: { + signedUrl: string; + byteLength?: number; + }) => ( +
), })); @@ -353,7 +363,7 @@ describe("ArtifactVisualizer", () => { }); }); - it("renders ParquetVisualizer for apacheparquet type", async () => { + it("renders ParquetVisualizer for apacheparquet type with the reported size", async () => { renderWithQuery( { await userEvent.click(screen.getByText("Preview")); await waitFor(() => { - expect(screen.getByTestId("parquet-visualizer")).toBeInTheDocument(); + expect(screen.getByTestId("parquet-visualizer")).toHaveAttribute( + "data-byte-length", + "1024", + ); + }); + }); + + it("renders ParquetVisualizer without a size when the artifact reports none", async () => { + renderWithQuery( + , + ); + + await userEvent.click(screen.getByText("Preview")); + + await waitFor(() => { + expect(screen.getByTestId("parquet-visualizer")).not.toHaveAttribute( + "data-byte-length", + ); }); }); }); diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactVisualizer.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactVisualizer.tsx index 0a7a6d84b8..59c13a2db4 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactVisualizer.tsx +++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactVisualizer.tsx @@ -197,6 +197,7 @@ const ArtifactVisualizer = ({ artifactId={artifact.id} type={normalizedType} isFullscreen={isFullscreen} + byteLength={artifactData?.total_size} /> )} diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ParquetVisualizer.test.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ParquetVisualizer.test.tsx index 16b53b5e70..2f690b8102 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ParquetVisualizer.test.tsx +++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ParquetVisualizer.test.tsx @@ -6,7 +6,6 @@ import { ErrorBoundary } from "react-error-boundary"; import { beforeEach, describe, expect, it, vi } from "vitest"; import useToastNotification from "@/hooks/useToastNotification"; -import { ArtifactFetchError } from "@/services/executionService"; import { PARQUET_PREVIEW_ROWS } from "./parquetUtils"; import ParquetVisualizer from "./ParquetVisualizer"; @@ -22,7 +21,6 @@ vi.mock("hyparquet", () => ({ parquetMetadata: vi.fn(), parquetMetadataAsync: vi.fn(), asyncBufferFromUrl: vi.fn(), - byteLengthFromUrl: vi.fn(), cachedAsyncBuffer: vi.fn((buffer) => buffer), toJson: vi.fn((value) => value), })); @@ -94,7 +92,6 @@ const { parquetMetadata, parquetMetadataAsync, asyncBufferFromUrl, - byteLengthFromUrl, } = await import("hyparquet"); const { downloadStringAsFile } = await import("@/utils/URL"); @@ -166,6 +163,14 @@ const mockWideDataset = (columnCount: number, numRows: number) => { ); }; +/** A `Range: bytes=0-0` probe answered with the total size, as object stores do. */ +const mockRangeProbe = (totalSize: number) => + vi.spyOn(globalThis, "fetch").mockResolvedValue({ + ok: true, + status: 206, + headers: new Headers({ "Content-Range": `bytes 0-0/${totalSize}` }), + } as unknown as Response); + beforeEach(() => { queryClient.clear(); lastDownloadFull = undefined; @@ -176,7 +181,7 @@ beforeEach(() => { vi.mocked(parquetMetadataAsync).mockReset(); vi.mocked(parquetReadObjects).mockReset(); vi.mocked(downloadStringAsFile).mockReset(); - vi.mocked(byteLengthFromUrl).mockResolvedValue(1024); + mockRangeProbe(1024); vi.mocked(asyncBufferFromUrl).mockResolvedValue({ byteLength: 1024, slice: vi.fn(), @@ -184,6 +189,56 @@ beforeEach(() => { }); describe("ParquetVisualizer", () => { + it("opens the file at the reported size without any size request", async () => { + const fetchSpy = mockRangeProbe(1024); + mockDataset(SCHEMA, 2); + + renderWithSuspense( + , + ); + + await waitFor(() => + expect(screen.getByTestId("table-visualizer")).toBeInTheDocument(), + ); + + expect(vi.mocked(asyncBufferFromUrl)).toHaveBeenCalledWith( + expect.objectContaining({ byteLength: 9446073 }), + ); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("probes the size with a ranged GET, never a HEAD, when none is reported", async () => { + const fetchSpy = mockRangeProbe(9446073); + mockDataset(SCHEMA, 2); + + renderWithSuspense( + , + ); + + await waitFor(() => + expect(screen.getByTestId("table-visualizer")).toBeInTheDocument(), + ); + + expect(vi.mocked(asyncBufferFromUrl)).toHaveBeenCalledWith( + expect.objectContaining({ byteLength: 9446073 }), + ); + expect(fetchSpy).toHaveBeenCalledWith( + "https://storage.example.com/unsized.parquet", + expect.objectContaining({ headers: { Range: "bytes=0-0" } }), + ); + const methods = fetchSpy.mock.calls.map(([, init]) => + (init?.method ?? "GET").toUpperCase(), + ); + expect(methods).not.toContain("HEAD"); + }); + it("reads via range requests, parses, and renders TableVisualizer with stats", async () => { mockDataset(SCHEMA, 2); @@ -381,6 +436,10 @@ describe("ParquetVisualizer", () => { expect(screen.queryByRole("button", { name: "Load max" })).toBeNull(); }); + /** + * A host that ignores `Range` and returns the whole object: the size probe + * sees 200 with no `Content-Range`, which is the no-range-support signal. + */ const mockFullDownload = (contentLength: string | null) => vi.spyOn(globalThis, "fetch").mockResolvedValue({ ok: true, @@ -390,10 +449,6 @@ describe("ParquetVisualizer", () => { } as unknown as Response); it("falls back to a full download when range requests are unavailable", async () => { - // Simulate a CORS/network failure on the range path (not an HTTP error). - vi.mocked(byteLengthFromUrl).mockRejectedValue( - new TypeError("Failed to fetch"), - ); mockFullDownload(String(1024)); vi.mocked(parquetMetadata).mockReturnValue({ schema: SCHEMA, @@ -418,9 +473,6 @@ describe("ParquetVisualizer", () => { }); it("errors instead of downloading a huge file when range requests are unavailable", async () => { - vi.mocked(byteLengthFromUrl).mockRejectedValue( - new TypeError("Failed to fetch"), - ); mockFullDownload(String(200 * 1024 * 1024)); renderWithSuspense( @@ -437,10 +489,6 @@ describe("ParquetVisualizer", () => { }); it("errors instead of downloading when Content-Length is absent", async () => { - vi.mocked(byteLengthFromUrl).mockRejectedValue( - new TypeError("Failed to fetch"), - ); - mockFullDownload(null); renderWithSuspense( @@ -514,9 +562,11 @@ describe("ParquetVisualizer", () => { }); it("surfaces artifact failures to the error boundary", async () => { - vi.mocked(byteLengthFromUrl).mockRejectedValue( - new ArtifactFetchError(500, "Server Error", "Failed to fetch artifact."), - ); + vi.spyOn(globalThis, "fetch").mockResolvedValue({ + ok: false, + status: 500, + statusText: "Server Error", + } as unknown as Response); renderWithSuspense( { const { data: base } = useSuspenseQuery({ - queryKey: ["artifact-parquet", signedUrl], + queryKey: ["artifact-parquet", signedUrl, byteLength], queryFn: async () => { - const opened = await openParquet(signedUrl); + const opened = await openParquet(signedUrl, byteLength); const columnCount = countColumns(opened.metadata); const { columns, rows } = await readParquetPreview( opened, diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/parquetUtils.ts b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/parquetUtils.ts index 869e45ebc8..152292b3ef 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/parquetUtils.ts +++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/parquetUtils.ts @@ -1,7 +1,6 @@ import { type AsyncBuffer, asyncBufferFromUrl, - byteLengthFromUrl, cachedAsyncBuffer, type FileMetaData, parquetMetadata, @@ -13,10 +12,7 @@ import { import { ArtifactFetchError } from "@/services/executionService"; -import { - fetchArtifactForHyparquet, - fetchArtifactOrThrow, -} from "./useArtifactFetch"; +import { fetchArtifactOrThrow } from "./useArtifactFetch"; import { type ArtifactCell, type ArtifactColumn, @@ -33,44 +29,67 @@ export interface OpenedParquet { metadata: FileMetaData; } -export async function openParquet(signedUrl: string): Promise { +export async function openParquet( + signedUrl: string, + byteLength?: number, +): Promise { try { - const byteLength = await byteLengthFromUrl( - signedUrl, - undefined, - fetchArtifactForHyparquet, - ); const source = cachedAsyncBuffer( await asyncBufferFromUrl({ url: signedUrl, - byteLength, - fetch: fetchArtifactForHyparquet, + byteLength: + typeof byteLength === "number" && byteLength > 0 + ? byteLength + : await byteLengthFromRangedGet(signedUrl), + fetch: fetchArtifactOrThrow, }), ); const metadata = await parquetMetadataAsync(source); return { source, metadata }; } catch (error) { if (error instanceof ArtifactFetchError) throw error; + return openByFullDownload(signedUrl); + } +} + +/** + * Signed object-store URLs reject HEAD — the method is part of the signature — + * and a bucket's CORS policy need not allow it either, so a blocked HEAD + * surfaces as an opaque network error rather than a 403 the caller could fall + * back from. A ranged GET is signed and needs no preflight, but reading the + * total size back requires the bucket to expose `Content-Range`; where it does + * not, this throws and the caller downloads the whole object instead. + */ +async function byteLengthFromRangedGet(url: string): Promise { + const controller = new AbortController(); + const response = await fetchArtifactOrThrow(url, { + headers: { Range: "bytes=0-0" }, + signal: controller.signal, + }); + const totalSize = response.headers + .get("Content-Range") + ?.match(/\/(\d+)$/)?.[1]; + controller.abort(); - const response = await fetchArtifactOrThrow(signedUrl); - const contentLengthHeader = response.headers.get("Content-Length"); - const contentLength = contentLengthHeader - ? Number(contentLengthHeader) - : NaN; - if ( - !Number.isFinite(contentLength) || - contentLength > MAX_VISUALIZABLE_SIZE_BYTES - ) { - throw new ArtifactFetchError( - 413, - "Payload Too Large", - "This parquet file is too large to preview without range-request support.", - ); - } - - const source = await response.arrayBuffer(); - return { source, metadata: parquetMetadata(source) }; + if (response.status !== 206 || !totalSize) { + throw new Error("Artifact host does not support range requests."); } + return Number(totalSize); +} + +async function openByFullDownload(signedUrl: string): Promise { + const response = await fetchArtifactOrThrow(signedUrl); + const contentLength = Number(response.headers.get("Content-Length")); + if (!contentLength || contentLength > MAX_VISUALIZABLE_SIZE_BYTES) { + throw new ArtifactFetchError( + 413, + "Payload Too Large", + "This parquet file is too large to preview.", + ); + } + + const source = await response.arrayBuffer(); + return { source, metadata: parquetMetadata(source) }; } export async function readParquetRows( diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/useArtifactFetch.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/useArtifactFetch.tsx index b23b226f49..c9e0728ba1 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/useArtifactFetch.tsx +++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/useArtifactFetch.tsx @@ -15,22 +15,6 @@ export const fetchArtifactOrThrow: typeof fetch = async (input, init) => { return response; }; -export const fetchArtifactForHyparquet: typeof fetch = async (input, init) => { - const response = await fetch(input, init); - if (response.ok) return response; - - const method = (init?.method ?? "GET").toUpperCase(); - if (method === "HEAD" && response.status === 403) { - return response; - } - - throw new ArtifactFetchError( - response.status, - response.statusText, - "Failed to fetch artifact.", - ); -}; - export function useArtifactFetch( queryKey: string, signedUrl: string, diff --git a/src/routes/PipelineRun/ArtifactPreview.test.tsx b/src/routes/PipelineRun/ArtifactPreview.test.tsx index 0cbeeb478b..765e0ee2c3 100644 --- a/src/routes/PipelineRun/ArtifactPreview.test.tsx +++ b/src/routes/PipelineRun/ArtifactPreview.test.tsx @@ -58,11 +58,13 @@ vi.mock( artifactId, name, isFullscreen, + byteLength, }: { type: string; artifactId: string; name: string; isFullscreen: boolean; + byteLength?: number; }) => (
), }; @@ -114,6 +117,7 @@ describe("ArtifactPreviewPage", () => { expect(preview).toHaveAttribute("data-artifact-id", "abc"); expect(preview).toHaveAttribute("data-name", "output"); expect(preview).toHaveAttribute("data-fullscreen", "true"); + expect(preview).toHaveAttribute("data-byte-length", "100"); }); }); diff --git a/src/routes/PipelineRun/ArtifactPreview.tsx b/src/routes/PipelineRun/ArtifactPreview.tsx index 3ca9712195..f3501a1f57 100644 --- a/src/routes/PipelineRun/ArtifactPreview.tsx +++ b/src/routes/PipelineRun/ArtifactPreview.tsx @@ -137,6 +137,7 @@ const ArtifactPreviewBody = ({ artifactId={artifactId} type={normalizedType} isFullscreen={true} + byteLength={artifactData?.total_size} />
diff --git a/src/routes/v2/pages/CompareView/components/ArtifactComparisonDialog.test.tsx b/src/routes/v2/pages/CompareView/components/ArtifactComparisonDialog.test.tsx new file mode 100644 index 0000000000..2787889762 --- /dev/null +++ b/src/routes/v2/pages/CompareView/components/ArtifactComparisonDialog.test.tsx @@ -0,0 +1,86 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import type { ArtifactNodeResponse } from "@/api/types.gen"; + +import { ArtifactComparisonDialog } from "./ArtifactComparisonDialog"; + +vi.mock( + "@/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/ArtifactPreviewContent", + () => ({ + PreviewContent: ({ + artifactId, + byteLength, + }: { + artifactId: string; + byteLength?: number; + }) => ( +
+ ), + PreviewSkeleton: () => null, + }), +); + +const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, +}); + +beforeEach(() => { + queryClient.clear(); +}); + +const artifact = ( + id: string, + artifactData?: ArtifactNodeResponse["artifact_data"], +): ArtifactNodeResponse => ({ id, artifact_data: artifactData }); + +const renderDialog = ( + artifactA: ArtifactNodeResponse, + artifactB: ArtifactNodeResponse, +) => + render( + + {}} + name="table" + labelA="run-a" + labelB="run-b" + artifactA={artifactA} + artifactB={artifactB} + typeA="Apache Parquet" + typeB="Apache Parquet" + /> + , + ); + +describe("ArtifactComparisonDialog", () => { + it("passes each side's reported size to its preview", () => { + renderDialog( + artifact("a1", { total_size: 1024, is_dir: false }), + artifact("b1", { total_size: 9446073, is_dir: false }), + ); + + const [paneA, paneB] = screen.getAllByTestId("preview-content"); + expect(paneA).toHaveAttribute("data-artifact-id", "a1"); + expect(paneA).toHaveAttribute("data-byte-length", "1024"); + expect(paneB).toHaveAttribute("data-artifact-id", "b1"); + expect(paneB).toHaveAttribute("data-byte-length", "9446073"); + }); + + it("omits the size when a side reports no artifact data", () => { + renderDialog( + artifact("a1"), + artifact("b1", { total_size: 1024, is_dir: false }), + ); + + const [paneA, paneB] = screen.getAllByTestId("preview-content"); + expect(paneA).not.toHaveAttribute("data-byte-length"); + expect(paneB).toHaveAttribute("data-byte-length", "1024"); + }); +}); diff --git a/src/routes/v2/pages/CompareView/components/ArtifactComparisonDialog.tsx b/src/routes/v2/pages/CompareView/components/ArtifactComparisonDialog.tsx index a17b60afae..177b408312 100644 --- a/src/routes/v2/pages/CompareView/components/ArtifactComparisonDialog.tsx +++ b/src/routes/v2/pages/CompareView/components/ArtifactComparisonDialog.tsx @@ -59,6 +59,8 @@ function ComparisonPane({ normalizeRawType(type ?? undefined), ); + const byteLength = artifact?.artifact_data?.total_size; + return ( )}