Update the usage task list to use clickhouse

Author: ericallam

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

@@ -1,9 +1,10 @@
-import { sqlDatabaseSchema } from "~/db.server";
+import { PrismaClientOrTransaction, sqlDatabaseSchema } from "~/db.server";
 import { env } from "~/env.server";
 import { getUsage, getUsageSeries } from "~/services/platform.v3.server";
 import { createTimeSeriesData } from "~/utils/graphs";
 import { BasePresenter } from "./basePresenter.server";
 import { DataPoint, linear } from "regression";
+import { clickhouseClient } from "~/services/clickhouseInstance.server";
 
 type Options = {
   organizationId: string;
@@ -103,27 +104,69 @@ export class UsagePresenter extends BasePresenter {
     });
 
     //usage by task
-    const tasks = this._replica.$queryRaw<TaskUsageItem[]>`
+    const tasks = await getTaskUsageByOrganization(
+      organizationId,
+      startOfMonth,
+      endOfMonth,
+      this._replica
+    );
+
+    return {
+      usage,
+      tasks,
+    };
+  }
+}
+
+async function getTaskUsageByOrganization(
+  organizationId: string,
+  startOfMonth: Date,
+  endOfMonth: Date,
+  replica: PrismaClientOrTransaction
+) {
+  if (clickhouseClient) {
+    const [queryError, tasks] = await clickhouseClient.taskRuns.getTaskUsageByOrganization({
+      startTime: startOfMonth.getTime(),
+      endTime: endOfMonth.getTime(),
+      organizationId,
+    });
+
+    if (queryError) {
+      throw queryError;
+    }
+
+    return tasks
+      .map((task) => ({
+        taskIdentifier: task.task_identifier,
+        runCount: Number(task.run_count),
+        averageDuration: Number(task.average_duration),
+        averageCost: Number(task.average_cost) + env.CENTS_PER_RUN / 100,
+        totalDuration: Number(task.total_duration),
+        totalCost: Number(task.total_cost) + Number(task.total_base_cost),
+      }))
+      .sort((a, b) => b.totalCost - a.totalCost);
+  } else {
+    return replica.$queryRaw<TaskUsageItem[]>`
     SELECT
-    tr."taskIdentifier",
-    COUNT(*) AS "runCount",
-    AVG(tr."usageDurationMs") AS "averageDuration",
-    SUM(tr."usageDurationMs") AS "totalDuration",
-    AVG(tr."costInCents") / 100.0 AS "averageCost",
-    SUM(tr."costInCents") / 100.0 AS "totalCost",
-    SUM(tr."baseCostInCents") / 100.0 AS "totalBaseCost"
-  FROM
-      ${sqlDatabaseSchema}."TaskRun" tr
-      JOIN ${sqlDatabaseSchema}."Project" pr ON pr.id = tr."projectId"
-      JOIN ${sqlDatabaseSchema}."Organization" org ON org.id = pr."organizationId"
-      JOIN ${sqlDatabaseSchema}."RuntimeEnvironment" env ON env."id" = tr."runtimeEnvironmentId"
-  WHERE
-      env.type <> 'DEVELOPMENT'
-      AND tr."createdAt" > ${startOfMonth}
-      AND tr."createdAt" < ${endOfMonth}
-      AND org.id = ${organizationId}
-  GROUP BY
-      tr."taskIdentifier";
+      tr."taskIdentifier",
+      COUNT(*) AS "runCount",
+      AVG(tr."usageDurationMs") AS "averageDuration",
+      SUM(tr."usageDurationMs") AS "totalDuration",
+      AVG(tr."costInCents") / 100.0 AS "averageCost",
+      SUM(tr."costInCents") / 100.0 AS "totalCost",
+      SUM(tr."baseCostInCents") / 100.0 AS "totalBaseCost"
+    FROM
+        ${sqlDatabaseSchema}."TaskRun" tr
+        JOIN ${sqlDatabaseSchema}."Project" pr ON pr.id = tr."projectId"
+        JOIN ${sqlDatabaseSchema}."Organization" org ON org.id = pr."organizationId"
+        JOIN ${sqlDatabaseSchema}."RuntimeEnvironment" env ON env."id" = tr."runtimeEnvironmentId"
+    WHERE
+        env.type <> 'DEVELOPMENT'
+        AND tr."createdAt" > ${startOfMonth}
+        AND tr."createdAt" < ${endOfMonth}
+        AND org.id = ${organizationId}
+    GROUP BY
+        tr."taskIdentifier";
   `.then((data) => {
       return data
         .map((item) => ({
@@ -136,10 +179,5 @@ export class UsagePresenter extends BasePresenter {
         }))
         .sort((a, b) => b.totalCost - a.totalCost);
     });
-
-    return {
-      usage,
-      tasks,
-    };
   }
 }

internal-packages/clickhouse/src/index.ts

@@ -9,6 +9,7 @@ import {
   getTaskActivityQueryBuilder,
   getCurrentRunningStats,
   getAverageDurations,
+  getTaskUsageByOrganization,
 } from "./taskRuns.js";
 import { Logger, type LogLevel } from "@trigger.dev/core/logger";
 import type { Agent as HttpAgent } from "http";
@@ -146,6 +147,7 @@ export class ClickHouse {
       getTaskActivity: getTaskActivityQueryBuilder(this.reader),
       getCurrentRunningStats: getCurrentRunningStats(this.reader),
       getAverageDurations: getAverageDurations(this.reader),
+      getTaskUsageByOrganization: getTaskUsageByOrganization(this.reader),
     };
   }
 }

internal-packages/clickhouse/src/taskRuns.ts

@@ -220,3 +220,49 @@ export function getAverageDurations(ch: ClickhouseReader, settings?: ClickHouseS
     settings,
   });
 }
+
+export const TaskUsageByOrganizationQueryResult = z.object({
+  task_identifier: z.string(),
+  run_count: z.number(),
+  average_duration: z.number(),
+  total_duration: z.number(),
+  average_cost: z.number(),
+  total_cost: z.number(),
+  total_base_cost: z.number(),
+});
+
+export const TaskUsageByOrganizationQueryParams = z.object({
+  startTime: z.number().int(),
+  endTime: z.number().int(),
+  organizationId: z.string(),
+});
+
+export function getTaskUsageByOrganization(ch: ClickhouseReader, settings?: ClickHouseSettings) {
+  return ch.query({
+    name: "getTaskUsageByOrganization",
+    query: `
+      SELECT
+      task_identifier,
+      count() AS run_count,
+      avg(usage_duration_ms) AS average_duration,
+      sum(usage_duration_ms) AS total_duration,
+      avg(cost_in_cents) / 100.0 AS average_cost,
+      sum(cost_in_cents) / 100.0 AS total_cost,
+      sum(base_cost_in_cents) / 100.0 AS total_base_cost
+  FROM trigger_dev.task_runs_v2 FINAL
+  WHERE
+      environment_type != 'DEVELOPMENT'
+      AND created_at >= fromUnixTimestamp64Milli({startTime: Int64})
+      AND created_at < fromUnixTimestamp64Milli({endTime: Int64})
+      AND organization_id = {organizationId: String}
+      AND _is_deleted = 0
+  GROUP BY
+      task_identifier
+  ORDER BY
+      total_cost DESC
+    `,
+    schema: TaskUsageByOrganizationQueryResult,
+    params: TaskUsageByOrganizationQueryParams,
+    settings,
+  });
+}