[sync] feat(v2): add computed task run-by-id command (#1465) (#2782)

Synced from teableio/teable-ee@48768e3

Co-authored-by: nichenqin <nichenqin@hotmail.com>

Author: tea-artist

packages/v2/adapter-table-repository-postgres/src/record/computed/RunComputedTaskByIdCommand.ts

@@ -0,0 +1,47 @@
+import { domainError, type DomainError, InternalCommand } from '@teable/v2-core';
+import { err, ok, type Result } from 'neverthrow';
+import { z } from 'zod';
+
+const runComputedTaskByIdInputSchema = z.object({
+  taskId: z.string().trim().min(1, 'Task ID is required'),
+  workerId: z.string().trim().min(1, 'Worker ID is required'),
+  allowProcessingTakeover: z.boolean().optional().default(true),
+});
+
+export type IRunComputedTaskByIdCommandInput = z.input<typeof runComputedTaskByIdInputSchema>;
+
+export type RunComputedTaskByIdResult = {
+  taskId: string;
+  workerId: string;
+  processed: true;
+};
+
+export class RunComputedTaskByIdCommand extends InternalCommand {
+  private constructor(
+    readonly taskId: string,
+    readonly workerId: string,
+    readonly allowProcessingTakeover: boolean
+  ) {
+    super();
+  }
+
+  static create(raw: unknown): Result<RunComputedTaskByIdCommand, DomainError> {
+    const parsed = runComputedTaskByIdInputSchema.safeParse(raw);
+    if (!parsed.success) {
+      return err(
+        domainError.validation({
+          message: 'Invalid RunComputedTaskByIdCommand input',
+          details: parsed.error.format(),
+        })
+      );
+    }
+
+    return ok(
+      new RunComputedTaskByIdCommand(
+        parsed.data.taskId,
+        parsed.data.workerId,
+        parsed.data.allowProcessingTakeover
+      )
+    );
+  }
+}

packages/v2/adapter-table-repository-postgres/src/record/computed/RunComputedTaskByIdHandler.ts

@@ -0,0 +1,52 @@
+import { domainError, type DomainError, type IExecutionContext } from '@teable/v2-core';
+import { inject, injectable } from '@teable/v2-di';
+import { err, ok, type Result } from 'neverthrow';
+
+import { v2RecordRepositoryPostgresTokens } from '../di/tokens';
+import type { ComputedUpdateWorker } from './worker/ComputedUpdateWorker';
+import { CommandHandler, type ICommandHandler } from '@teable/v2-core';
+import {
+  RunComputedTaskByIdCommand,
+  type RunComputedTaskByIdResult,
+} from './RunComputedTaskByIdCommand';
+
+@CommandHandler(RunComputedTaskByIdCommand)
+@injectable()
+export class RunComputedTaskByIdHandler
+  implements ICommandHandler<RunComputedTaskByIdCommand, RunComputedTaskByIdResult>
+{
+  constructor(
+    @inject(v2RecordRepositoryPostgresTokens.computedUpdateWorker)
+    private readonly worker: ComputedUpdateWorker
+  ) {}
+
+  async handle(
+    context: IExecutionContext,
+    command: RunComputedTaskByIdCommand
+  ): Promise<Result<RunComputedTaskByIdResult, DomainError>> {
+    const result = await this.worker.runTaskById({
+      taskId: command.taskId,
+      workerId: command.workerId,
+      actorId: context.actorId,
+      tracer: context.tracer,
+      requestId: context.requestId,
+      allowProcessingTakeover: command.allowProcessingTakeover,
+    });
+    if (result.isErr()) return err(result.error);
+
+    if (!result.value) {
+      return err(
+        domainError.notFound({
+          code: 'computed_task.not_retryable',
+          message: `Computed task not found or not retryable: ${command.taskId}`,
+        })
+      );
+    }
+
+    return ok({
+      taskId: command.taskId,
+      workerId: command.workerId,
+      processed: true,
+    });
+  }
+}

packages/v2/adapter-table-repository-postgres/src/record/computed/index.ts

@@ -6,6 +6,8 @@ export * from './ComputedUpdateLock';
 export * from './ComputedUpdatePlanner';
 export * from './ComputedUpdateRun';
 export * from './ExternalComputedRefreshService';
+export * from './RunComputedTaskByIdCommand';
+export * from './RunComputedTaskByIdHandler';
 export * from './UserRenamePropagationService';
 export * from './FieldDependencyGraph';
 export * from './UpdateFromSelectBuilder';

packages/v2/adapter-table-repository-postgres/src/record/computed/outbox/ComputedUpdateOutbox.ts

@@ -30,6 +30,7 @@ import { defaultComputedUpdateOutboxConfig } from './IComputedUpdateOutbox';
 import type {
   IComputedUpdateOutbox,
   ClaimBatchParams,
+  ClaimByIdParams,
   ComputedUpdateOutboxConfig,
   AnyOutboxItem,
   FieldBackfillOutboxItem,
@@ -438,6 +439,67 @@ export class ComputedUpdateOutbox implements IComputedUpdateOutbox {
     }
   }
 
+  async claimById(
+    params: ClaimByIdParams,
+    context?: IExecutionContext
+  ): Promise<Result<AnyOutboxItem | null, DomainError>> {
+    const span = context?.tracer?.startSpan('teable.outbox.claimById', {
+      'outbox.taskId': params.taskId,
+      'outbox.workerId': params.workerId,
+    });
+
+    const executeClaim = async (): Promise<Result<AnyOutboxItem | null, DomainError>> => {
+      const now = params.now ?? new Date();
+      const db = resolvePostgresDbOrTx(this.db, context) as unknown as Kysely<DynamicDB>;
+      const retryableStatuses = params.allowProcessingTakeover
+        ? [DEFAULT_STATUS, 'processing']
+        : [DEFAULT_STATUS];
+
+      return runInTransaction(db, context, async (trx) => {
+        const row = await trx
+          .selectFrom(OUTBOX_TABLE)
+          .selectAll()
+          .where('id', '=', params.taskId)
+          .where('status', 'in', retryableStatuses)
+          .forUpdate()
+          .executeTakeFirst();
+
+        if (!row) return ok(null);
+
+        await trx
+          .updateTable(OUTBOX_TABLE)
+          .set({
+            status: 'processing',
+            locked_at: now,
+            locked_by: params.workerId,
+            updated_at: now,
+          })
+          .where('id', '=', params.taskId)
+          .execute();
+
+        const seedMap = await this.loadSeedRecords(trx, [row]);
+        const claimedRow = {
+          ...row,
+          status: 'processing',
+          locked_at: now,
+          locked_by: params.workerId,
+          updated_at: now,
+        };
+
+        return ok(toAnyOutboxItem(claimedRow, seedMap.get(String(row.id)) ?? []));
+      });
+    };
+
+    try {
+      if (span && context?.tracer) {
+        return await context.tracer.withSpan(span, executeClaim);
+      }
+      return await executeClaim();
+    } finally {
+      span?.end();
+    }
+  }
+
   async markDone(taskId: string, context?: IExecutionContext): Promise<Result<void, DomainError>> {
     const span = context?.tracer?.startSpan('teable.outbox.markDone', {
       'outbox.taskId': taskId,

packages/v2/adapter-table-repository-postgres/src/record/computed/outbox/IComputedUpdateOutbox.ts

@@ -32,6 +32,13 @@ export type ClaimBatchParams = {
   now?: Date;
 };
 
+export type ClaimByIdParams = {
+  taskId: string;
+  workerId: string;
+  now?: Date;
+  allowProcessingTakeover?: boolean;
+};
+
 /**
  * Outbox item for field backfill tasks.
  */
@@ -113,6 +120,11 @@ export interface IComputedUpdateOutbox {
     context?: IExecutionContext
   ): Promise<Result<ReadonlyArray<AnyOutboxItem>, DomainError>>;
 
+  claimById(
+    params: ClaimByIdParams,
+    context?: IExecutionContext
+  ): Promise<Result<AnyOutboxItem | null, DomainError>>;
+
   markDone(taskId: string, context?: IExecutionContext): Promise<Result<void, DomainError>>;
 
   markFailed(

packages/v2/adapter-table-repository-postgres/src/record/computed/strategies/HybridWithOutboxStrategy.spec.ts

@@ -128,6 +128,7 @@ const createOutboxStub = () => {
     enqueueSeedTask,
     enqueueFieldBackfill,
     claimBatch: async () => ok([]),
+    claimById: async () => ok(null),
     markDone: async () => ok(undefined),
     markFailed: async () => ok(undefined),
   };

packages/v2/adapter-table-repository-postgres/src/record/computed/worker/ComputedUpdateWorker.spec.ts

@@ -105,6 +105,7 @@ describe('ComputedUpdateWorker', () => {
         enqueueSeedTask: vi.fn(),
         enqueueFieldBackfill: vi.fn(),
         claimBatch: vi.fn().mockResolvedValue(ok([])),
+        claimById: vi.fn(),
         markDone: vi.fn(),
         markFailed: vi.fn(),
       };
@@ -142,6 +143,7 @@ describe('ComputedUpdateWorker', () => {
         enqueueSeedTask: vi.fn(),
         enqueueFieldBackfill: vi.fn(),
         claimBatch: vi.fn().mockResolvedValue(ok([task])),
+        claimById: vi.fn(),
         markDone: vi.fn(),
         markFailed,
       };
@@ -191,6 +193,7 @@ describe('ComputedUpdateWorker', () => {
         enqueueSeedTask: vi.fn(),
         enqueueFieldBackfill: vi.fn(),
         claimBatch: vi.fn().mockResolvedValue(ok([task])),
+        claimById: vi.fn(),
         markDone,
         markFailed: vi.fn(),
       };
@@ -242,6 +245,7 @@ describe('ComputedUpdateWorker', () => {
         enqueueSeedTask: vi.fn(),
         enqueueFieldBackfill: vi.fn(),
         claimBatch: vi.fn().mockResolvedValue(ok([task1, task2, task3])),
+        claimById: vi.fn(),
         markDone,
         markFailed: vi.fn(),
       };
@@ -305,6 +309,7 @@ describe('ComputedUpdateWorker', () => {
         enqueueSeedTask: vi.fn(),
         enqueueFieldBackfill: vi.fn(),
         claimBatch: vi.fn().mockResolvedValue(ok([task])),
+        claimById: vi.fn(),
         markDone,
         markFailed: vi.fn(),
       };
@@ -369,6 +374,7 @@ describe('ComputedUpdateWorker', () => {
         enqueueSeedTask: vi.fn(),
         enqueueFieldBackfill: vi.fn(),
         claimBatch: vi.fn().mockResolvedValue(ok([task])),
+        claimById: vi.fn(),
         markDone: vi.fn(),
         markFailed: vi.fn().mockResolvedValue(ok(undefined)),
       };
@@ -425,6 +431,7 @@ describe('ComputedUpdateWorker', () => {
         enqueueSeedTask: vi.fn(),
         enqueueFieldBackfill: vi.fn(),
         claimBatch: vi.fn().mockResolvedValue(ok([task])),
+        claimById: vi.fn(),
         markDone: vi.fn(),
         markFailed,
       };
@@ -470,4 +477,92 @@ describe('ComputedUpdateWorker', () => {
       );
     });
   });
+
+  describe('runTaskById', () => {
+    it('claims and processes the specified task id', async () => {
+      const task = createMockTask();
+      const markDone = vi.fn().mockResolvedValue(ok(undefined));
+      const claimById = vi.fn().mockResolvedValue(ok(task));
+
+      const outbox: IComputedUpdateOutbox = {
+        enqueueOrMerge: vi.fn(),
+        enqueueSeedTask: vi.fn(),
+        enqueueFieldBackfill: vi.fn(),
+        claimBatch: vi.fn().mockResolvedValue(ok([])),
+        claimById,
+        markDone,
+        markFailed: vi.fn(),
+      };
+
+      const updater = createUpdaterStub({
+        execute: vi.fn().mockResolvedValue(ok({ changesByStep: [] })),
+        collectDirtySeedGroups: vi.fn().mockResolvedValue(ok([])),
+      });
+
+      const planner = {
+        planStage: vi.fn().mockResolvedValue(ok({ steps: [], edges: [] })),
+      } as unknown as ComputedUpdatePlanner;
+
+      const worker = new ComputedUpdateWorker(
+        outbox,
+        updater,
+        planner,
+        createUnitOfWork(),
+        createLogger(),
+        createHasher(),
+        createTableRepository(),
+        createBackfillService(),
+        createEventBus()
+      );
+
+      const result = await worker.runTaskById({
+        taskId: task.id,
+        workerId: 'manual-worker',
+      });
+
+      expect(result.isOk()).toBe(true);
+      expect(result._unsafeUnwrap()).toBe(true);
+      expect(claimById).toHaveBeenCalledWith(
+        expect.objectContaining({
+          taskId: task.id,
+          workerId: 'manual-worker',
+          allowProcessingTakeover: true,
+        }),
+        expect.anything()
+      );
+      expect(markDone).toHaveBeenCalledWith(task.id, expect.anything());
+    });
+
+    it('returns false when the task cannot be claimed by id', async () => {
+      const outbox: IComputedUpdateOutbox = {
+        enqueueOrMerge: vi.fn(),
+        enqueueSeedTask: vi.fn(),
+        enqueueFieldBackfill: vi.fn(),
+        claimBatch: vi.fn().mockResolvedValue(ok([])),
+        claimById: vi.fn().mockResolvedValue(ok(null)),
+        markDone: vi.fn(),
+        markFailed: vi.fn(),
+      };
+
+      const worker = new ComputedUpdateWorker(
+        outbox,
+        createUpdaterStub(),
+        {} as ComputedUpdatePlanner,
+        createUnitOfWork(),
+        createLogger(),
+        createHasher(),
+        createTableRepository(),
+        createBackfillService(),
+        createEventBus()
+      );
+
+      const result = await worker.runTaskById({
+        taskId: 'cuo-missing',
+        workerId: 'manual-worker',
+      });
+
+      expect(result.isOk()).toBe(true);
+      expect(result._unsafeUnwrap()).toBe(false);
+    });
+  });
 });