chunzhi/comfyui-deploy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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
This commit is contained in:
BennyKok 2024-01-01 23:13:01 +08:00
parent 75776e6d8f
commit c66de45522
19 changed files with 907 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

BIN
web/bun.lockb
View File

Binary file not shown.

2
web/drizzle/0016_overrated_cable.sql Normal file
View File

@ -0,0 +1,2 @@
ALTER TABLE "comfyui_deploy"."machines" ADD COLUMN "org_id" text;--> statement-breakpoint
ALTER TABLE "comfyui_deploy"."workflows" ADD COLUMN "org_id" text;

675
web/drizzle/meta/0016_snapshot.json Normal file
View File

@ -0,0 +1,675 @@
{
"id": "89f22940-a898-4086-8a65-5f9c40fc7f0d",
"prevId": "92bef822-0089-48aa-8f43-5bba40cdce2e",
"version": "5",
"dialect": "pg",
"tables": {
"api_keys": {
"name": "api_keys",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"key": {
"name": "key",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"org_id": {
"name": "org_id",
"type": "text",
"primaryKey": false,
"notNull": false
},
"revoked": {
"name": "revoked",
"type": "boolean",
"primaryKey": false,
"notNull": true,
"default": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"api_keys_user_id_users_id_fk": {
"name": "api_keys_user_id_users_id_fk",
"tableFrom": "api_keys",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {
"api_keys_key_unique": {
"name": "api_keys_key_unique",
"nullsNotDistinct": false,
"columns": [
"key"
]
}
}
},
"deployments": {
"name": "deployments",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"workflow_version_id": {
"name": "workflow_version_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"workflow_id": {
"name": "workflow_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"machine_id": {
"name": "machine_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"environment": {
"name": "environment",
"type": "deployment_environment",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"deployments_user_id_users_id_fk": {
"name": "deployments_user_id_users_id_fk",
"tableFrom": "deployments",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
},
"deployments_workflow_version_id_workflow_versions_id_fk": {
"name": "deployments_workflow_version_id_workflow_versions_id_fk",
"tableFrom": "deployments",
"tableTo": "workflow_versions",
"columnsFrom": [
"workflow_version_id"
],
"columnsTo": [
"id"
],
"onDelete": "no action",
"onUpdate": "no action"
},
"deployments_workflow_id_workflows_id_fk": {
"name": "deployments_workflow_id_workflows_id_fk",
"tableFrom": "deployments",
"tableTo": "workflows",
"columnsFrom": [
"workflow_id"
],
"columnsTo": [
"id"
],
"onDelete": "no action",
"onUpdate": "no action"
},
"deployments_machine_id_machines_id_fk": {
"name": "deployments_machine_id_machines_id_fk",
"tableFrom": "deployments",
"tableTo": "machines",
"columnsFrom": [
"machine_id"
],
"columnsTo": [
"id"
],
"onDelete": "no action",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"machines": {
"name": "machines",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"org_id": {
"name": "org_id",
"type": "text",
"primaryKey": false,
"notNull": false
},
"endpoint": {
"name": "endpoint",
"type": "text",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"disabled": {
"name": "disabled",
"type": "boolean",
"primaryKey": false,
"notNull": true,
"default": false
},
"auth_token": {
"name": "auth_token",
"type": "text",
"primaryKey": false,
"notNull": false
},
"type": {
"name": "type",
"type": "machine_type",
"primaryKey": false,
"notNull": true,
"default": "'classic'"
}
},
"indexes": {},
"foreignKeys": {
"machines_user_id_users_id_fk": {
"name": "machines_user_id_users_id_fk",
"tableFrom": "machines",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"users": {
"name": "users",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "text",
"primaryKey": true,
"notNull": true
},
"username": {
"name": "username",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": false,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": false,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflow_run_outputs": {
"name": "workflow_run_outputs",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"run_id": {
"name": "run_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"data": {
"name": "data",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflow_run_outputs_run_id_workflow_runs_id_fk": {
"name": "workflow_run_outputs_run_id_workflow_runs_id_fk",
"tableFrom": "workflow_run_outputs",
"tableTo": "workflow_runs",
"columnsFrom": [
"run_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflow_runs": {
"name": "workflow_runs",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"workflow_version_id": {
"name": "workflow_version_id",
"type": "uuid",
"primaryKey": false,
"notNull": false
},
"workflow_inputs": {
"name": "workflow_inputs",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"workflow_id": {
"name": "workflow_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"machine_id": {
"name": "machine_id",
"type": "uuid",
"primaryKey": false,
"notNull": false
},
"origin": {
"name": "origin",
"type": "workflow_run_origin",
"primaryKey": false,
"notNull": true,
"default": "'api'"
},
"status": {
"name": "status",
"type": "workflow_run_status",
"primaryKey": false,
"notNull": true,
"default": "'not-started'"
},
"ended_at": {
"name": "ended_at",
"type": "timestamp",
"primaryKey": false,
"notNull": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflow_runs_workflow_version_id_workflow_versions_id_fk": {
"name": "workflow_runs_workflow_version_id_workflow_versions_id_fk",
"tableFrom": "workflow_runs",
"tableTo": "workflow_versions",
"columnsFrom": [
"workflow_version_id"
],
"columnsTo": [
"id"
],
"onDelete": "set null",
"onUpdate": "no action"
},
"workflow_runs_workflow_id_workflows_id_fk": {
"name": "workflow_runs_workflow_id_workflows_id_fk",
"tableFrom": "workflow_runs",
"tableTo": "workflows",
"columnsFrom": [
"workflow_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
},
"workflow_runs_machine_id_machines_id_fk": {
"name": "workflow_runs_machine_id_machines_id_fk",
"tableFrom": "workflow_runs",
"tableTo": "machines",
"columnsFrom": [
"machine_id"
],
"columnsTo": [
"id"
],
"onDelete": "set null",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflows": {
"name": "workflows",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"org_id": {
"name": "org_id",
"type": "text",
"primaryKey": false,
"notNull": false
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflows_user_id_users_id_fk": {
"name": "workflows_user_id_users_id_fk",
"tableFrom": "workflows",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflow_versions": {
"name": "workflow_versions",
"schema": "comfyui_deploy",
"columns": {
"workflow_id": {
"name": "workflow_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"workflow": {
"name": "workflow",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"workflow_api": {
"name": "workflow_api",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"version": {
"name": "version",
"type": "integer",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflow_versions_workflow_id_workflows_id_fk": {
"name": "workflow_versions_workflow_id_workflows_id_fk",
"tableFrom": "workflow_versions",
"tableTo": "workflows",
"columnsFrom": [
"workflow_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
}
},
"enums": {
"deployment_environment": {
"name": "deployment_environment",
"values": {
"staging": "staging",
"production": "production"
}
},
"machine_type": {
"name": "machine_type",
"values": {
"classic": "classic",
"runpod-serverless": "runpod-serverless"
}
},
"workflow_run_origin": {
"name": "workflow_run_origin",
"values": {
"manual": "manual",
"api": "api"
}
},
"workflow_run_status": {
"name": "workflow_run_status",
"values": {
"not-started": "not-started",
"running": "running",
"uploading": "uploading",
"success": "success",
"failed": "failed"
}
}
},
"schemas": {
"comfyui_deploy": "comfyui_deploy"
},
"_meta": {
"schemas": {},
"tables": {},
"columns": {}
}
}

7
web/drizzle/meta/_journal.json
View File

@ -113,6 +113,13 @@
"when": 1703409502387,
"tag": "0015_simple_killmonger",
"breakpoints": true
},
{
"idx": 16,
"version": "5",
"when": 1704092053001,
"tag": "0016_overrated_cable",
"breakpoints": true
}
]
}

1
web/package.json
View File

@ -61,6 +61,7 @@
"lucide-react": "^0.294.0",
"mdast-util-to-string": "^4.0.0",
"mdx-annotations": "^0.1.4",
"mitata": "^0.1.6",
"nanoid": "^5.0.4",
"next": "14.0.3",
"next-plausible": "^3.12.0",

78
web/src/app/(app)/api/run/route.ts
View File

@ -25,26 +25,57 @@ async function checkToken(request: Request) {
const token = request.headers.get("Authorization")?.split(" ")?.[1]; // Assuming token is sent as "Bearer your_token"
const userData = token ? parseJWT(token) : undefined;
if (!userData || token === undefined) {
return new NextResponse("Invalid or expired token", {
status: 401,
});
return {
error: new NextResponse("Invalid or expired token", {
status: 401,
}),
};
} else {
const revokedKey = await isKeyRevoked(token);
if (revokedKey)
return new NextResponse("Revoked token", {
status: 401,
});
return {
error: new NextResponse("Revoked token", {
status: 401,
}),
};
}
return {
data: userData,
};
}
export async function GET(request: Request) {
const invalidRequest = await checkToken(request);
if (invalidRequest) return invalidRequest;
const apiKeyTokenData = await checkToken(request);
if (apiKeyTokenData.error) return apiKeyTokenData.error;
const [data, error] = await parseDataSafe(Request2, request);
if (!data || error) return error;
const run = await getRunsData(data.run_id);
// return NextResponse.json(
// await db
// .select()
// .from(workflowTable)
// .innerJoin(
// workflowRunsTable,
// eq(workflowTable.id, workflowRunsTable.workflow_id)
// )
// .where(
// and(
// eq(workflowTable.id, workflowRunsTable.workflow_id),
// apiKeyTokenData.data.org_id
// ? eq(workflowTable.org_id, apiKeyTokenData.data.org_id)
// : eq(workflowTable.user_id, apiKeyTokenData.data.user_id!)
// )
// ),
// {
// status: 200,
// }
// );
const run = await getRunsData(apiKeyTokenData.data, data.run_id);
if (!run) return new NextResponse("Run not found", { status: 404 });
if (run?.status === "success" && run?.outputs?.length > 0) {
for (let i = 0; i < run.outputs.length; i++) {
@ -74,8 +105,8 @@ export async function GET(request: Request) {
}
export async function POST(request: Request) {
const invalidRequest = await checkToken(request);
if (invalidRequest) return invalidRequest;
const apiKeyTokenData = await checkToken(request);
if (apiKeyTokenData.error) return apiKeyTokenData.error;
const [data, error] = await parseDataSafe(Request, request);
if (!data || error) return error;
@ -87,16 +118,31 @@ export async function POST(request: Request) {
try {
const deploymentData = await db.query.deploymentsTable.findFirst({
where: eq(deploymentsTable.id, deployment_id),
with: {
machine: true,
version: {
with: {
workflow: {
columns: {
org_id: true,
user_id: true,
},
},
},
},
},
});
if (!deploymentData) throw new Error("Deployment not found");
const run_id = await createRun(
const run_id = await createRun({
origin,
deploymentData.workflow_version_id,
deploymentData.machine_id,
inputs
);
workflow_version_id: deploymentData.version,
machine_id: deploymentData.machine,
inputs,
isManualRun: false,
apiUser: apiKeyTokenData.data,
});
if ("error" in run_id) throw new Error(run_id.error);

1
web/src/app/(app)/api/upload/route.ts
View File

@ -79,6 +79,7 @@ export async function POST(request: Request) {
.values({
user_id,
name: workflow_name,
org_id: org_id,
})
.returning();

13
web/src/app/(app)/layout.tsx
View File

@ -1,8 +1,8 @@
import "./globals.css";
import { NavbarRight } from "@/components/NavbarRight";
import { NavbarMenu } from "@/components/NavbarMenu";
import { Button } from "@/components/ui/button";
import { TooltipProvider } from "@/components/ui/tooltip";
import { ClerkProvider, UserButton } from "@clerk/nextjs";
import { ClerkProvider, OrganizationSwitcher, UserButton } from "@clerk/nextjs";
import { Github } from "lucide-react";
import type { Metadata } from "next";
import meta from "next-gen/config";
@ -54,9 +54,16 @@ export default function RootLayout({
>
{meta.name}
</a>
<NavbarRight />
<OrganizationSwitcher
appearance={{
elements: {
rootBox: "flex items-center justify-center",
},
}}
/>
</div>
<div className="flex flex-row items-center gap-2">
<NavbarMenu />
<Button
asChild
variant="link"

9
web/src/app/(app)/machines/page.tsx
View File

@ -2,14 +2,14 @@ import { MachineList } from "@/components/MachineList";
import { db } from "@/db/db";
import { machinesTable } from "@/db/schema";
import { auth } from "@clerk/nextjs";
import { desc, eq } from "drizzle-orm";
import { desc, eq, isNull, and } from "drizzle-orm";
export default function Page() {
return <MachineListServer />;
}
async function MachineListServer() {
const { userId } = await auth();
const { userId, orgId } = await auth();
if (!userId) {
return <div>No auth</div>;
@ -17,7 +17,10 @@ async function MachineListServer() {
const machines = await db.query.machinesTable.findMany({
orderBy: desc(machinesTable.updated_at),
where: eq(machinesTable.user_id, userId),
where:
orgId != undefined
? eq(machinesTable.org_id, orgId)
: and(eq(machinesTable.user_id, userId), isNull(machinesTable.org_id)),
});
return (

9
web/src/app/(app)/workflows/page.tsx
View File

@ -2,14 +2,14 @@ import { WorkflowList } from "@/components/WorkflowList";
import { db } from "@/db/db";
import { usersTable, workflowTable, workflowVersionTable } from "@/db/schema";
import { auth, clerkClient } from "@clerk/nextjs";
import { desc, eq } from "drizzle-orm";
import { and, desc, eq, isNull } from "drizzle-orm";
export default function Home() {
return <WorkflowServer />;
}
async function WorkflowServer() {
const { userId } = await auth();
const { userId, orgId } = await auth();
if (!userId) {
return <div>No auth</div>;
@ -36,7 +36,10 @@ async function WorkflowServer() {
},
},
orderBy: desc(workflowTable.updated_at),
where: eq(workflowTable.user_id, userId),
where:
orgId != undefined
? eq(workflowTable.org_id, orgId)
: and(eq(workflowTable.user_id, userId), isNull(workflowTable.org_id)),
});
return (

8
web/src/components/NavbarRight.tsx → web/src/components/NavbarMenu.tsx
View File

@ -12,7 +12,7 @@ import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { usePathname } from "next/navigation";
import { useRouter } from "next/navigation";
export function NavbarRight() {
export function NavbarMenu() {
const pathnames = usePathname();
const pathname = `/${pathnames.split("/")[1]}`;
@ -34,10 +34,11 @@ export function NavbarRight() {
];
return (
<div>
<div className="mr-2">
{/* <div className="w-full h-full absolute inset-x-0 top-0 flex items-center justify-center pointer-events-none"> */}
<Tabs
defaultValue={pathname}
className="w-[300px] hidden lg:flex"
className="w-[300px] hidden lg:flex pointer-events-auto"
onValueChange={(value) => {
router.push(value);
}}
@ -50,6 +51,7 @@ export function NavbarRight() {
))}
</TabsList>
</Tabs>
{/* </div> */}
<div className="w-[100px] flex lg:hidden">
<Select

2
web/src/components/VersionDetails.tsx
View File

@ -29,7 +29,7 @@ export function VersionDetails({
return (
<div className="mt-4">
Workflow Details
Workflow Inputs
<div className="border rounded-lg p-2">
{inputs && inputs.length > 0 ? (
<div className="flex flex-col gap-2">

8
web/src/components/VersionSelect.tsx
View File

@ -165,7 +165,13 @@ export function RunWorkflowButton({
try {
const origin = window.location.origin;
await callServerPromise(
createRun(origin, workflow_version_id, machine, val, true)
createRun({
origin,
workflow_version_id,
machine_id: machine,
inputs: val,
isManualRun: true,
})
);
// console.log(res.json());
setIsLoading(false);

7
web/src/db/schema.ts
View File

@ -29,6 +29,7 @@ export const workflowTable = dbSchema.table("workflows", {
onDelete: "cascade",
})
.notNull(),
org_id: text("org_id"),
name: text("name").notNull(),
created_at: timestamp("created_at").defaultNow().notNull(),
updated_at: timestamp("updated_at").defaultNow().notNull(),
@ -144,6 +145,10 @@ export const workflowRunRelations = relations(
references: [workflowVersionTable.id],
}),
outputs: many(workflowRunOutputs),
workflow: one(workflowTable, {
fields: [workflowRunsTable.workflow_id],
references: [workflowTable.id],
}),
})
);
@ -180,6 +185,7 @@ export const machinesTable = dbSchema.table("machines", {
})
.notNull(),
name: text("name").notNull(),
org_id: text("org_id"),
endpoint: text("endpoint").notNull(),
created_at: timestamp("created_at").defaultNow().notNull(),
updated_at: timestamp("updated_at").defaultNow().notNull(),
@ -244,3 +250,4 @@ export const apiKeyTable = dbSchema.table("api_keys", {
export type UserType = InferSelectModel<typeof usersTable>;
export type WorkflowType = InferSelectModel<typeof workflowTable>;
export type MachineType = InferSelectModel<typeof machinesTable>;
export type WorkflowVersionType = InferSelectModel<typeof workflowVersionTable>;

2
web/src/server/APIKeyBodyRequest.ts
View File

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

73
web/src/server/createRun.ts
View File

@ -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();

10
web/src/server/curdApiKeys.ts
View File

@ -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),
});
}

3
web/src/server/curdMachine.ts
View File

@ -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");

56
web/src/server/getRunsOutput.tsx
View File

@ -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;
}