Use clickhouse in task list aggregation queries instead of pg (keep pg for self-hosters)

Author: ericallam

apps/webapp/app/env.server.ts

@@ -741,6 +741,14 @@ const EnvironmentSchema = z.object({
   RUN_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS: z.coerce.number().int().optional(),
   RUN_REPLICATION_MAX_OPEN_CONNECTIONS: z.coerce.number().int().default(10),
 
+  // Clickhouse
+  CLICKHOUSE_URL: z.string().optional(),
+  CLICKHOUSE_KEEP_ALIVE_ENABLED: z.string().default("1"),
+  CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS: z.coerce.number().int().optional(),
+  CLICKHOUSE_MAX_OPEN_CONNECTIONS: z.coerce.number().int().default(10),
+  CLICKHOUSE_LOG_LEVEL: z.enum(["log", "error", "warn", "info", "debug"]).default("info"),
+  CLICKHOUSE_COMPRESSION_REQUEST: z.string().default("1"),
+
   // Bootstrap
   TRIGGER_BOOTSTRAP_ENABLED: z.string().default("0"),
   TRIGGER_BOOTSTRAP_WORKER_GROUP_NAME: z.string().optional(),

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

@@ -1,15 +1,14 @@
 import { CURRENT_DEPLOYMENT_LABEL } from "@trigger.dev/core/v3/isomorphic";
+import { PrismaClientOrTransaction, type TaskTriggerSource } from "@trigger.dev/database";
+import { $replica, sqlDatabaseSchema } from "~/db.server";
+import { clickhouseClient } from "~/services/clickhouseInstance.server";
 import {
-  Prisma,
-  type TaskRunStatus as DBTaskRunStatus,
-  type TaskRunStatus as TaskRunStatusType,
-  type TaskTriggerSource,
-} from "@trigger.dev/database";
-import { QUEUED_STATUSES } from "~/components/runs/v3/TaskRunStatus";
-import { TaskRunStatus } from "~/database-types";
-import { sqlDatabaseSchema } from "~/db.server";
-import { logger } from "~/services/logger.server";
-import { BasePresenter } from "./basePresenter.server";
+  ClickHouseEnvironmentMetricsRepository,
+  DailyTaskActivity,
+  EnvironmentMetricsRepository,
+  PostgrestEnvironmentMetricsRepository,
+} from "~/services/environmentMetricsRepository.server";
+import { singleton } from "~/utils/singleton";
 
 export type TaskListItem = {
   slug: string;
@@ -18,11 +17,14 @@ export type TaskListItem = {
   triggerSource: TaskTriggerSource;
 };
 
-type Return = Awaited<ReturnType<TaskListPresenter["call"]>>;
+export type TaskActivity = DailyTaskActivity[string];
 
-export type TaskActivity = Awaited<Return["activity"]>[string];
+export class TaskListPresenter {
+  constructor(
+    private readonly environmentMetricsRepository: EnvironmentMetricsRepository,
+    private readonly _replica: PrismaClientOrTransaction
+  ) {}
 
-export class TaskListPresenter extends BasePresenter {
   public async call({ environmentId, projectId }: { environmentId: string; projectId: string }) {
     const tasks = await this._replica.$queryRaw<
       {
@@ -56,187 +58,57 @@ export class TaskListPresenter extends BasePresenter {
     //then get the activity for each task
     const activity = this.#getActivity(
       tasks.map((t) => t.slug),
-      projectId,
       environmentId
     );
 
     const runningStats = this.#getRunningStats(
       tasks.map((t) => t.slug),
-      projectId,
       environmentId
     );
 
     const durations = this.#getAverageDurations(
       tasks.map((t) => t.slug),
-      projectId,
       environmentId
     );
 
     return { tasks, activity, runningStats, durations };
   }
 
-  async #getActivity(tasks: string[], projectId: string, environmentId: string) {
-    if (tasks.length === 0) {
-      return {};
-    }
-
-    const activity = await this._replica.$queryRaw<
-      {
-        taskIdentifier: string;
-        status: TaskRunStatusType;
-        day: Date;
-        count: BigInt;
-      }[]
-    >`
-    SELECT
-    tr."taskIdentifier",
-    tr."status",
-    DATE(tr."createdAt") as day,
-    COUNT(*)
-  FROM
-    ${sqlDatabaseSchema}."TaskRun" as tr
-  WHERE
-    tr."taskIdentifier" IN (${Prisma.join(tasks)})
-    AND tr."projectId" = ${projectId}
-    AND tr."runtimeEnvironmentId" = ${environmentId}
-    AND tr."createdAt" >= (current_date - interval '6 days')
-  GROUP BY
-    tr."taskIdentifier",
-    tr."status",
-    day
-  ORDER BY
-    tr."taskIdentifier" ASC,
-    day ASC,
-    tr."status" ASC;`;
-
-    //today with no time
-    const today = new Date();
-    today.setUTCHours(0, 0, 0, 0);
-
-    return activity.reduce((acc, a) => {
-      let existingTask = acc[a.taskIdentifier];
-
-      if (!existingTask) {
-        existingTask = [];
-        //populate the array with the past 7 days
-        for (let i = 6; i >= 0; i--) {
-          const day = new Date(today);
-          day.setUTCDate(today.getDate() - i);
-          day.setUTCHours(0, 0, 0, 0);
-
-          existingTask.push({
-            day: day.toISOString(),
-            [TaskRunStatus.COMPLETED_SUCCESSFULLY]: 0,
-          } as { day: string } & Record<TaskRunStatusType, number>);
-        }
-
-        acc[a.taskIdentifier] = existingTask;
-      }
-
-      const dayString = a.day.toISOString();
-      const day = existingTask.find((d) => d.day === dayString);
-
-      if (!day) {
-        logger.warn(`Day not found for TaskRun`, {
-          day: dayString,
-          taskIdentifier: a.taskIdentifier,
-          existingTask,
-        });
-        return acc;
-      }
-
-      day[a.status] = Number(a.count);
-
-      return acc;
-    }, {} as Record<string, ({ day: string } & Record<TaskRunStatusType, number>)[]>);
+  async #getActivity(tasks: string[], environmentId: string) {
+    return this.environmentMetricsRepository.getDailyTaskActivity({
+      environmentId,
+      days: 6,
+      tasks,
+    });
   }
 
-  async #getRunningStats(tasks: string[], projectId: string, environmentId: string) {
-    if (tasks.length === 0) {
-      return {};
-    }
-
-    const stats = await this._replica.$queryRaw<
-      {
-        taskIdentifier: string;
-        status: DBTaskRunStatus;
-        count: BigInt;
-      }[]
-    >`
-    SELECT
-    tr."taskIdentifier",
-    tr.status,
-    COUNT(*)
-  FROM
-    ${sqlDatabaseSchema}."TaskRun" as tr
-  WHERE
-    tr."taskIdentifier" IN (${Prisma.join(tasks)})
-    AND tr."projectId" = ${projectId}
-    AND tr."runtimeEnvironmentId" = ${environmentId}
-    AND tr."status" = ANY(ARRAY[${Prisma.join([
-      ...QUEUED_STATUSES,
-      "EXECUTING",
-    ])}]::\"TaskRunStatus\"[])
-  GROUP BY
-    tr."taskIdentifier",
-    tr.status
-  ORDER BY
-    tr."taskIdentifier" ASC`;
-
-    //create an object combining the queued and concurrency counts
-    const result: Record<string, { queued: number; running: number }> = {};
-    for (const task of tasks) {
-      const queued = stats.filter(
-        (q) => q.taskIdentifier === task && QUEUED_STATUSES.includes(q.status)
-      );
-      const queuedCount =
-        queued.length === 0
-          ? 0
-          : queued.reduce((acc, q) => {
-              return acc + Number(q.count);
-            }, 0);
-
-      const running = stats.filter((r) => r.taskIdentifier === task && r.status === "EXECUTING");
-      const runningCount =
-        running.length === 0
-          ? 0
-          : running.reduce((acc, r) => {
-              return acc + Number(r.count);
-            }, 0);
+  async #getRunningStats(tasks: string[], environmentId: string) {
+    return this.environmentMetricsRepository.getCurrentRunningStats({
+      environmentId,
+      days: 6,
+      tasks,
+    });
+  }
 
-      result[task] = {
-        queued: queuedCount,
-        running: runningCount,
-      };
-    }
-    return result;
+  async #getAverageDurations(tasks: string[], environmentId: string) {
+    return this.environmentMetricsRepository.getAverageDurations({
+      environmentId,
+      days: 6,
+      tasks,
+    });
   }
+}
 
-  async #getAverageDurations(tasks: string[], projectId: string, environmentId: string) {
-    if (tasks.length === 0) {
-      return {};
-    }
+export const taskListPresenter = singleton("taskListPresenter", setupTaskListPresenter);
 
-    const durations = await this._replica.$queryRaw<
-      {
-        taskIdentifier: string;
-        duration: Number;
-      }[]
-    >`
-    SELECT
-      tr."taskIdentifier",
-      AVG(EXTRACT(EPOCH FROM (tr."updatedAt" - COALESCE(tr."startedAt", tr."lockedAt")))) as duration
-      FROM
-      ${sqlDatabaseSchema}."TaskRun" as tr
-    WHERE
-      tr."taskIdentifier" IN (${Prisma.join(tasks)})
-      AND tr."projectId" = ${projectId}
-      AND tr."runtimeEnvironmentId" = ${environmentId}
-      AND tr."createdAt" >= (current_date - interval '6 days')
-      AND tr."status" IN ('COMPLETED_SUCCESSFULLY', 'COMPLETED_WITH_ERRORS')
-    GROUP BY
-      tr."taskIdentifier";`;
+function setupTaskListPresenter() {
+  const environmentMetricsRepository = clickhouseClient
+    ? new ClickHouseEnvironmentMetricsRepository({
+        clickhouse: clickhouseClient,
+      })
+    : new PostgrestEnvironmentMetricsRepository({
+        prisma: $replica,
+      });
 
-    return Object.fromEntries(durations.map((s) => [s.taskIdentifier, Number(s.duration)]));
-  }
+  return new TaskListPresenter(environmentMetricsRepository, $replica);
 }

apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam._index/route.tsx

@@ -75,6 +75,7 @@ import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server";
 import {
   type TaskActivity,
   type TaskListItem,
+  taskListPresenter,
   TaskListPresenter,
 } from "~/presenters/v3/TaskListPresenter.server";
 import {
@@ -123,8 +124,7 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
   }
 
   try {
-    const presenter = new TaskListPresenter();
-    const { tasks, activity, runningStats, durations } = await presenter.call({
+    const { tasks, activity, runningStats, durations } = await taskListPresenter.call({
       environmentId: environment.id,
       projectId: project.id,
     });

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

@@ -0,0 +1,30 @@
+import { ClickHouse } from "@internal/clickhouse";
+import { env } from "~/env.server";
+import { singleton } from "~/utils/singleton";
+
+export const clickhouseClient = singleton("clickhouseClient", initializeClickhouseClient);
+
+function initializeClickhouseClient() {
+  if (!env.CLICKHOUSE_URL) {
+    console.log("🗃️  Clickhouse service not enabled");
+    return;
+  }
+
+  console.log("🗃️  Clickhouse service enabled");
+
+  const clickhouse = new ClickHouse({
+    url: env.CLICKHOUSE_URL,
+    name: "clickhouse-instance",
+    keepAlive: {
+      enabled: env.CLICKHOUSE_KEEP_ALIVE_ENABLED === "1",
+      idleSocketTtl: env.CLICKHOUSE_KEEP_ALIVE_IDLE_SOCKET_TTL_MS,
+    },
+    logLevel: env.CLICKHOUSE_LOG_LEVEL,
+    compression: {
+      request: true,
+    },
+    maxOpenConnections: env.CLICKHOUSE_MAX_OPEN_CONNECTIONS,
+  });
+
+  return clickhouse;
+}