From ee685077554b45d4b9f97dea97cb0aa00ff3b738 Mon Sep 17 00:00:00 2001
From: BennyKok <itechbenny@gmail.com>
Date: Wed, 20 Dec 2023 23:14:01 +0800
Subject: [PATCH] feat(web): add custom file download button for runs output

---
 web/src/app/api/file-upload/route.ts |  3 ++-
 web/src/app/api/update-run/route.ts  |  4 ++--
 web/src/app/api/view/route.ts        | 12 ++++-------
 web/src/components/LoadingIcon.tsx   |  2 +-
 web/src/components/OutputRender.tsx  | 31 ++++++++++++++++++++++++++++
 web/src/components/RunDisplay.tsx    | 15 +-------------
 web/src/components/RunInputs.tsx     |  2 +-
 web/src/components/RunOutputs.tsx    |  2 +-
 web/src/components/VersionSelect.tsx |  4 ++--
 web/src/server/getFileDownloadUrl.ts |  9 ++++++++
 10 files changed, 54 insertions(+), 30 deletions(-)
 create mode 100644 web/src/components/OutputRender.tsx
 create mode 100644 web/src/server/getFileDownloadUrl.ts

diff --git a/web/src/app/api/file-upload/route.ts b/web/src/app/api/file-upload/route.ts
index 0d96918..62a0243 100644
--- a/web/src/app/api/file-upload/route.ts
+++ b/web/src/app/api/file-upload/route.ts
@@ -6,12 +6,13 @@ import { z } from "zod";
 const Request = z.object({
   file_name: z.string(),
   run_id: z.string(),
-  type: z.enum(["image/png", "image/jpeg"]),
+  type: z.string(),
 });
 
 export const dynamic = "force-dynamic";
 
 export async function GET(request: Request) {
+  console.log("yo hello");
   const [data, error] = await parseDataSafe(Request, request);
   if (!data || error) return error;
 
diff --git a/web/src/app/api/update-run/route.ts b/web/src/app/api/update-run/route.ts
index d928bcb..b4a2f0f 100644
--- a/web/src/app/api/update-run/route.ts
+++ b/web/src/app/api/update-run/route.ts
@@ -17,7 +17,7 @@ export async function POST(request: Request) {
 
   const { run_id, status, output_data } = data;
 
-  console.log(run_id, status, output_data);
+  // console.log(run_id, status, output_data);
 
   if (output_data) {
     const workflow_run_output = await db.insert(workflowRunOutputs).values({
@@ -25,7 +25,7 @@ export async function POST(request: Request) {
       data: output_data,
     });
   } else if (status) {
-    console.log("status", status);
+    // console.log("status", status);
     const workflow_run = await db
       .update(workflowRunsTable)
       .set({
diff --git a/web/src/app/api/view/route.ts b/web/src/app/api/view/route.ts
index d3ca8ac..2a3e784 100644
--- a/web/src/app/api/view/route.ts
+++ b/web/src/app/api/view/route.ts
@@ -1,12 +1,8 @@
-import { replaceCDNUrl } from "@/server/resource";
 import { NextResponse, type NextRequest } from "next/server";
+import { getFileDownloadUrl } from "../../../server/getFileDownloadUrl";
 
 export async function GET(request: NextRequest) {
   const file = new URL(request.url).searchParams.get("file");
-  // console.log(`${process.env.SPACES_ENDPOINT}/comfyui-deploy/${file}`);
-  return NextResponse.redirect(
-    replaceCDNUrl(
-      `${process.env.SPACES_ENDPOINT}/${process.env.SPACES_BUCKET}/${file}`
-    )
-  );
-}
+  if (!file) return NextResponse.redirect("/");
+  return NextResponse.redirect(getFileDownloadUrl(file));
+}
\ No newline at end of file
diff --git a/web/src/components/LoadingIcon.tsx b/web/src/components/LoadingIcon.tsx
index 65f0078..b925661 100644
--- a/web/src/components/LoadingIcon.tsx
+++ b/web/src/components/LoadingIcon.tsx
@@ -3,5 +3,5 @@ import { LoaderIcon } from "lucide-react";
 import * as React from "react";
 
 export function LoadingIcon() {
-  return <LoaderIcon size={14} className="ml-2 animate-spin" />;
+  return <LoaderIcon size={14} className="animate-spin" />;
 }
diff --git a/web/src/components/OutputRender.tsx b/web/src/components/OutputRender.tsx
new file mode 100644
index 0000000..b0b652f
--- /dev/null
+++ b/web/src/components/OutputRender.tsx
@@ -0,0 +1,31 @@
+'use client'
+
+import { Download } from "lucide-react";
+import { Button } from "@/components/ui/button";
+import { getFileDownloadUrl } from "@/server/getFileDownloadUrl";
+
+
+export function OutputRender(props: { run_id: string; filename: string; }) {
+  if (props.filename.endsWith(".png")) {
+    return (
+      <img
+        className="max-w-[200px]"
+        alt={props.filename}
+        src={`/api/view?file=${encodeURIComponent(
+          `outputs/runs/${props.run_id}/${props.filename}`
+        )}`} />
+    );
+  } else {
+    return <Button className="gap-2" onClick={async () => {
+      const url = await getFileDownloadUrl(`outputs/runs/${props.run_id}/${props.filename}`);
+
+      const a = document.createElement('a');
+      a.href = url;
+      a.download = props.filename;
+      a.target = '_blank'; // Add this line
+      document.body.appendChild(a);
+      a.click();
+      document.body.removeChild(a);
+    }}>Download <Download size={14} /></Button>;
+  }
+}
diff --git a/web/src/components/RunDisplay.tsx b/web/src/components/RunDisplay.tsx
index 8aa9613..18ce7a6 100644
--- a/web/src/components/RunDisplay.tsx
+++ b/web/src/components/RunDisplay.tsx
@@ -1,3 +1,4 @@
+
 import { RunInputs } from "@/components/RunInputs";
 import { LiveStatus } from "./LiveStatus";
 import { RunOutputs } from "@/components/RunOutputs";
@@ -51,17 +52,3 @@ export async function RunDisplay({
     </Dialog>
   );
 }
-
-export function OutputRender(props: { run_id: string; filename: string }) {
-  if (props.filename.endsWith(".png")) {
-    return (
-      <img
-        className="max-w-[200px]"
-        alt={props.filename}
-        src={`/api/view?file=${encodeURIComponent(
-          `outputs/runs/${props.run_id}/${props.filename}`
-        )}`}
-      />
-    );
-  }
-}
diff --git a/web/src/components/RunInputs.tsx b/web/src/components/RunInputs.tsx
index 97695f1..7e24f42 100644
--- a/web/src/components/RunInputs.tsx
+++ b/web/src/components/RunInputs.tsx
@@ -1,4 +1,4 @@
-import { OutputRender } from "./RunDisplay";
+import { OutputRender } from "./OutputRender";
 import { CodeBlock } from "@/components/CodeBlock";
 import {
   Table,
diff --git a/web/src/components/RunOutputs.tsx b/web/src/components/RunOutputs.tsx
index 2c78d35..ae61ff6 100644
--- a/web/src/components/RunOutputs.tsx
+++ b/web/src/components/RunOutputs.tsx
@@ -1,4 +1,4 @@
-import { OutputRender } from "./RunDisplay";
+import { OutputRender } from "./OutputRender";
 import { CodeBlock } from "@/components/CodeBlock";
 import {
   Table,
diff --git a/web/src/components/VersionSelect.tsx b/web/src/components/VersionSelect.tsx
index 188c87f..30f80ab 100644
--- a/web/src/components/VersionSelect.tsx
+++ b/web/src/components/VersionSelect.tsx
@@ -128,7 +128,7 @@ export function RunWorkflowButton({
         }
       }}
     >
-      Run <Play size={14} /> {isLoading && <LoadingIcon />}
+      Run  {isLoading ?  <LoadingIcon /> : <Play size={14} />}
     </Button>
   );
 }
@@ -155,7 +155,7 @@ export function CreateDeploymentButton({
     <DropdownMenu>
       <DropdownMenuTrigger asChild>
         <Button className="gap-2" disabled={isLoading} variant="outline">
-          Deploy <MoreVertical size={14} /> {isLoading && <LoadingIcon />}
+          Deploy  {isLoading ? <LoadingIcon /> : <MoreVertical size={14} /> }
         </Button>
       </DropdownMenuTrigger>
       <DropdownMenuContent className="w-56">
diff --git a/web/src/server/getFileDownloadUrl.ts b/web/src/server/getFileDownloadUrl.ts
new file mode 100644
index 0000000..32ac026
--- /dev/null
+++ b/web/src/server/getFileDownloadUrl.ts
@@ -0,0 +1,9 @@
+'use server'
+
+import { replaceCDNUrl } from "@/server/resource";
+
+export async function getFileDownloadUrl(file: string) {
+  return replaceCDNUrl(
+    `${process.env.SPACES_ENDPOINT}/${process.env.SPACES_BUCKET}/${file}`
+  );
+}