From ad977f56f36db97f3bbfe56c7a6a033ac09389b1 Mon Sep 17 00:00:00 2001
From: BennyKok <itechbenny@gmail.com>
Date: Thu, 18 Jan 2024 19:50:41 +0800
Subject: [PATCH] feat: add clone machine code

---
 web/src/app/(app)/share/[share_id]/page.tsx | 30 +++++++++-----
 web/src/components/ButtonActionLoader.tsx   | 43 +++++++++++++++++++++
 web/src/server/curdDeploments.ts            | 33 ++++++++++++++++
 3 files changed, 96 insertions(+), 10 deletions(-)

diff --git a/web/src/app/(app)/share/[share_id]/page.tsx b/web/src/app/(app)/share/[share_id]/page.tsx
index ae8de65..5402aca 100644
--- a/web/src/app/(app)/share/[share_id]/page.tsx
+++ b/web/src/app/(app)/share/[share_id]/page.tsx
@@ -1,9 +1,8 @@
-import { ButtonAction } from "@/components/ButtonActionLoader";
+import { ButtonActionMenu } from "@/components/ButtonActionLoader";
 import {
   PublicRunOutputs,
   RunWorkflowInline,
 } from "@/components/VersionSelect";
-import { Button } from "@/components/ui/button";
 import {
   Card,
   CardContent,
@@ -16,7 +15,11 @@ import { usersTable } from "@/db/schema";
 import { getInputsFromWorkflow } from "@/lib/getInputsFromWorkflow";
 import { getRelativeTime } from "@/lib/getRelativeTime";
 import { setInitialUserData } from "@/lib/setInitialUserData";
-import { cloneWorkflow, findSharedDeployment } from "@/server/curdDeploments";
+import {
+  cloneMachine,
+  cloneWorkflow,
+  findSharedDeployment,
+} from "@/server/curdDeploments";
 import { auth, clerkClient } from "@clerk/nextjs/server";
 import { eq } from "drizzle-orm";
 import { redirect } from "next/navigation";
@@ -63,13 +66,20 @@ export default async function Page({
               {" / "}
               {sharedDeployment.workflow.name}
             </div>
-            <Button asChild className="gap-2" variant="outline" type="submit">
-              <ButtonAction
-                action={cloneWorkflow.bind(null, sharedDeployment.id)}
-              >
-                Clone
-              </ButtonAction>
-            </Button>
+
+            <ButtonActionMenu
+              title="Clone"
+              actions={[
+                {
+                  title: "Workflow",
+                  action: cloneWorkflow.bind(null, sharedDeployment.id),
+                },
+                {
+                  title: "Machine",
+                  action: cloneMachine.bind(null, sharedDeployment.id),
+                },
+              ]}
+            />
           </CardTitle>
           <CardDescription suppressHydrationWarning={true}>
             {getRelativeTime(sharedDeployment?.updated_at)}
diff --git a/web/src/components/ButtonActionLoader.tsx b/web/src/components/ButtonActionLoader.tsx
index df703c7..331a3bb 100644
--- a/web/src/components/ButtonActionLoader.tsx
+++ b/web/src/components/ButtonActionLoader.tsx
@@ -2,6 +2,14 @@
 
 import { LoadingIcon } from "@/components/LoadingIcon";
 import { callServerPromise } from "@/components/callServerPromise";
+import { Button } from "@/components/ui/button";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import { MoreVertical } from "lucide-react";
 import { useRouter } from "next/navigation";
 import { useState } from "react";
 
@@ -33,3 +41,38 @@ export function ButtonAction({
     </button>
   );
 }
+
+export function ButtonActionMenu(props: {
+  title?: string;
+  actions: {
+    title: string;
+    action: () => Promise<any>;
+  }[];
+}) {
+  const [isLoading, setIsLoading] = useState(false);
+
+  return (
+    <DropdownMenu>
+      <DropdownMenuTrigger asChild>
+        <Button className="gap-2" variant="outline" disabled={isLoading}>
+          {props.title}
+          {isLoading ? <LoadingIcon /> : <MoreVertical size={14} />}
+        </Button>
+      </DropdownMenuTrigger>
+      <DropdownMenuContent className="w-56">
+        {props.actions.map((action) => (
+          <DropdownMenuItem
+            key={action.title}
+            onClick={async () => {
+              setIsLoading(true);
+              await callServerPromise(action.action());
+              setIsLoading(false);
+            }}
+          >
+            {action.title}
+          </DropdownMenuItem>
+        ))}
+      </DropdownMenuContent>
+    </DropdownMenu>
+  );
+}
diff --git a/web/src/server/curdDeploments.ts b/web/src/server/curdDeploments.ts
index 0a0d2f4..4724f54 100644
--- a/web/src/server/curdDeploments.ts
+++ b/web/src/server/curdDeploments.ts
@@ -4,6 +4,7 @@ import { db } from "@/db/db";
 import type { DeploymentType } from "@/db/schema";
 import { deploymentsTable, workflowTable } from "@/db/schema";
 import { createNewWorkflow } from "@/server/createNewWorkflow";
+import { addCustomMachine } from "@/server/curdMachine";
 import { withServerPromise } from "@/server/withServerPromise";
 import { auth } from "@clerk/nextjs";
 import { and, eq, isNull } from "drizzle-orm";
@@ -162,3 +163,35 @@ export const cloneWorkflow = withServerPromise(
     };
   }
 );
+
+export const cloneMachine = withServerPromise(async (deployment_id: string) => {
+  const deployment = await db.query.deploymentsTable.findFirst({
+    where: and(
+      eq(deploymentsTable.environment, "public-share"),
+      eq(deploymentsTable.id, deployment_id)
+    ),
+    with: {
+      machine: true,
+    },
+  });
+
+  if (!deployment) throw new Error("No deployment found");
+  if (deployment.machine.type !== "comfy-deploy-serverless")
+    throw new Error("Can only clone comfy-deploy-serverlesss");
+
+  const { userId, orgId } = auth();
+
+  if (!userId) throw new Error("No user id");
+
+  await addCustomMachine({
+    gpu: deployment.machine.gpu,
+    models: deployment.machine.models,
+    snapshot: deployment.machine.snapshot,
+    name: `${deployment.machine.name} (Cloned)`,
+    type: "comfy-deploy-serverless",
+  });
+
+  return {
+    message: "Successfully cloned workflow",
+  };
+});