made accessing run and payload fields more type safe

Author: ericallam

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

@@ -1,5 +1,5 @@
 import type { ClickHouse, TaskRunInsertArray, PayloadInsertArray } from "@internal/clickhouse";
-import { TASK_RUN_INDEX, PAYLOAD_INDEX } from "@internal/clickhouse";
+import { getTaskRunField, getPayloadField } from "@internal/clickhouse";
 import { type RedisOptions } from "@internal/redis";
 import {
   LogicalReplicationClient,
@@ -580,28 +580,28 @@ export class RunsReplicationService {
         // batch inserts in clickhouse are more performant if the items
         // are pre-sorted by the primary key
         .sort((a, b) => {
-          const aOrgId = a[TASK_RUN_INDEX.organization_id] as string;
-          const bOrgId = b[TASK_RUN_INDEX.organization_id] as string;
+          const aOrgId = getTaskRunField(a, "organization_id");
+          const bOrgId = getTaskRunField(b, "organization_id");
           if (aOrgId !== bOrgId) {
             return aOrgId < bOrgId ? -1 : 1;
           }
-          const aProjId = a[TASK_RUN_INDEX.project_id] as string;
-          const bProjId = b[TASK_RUN_INDEX.project_id] as string;
+          const aProjId = getTaskRunField(a, "project_id");
+          const bProjId = getTaskRunField(b, "project_id");
           if (aProjId !== bProjId) {
             return aProjId < bProjId ? -1 : 1;
           }
-          const aEnvId = a[TASK_RUN_INDEX.environment_id] as string;
-          const bEnvId = b[TASK_RUN_INDEX.environment_id] as string;
+          const aEnvId = getTaskRunField(a, "environment_id");
+          const bEnvId = getTaskRunField(b, "environment_id");
           if (aEnvId !== bEnvId) {
             return aEnvId < bEnvId ? -1 : 1;
           }
-          const aCreatedAt = a[TASK_RUN_INDEX.created_at] as number;
-          const bCreatedAt = b[TASK_RUN_INDEX.created_at] as number;
+          const aCreatedAt = getTaskRunField(a, "created_at");
+          const bCreatedAt = getTaskRunField(b, "created_at");
           if (aCreatedAt !== bCreatedAt) {
             return aCreatedAt - bCreatedAt;
           }
-          const aRunId = a[TASK_RUN_INDEX.run_id] as string;
-          const bRunId = b[TASK_RUN_INDEX.run_id] as string;
+          const aRunId = getTaskRunField(a, "run_id");
+          const bRunId = getTaskRunField(b, "run_id");
           if (aRunId === bRunId) return 0;
           return aRunId < bRunId ? -1 : 1;
         });
@@ -612,8 +612,8 @@ export class RunsReplicationService {
         // batch inserts in clickhouse are more performant if the items
         // are pre-sorted by the primary key
         .sort((a, b) => {
-          const aRunId = a[PAYLOAD_INDEX.run_id] as string;
-          const bRunId = b[PAYLOAD_INDEX.run_id] as string;
+          const aRunId = getPayloadField(a, "run_id");
+          const bRunId = getPayloadField(b, "run_id");
           if (aRunId === bRunId) return 0;
           return aRunId < bRunId ? -1 : 1;
         });

apps/webapp/test/runsReplicationService.part2.test.ts

@@ -1,4 +1,4 @@
-import { ClickHouse, TASK_RUN_INDEX, PAYLOAD_INDEX } from "@internal/clickhouse";
+import { ClickHouse, getTaskRunField, getPayloadField } from "@internal/clickhouse";
 import { containerTest } from "@internal/testcontainers";
 import { Logger } from "@trigger.dev/core/logger";
 import { readFile } from "node:fs/promises";
@@ -889,13 +889,13 @@ describe("RunsReplicationService (part 2/2)", () => {
       await setTimeout(1000);
 
       expect(batchFlushedEvents?.[0].taskRunInserts).toHaveLength(2);
-      // Use TASK_RUN_INDEX for type-safe array access
-      expect(batchFlushedEvents?.[0].taskRunInserts[0][TASK_RUN_INDEX.run_id]).toEqual(run.id);
-      expect(batchFlushedEvents?.[0].taskRunInserts[0][TASK_RUN_INDEX.status]).toEqual(
+      // Use getTaskRunField for type-safe array access
+      expect(getTaskRunField(batchFlushedEvents![0].taskRunInserts[0], "run_id")).toEqual(run.id);
+      expect(getTaskRunField(batchFlushedEvents![0].taskRunInserts[0], "status")).toEqual(
         "PENDING_VERSION"
       );
-      expect(batchFlushedEvents?.[0].taskRunInserts[1][TASK_RUN_INDEX.run_id]).toEqual(run.id);
-      expect(batchFlushedEvents?.[0].taskRunInserts[1][TASK_RUN_INDEX.status]).toEqual(
+      expect(getTaskRunField(batchFlushedEvents![0].taskRunInserts[1], "run_id")).toEqual(run.id);
+      expect(getTaskRunField(batchFlushedEvents![0].taskRunInserts[1], "status")).toEqual(
         "COMPLETED_SUCCESSFULLY"
       );
 
@@ -1063,22 +1063,22 @@ describe("RunsReplicationService (part 2/2)", () => {
 
       // Verify sorting order: organization_id, project_id, environment_id, created_at, run_id
       for (let i = 1; i < batchFlushedEvents[0]?.taskRunInserts.length; i++) {
-        const prev = batchFlushedEvents[0]?.taskRunInserts[i - 1];
-        const curr = batchFlushedEvents[0]?.taskRunInserts[i];
+        const prev = batchFlushedEvents[0]!.taskRunInserts[i - 1];
+        const curr = batchFlushedEvents[0]!.taskRunInserts[i];
 
         const prevKey = [
-          prev[TASK_RUN_INDEX.organization_id],
-          prev[TASK_RUN_INDEX.project_id],
-          prev[TASK_RUN_INDEX.environment_id],
-          prev[TASK_RUN_INDEX.created_at],
-          prev[TASK_RUN_INDEX.run_id],
+          getTaskRunField(prev, "organization_id"),
+          getTaskRunField(prev, "project_id"),
+          getTaskRunField(prev, "environment_id"),
+          getTaskRunField(prev, "created_at"),
+          getTaskRunField(prev, "run_id"),
         ];
         const currKey = [
-          curr[TASK_RUN_INDEX.organization_id],
-          curr[TASK_RUN_INDEX.project_id],
-          curr[TASK_RUN_INDEX.environment_id],
-          curr[TASK_RUN_INDEX.created_at],
-          curr[TASK_RUN_INDEX.run_id],
+          getTaskRunField(curr, "organization_id"),
+          getTaskRunField(curr, "project_id"),
+          getTaskRunField(curr, "environment_id"),
+          getTaskRunField(curr, "created_at"),
+          getTaskRunField(curr, "run_id"),
         ];
 
         const keysAreEqual = prevKey.every((val, idx) => val === currKey[idx]);
@@ -1106,9 +1106,9 @@ describe("RunsReplicationService (part 2/2)", () => {
 
       // Verify payloadInserts are also sorted by run_id
       for (let i = 1; i < batchFlushedEvents[0]?.payloadInserts.length; i++) {
-        const prev = batchFlushedEvents[0]?.payloadInserts[i - 1];
-        const curr = batchFlushedEvents[0]?.payloadInserts[i];
-        expect(prev[PAYLOAD_INDEX.run_id] <= curr[PAYLOAD_INDEX.run_id]).toBeTruthy();
+        const prev = batchFlushedEvents[0]!.payloadInserts[i - 1];
+        const curr = batchFlushedEvents[0]!.payloadInserts[i];
+        expect(getPayloadField(prev, "run_id") <= getPayloadField(curr, "run_id")).toBeTruthy();
       }
 
       await runsReplicationService.stop();

internal-packages/clickhouse/src/client/client.ts

@@ -708,7 +708,7 @@ export class ClickhouseClient implements ClickhouseReader, ClickhouseWriter {
             table: req.table,
           });
 
-          recordSpanError(span, clickhouseError);
+          recordClickhouseError(span, clickhouseError);
           return [new InsertError(clickhouseError.message), null];
         }
 

internal-packages/clickhouse/src/index.ts

@@ -31,12 +31,14 @@ export type * from "./taskRuns.js";
 export type * from "./taskEvents.js";
 export type * from "./client/queryBuilder.js";
 
-// Re-export column constants and indices for type-safe array access
+// Re-export column constants, indices, and type-safe accessors
 export {
   TASK_RUN_COLUMNS,
   TASK_RUN_INDEX,
   PAYLOAD_COLUMNS,
   PAYLOAD_INDEX,
+  getTaskRunField,
+  getPayloadField,
 } from "./taskRuns.js";
 
 // TSQL query execution

internal-packages/clickhouse/src/taskRuns.ts

@@ -109,6 +109,73 @@ export const TASK_RUN_INDEX = Object.fromEntries(
   TASK_RUN_COLUMNS.map((col, idx) => [col, idx])
 ) as { readonly [K in TaskRunColumnName]: number };
 
+/**
+ * Type mapping from column name to its type in TaskRunInsertArray.
+ * This enables type-safe field access without manual casting.
+ */
+export type TaskRunFieldTypes = {
+  environment_id: string;
+  organization_id: string;
+  project_id: string;
+  run_id: string;
+  updated_at: number;
+  created_at: number;
+  status: string;
+  environment_type: string;
+  friendly_id: string;
+  attempt: number;
+  engine: string;
+  task_identifier: string;
+  queue: string;
+  schedule_id: string;
+  batch_id: string;
+  completed_at: number | null;
+  started_at: number | null;
+  executed_at: number | null;
+  delay_until: number | null;
+  queued_at: number | null;
+  expired_at: number | null;
+  usage_duration_ms: number;
+  cost_in_cents: number;
+  base_cost_in_cents: number;
+  output: { data: unknown };
+  error: { data: unknown };
+  tags: string[];
+  task_version: string;
+  sdk_version: string;
+  cli_version: string;
+  machine_preset: string;
+  root_run_id: string;
+  parent_run_id: string;
+  depth: number;
+  span_id: string;
+  trace_id: string;
+  idempotency_key: string;
+  expiration_ttl: string;
+  is_test: boolean;
+  _version: string;
+  _is_deleted: number;
+  concurrency_key: string;
+  bulk_action_group_ids: string[];
+  worker_queue: string;
+  max_duration_in_seconds: number | null;
+};
+
+/**
+ * Type-safe accessor for TaskRunInsertArray fields.
+ * Returns the correct type for each field without manual casting.
+ *
+ * @example
+ * const orgId = getTaskRunField(run, "organization_id"); // type: string
+ * const createdAt = getTaskRunField(run, "created_at"); // type: number
+ */
+export function getTaskRunField<K extends TaskRunColumnName>(
+  run: TaskRunInsertArray,
+  field: K
+): TaskRunFieldTypes[K] {
+  return run[TASK_RUN_INDEX[field]] as TaskRunFieldTypes[K];
+}
+
 export function insertTaskRunsCompactArrays(ch: ClickhouseWriter, settings?: ClickHouseSettings) {
   return ch.insertCompactRaw({
     name: "insertTaskRunsCompactArrays",
@@ -149,9 +216,29 @@ export const PAYLOAD_COLUMNS = ["run_id", "created_at", "payload"] as const;
 export type PayloadColumnName = (typeof PAYLOAD_COLUMNS)[number];
 
 // Type-safe column indices generated from PAYLOAD_COLUMNS
-export const PAYLOAD_INDEX = Object.fromEntries(
-  PAYLOAD_COLUMNS.map((col, idx) => [col, idx])
-) as { readonly [K in PayloadColumnName]: number };
+export const PAYLOAD_INDEX = Object.fromEntries(PAYLOAD_COLUMNS.map((col, idx) => [col, idx])) as {
+  readonly [K in PayloadColumnName]: number;
+};
+
+/**
+ * Type mapping from column name to its type in PayloadInsertArray.
+ */
+export type PayloadFieldTypes = {
+  run_id: string;
+  created_at: number;
+  payload: { data: unknown };
+};
+
+/**
+ * Type-safe accessor for PayloadInsertArray fields.
+ * Returns the correct type for each field without manual casting.
+ */
+export function getPayloadField<K extends PayloadColumnName>(
+  payload: PayloadInsertArray,
+  field: K
+): PayloadFieldTypes[K] {
+  return payload[PAYLOAD_INDEX[field]] as PayloadFieldTypes[K];
+}
 
 /**
  * Type-safe tuple representing a task run insert array.
@@ -209,11 +296,7 @@ export type TaskRunInsertArray = [
  * Type-safe tuple representing a payload insert array.
  * Order matches PAYLOAD_COLUMNS exactly.
  */
-export type PayloadInsertArray = [
-  run_id: string,
-  created_at: number,
-  payload: { data: unknown },
-];
+export type PayloadInsertArray = [run_id: string, created_at: number, payload: { data: unknown }];
 
 export function insertRawTaskRunPayloadsCompactArrays(
   ch: ClickhouseWriter,