feat(all): add organisation, api updates with organisation check

Now calling run endpoints with GET, POST will check against the API key, whether there is org_id or not, if the operation workflow doesnt match with the user org or user id, will return with workflow not found, run not found

web/src/server/APIKeyBodyRequest.ts

@@ -5,3 +5,5 @@ export const APIKeyBodyRequest = z.object({
   org_id: z.string().optional(),
   iat: z.number(),
 });
+
+export type APIKeyUserType = z.infer<typeof APIKeyBodyRequest>;

web/src/server/createRun.ts

@@ -2,7 +2,9 @@
 
 import { withServerPromise } from "./withServerPromise";
 import { db } from "@/db/db";
+import type { MachineType, WorkflowVersionType } from "@/db/schema";
 import { machinesTable, workflowRunsTable } from "@/db/schema";
+import type { APIKeyUserType } from "@/server/APIKeyBodyRequest";
 import { ComfyAPI_Run } from "@/types/ComfyAPI_Run";
 import { and, eq } from "drizzle-orm";
 import { revalidatePath } from "next/cache";
@@ -10,34 +12,67 @@ import "server-only";
 import { v4 } from "uuid";
 
 export const createRun = withServerPromise(
-  async (
-    origin: string,
-    workflow_version_id: string,
-    machine_id: string,
-    inputs?: Record<string, string>,
-    isManualRun?: boolean
-  ) => {
-    const machine = await db.query.machinesTable.findFirst({
-      where: and(
-        eq(machinesTable.id, machine_id),
-        eq(machinesTable.disabled, false)
-      ),
-    });
+  async ({
+    origin,
+    workflow_version_id,
+    machine_id,
+    inputs,
+    isManualRun,
+    apiUser,
+  }: {
+    origin: string;
+    workflow_version_id: string | WorkflowVersionType;
+    machine_id: string | MachineType;
+    inputs?: Record<string, string>;
+    isManualRun?: boolean;
+    apiUser?: APIKeyUserType;
+  }) => {
+    const machine =
+      typeof machine_id === "string"
+        ? await db.query.machinesTable.findFirst({
+            where: and(
+              eq(machinesTable.id, machine_id),
+              eq(machinesTable.disabled, false)
+            ),
+          })
+        : machine_id;
 
     if (!machine) {
       throw new Error("Machine not found");
     }
 
-    const workflow_version_data = await db.query.workflowVersionTable.findFirst(
-      {
-        where: eq(workflowRunsTable.id, workflow_version_id),
-      }
-    );
+    const workflow_version_data =
+      typeof workflow_version_id === "string"
+        ? await db.query.workflowVersionTable.findFirst({
+            where: eq(workflowRunsTable.id, workflow_version_id),
+            with: {
+              workflow: {
+                columns: {
+                  org_id: true,
+                  user_id: true,
+                },
+              },
+            },
+          })
+        : workflow_version_id;
 
     if (!workflow_version_data) {
       throw new Error("Workflow version not found");
     }
 
+    if (apiUser)
+      if (apiUser.org_id) {
+        // is org api call, check org only
+        if (apiUser.org_id != workflow_version_data.workflow.org_id) {
+          throw new Error("Workflow not found");
+        }
+      } else {
+        // is user api call, check user only
+        if (apiUser.user_id != workflow_version_data.workflow.user_id) {
+          throw new Error("Workflow not found");
+        }
+      }
+
     const workflow_api = workflow_version_data.workflow_api;
 
     // Replace the inputs
@@ -68,7 +103,7 @@ export const createRun = withServerPromise(
         workflow_id: workflow_version_data.workflow_id,
         workflow_version_id: workflow_version_data.id,
         workflow_inputs: inputs,
-        machine_id,
+        machine_id: machine.id,
         origin: isManualRun ? "manual" : "api",
       })
       .returning();

web/src/server/curdApiKeys.ts

@@ -3,7 +3,7 @@
 import { db } from "@/db/db";
 import { apiKeyTable } from "@/db/schema";
 import { auth } from "@clerk/nextjs";
-import { and, desc, eq } from "drizzle-orm";
+import { and, desc, eq, isNull } from "drizzle-orm";
 import jwt from "jsonwebtoken";
 import { revalidatePath } from "next/cache";
 
@@ -29,7 +29,7 @@ export async function addNewAPIKey(name: string) {
   if (orgId) {
     token = jwt.sign(
       { user_id: userId, org_id: orgId },
-      process.env.JWT_SECRET!,
+      process.env.JWT_SECRET!
     );
   } else {
     token = jwt.sign({ user_id: userId }, process.env.JWT_SECRET!);
@@ -90,7 +90,11 @@ export async function getAPIKeys() {
     });
   } else {
     return await db.query.apiKeyTable.findMany({
-      where: and(eq(apiKeyTable.user_id, userId), eq(apiKeyTable.revoked, false)),
+      where: and(
+        eq(apiKeyTable.user_id, userId),
+        isNull(apiKeyTable.org_id),
+        eq(apiKeyTable.revoked, false)
+      ),
       orderBy: desc(apiKeyTable.created_at),
     });
   }
@@ -102,4 +106,4 @@ export async function isKeyRevoked(key: string) {
   });
 
   return revokedKey !== undefined;
-}
+}

web/src/server/curdMachine.ts

@@ -24,11 +24,12 @@ export async function getMachines() {
 
 export const addMachine = withServerPromise(
   async (data: z.infer<typeof addMachineSchema>) => {
-    const { userId } = auth();
+    const { userId, orgId } = auth();
     if (!userId) return { error: "No user id" };
     // console.log(name, endpoint);
     await db.insert(machinesTable).values({
       ...data,
+      org_id: orgId,
       user_id: userId,
     });
     revalidatePath("/machines");

web/src/server/getRunsOutput.tsx

@@ -3,7 +3,8 @@
 import { RunOutputs } from "@/components/RunOutputs";
 import { db } from "@/db/db";
 import { workflowRunOutputs, workflowRunsTable } from "@/db/schema";
-import { eq } from "drizzle-orm";
+import type { APIKeyUserType } from "@/server/APIKeyBodyRequest";
+import { and, eq } from "drizzle-orm";
 
 export async function getRunsOutputDisplay(run_id: string) {
   return <RunOutputs run_id={run_id} />;
@@ -17,11 +18,38 @@ export async function getRunsOutput(run_id: string) {
     .where(eq(workflowRunOutputs.run_id, run_id));
 }
 
-export async function getRunsData(run_id: string) {
-  // throw new Error("Not implemented");
-  return await db.query.workflowRunsTable.findFirst({
-    where: eq(workflowRunsTable.id, run_id),
+export async function getRunsData(user: APIKeyUserType, run_id: string) {
+  const data = await db.query.workflowRunsTable.findFirst({
+    where: and(
+      eq(workflowRunsTable.id, run_id)
+      // inArray(
+      //   workflowRunsTable.workflow_id,
+      //   db
+      //     .select({
+      //       id: workflowTable.id,
+      //     })
+      //     .from(workflowTable)
+      //     .innerJoin(
+      //       workflowRunsTable,
+      //       eq(workflowTable.id, workflowRunsTable.workflow_id)
+      //     )
+      //     .where(
+      //       and(
+      //         eq(workflowTable.id, workflowRunsTable.workflow_id),
+      //         user.org_id
+      //           ? eq(workflowTable.org_id, user.org_id)
+      //           : eq(workflowTable.user_id, user.user_id!)
+      //       )
+      //     )
+      // )
+    ),
     with: {
+      workflow: {
+        columns: {
+          org_id: true,
+          user_id: true,
+        },
+      },
       outputs: {
         columns: {
           data: true,
@@ -29,4 +57,22 @@ export async function getRunsData(run_id: string) {
       },
     },
   });
+
+  if (!data) {
+    return null;
+  }
+
+  if (user.org_id) {
+    // is org api call, check org only
+    if (data.workflow.org_id != user.org_id) {
+      return null;
+    }
+  } else {
+    // is user api call, check user only
+    if (data.workflow.user_id != user.user_id) {
+      return null;
+    }
+  }
+
+  return data;
 }