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
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

Binary file not shown.

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;

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": {}
}
}

View File

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

View File

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

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 token = request.headers.get("Authorization")?.split(" ")?.[1]; // Assuming token is sent as "Bearer your_token"
const userData = token ? parseJWT(token) : undefined; const userData = token ? parseJWT(token) : undefined;
if (!userData || token === undefined) { if (!userData || token === undefined) {
return new NextResponse("Invalid or expired token", { return {
status: 401, error: new NextResponse("Invalid or expired token", {
}); status: 401,
}),
};
} else { } else {
const revokedKey = await isKeyRevoked(token); const revokedKey = await isKeyRevoked(token);
if (revokedKey) if (revokedKey)
return new NextResponse("Revoked token", { return {
status: 401, error: new NextResponse("Revoked token", {
}); status: 401,
}),
};
} }
return {
data: userData,
};
} }
export async function GET(request: Request) { export async function GET(request: Request) {
const invalidRequest = await checkToken(request); const apiKeyTokenData = await checkToken(request);
if (invalidRequest) return invalidRequest; if (apiKeyTokenData.error) return apiKeyTokenData.error;
const [data, error] = await parseDataSafe(Request2, request); const [data, error] = await parseDataSafe(Request2, request);
if (!data || error) return error; 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) { if (run?.status === "success" && run?.outputs?.length > 0) {
for (let i = 0; i < run.outputs.length; i++) { 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) { export async function POST(request: Request) {
const invalidRequest = await checkToken(request); const apiKeyTokenData = await checkToken(request);
if (invalidRequest) return invalidRequest; if (apiKeyTokenData.error) return apiKeyTokenData.error;
const [data, error] = await parseDataSafe(Request, request); const [data, error] = await parseDataSafe(Request, request);
if (!data || error) return error; if (!data || error) return error;
@ -87,16 +118,31 @@ export async function POST(request: Request) {
try { try {
const deploymentData = await db.query.deploymentsTable.findFirst({ const deploymentData = await db.query.deploymentsTable.findFirst({
where: eq(deploymentsTable.id, deployment_id), 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"); if (!deploymentData) throw new Error("Deployment not found");
const run_id = await createRun( const run_id = await createRun({
origin, origin,
deploymentData.workflow_version_id, workflow_version_id: deploymentData.version,
deploymentData.machine_id, machine_id: deploymentData.machine,
inputs inputs,
); isManualRun: false,
apiUser: apiKeyTokenData.data,
});
if ("error" in run_id) throw new Error(run_id.error); if ("error" in run_id) throw new Error(run_id.error);

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -165,7 +165,13 @@ export function RunWorkflowButton({
try { try {
const origin = window.location.origin; const origin = window.location.origin;
await callServerPromise( 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()); // console.log(res.json());
setIsLoading(false); setIsLoading(false);

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -3,7 +3,8 @@
import { RunOutputs } from "@/components/RunOutputs"; import { RunOutputs } from "@/components/RunOutputs";
import { db } from "@/db/db"; import { db } from "@/db/db";
import { workflowRunOutputs, workflowRunsTable } from "@/db/schema"; 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) { export async function getRunsOutputDisplay(run_id: string) {
return <RunOutputs run_id={run_id} />; return <RunOutputs run_id={run_id} />;
@ -17,11 +18,38 @@ export async function getRunsOutput(run_id: string) {
.where(eq(workflowRunOutputs.run_id, run_id)); .where(eq(workflowRunOutputs.run_id, run_id));
} }
export async function getRunsData(run_id: string) { export async function getRunsData(user: APIKeyUserType, run_id: string) {
// throw new Error("Not implemented"); const data = await db.query.workflowRunsTable.findFirst({
return await db.query.workflowRunsTable.findFirst({ where: and(
where: eq(workflowRunsTable.id, run_id), 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: { with: {
workflow: {
columns: {
org_id: true,
user_id: true,
},
},
outputs: { outputs: {
columns: { columns: {
data: true, 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;
} }