feat(server): copilot job models (#10732)

packages/backend/server/src/models/common/copilot.ts

@@ -0,0 +1,12 @@
+import { AiJobStatus, AiJobType } from '@prisma/client';
+import type { JsonValue } from '@prisma/client/runtime/library';
+
+export interface CopilotJob {
+  id?: string;
+  workspaceId: string;
+  blobId: string;
+  createdBy?: string;
+  type: AiJobType;
+  status?: AiJobStatus;
+  payload?: JsonValue;
+}

packages/backend/server/src/models/copilot-job.ts

@@ -0,0 +1,151 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+import { AiJobStatus, AiJobType } from '@prisma/client';
+import type { ZodType } from 'zod';
+
+import { BaseModel } from './base';
+import { CopilotJob } from './common/copilot';
+
+type CreateCopilotJobInput = Omit<CopilotJob, 'id' | 'status' | 'payload'>;
+type UpdateCopilotJobInput = Pick<CopilotJob, 'status' | 'payload'>;
+
+/**
+ * Copilot Job Model
+ */
+@Injectable()
+export class CopilotJobModel extends BaseModel {
+  async create(job: CreateCopilotJobInput) {
+    const row = await this.db.aiJobs.create({
+      data: {
+        workspaceId: job.workspaceId,
+        blobId: job.blobId,
+        createdBy: job.createdBy,
+        type: job.type,
+        status: AiJobStatus.pending,
+        payload: {},
+      },
+      select: {
+        id: true,
+        status: true,
+      },
+    });
+    return row;
+  }
+
+  async has(workspaceId: string, blobId: string) {
+    const row = await this.db.aiJobs.findFirst({
+      where: {
+        workspaceId,
+        blobId,
+      },
+    });
+    return !!row;
+  }
+
+  async update(jobId: string, data: UpdateCopilotJobInput) {
+    const ret = await this.db.aiJobs.updateMany({
+      where: {
+        id: jobId,
+      },
+      data: {
+        status: data.status || undefined,
+        payload: data.payload || undefined,
+      },
+    });
+    return ret.count > 0;
+  }
+
+  @Transactional()
+  async claim(jobId: string, userId: string) {
+    const job = await this.get(jobId);
+
+    if (
+      job &&
+      job.createdBy === userId &&
+      job.status === AiJobStatus.finished
+    ) {
+      await this.update(jobId, { status: AiJobStatus.claimed });
+    }
+
+    const ret = await this.db.aiJobs.findFirst({
+      where: { id: jobId },
+      select: { status: true },
+    });
+    return ret?.status;
+  }
+
+  async getWithUser(
+    userId: string,
+    workspaceId: string,
+    jobId?: string,
+    type?: AiJobType
+  ) {
+    const row = await this.db.aiJobs.findFirst({
+      where: {
+        id: jobId,
+        workspaceId,
+        type,
+        OR: [
+          {
+            createdBy: userId,
+            status: { in: [AiJobStatus.finished, AiJobStatus.claimed] },
+          },
+          { createdBy: { not: userId }, status: AiJobStatus.claimed },
+        ],
+      },
+    });
+
+    if (!row) {
+      return null;
+    }
+
+    return {
+      id: row.id,
+      workspaceId: row.workspaceId,
+      blobId: row.blobId,
+      createdBy: row.createdBy || undefined,
+      type: row.type,
+      status: row.status,
+      payload: row.payload,
+    };
+  }
+
+  async get(jobId: string): Promise<CopilotJob | null> {
+    const row = await this.db.aiJobs.findFirst({
+      where: {
+        id: jobId,
+      },
+    });
+
+    if (!row) {
+      return null;
+    }
+
+    return {
+      id: row.id,
+      workspaceId: row.workspaceId,
+      blobId: row.blobId,
+      createdBy: row.createdBy || undefined,
+      type: row.type,
+      status: row.status,
+      payload: row.payload,
+    };
+  }
+
+  async getPayload<
+    C extends ZodType<any>,
+    O = C extends ZodType<infer T> ? T : never,
+  >(jobId: string, schema: C): Promise<O> {
+    const row = await this.db.aiJobs.findUnique({
+      where: {
+        id: jobId,
+      },
+      select: {
+        payload: true,
+      },
+    });
+
+    const ret = schema.safeParse(row?.payload);
+    return ret.success ? ret.data : ({} as O);
+  }
+}

packages/backend/server/src/models/index.ts

@@ -7,6 +7,7 @@ import {
 import { ModuleRef } from '@nestjs/core';
 
 import { ApplyType } from '../base';
+import { CopilotJobModel } from './copilot-job';
 import { DocModel } from './doc';
 import { DocUserModel } from './doc-user';
 import { FeatureModel } from './feature';
@@ -38,6 +39,7 @@ const MODELS = {
   history: HistoryModel,
   notification: NotificationModel,
   settings: SettingsModel,
+  copilotJob: CopilotJobModel,
 };
 
 type ModelsType = {
@@ -90,6 +92,7 @@ const ModelsSymbolProvider: ExistingProvider = {
 export class ModelsModule {}
 
 export * from './common';
+export * from './copilot-job';
 export * from './doc';
 export * from './doc-user';
 export * from './feature';