BaseService tx support, minor improvements

Author: matt-aitken

apps/webapp/app/v3/failedTaskRun.server.ts

@@ -41,7 +41,6 @@ export class FailedTaskRunService extends BaseService {
 
     const finalizeService = new FinalizeTaskRunService();
     await finalizeService.call({
-      tx: this._prisma,
       id: taskRun.id,
       status: "SYSTEM_FAILURE",
       completedAt: new Date(),

apps/webapp/app/v3/services/cancelAttempt.server.ts

@@ -50,8 +50,6 @@ export class CancelAttemptService extends BaseService {
         return;
       }
 
-      const finalizeService = new FinalizeTaskRunService();
-
       await $transaction(this._prisma, async (tx) => {
         await tx.taskRunAttempt.update({
           where: {
@@ -63,8 +61,8 @@ export class CancelAttemptService extends BaseService {
           },
         });
 
+        const finalizeService = new FinalizeTaskRunService(tx);
         await finalizeService.call({
-          tx,
           id: taskRunId,
           status: isCancellableRunStatus(taskRunAttempt.taskRun.status) ? "INTERRUPTED" : undefined,
           completedAt: isCancellableRunStatus(taskRunAttempt.taskRun.status)

apps/webapp/app/v3/services/cancelTaskRun.server.ts

@@ -49,7 +49,6 @@ export class CancelTaskRunService extends BaseService {
 
     const finalizeService = new FinalizeTaskRunService();
     const cancelledTaskRun = await finalizeService.call({
-      tx: this._prisma,
       id: taskRun.id,
       status: "CANCELED",
       completedAt: opts.cancelledAt,

apps/webapp/app/v3/services/completeAttempt.server.ts

@@ -51,7 +51,7 @@ export class CompleteAttemptService extends BaseService {
         id: execution.attempt.id,
       });
 
-      const run = await this._prisma.taskRun.findUnique({
+      const run = await this._prisma.taskRun.findFirst({
         where: {
           friendlyId: execution.run.id,
         },
@@ -70,7 +70,6 @@ export class CompleteAttemptService extends BaseService {
 
       const finalizeService = new FinalizeTaskRunService();
       await finalizeService.call({
-        tx: this._prisma,
         id: run.id,
         status: "SYSTEM_FAILURE",
         completedAt: new Date(),
@@ -111,8 +110,6 @@ export class CompleteAttemptService extends BaseService {
     taskRunAttempt: NonNullable<FoundAttempt>,
     env?: AuthenticatedEnvironment
   ): Promise<"COMPLETED"> {
-    const finalizeService = new FinalizeTaskRunService();
-
     await $transaction(this._prisma, async (tx) => {
       await tx.taskRunAttempt.update({
         where: { id: taskRunAttempt.id },
@@ -125,9 +122,9 @@ export class CompleteAttemptService extends BaseService {
         },
       });
 
+      const finalizeService = new FinalizeTaskRunService(tx);
       await finalizeService.call({
-        tx,
-        id: taskRunAttempt.taskRun.id,
+        id: taskRunAttempt.taskRunId,
         status: "COMPLETED_SUCCESSFULLY",
         completedAt: new Date(),
       });
@@ -271,7 +268,6 @@ export class CompleteAttemptService extends BaseService {
 
         const finalizeService = new FinalizeTaskRunService();
         await finalizeService.call({
-          tx: this._prisma,
           id: taskRunAttempt.taskRunId,
           status: "SYSTEM_FAILURE",
           completedAt: new Date(),
@@ -311,7 +307,6 @@ export class CompleteAttemptService extends BaseService {
       ) {
         const finalizeService = new FinalizeTaskRunService();
         await finalizeService.call({
-          tx: this._prisma,
           id: taskRunAttempt.taskRunId,
           status: "SYSTEM_FAILURE",
           completedAt: new Date(),
@@ -343,7 +338,6 @@ export class CompleteAttemptService extends BaseService {
       } else {
         const finalizeService = new FinalizeTaskRunService();
         await finalizeService.call({
-          tx: this._prisma,
           id: taskRunAttempt.taskRunId,
           status: "COMPLETED_WITH_ERRORS",
           completedAt: new Date(),

apps/webapp/app/v3/services/crashTaskRun.server.ts

@@ -46,7 +46,6 @@ export class CrashTaskRunService extends BaseService {
 
     const finalizeService = new FinalizeTaskRunService();
     const crashedTaskRun = await finalizeService.call({
-      tx: this._prisma,
       id: taskRun.id,
       status: "CRASHED",
       completedAt: new Date(),

apps/webapp/app/v3/services/expireEnqueuedRun.server.ts

@@ -41,7 +41,6 @@ export class ExpireEnqueuedRunService extends BaseService {
 
     const finalizeService = new FinalizeTaskRunService();
     await finalizeService.call({
-      tx: this._prisma,
       id: run.id,
       status: "EXPIRED",
       expiredAt: new Date(),

apps/webapp/app/v3/services/finalizeTaskRun.server.ts

@@ -1,12 +1,10 @@
-import { type Prisma, type TaskRun, type TaskRunStatus } from "@trigger.dev/database";
-import { type PrismaClientOrTransaction } from "~/db.server";
+import { type Prisma, type TaskRun } from "@trigger.dev/database";
+import { type FINISHED_STATUSES } from "~/components/runs/v3/TaskRunStatus";
+import { logger } from "~/services/logger.server";
 import { marqs } from "~/v3/marqs/index.server";
 import { BaseService } from "./baseService.server";
-import { logger } from "~/services/logger.server";
-import { type FINISHED_STATUSES } from "~/components/runs/v3/TaskRunStatus";
 
 type BaseInput = {
-  tx: PrismaClientOrTransaction;
   id: string;
   status?: FINISHED_STATUSES;
   expiredAt?: Date;
@@ -27,7 +25,6 @@ type Output<T extends Prisma.TaskRunInclude | undefined> = T extends Prisma.Task
 
 export class FinalizeTaskRunService extends BaseService {
   public async call<T extends Prisma.TaskRunInclude | undefined>({
-    tx,
     id,
     status,
     expiredAt,
@@ -51,7 +48,7 @@ export class FinalizeTaskRunService extends BaseService {
       completedAt,
     });
 
-    const run = await tx.taskRun.update({
+    const run = await this._prisma.taskRun.update({
       where: { id },
       data: { status, expiredAt, completedAt },
       ...(include ? { include } : {}),