v3: fix raw queries for custom schemas (#1033)

* add custom validation to db url env vars

* extract schema from db url and use in all raw queries

* use qualified names in scheduling raw queries

* cook a few raw queries

* Added missing raw query schema specifier to DeploymentListPresenter

---------

Co-authored-by: Matt Aitken <[email protected]>

Author: nicktrn

apps/webapp/app/db.server.ts

@@ -4,6 +4,7 @@ import { z } from "zod";
 import { logger } from "./services/logger.server";
 import { env } from "./env.server";
 import { singleton } from "./utils/singleton";
+import { isValidDatabaseUrl } from "./utils/db";
 
 export type PrismaTransactionClient = Omit<
   PrismaClient,
@@ -138,3 +139,23 @@ export type { PrismaClient } from "@trigger.dev/database";
 export const PrismaErrorSchema = z.object({
   code: z.string(),
 });
+
+function getDatabaseSchema() {
+  if (!isValidDatabaseUrl(env.DATABASE_URL)) {
+    throw new Error("Invalid Database URL");
+  }
+
+  const databaseUrl = new URL(env.DATABASE_URL);
+  const schemaFromSearchParam = databaseUrl.searchParams.get("schema");
+
+  if (!schemaFromSearchParam) {
+    console.debug("❗ database schema unspecified, will default to `public` schema");
+    return "public";
+  }
+
+  return schemaFromSearchParam;
+}
+
+export const DATABASE_SCHEMA = singleton("DATABASE_SCHEMA", getDatabaseSchema);
+
+export const sqlDatabaseSchema = Prisma.sql([`${DATABASE_SCHEMA}`]);

apps/webapp/app/env.server.ts

@@ -1,13 +1,24 @@
 import { z } from "zod";
 import { SecretStoreOptionsSchema } from "./services/secrets/secretStore.server";
 import { isValidRegex } from "./utils/regex";
+import { isValidDatabaseUrl } from "./utils/db";
 
 const EnvironmentSchema = z.object({
   NODE_ENV: z.union([z.literal("development"), z.literal("production"), z.literal("test")]),
-  DATABASE_URL: z.string(),
+  DATABASE_URL: z
+    .string()
+    .refine(
+      isValidDatabaseUrl,
+      "DATABASE_URL is invalid, for details please check the additional output above this message."
+    ),
   DATABASE_CONNECTION_LIMIT: z.coerce.number().int().default(10),
   DATABASE_POOL_TIMEOUT: z.coerce.number().int().default(60),
-  DIRECT_URL: z.string(),
+  DIRECT_URL: z
+    .string()
+    .refine(
+      isValidDatabaseUrl,
+      "DIRECT_URL is invalid, for details please check the additional output above this message."
+    ),
   SESSION_SECRET: z.string(),
   MAGIC_LINK_SECRET: z.string(),
   ENCRYPTION_KEY: z.string(),

apps/webapp/app/presenters/JobListPresenter.server.ts

@@ -3,7 +3,7 @@ import {
   DisplayPropertySchema,
   EventSpecificationSchema,
 } from "@trigger.dev/core";
-import { PrismaClient, Prisma, prisma } from "~/db.server";
+import { PrismaClient, Prisma, prisma, sqlDatabaseSchema } from "~/db.server";
 import { Organization } from "~/models/organization.server";
 import { Project } from "~/models/project.server";
 import { User } from "~/models/user.server";
@@ -122,7 +122,7 @@ export class JobListPresenter {
               "jobId",
               ROW_NUMBER() OVER(PARTITION BY "jobId" ORDER BY "createdAt" DESC) as rn
           FROM 
-              "JobRun" 
+              ${sqlDatabaseSchema}."JobRun" 
           WHERE 
               "jobId" IN (${Prisma.join(jobs.map((j) => j.id))})
       ) t

apps/webapp/app/presenters/OrgUsagePresenter.server.ts

@@ -1,5 +1,5 @@
 import { estimate } from "@trigger.dev/billing";
-import { PrismaClient, prisma } from "~/db.server";
+import { sqlDatabaseSchema, PrismaClient, prisma } from "~/db.server";
 import { featuresForRequest } from "~/features.server";
 import { BillingService } from "~/services/billing.server";
 
@@ -53,7 +53,7 @@ export class OrgUsagePresenter {
         month: string;
         count: number;
       }[]
-    >`SELECT TO_CHAR("createdAt", 'YYYY-MM') as month, COUNT(*) as count FROM "JobRun" WHERE "organizationId" = ${organization.id} AND "createdAt" >= NOW() - INTERVAL '6 months' AND "internal" = FALSE GROUP BY month ORDER BY month ASC`;
+    >`SELECT TO_CHAR("createdAt", 'YYYY-MM') as month, COUNT(*) as count FROM ${sqlDatabaseSchema}."JobRun" WHERE "organizationId" = ${organization.id} AND "createdAt" >= NOW() - INTERVAL '6 months' AND "internal" = FALSE GROUP BY month ORDER BY month ASC`;
 
     const hasMonthlyRunData = monthlyRunsDataRaw.length > 0;
     const monthlyRunsData = monthlyRunsDataRaw.map((obj) => ({
@@ -117,7 +117,7 @@ export class OrgUsagePresenter {
 
     const dailyRunsRawData = await this.#prismaClient.$queryRaw<
       { day: Date; runs: BigInt }[]
-    >`SELECT date_trunc('day', "createdAt") as day, COUNT(*) as runs FROM "JobRun" WHERE "organizationId" = ${organization.id} AND "createdAt" >= NOW() - INTERVAL '30 days' AND "internal" = FALSE GROUP BY day`;
+    >`SELECT date_trunc('day', "createdAt") as day, COUNT(*) as runs FROM ${sqlDatabaseSchema}."JobRun" WHERE "organizationId" = ${organization.id} AND "createdAt" >= NOW() - INTERVAL '30 days' AND "internal" = FALSE GROUP BY day`;
 
     const hasDailyRunsData = dailyRunsRawData.length > 0;
     const dailyRunsDataFilledIn = fillInMissingDailyRuns(ThirtyDaysAgo, 31, dailyRunsRawData);

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

@@ -1,5 +1,5 @@
 import { WorkerDeploymentStatus } from "@trigger.dev/database";
-import { PrismaClient, prisma } from "~/db.server";
+import { sqlDatabaseSchema, PrismaClient, prisma } from "~/db.server";
 import { Organization } from "~/models/organization.server";
 import { Project } from "~/models/project.server";
 import { User } from "~/models/user.server";
@@ -97,7 +97,7 @@ export class DeploymentListPresenter {
   wd."id", 
   wd."shortCode", 
   wd."version", 
-  (SELECT COUNT(*) FROM "BackgroundWorkerTask" WHERE "BackgroundWorkerTask"."workerId" = wd."workerId") AS "tasksCount", 
+  (SELECT COUNT(*) FROM ${sqlDatabaseSchema}."BackgroundWorkerTask" WHERE "BackgroundWorkerTask"."workerId" = wd."workerId") AS "tasksCount", 
   wd."environmentId", 
   wd."status", 
   u."id" AS "userId", 
@@ -106,9 +106,9 @@ export class DeploymentListPresenter {
   u."avatarUrl" AS "userAvatarUrl", 
   wd."deployedAt"
 FROM 
-  "WorkerDeployment" as wd
+  ${sqlDatabaseSchema}."WorkerDeployment" as wd
 INNER JOIN 
-  "User" as u ON wd."triggeredById" = u."id" 
+  ${sqlDatabaseSchema}."User" as u ON wd."triggeredById" = u."id" 
 WHERE 
   wd."projectId" = ${project.id}
 ORDER BY 

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

@@ -58,12 +58,13 @@ export class EditSchedulePresenter {
       },
     });
 
-    const possibleTasks = await this.#prismaClient.$queryRaw<{ slug: string }[]>`
-    SELECT DISTINCT(slug)
-    FROM "BackgroundWorkerTask"
-    WHERE "projectId" = ${project.id} 
-    AND "triggerSource" = 'SCHEDULED';
-    `;
+    const possibleTasks = await this.#prismaClient.backgroundWorkerTask.findMany({
+      distinct: ["slug"],
+      where: {
+        projectId: project.id,
+        triggerSource: "SCHEDULED",
+      },
+    });
 
     const possibleEnvironments = project.environments.map((environment) => {
       let userName: undefined | string;

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

@@ -1,6 +1,6 @@
 import { Prisma, TaskRunStatus } from "@trigger.dev/database";
 import { Direction } from "~/components/runs/RunStatuses";
-import { PrismaClient, prisma } from "~/db.server";
+import { sqlDatabaseSchema, PrismaClient, prisma } from "~/db.server";
 import { getUsername } from "~/utils/username";
 import { CANCELLABLE_STATUSES } from "~/v3/services/cancelTaskRun.server";
 
@@ -85,11 +85,12 @@ export class RunListPresenter {
     });
 
     //get all possible tasks
-    const possibleTasks = await this.#prismaClient.$queryRaw<{ slug: string }[]>`
-    SELECT DISTINCT(slug)
-    FROM "BackgroundWorkerTask"
-    WHERE "projectId" = ${project.id};
-    `;
+    const possibleTasks = await this.#prismaClient.backgroundWorkerTask.findMany({
+      distinct: ["slug"],
+      where: {
+        projectId: project.id,
+      },
+    });
 
     //get the runs
     let runs = await this.#prismaClient.$queryRaw<
@@ -122,15 +123,15 @@ export class RunListPresenter {
     tr."isTest" AS "isTest",
     COUNT(tra.id) AS attempts
   FROM
-    "TaskRun" tr
+    ${sqlDatabaseSchema}."TaskRun" tr
   LEFT JOIN
     (
       SELECT *,
         ROW_NUMBER() OVER (PARTITION BY "taskRunId" ORDER BY "createdAt" DESC) rn
-      FROM "TaskRunAttempt"
+      FROM ${sqlDatabaseSchema}."TaskRunAttempt"
     ) tra ON tr.id = tra."taskRunId" AND tra.rn = 1
   LEFT JOIN
-    "BackgroundWorker" bw ON tra."backgroundWorkerId" = bw.id
+    ${sqlDatabaseSchema}."BackgroundWorker" bw ON tra."backgroundWorkerId" = bw.id
   WHERE
       -- project
       tr."projectId" = ${project.id}

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

@@ -1,6 +1,6 @@
 import { Prisma, RuntimeEnvironmentType } from "@trigger.dev/database";
 import { ScheduleListFilters } from "~/components/runs/v3/ScheduleFilters";
-import { PrismaClient, prisma } from "~/db.server";
+import { PrismaClient, prisma, sqlDatabaseSchema } from "~/db.server";
 import { getUsername } from "~/utils/username";
 import { calculateNextScheduledTimestamp } from "~/v3/utils/calculateNextSchedule.server";
 
@@ -81,12 +81,13 @@ export class ScheduleListPresenter {
     });
 
     //get all possible scheduled tasks
-    const possibleTasks = await this.#prismaClient.$queryRaw<{ slug: string }[]>`
-    SELECT DISTINCT(slug)
-    FROM "BackgroundWorkerTask"
-    WHERE "projectId" = ${project.id}
-    AND "triggerSource" = 'SCHEDULED';
-    `;
+    const possibleTasks = await this.#prismaClient.backgroundWorkerTask.findMany({
+      distinct: ["slug"],
+      where: {
+        projectId: project.id,
+        triggerSource: "SCHEDULED",
+      },
+    });
 
     //do this here to protect against SQL injection
     search = search && search !== "" ? `%${search}%` : undefined;
@@ -201,11 +202,11 @@ export class ScheduleListPresenter {
     SELECT t."scheduleId", t."createdAt"
     FROM (
       SELECT "scheduleId", MAX("createdAt") as "LatestRun"
-      FROM "TaskRun"
+      FROM ${sqlDatabaseSchema}."TaskRun"
       WHERE "scheduleId" IN (${Prisma.join(rawSchedules.map((s) => s.id))})
       GROUP BY "scheduleId"
     ) r
-    JOIN "TaskRun" t
+    JOIN ${sqlDatabaseSchema}."TaskRun" t
     ON t."scheduleId" = r."scheduleId" AND t."createdAt" = r."LatestRun";`
         : [];
 

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

@@ -1,5 +1,5 @@
 import { Prisma, TaskRunStatus, TaskTriggerSource } from "@trigger.dev/database";
-import { PrismaClient, prisma } from "~/db.server";
+import { PrismaClient, prisma, sqlDatabaseSchema } from "~/db.server";
 import { Organization } from "~/models/organization.server";
 import { Project } from "~/models/project.server";
 import { User } from "~/models/user.server";
@@ -73,7 +73,7 @@ export class TaskListPresenter {
       bwt."createdAt",
       bwt."triggerSource"
     FROM
-      "BackgroundWorkerTask" as bwt
+      ${sqlDatabaseSchema}."BackgroundWorkerTask" as bwt
     WHERE bwt."projectId" = ${project.id}
     ORDER BY
       bwt.slug,
@@ -101,7 +101,7 @@ export class TaskListPresenter {
           "lockedById",
           ROW_NUMBER() OVER (PARTITION BY "lockedById" ORDER BY "updatedAt" DESC) AS rn
         FROM
-          "TaskRun"
+          ${sqlDatabaseSchema}."TaskRun"
         WHERE
           "lockedById" IN(${Prisma.join(tasks.map((t) => t.id))})
             ) t

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

@@ -1,5 +1,5 @@
 import { TaskTriggerSource } from "@trigger.dev/database";
-import { PrismaClient, prisma } from "~/db.server";
+import { sqlDatabaseSchema, PrismaClient, prisma } from "~/db.server";
 import { TestSearchParams } from "~/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.test/route";
 import { sortEnvironments } from "~/services/environmentSort.server";
 import { createSearchParams } from "~/utils/searchParams";
@@ -97,14 +97,14 @@ export class TestPresenter {
             bw.*,
             ROW_NUMBER() OVER(ORDER BY string_to_array(bw.version, '.')::int[] DESC) AS rn
       FROM 
-            "BackgroundWorker" bw
+            ${sqlDatabaseSchema}."BackgroundWorker" bw
       WHERE "runtimeEnvironmentId" = ${matchingEnvironment.id}
     ),
     latest_workers AS (SELECT * FROM workers WHERE rn = 1)
-    SELECT "BackgroundWorkerTask".id, version, slug as "taskIdentifier", "filePath", "exportName", "BackgroundWorkerTask"."friendlyId", "BackgroundWorkerTask"."triggerSource"
+    SELECT bwt.id, version, slug as "taskIdentifier", "filePath", "exportName", bwt."friendlyId" 
     FROM latest_workers
-    JOIN "BackgroundWorkerTask" ON "BackgroundWorkerTask"."workerId" = latest_workers.id
-    ORDER BY "BackgroundWorkerTask"."exportName" ASC;
+    JOIN ${sqlDatabaseSchema}."BackgroundWorkerTask" bwt ON bwt."workerId" = latest_workers.id
+    ORDER BY bwt."exportName" ASC;
     `;
 
     return {

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

@@ -5,7 +5,7 @@ import {
   TaskRunStatus,
   TaskTriggerSource,
 } from "@trigger.dev/database";
-import { PrismaClient, prisma } from "~/db.server";
+import { sqlDatabaseSchema, PrismaClient, prisma } from "~/db.server";
 import { getUsername } from "~/utils/username";
 
 type TestTaskOptions = {
@@ -107,9 +107,9 @@ export class TestTaskPresenter {
       SELECT 
           tr.* 
       FROM 
-          "TaskRun" as tr
+          ${sqlDatabaseSchema}."TaskRun" as tr
       JOIN
-          "BackgroundWorkerTask" as bwt
+          ${sqlDatabaseSchema}."BackgroundWorkerTask" as bwt
       ON
           tr."taskIdentifier" = bwt.slug
       WHERE

apps/webapp/app/utils/db.ts

@@ -0,0 +1,18 @@
+export function isValidDatabaseUrl(url: string) {
+  try {
+    const databaseUrl = new URL(url);
+    const schemaFromSearchParam = databaseUrl.searchParams.get("schema");
+
+    if (schemaFromSearchParam === "") {
+      console.error(
+        "Invalid Database URL: The schema search param can't have an empty value. To use the `public` schema, either omit the schema param entirely or specify it in full: `?schema=public`"
+      );
+      return false;
+    }
+
+    return true;
+  } catch (err) {
+    console.error(err);
+    return false;
+  }
+}