From 127b0df7323c77b7acb931c40d6361e05705c335 Mon Sep 17 00:00:00 2001
From: Camiel van Schoonhoven <camiel.vanschoonhoven@shopify.com>
Date: Mon, 23 Mar 2026 12:26:23 -0700
Subject: [PATCH] chore: Simplify IOCell

---
 .../ArtifactVisualizer/ArtifactVisualizer.tsx |  91 +++++++++--
 .../ArtifactVisualizer/CsvVisualizer.tsx      |  20 ++-
 .../ArtifactVisualizer/JsonVisualizer.tsx     |  18 +-
 .../ArtifactVisualizer/TableVisualizer.tsx    |   4 +-
 .../TaskOverview/IOSection/IOCell/IOCell.tsx  | 154 +++++-------------
 5 files changed, 145 insertions(+), 142 deletions(-)

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 3221d2906..4098f852d 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
@@ -37,19 +37,19 @@ const VISUALIZABLE_TYPES = new Set([
   "apacheparquet",
 ]);
 
-interface ArtifactVisualizerProps {
+type ArtifactVisualizerProps = {
   artifact: ArtifactNodeResponse;
   name: string;
   type: string;
-}
+  value?: string;
+};
 
 const ArtifactVisualizer = ({
   artifact,
   name,
   type,
+  value,
 }: ArtifactVisualizerProps) => {
-  const { backendUrl } = useBackend();
-
   const [isFullscreen, setIsFullscreen] = useState(false);
 
   const normalizedType = type?.toLowerCase().replace(/\s/g, "") ?? "text";
@@ -67,10 +67,20 @@ const ArtifactVisualizer = ({
   return (
     <Dialog onOpenChange={handleOpenChange}>
       <DialogTrigger asChild>
-        <Button variant="ghost" size="xs">
-          <Icon name="Eye" />
-          Preview
-        </Button>
+        {value ? (
+          <Button variant="ghost" size="xs">
+            <Icon
+              name="Maximize2"
+              size="xs"
+              className="text-muted-foreground"
+            />
+          </Button>
+        ) : (
+          <Button variant="ghost" size="xs">
+            <Icon name="Eye" />
+            Preview
+          </Button>
+        )}
       </DialogTrigger>
       <DialogContent
         className={cn(
@@ -110,23 +120,69 @@ const ArtifactVisualizer = ({
           </DialogDescription>
         </DialogHeader>
         <div className="overflow-auto flex-1 min-h-0">
-          <PreviewContent
-            artifactId={artifact.id}
-            type={normalizedType}
-            backendUrl={backendUrl}
-            name={name}
-            isFullscreen={isFullscreen}
-          />
+          {value ? (
+            <InlineContent
+              name={name}
+              value={value}
+              type={normalizedType}
+              isFullscreen={isFullscreen}
+            />
+          ) : (
+            <PreviewContent
+              name={name}
+              artifactId={artifact.id}
+              type={normalizedType}
+              isFullscreen={isFullscreen}
+            />
+          )}
         </div>
       </DialogContent>
     </Dialog>
   );
 };
 
+interface InlineContentProps {
+  type: string;
+  name: string;
+  value: string;
+  isFullscreen: boolean;
+}
+
+const InlineContent = ({
+  type,
+  name,
+  value,
+  isFullscreen,
+}: InlineContentProps) => {
+  switch (type) {
+    case "csv":
+      return (
+        <CsvVisualizer
+          value={value}
+          delimiter=","
+          isFullscreen={isFullscreen}
+        />
+      );
+    case "tsv":
+      return (
+        <CsvVisualizer
+          value={value}
+          delimiter={"\t"}
+          isFullscreen={isFullscreen}
+        />
+      );
+    case "jsonobject":
+    case "jsonarray":
+      return <JsonVisualizer value={value} name={name} />;
+    case "text":
+    default:
+      return <TextVisualizer value={value} />;
+  }
+};
+
 interface PreviewContentProps {
   artifactId: string;
   type: string;
-  backendUrl: string;
   name: string;
   isFullscreen: boolean;
 }
@@ -134,10 +190,11 @@ interface PreviewContentProps {
 const PreviewContent = ({
   artifactId,
   type,
-  backendUrl,
   name,
   isFullscreen,
 }: PreviewContentProps) => {
+  const { backendUrl } = useBackend();
+
   const { data, isLoading, error } = useQuery({
     queryKey: ["artifact-signed-url", artifactId],
     queryFn: () => getArtifactSignedUrl(artifactId, backendUrl),
diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/CsvVisualizer.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/CsvVisualizer.tsx
index 0a8be3582..b88c9d4dd 100644
--- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/CsvVisualizer.tsx
+++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/CsvVisualizer.tsx
@@ -7,20 +7,25 @@ import { TWENTY_FOUR_HOURS_IN_MS } from "@/utils/constants";
 import TableVisualizer from "./TableVisualizer";
 import { parseCsvPreview } from "./utils";
 
-interface CsvVisualizerProps {
-  signedUrl: string;
+type CsvVisualizerProps = {
   delimiter: string;
   isFullscreen: boolean;
-}
+} & (
+  | { value: string; signedUrl?: never }
+  | { value?: never; signedUrl: string }
+);
 
 const CsvVisualizer = ({
-  signedUrl,
   delimiter,
   isFullscreen,
+  value,
+  signedUrl,
 }: CsvVisualizerProps) => {
   const { data, isLoading, error } = useQuery({
     queryKey: ["artifact-csv", signedUrl, delimiter],
     queryFn: async () => {
+      if (!signedUrl) return null;
+
       const response = await fetch(signedUrl);
       if (!response.ok) {
         throw new Error(`(${response.status}) Failed to fetch artifact.`);
@@ -31,8 +36,11 @@ const CsvVisualizer = ({
     },
     staleTime: TWENTY_FOUR_HOURS_IN_MS,
     retry: false,
+    enabled: !!signedUrl,
   });
 
+  const parsed = value ? parseCsvPreview(value, delimiter) : data;
+
   if (isLoading) return <Spinner />;
 
   if (error) {
@@ -43,7 +51,7 @@ const CsvVisualizer = ({
     );
   }
 
-  if (!data || data.headers.length === 0) {
+  if (!parsed || parsed.headers.length === 0) {
     return (
       <Paragraph tone="subdued" size="xs">
         No data
@@ -53,7 +61,7 @@ const CsvVisualizer = ({
 
   return (
     <TableVisualizer
-      data={data}
+      data={parsed}
       signedUrl={signedUrl}
       isFullscreen={isFullscreen}
     />
diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/JsonVisualizer.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/JsonVisualizer.tsx
index 0cc64517e..dc41be113 100644
--- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/JsonVisualizer.tsx
+++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/JsonVisualizer.tsx
@@ -6,15 +6,19 @@ import { TWENTY_FOUR_HOURS_IN_MS } from "@/utils/constants";
 
 import IOCodeViewer from "../IOCodeViewer";
 
-interface JsonVisualizerProps {
-  signedUrl: string;
+type JsonVisualizerProps = {
   name: string;
-}
+} & (
+  | { value: string; signedUrl?: never }
+  | { value?: never; signedUrl: string }
+);
 
-const JsonVisualizer = ({ signedUrl, name }: JsonVisualizerProps) => {
+const JsonVisualizer = ({ name, value, signedUrl }: JsonVisualizerProps) => {
   const { data, isLoading, error } = useQuery({
     queryKey: ["artifact-json", signedUrl],
     queryFn: async () => {
+      if (!signedUrl) return null;
+
       const response = await fetch(signedUrl);
       if (!response.ok) {
         throw new Error(`(${response.status}) Failed to fetch artifact.`);
@@ -24,6 +28,7 @@ const JsonVisualizer = ({ signedUrl, name }: JsonVisualizerProps) => {
     },
     staleTime: TWENTY_FOUR_HOURS_IN_MS,
     retry: false,
+    enabled: !!signedUrl,
   });
 
   if (isLoading) return <Spinner />;
@@ -36,9 +41,10 @@ const JsonVisualizer = ({ signedUrl, name }: JsonVisualizerProps) => {
     );
   }
 
-  if (!data) return null;
+  const content = value ?? data;
+  if (!content) return null;
 
-  return <IOCodeViewer title={name} value={data} />;
+  return <IOCodeViewer title={name} value={content} />;
 };
 
 export default JsonVisualizer;
diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/TableVisualizer.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/TableVisualizer.tsx
index 64aaa7ff8..e98371eda 100644
--- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/TableVisualizer.tsx
+++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/ArtifactVisualizer/TableVisualizer.tsx
@@ -18,7 +18,7 @@ import {
 
 interface TableVisualizerProps {
   data: ArtifactTableData;
-  signedUrl: string;
+  signedUrl?: string;
   isFullscreen: boolean;
 }
 
@@ -42,7 +42,7 @@ const TableVisualizer = ({
             ? `Showing all ${displayedRows.length} rows`
             : `Showing first ${displayedRows.length} rows`}
         </Paragraph>
-        {!isShowingAllRows && (
+        {!isShowingAllRows && signedUrl && (
           <Link
             href={signedUrl}
             target="_blank"
diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/IOCell.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/IOCell.tsx
index 18d7c6acd..c212c422b 100644
--- a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/IOCell.tsx
+++ b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskOverview/IOSection/IOCell/IOCell.tsx
@@ -1,26 +1,11 @@
-import { useState } from "react";
-
-import type {
-  ArtifactDataResponse,
-  ArtifactNodeResponse,
-} from "@/api/types.gen";
+import type { ArtifactNodeResponse } from "@/api/types.gen";
 import { CopyText } from "@/components/shared/CopyText/CopyText";
-import { Button } from "@/components/ui/button";
-import {
-  Dialog,
-  DialogContent,
-  DialogDescription,
-  DialogFooter,
-  DialogTitle,
-} from "@/components/ui/dialog";
-import { Icon } from "@/components/ui/icon";
 import { BlockStack, InlineStack } from "@/components/ui/layout";
 import { Text } from "@/components/ui/typography";
 import { formatBytes } from "@/utils/string";
 
 import ArtifactURI from "./ArtifactURI";
 import ArtifactVisualizer from "./ArtifactVisualizer/ArtifactVisualizer";
-import TextVisualizer from "./ArtifactVisualizer/TextVisualizer";
 
 interface IOCellProps {
   name: string;
@@ -34,82 +19,59 @@ const IOCell = ({ name, type, artifact }: IOCellProps) => {
   const hasInlineValue = canShowInlineValue(inlineValue);
   const hasDetails = Boolean(artifactData?.uri || hasInlineValue);
 
-  const [isOpen, setIsOpen] = useState(false);
-
   const artifactType =
     type ?? artifact?.type_name ?? (artifactData?.is_dir ? "Directory" : "Any");
 
   return (
-    <>
-      <BlockStack gap="1" className="w-full p-2 border bg-white rounded-md">
-        <InlineStack
-          gap="2"
-          align="space-between"
-          blockAlign="center"
-          className="w-full"
-        >
-          <InlineStack gap="4" blockAlign="center">
-            <CopyText compact size="sm">
-              {name}
-            </CopyText>
+    <BlockStack gap="1" className="w-full p-2 border bg-white rounded-md">
+      <InlineStack
+        gap="2"
+        align="space-between"
+        blockAlign="center"
+        className="w-full"
+      >
+        <InlineStack gap="4" blockAlign="center">
+          <CopyText compact size="sm">
+            {name}
+          </CopyText>
 
-            <InlineStack gap="2" blockAlign="center">
-              <Text size="xs" tone="subdued">
-                {artifactType}
-              </Text>
+          <InlineStack gap="2" blockAlign="center">
+            <Text size="xs" tone="subdued">
+              {artifactType}
+            </Text>
 
-              {artifactData?.total_size && (
-                <Text size="xs" tone="subdued" font="mono">
-                  ({formatBytes(artifactData.total_size)})
-                </Text>
-              )}
-            </InlineStack>
+            {artifactData?.total_size && (
+              <Text size="xs" tone="subdued" font="mono">
+                ({formatBytes(artifactData.total_size)})
+              </Text>
+            )}
           </InlineStack>
+        </InlineStack>
 
-          {artifact && hasDetails && !hasInlineValue && !!type && (
-            <ArtifactVisualizer artifact={artifact} name={name} type={type} />
-          )}
+        {artifact && (hasDetails || hasInlineValue) && (
+          <ArtifactVisualizer
+            artifact={artifact}
+            name={name}
+            type={type ?? "text"}
+            value={hasInlineValue ? inlineValue : undefined}
+          />
+        )}
+      </InlineStack>
 
-          {hasInlineValue && (
-            <InlineStack gap="1" wrap="nowrap" blockAlign="center">
-              <CopyText
-                size="xs"
-                compact
-                className="font-mono text-success line-clamp-2 break-all"
-              >
-                {inlineValue}
-              </CopyText>
+      {hasInlineValue && (
+        <CopyText
+          size="xs"
+          compact
+          className="font-mono text-success line-clamp-2 break-all"
+        >
+          {inlineValue}
+        </CopyText>
+      )}
 
-              {inlineValue.length > 16 && (
-                <Button
-                  variant="ghost"
-                  size="xs"
-                  onClick={() => setIsOpen(true)}
-                >
-                  <Icon
-                    name="Maximize2"
-                    size="xs"
-                    className="text-muted-foreground"
-                  />
-                </Button>
-              )}
-            </InlineStack>
-          )}
-        </InlineStack>
-        {artifactData?.uri && (
-          <ArtifactURI uri={artifactData.uri} isDir={artifactData.is_dir} />
-        )}
-      </BlockStack>
-      {inlineValue && (
-        <TextPreview
-          open={isOpen}
-          title={name}
-          value={inlineValue}
-          artifactData={artifactData}
-          onClose={() => setIsOpen(false)}
-        />
+      {artifactData?.uri && (
+        <ArtifactURI uri={artifactData.uri} isDir={artifactData.is_dir} />
       )}
-    </>
+    </BlockStack>
   );
 };
 
@@ -126,33 +88,3 @@ const canShowInlineValue = (
   }
   return false;
 };
-
-const TextPreview = ({
-  open,
-  title,
-  value,
-  artifactData,
-  onClose,
-}: {
-  open: boolean;
-  title: string;
-  value: string;
-  artifactData?: ArtifactDataResponse;
-  onClose: () => void;
-}) => {
-  return (
-    <Dialog open={open} onOpenChange={onClose}>
-      <DialogContent className="max-h-9/10 flex flex-col">
-        <DialogTitle>{title}</DialogTitle>
-        <DialogDescription>Complete artifact</DialogDescription>
-        {artifactData?.uri && (
-          <ArtifactURI uri={artifactData.uri} isDir={artifactData.is_dir} />
-        )}
-        <TextVisualizer value={value} />
-        <DialogFooter>
-          <Button onClick={onClose}>Close</Button>
-        </DialogFooter>
-      </DialogContent>
-    </Dialog>
-  );
-};