db efficiency improvements

apps/webapp/app/db.server.ts

@@ -13,6 +13,8 @@ import { logger } from "./services/logger.server";
 import { isValidDatabaseUrl } from "./utils/db";
 import { singleton } from "./utils/singleton";
 import { $transaction as transac } from "@trigger.dev/database";
+import { startActiveSpan } from "./v3/tracer.server";
+import { Span } from "@opentelemetry/api";
 
 export type {
   PrismaTransactionClient,
@@ -21,25 +23,76 @@ export type {
   PrismaReplicaClient,
 };
 
+export async function $transaction<R>(
+  prisma: PrismaClientOrTransaction,
+  name: string,
+  fn: (prisma: PrismaTransactionClient, span?: Span) => Promise<R>,
+  options?: PrismaTransactionOptions
+): Promise<R | undefined>;
 export async function $transaction<R>(
   prisma: PrismaClientOrTransaction,
   fn: (prisma: PrismaTransactionClient) => Promise<R>,
   options?: PrismaTransactionOptions
+): Promise<R | undefined>;
+export async function $transaction<R>(
+  prisma: PrismaClientOrTransaction,
+  fnOrName: ((prisma: PrismaTransactionClient) => Promise<R>) | string,
+  fnOrOptions?: ((prisma: PrismaTransactionClient) => Promise<R>) | PrismaTransactionOptions,
+  options?: PrismaTransactionOptions
 ): Promise<R | undefined> {
-  return transac(
-    prisma,
-    fn,
-    (error) => {
-      logger.error("prisma.$transaction error", {
-        code: error.code,
-        meta: error.meta,
-        stack: error.stack,
-        message: error.message,
-        name: error.name,
-      });
-    },
-    options
-  );
+  if (typeof fnOrName === "string") {
+    return await startActiveSpan(fnOrName, async (span) => {
+      span.setAttribute("$transaction", true);
+
+      if (options?.isolationLevel) {
+        span.setAttribute("isolation_level", options.isolationLevel);
+      }
+
+      if (options?.timeout) {
+        span.setAttribute("timeout", options.timeout);
+      }
+
+      if (options?.maxWait) {
+        span.setAttribute("max_wait", options.maxWait);
+      }
+
+      if (options?.swallowPrismaErrors) {
+        span.setAttribute("swallow_prisma_errors", options.swallowPrismaErrors);
+      }
+
+      const fn = fnOrOptions as (prisma: PrismaTransactionClient, span: Span) => Promise<R>;
+
+      return transac(
+        prisma,
+        (client) => fn(client, span),
+        (error) => {
+          logger.error("prisma.$transaction error", {
+            code: error.code,
+            meta: error.meta,
+            stack: error.stack,
+            message: error.message,
+            name: error.name,
+          });
+        },
+        options
+      );
+    });
+  } else {
+    return transac(
+      prisma,
+      fnOrName,
+      (error) => {
+        logger.error("prisma.$transaction error", {
+          code: error.code,
+          meta: error.meta,
+          stack: error.stack,
+          message: error.message,
+          name: error.name,
+        });
+      },
+      typeof fnOrOptions === "function" ? undefined : fnOrOptions
+    );
+  }
 }
 
 export { Prisma };

apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts

@@ -9,7 +9,7 @@ export class ApiBatchResultsPresenter extends BasePresenter {
     env: AuthenticatedEnvironment
   ): Promise<BatchTaskRunExecutionResult | undefined> {
     return this.traceWithEnv("call", env, async (span) => {
-      const batchRun = await this._prisma.batchTaskRun.findUnique({
+      const batchRun = await this._prisma.batchTaskRun.findFirst({
         where: {
           friendlyId,
           runtimeEnvironmentId: env.id,

apps/webapp/app/presenters/v3/ApiRunResultPresenter.server.ts

@@ -9,7 +9,7 @@ export class ApiRunResultPresenter extends BasePresenter {
     env: AuthenticatedEnvironment
   ): Promise<TaskRunExecutionResult | undefined> {
     return this.traceWithEnv("call", env, async (span) => {
-      const taskRun = await this._prisma.taskRun.findUnique({
+      const taskRun = await this._prisma.taskRun.findFirst({
         where: {
           friendlyId,
           runtimeEnvironmentId: env.id,

apps/webapp/app/presenters/v3/DeploymentPresenter.server.ts

@@ -53,12 +53,10 @@ export class DeploymentPresenter {
       },
     });
 
-    const deployment = await this.#prismaClient.workerDeployment.findUniqueOrThrow({
+    const deployment = await this.#prismaClient.workerDeployment.findFirstOrThrow({
       where: {
-        projectId_shortCode: {
-          projectId: project.id,
-          shortCode: deploymentShortCode,
-        },
+        projectId: project.id,
+        shortCode: deploymentShortCode,
       },
       select: {
         id: true,

apps/webapp/app/presenters/v3/EnvironmentVariablesPresenter.server.ts

@@ -15,7 +15,7 @@ export class EnvironmentVariablesPresenter {
   }
 
   public async call({ userId, projectSlug }: { userId: User["id"]; projectSlug: Project["slug"] }) {
-    const project = await this.#prismaClient.project.findUnique({
+    const project = await this.#prismaClient.project.findFirst({
       select: {
         id: true,
       },

apps/webapp/app/presenters/v3/NewAlertChannelPresenter.server.ts

@@ -8,7 +8,7 @@ import { WebClient } from "@slack/web-api";
 
 export class NewAlertChannelPresenter extends BasePresenter {
   public async call(projectId: string) {
-    const project = await this._prisma.project.findUniqueOrThrow({
+    const project = await this._prisma.project.findFirstOrThrow({
       where: {
         id: projectId,
       },

apps/webapp/app/presenters/v3/RunListPresenter.server.ts

@@ -131,7 +131,7 @@ export class RunListPresenter extends BasePresenter {
 
     //bulk id
     if (bulkId) {
-      const bulkAction = await this._replica.bulkActionGroup.findUnique({
+      const bulkAction = await this._replica.bulkActionGroup.findFirst({
         select: {
           items: {
             select: {

apps/webapp/app/presenters/v3/RunStreamPresenter.server.ts

@@ -21,7 +21,7 @@ export class RunStreamPresenter {
     request: Request;
     runFriendlyId: TaskRun["friendlyId"];
   }) {
-    const run = await this.#prismaClient.taskRun.findUnique({
+    const run = await this.#prismaClient.taskRun.findFirst({
       where: {
         friendlyId: runFriendlyId,
       },

apps/webapp/app/presenters/v3/SpanPresenter.server.ts

@@ -29,7 +29,7 @@ export class SpanPresenter extends BasePresenter {
     spanId: string;
     runFriendlyId: string;
   }) {
-    const project = await this._replica.project.findUnique({
+    const project = await this._replica.project.findFirst({
       where: {
         slug: projectSlug,
       },

apps/webapp/app/presenters/v3/TasksStreamPresenter.server.ts

@@ -33,7 +33,7 @@ export class TasksStreamPresenter {
     projectSlug: string;
     userId: string;
   }) {
-    const project = await this.#prismaClient.project.findUnique({
+    const project = await this.#prismaClient.project.findFirst({
       where: {
         slug: projectSlug,
         organization: {

apps/webapp/app/routes/api.v1.deployments.$deploymentId.ts

@@ -27,7 +27,7 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
 
   const { deploymentId } = parsedParams.data;
 
-  const deployment = await prisma.workerDeployment.findUnique({
+  const deployment = await prisma.workerDeployment.findFirst({
     where: {
       friendlyId: deploymentId,
       environmentId: authenticatedEnv.id,

apps/webapp/app/routes/api.v1.projects.$projectRef.envvars.ts

@@ -26,7 +26,7 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
 
   const { projectRef } = parsedParams.data;
 
-  const project = await prisma.project.findUnique({
+  const project = await prisma.project.findFirst({
     where: {
       externalRef: projectRef,
       environments: {

apps/webapp/app/routes/api.v1.projects.$projectRef.ts

@@ -27,7 +27,7 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
 
   const { projectRef } = parsedParams.data;
 
-  const project = await prisma.project.findUnique({
+  const project = await prisma.project.findFirst({
     where: {
       externalRef: projectRef,
       organization: {

apps/webapp/app/services/autoIncrementCounter.server.ts

@@ -1,11 +1,5 @@
 import Redis, { RedisOptions } from "ioredis";
-import {
-  $transaction,
-  Prisma,
-  PrismaClientOrTransaction,
-  PrismaTransactionOptions,
-  prisma,
-} from "~/db.server";
+import { Prisma, PrismaClientOrTransaction, PrismaTransactionOptions, prisma } from "~/db.server";
 import { env } from "~/env.server";
 import { singleton } from "~/utils/singleton";
 

apps/webapp/app/v3/environmentVariables/environmentVariablesRepository.server.ts

@@ -57,7 +57,7 @@ export class EnvironmentVariablesRepository implements Repository {
       }[];
     }
   ): Promise<CreateResult> {
-    const project = await this.prismaClient.project.findUnique({
+    const project = await this.prismaClient.project.findFirst({
       where: {
         id: projectId,
         deletedAt: null,
@@ -136,7 +136,7 @@ export class EnvironmentVariablesRepository implements Repository {
 
     try {
       for (const variable of values) {
-        const result = await $transaction(this.prismaClient, async (tx) => {
+        const result = await $transaction(this.prismaClient, "create env var", async (tx) => {
           const environmentVariable = await tx.environmentVariable.upsert({
             where: {
               projectId_key: {
@@ -227,7 +227,7 @@ export class EnvironmentVariablesRepository implements Repository {
       keepEmptyValues?: boolean;
     }
   ): Promise<Result> {
-    const project = await this.prismaClient.project.findUnique({
+    const project = await this.prismaClient.project.findFirst({
       where: {
         id: projectId,
         deletedAt: null,
@@ -266,7 +266,7 @@ export class EnvironmentVariablesRepository implements Repository {
       }
     }
 
-    const environmentVariable = await this.prismaClient.environmentVariable.findUnique({
+    const environmentVariable = await this.prismaClient.environmentVariable.findFirst({
       select: {
         id: true,
         key: true,
@@ -280,20 +280,18 @@ export class EnvironmentVariablesRepository implements Repository {
     }
 
     try {
-      await $transaction(this.prismaClient, async (tx) => {
+      await $transaction(this.prismaClient, "edit env var", async (tx) => {
         const secretStore = getSecretStore("DATABASE", {
           prismaClient: tx,
         });
 
         //create the secret values and references
         for (const value of values) {
           const key = secretKey(projectId, value.environmentId, environmentVariable.key);
-          const existingValue = await tx.environmentVariableValue.findUnique({
+          const existingValue = await tx.environmentVariableValue.findFirst({
             where: {
-              variableId_environmentId: {
-                variableId: environmentVariable.id,
-                environmentId: value.environmentId,
-              },
+              variableId: environmentVariable.id,
+              environmentId: value.environmentId,
             },
           });
 
@@ -356,7 +354,7 @@ export class EnvironmentVariablesRepository implements Repository {
   }
 
   async getProject(projectId: string): Promise<ProjectEnvironmentVariable[]> {
-    const project = await this.prismaClient.project.findUnique({
+    const project = await this.prismaClient.project.findFirst({
       where: {
         id: projectId,
         deletedAt: null,
@@ -429,7 +427,7 @@ export class EnvironmentVariablesRepository implements Repository {
   }
 
   async getEnvironment(projectId: string, environmentId: string): Promise<EnvironmentVariable[]> {
-    const project = await this.prismaClient.project.findUnique({
+    const project = await this.prismaClient.project.findFirst({
       where: {
         id: projectId,
         deletedAt: null,
@@ -483,7 +481,7 @@ export class EnvironmentVariablesRepository implements Repository {
   }
 
   async delete(projectId: string, options: DeleteEnvironmentVariable): Promise<Result> {
-    const project = await this.prismaClient.project.findUnique({
+    const project = await this.prismaClient.project.findFirst({
       where: {
         id: projectId,
         deletedAt: null,
@@ -501,7 +499,7 @@ export class EnvironmentVariablesRepository implements Repository {
       return { success: false as const, error: "Project not found" };
     }
 
-    const environmentVariable = await this.prismaClient.environmentVariable.findUnique({
+    const environmentVariable = await this.prismaClient.environmentVariable.findFirst({
       select: {
         id: true,
         key: true,
@@ -526,7 +524,7 @@ export class EnvironmentVariablesRepository implements Repository {
     }
 
     try {
-      await $transaction(this.prismaClient, async (tx) => {
+      await $transaction(this.prismaClient, "delete env var", async (tx) => {
         await tx.environmentVariable.delete({
           where: {
             id: options.id,
@@ -564,7 +562,7 @@ export class EnvironmentVariablesRepository implements Repository {
   }
 
   async deleteValue(projectId: string, options: DeleteEnvironmentVariableValue): Promise<Result> {
-    const project = await this.prismaClient.project.findUnique({
+    const project = await this.prismaClient.project.findFirst({
       where: {
         id: projectId,
         deletedAt: null,
@@ -582,7 +580,7 @@ export class EnvironmentVariablesRepository implements Repository {
       return { success: false as const, error: "Project not found" };
     }
 
-    const environmentVariable = await this.prismaClient.environmentVariable.findUnique({
+    const environmentVariable = await this.prismaClient.environmentVariable.findFirst({
       select: {
         id: true,
         key: true,
@@ -619,7 +617,7 @@ export class EnvironmentVariablesRepository implements Repository {
     }
 
     try {
-      await $transaction(this.prismaClient, async (tx) => {
+      await $transaction(this.prismaClient, "delete env var value", async (tx) => {
         const secretStore = getSecretStore("DATABASE", {
           prismaClient: tx,
         });

apps/webapp/app/v3/marqs/devQueueConsumer.server.ts

@@ -81,7 +81,7 @@ export class DevQueueConsumer {
   }
 
   public async registerBackgroundWorker(id: string, inProgressRuns: string[] = []) {
-    const backgroundWorker = await prisma.backgroundWorker.findUnique({
+    const backgroundWorker = await prisma.backgroundWorker.findFirst({
       where: { friendlyId: id, runtimeEnvironmentId: this.env.id },
       include: {
         tasks: true,
@@ -170,7 +170,7 @@ export class DevQueueConsumer {
   public async taskHeartbeat(workerId: string, id: string) {
     logger.debug("[DevQueueConsumer] taskHeartbeat()", { id });
 
-    const taskRunAttempt = await prisma.taskRunAttempt.findUnique({
+    const taskRunAttempt = await prisma.taskRunAttempt.findFirst({
       where: { friendlyId: id },
     });
 

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

@@ -11,7 +11,7 @@ import { TaskRunErrorCodes } from "@trigger.dev/core/v3";
 
 export class RequeueTaskRunService extends BaseService {
   public async call(runId: string) {
-    const taskRun = await this._prisma.taskRun.findUnique({
+    const taskRun = await this._prisma.taskRun.findFirst({
       where: {
         id: runId,
       },