feat: support run document index tassk in background (#130)

* feat: support run document index tassk in background * fix
pingcap · May 11, 2024 · 6ac0c8d · 6ac0c8d
1 parent 477b63e
commit 6ac0c8d
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 5 deletions.
diff --git a/ddl/0-initial-schema.sql b/ddl/0-initial-schema.sql
@@ -75,7 +75,7 @@ CREATE TABLE `document_index_task` (
     index_id         INT            NOT NULL,
     document_id      INT            NOT NULL,
     type             ENUM('CREATE_INDEX', 'REINDEX') NOT NULL ,
-    status           ENUM('CREATED', 'PENDING', 'INDEXING', 'SUCCEED', 'FAILED') NOT NULL ,
+    status           ENUM('CREATED', 'PENDING', 'INDEXING', 'SUCCEED', 'FAILED', 'CANCELED') NOT NULL ,
     info             JSON           NULL NOT NULL ,
     message          TEXT           NULL,
     created_at       DATETIME       NOT NULL DEFAULT CURRENT_TIMESTAMP,

diff --git a/src/app/api/v1/documents/index/route.ts b/src/app/api/v1/documents/index/route.ts
@@ -8,14 +8,15 @@ const IndexDocumentsOptionsSchema = z.object({
     .int()
     .array()
     .min(1, 'Must provide at least one document'),
-  indexName: z.string()
+  indexName: z.string(),
+  runInBackground: z.boolean().default(false)
 });
 
 export const POST = defineHandler({
   auth: 'admin',
   body: IndexDocumentsOptionsSchema
 },  async ({ body}) => {
-  const { documentIds, indexName } = body;
+  const { documentIds, indexName, runInBackground } = body;
 
   const index = await getIndexByNameOrThrow(indexName);
   const documentIdStr = documentIds.map((id) => `#${id}`).join(', ')
@@ -29,6 +30,12 @@ export const POST = defineHandler({
   const taskIdStr = taskIds.map((id) => `#${id}`).join(', ')
   console.log(`Create document index tasks ${taskIdStr}.`);
 
+  if (runInBackground) {
+    return {
+      taskIds
+    }
+  }
+
   // Execute document index tasks.
   const results = await Promise.allSettled(
     taskIds.map(taskId => service.runDocumentIndexTask(taskId))
@@ -48,6 +55,7 @@ export const POST = defineHandler({
   });
 
   return {
+    taskIds,
     succeed,
     failed
   }

diff --git a/src/core/db/schema.d.ts b/src/core/db/schema.d.ts
@@ -97,7 +97,7 @@ export interface DocumentIndexTask {
   info: Json;
   message: string | null;
   started_at: Date | null;
-  status: "CREATED" | "FAILED" | "INDEXING" | "PENDING" | "SUCCEED";
+  status: "CANCELED" | "CREATED" | "FAILED" | "INDEXING" | "PENDING" | "SUCCEED";
   type: "CREATE_INDEX" | "REINDEX";
 }
 

diff --git a/src/lib/knowledge-graph/client.ts b/src/lib/knowledge-graph/client.ts
@@ -55,7 +55,7 @@ export class KnowledgeGraphClient {
         },
         body: JSON.stringify(doc)
       });
-      if (res.ok) {
+      if (!res.ok) {
         throw new Error(`Failed to call build knowledge graph index API for doc (url: ${doc.uri}): ${res.statusText}`);
       }
       const data = await res.json();