more clickhouse events stuff

Author: ericallam

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

@@ -3,7 +3,7 @@ import { createTreeFromFlatItems, flattenTree } from "~/components/primitives/Tr
 import { prisma, type PrismaClient } from "~/db.server";
 import { createTimelineSpanEventsFromSpanEvents } from "~/utils/timelineSpanEvents";
 import { getUsername } from "~/utils/username";
-import { eventRepository } from "~/v3/eventRepository.server";
+import { resolveEventRepositoryForStore } from "~/v3/eventRepository";
 import { getTaskEventStoreTableForRun } from "~/v3/taskEventStore.server";
 import { isFinalRunStatus } from "~/v3/taskStatus";
 
@@ -140,9 +140,12 @@ export class RunPresenter {
       };
     }
 
+    const eventRepository = resolveEventRepositoryForStore(run.taskEventStore);
+
     // get the events
     const traceSummary = await eventRepository.getTraceSummary(
       getTaskEventStoreTableForRun(run),
+      run.runtimeEnvironment.id,
       run.traceId,
       run.rootTaskRun?.createdAt ?? run.createdAt,
       run.completedAt ?? undefined,

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

@@ -3,7 +3,7 @@ import { eventStream } from "remix-utils/sse/server";
 import { PrismaClient, prisma } from "~/db.server";
 import { logger } from "~/services/logger.server";
 import { throttle } from "~/utils/throttle";
-import { eventRepository } from "~/v3/eventRepository.server";
+import { tracePubSub } from "~/v3/services/tracePubSub.server";
 
 const pingInterval = 1000;
 
@@ -41,7 +41,7 @@ export class RunStreamPresenter {
 
     let pinger: NodeJS.Timeout | undefined = undefined;
 
-    const { unsubscribe, eventEmitter } = await eventRepository.subscribeToTrace(run.traceId);
+    const { unsubscribe, eventEmitter } = await tracePubSub.subscribeToTrace(run.traceId);
 
     return eventStream(request.signal, (send, close) => {
       const safeSend = (args: { event?: string; data: string }) => {

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

@@ -10,21 +10,23 @@ import {
 import { AttemptId, getMaxDuration, parseTraceparent } from "@trigger.dev/core/v3/isomorphic";
 import { RUNNING_STATUSES } from "~/components/runs/v3/TaskRunStatus";
 import { logger } from "~/services/logger.server";
-import { eventRepository, rehydrateAttribute } from "~/v3/eventRepository.server";
+import { rehydrateAttribute } from "~/v3/eventRepository/eventRepository.server";
 import { machinePresetFromRun } from "~/v3/machinePresets.server";
 import { getTaskEventStoreTableForRun, type TaskEventStoreTable } from "~/v3/taskEventStore.server";
 import { isFailedRunStatus, isFinalRunStatus } from "~/v3/taskStatus";
 import { BasePresenter } from "./basePresenter.server";
 import { WaitpointPresenter } from "./WaitpointPresenter.server";
 import { engine } from "~/v3/runEngine.server";
+import { resolveEventRepositoryForStore } from "~/v3/eventRepository";
+import { IEventRepository, SpanDetail } from "~/v3/eventRepository/eventRepository.types";
 
 type Result = Awaited<ReturnType<SpanPresenter["call"]>>;
 export type Span = NonNullable<NonNullable<Result>["span"]>;
 export type SpanRun = NonNullable<NonNullable<Result>["run"]>;
 type FindRunResult = NonNullable<
   Awaited<ReturnType<InstanceType<typeof SpanPresenter>["findRun"]>>
 >;
-type GetSpanResult = NonNullable<Awaited<ReturnType<(typeof eventRepository)["getSpan"]>>>;
+type GetSpanResult = SpanDetail;
 
 export class SpanPresenter extends BasePresenter {
   public async call({
@@ -74,14 +76,20 @@ export class SpanPresenter extends BasePresenter {
       return;
     }
 
+    const { traceId } = parentRun;
+
+    const eventRepository = resolveEventRepositoryForStore(parentRun.taskEventStore);
+
     const eventStore = getTaskEventStoreTableForRun(parentRun);
 
     const run = await this.getRun({
       eventStore,
-      environmentId: parentRun.runtimeEnvironmentId,
+      traceId,
+      eventRepository,
       spanId,
       createdAt: parentRun.createdAt,
       completedAt: parentRun.completedAt,
+      environmentId: parentRun.runtimeEnvironmentId,
     });
     if (run) {
       return {
@@ -93,10 +101,12 @@ export class SpanPresenter extends BasePresenter {
     const span = await this.#getSpan({
       eventStore,
       spanId,
+      traceId,
       environmentId: parentRun.runtimeEnvironmentId,
       projectId: parentRun.projectId,
       createdAt: parentRun.createdAt,
       completedAt: parentRun.completedAt,
+      eventRepository,
     });
 
     if (!span) {
@@ -112,29 +122,30 @@ export class SpanPresenter extends BasePresenter {
   async getRun({
     eventStore,
     environmentId,
+    traceId,
+    eventRepository,
     spanId,
     createdAt,
     completedAt,
   }: {
     eventStore: TaskEventStoreTable;
     environmentId: string;
+    traceId: string;
+    eventRepository: IEventRepository;
     spanId: string;
     createdAt: Date;
     completedAt: Date | null;
   }) {
-    const span = await eventRepository.getSpan({
-      storeTable: eventStore,
+    const originalRunId = await eventRepository.getSpanOriginalRunId(
+      eventStore,
+      environmentId,
       spanId,
       environmentId,
-      startCreatedAt: createdAt,
-      endCreatedAt: completedAt ?? undefined,
-    });
-
-    if (!span) {
-      return;
-    }
+      createdAt,
+      completedAt ?? undefined
+    );
 
-    const run = await this.findRun({ span, spanId });
+    const run = await this.findRun({ originalRunId, spanId, environmentId });
 
     if (!run) {
       return;
@@ -259,7 +270,7 @@ export class SpanPresenter extends BasePresenter {
       workerQueue: run.workerQueue,
       traceId: run.traceId,
       spanId: run.spanId,
-      isCached: !!span.originalRun,
+      isCached: !!originalRunId,
       machinePreset: machine?.name,
       externalTraceId,
     };
@@ -294,7 +305,15 @@ export class SpanPresenter extends BasePresenter {
     };
   }
 
-  async findRun({ span, spanId }: { span: GetSpanResult; spanId: string }) {
+  async findRun({
+    originalRunId,
+    spanId,
+    environmentId,
+  }: {
+    originalRunId?: string;
+    spanId: string;
+    environmentId: string;
+  }) {
     const run = await this._replica.taskRun.findFirst({
       select: {
         id: true,
@@ -404,12 +423,14 @@ export class SpanPresenter extends BasePresenter {
           },
         },
       },
-      where: span.originalRun
+      where: originalRunId
         ? {
-            friendlyId: span.originalRun,
+            friendlyId: originalRunId,
+            runtimeEnvironmentId: environmentId,
           }
         : {
             spanId,
+            runtimeEnvironmentId: environmentId,
           },
     });
 
@@ -418,27 +439,32 @@ export class SpanPresenter extends BasePresenter {
 
   async #getSpan({
     eventStore,
+    eventRepository,
+    traceId,
     spanId,
     environmentId,
     projectId,
     createdAt,
     completedAt,
   }: {
+    eventRepository: IEventRepository;
+    traceId: string;
     spanId: string;
     environmentId: string;
     projectId: string;
     eventStore: TaskEventStoreTable;
     createdAt: Date;
     completedAt: Date | null;
   }) {
-    const span = await eventRepository.getSpan({
-      storeTable: eventStore,
+    const span = await eventRepository.getSpan(
+      eventStore,
+      environmentId,
       spanId,
       environmentId,
-      startCreatedAt: createdAt,
-      endCreatedAt: completedAt ?? undefined,
-      options: { includeDebugLogs: true },
-    });
+      createdAt,
+      completedAt ?? undefined,
+      { includeDebugLogs: true }
+    );
 
     if (!span) {
       return;

apps/webapp/app/routes/api.v1.runs.$runId.events.ts

@@ -2,7 +2,7 @@ import { json } from "@remix-run/server-runtime";
 import { z } from "zod";
 import { getTaskEventStoreTableForRun } from "~/v3/taskEventStore.server";
 import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server";
-import { eventRepository } from "~/v3/eventRepository.server";
+import { eventRepository } from "~/v3/eventRepository/eventRepository.server";
 import { ApiRetrieveRunPresenter } from "~/presenters/v3/ApiRetrieveRunPresenter.server";
 
 const ParamsSchema = z.object({
@@ -30,9 +30,10 @@ export const loader = createLoaderApiRoute(
       superScopes: ["read:runs", "read:all", "admin"],
     },
   },
-  async ({ resource: run }) => {
+  async ({ resource: run, authentication }) => {
     const runEvents = await eventRepository.getRunEvents(
       getTaskEventStoreTableForRun(run),
+      authentication.environment.id,
       run.friendlyId,
       run.createdAt,
       run.completedAt ?? undefined

apps/webapp/app/routes/api.v1.runs.$runId.trace.ts

@@ -3,7 +3,7 @@ import { BatchId } from "@trigger.dev/core/v3/isomorphic";
 import { z } from "zod";
 import { $replica } from "~/db.server";
 import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server";
-import { eventRepository } from "~/v3/eventRepository.server";
+import { eventRepository } from "~/v3/eventRepository/eventRepository.server";
 import { getTaskEventStoreTableForRun } from "~/v3/taskEventStore.server";
 
 const ParamsSchema = z.object({
@@ -35,9 +35,10 @@ export const loader = createLoaderApiRoute(
       superScopes: ["read:runs", "read:all", "admin"],
     },
   },
-  async ({ resource: run }) => {
+  async ({ resource: run, authentication }) => {
     const traceSummary = await eventRepository.getTraceDetailedSummary(
       getTaskEventStoreTableForRun(run),
+      authentication.environment.id,
       run.traceId,
       run.createdAt,
       run.completedAt ?? undefined

apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.logs.debug.ts

@@ -9,7 +9,7 @@ import { z } from "zod";
 import { prisma } from "~/db.server";
 import { logger } from "~/services/logger.server";
 import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
-import { recordRunDebugLog } from "~/v3/eventRepository.server";
+import { recordRunDebugLog } from "~/v3/eventRepository/eventRepository.server";
 
 // const { action } = createActionApiRoute(
 //   {

apps/webapp/app/routes/engine.v1.dev.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.start.ts

@@ -1,5 +1,5 @@
 import { json, TypedResponse } from "@remix-run/server-runtime";
-import { MachinePreset } from "@trigger.dev/core/v3";
+import { MachinePreset, SemanticInternalAttributes } from "@trigger.dev/core/v3";
 import { RunId, SnapshotId } from "@trigger.dev/core/v3/isomorphic";
 import {
   WorkerApiRunAttemptStartRequestBody,
@@ -57,7 +57,8 @@ const { action } = createActionApiRoute(
       const envVars = await getEnvVars(
         authentication.environment,
         engineResult.run.id,
-        engineResult.execution.machine ?? defaultMachinePreset
+        engineResult.execution.machine ?? defaultMachinePreset,
+        engineResult.run.taskEventStore
       );
 
       return json({
@@ -77,7 +78,8 @@ const { action } = createActionApiRoute(
 async function getEnvVars(
   environment: RuntimeEnvironment,
   runId: string,
-  machinePreset: MachinePreset
+  machinePreset: MachinePreset,
+  taskEventStore?: string
 ): Promise<Record<string, string>> {
   const variables = await resolveVariablesForEnvironment(environment);
 
@@ -94,6 +96,19 @@ async function getEnvVars(
     ]
   );
 
+  if (taskEventStore) {
+    const resourceAttributes = JSON.stringify({
+      [SemanticInternalAttributes.TASK_EVENT_STORE]: taskEventStore,
+    });
+
+    variables.push(
+      ...[
+        { key: "OTEL_RESOURCE_ATTRIBUTES", value: resourceAttributes },
+        { key: "TRIGGER_OTEL_RESOURCE_ATTRIBUTES", value: resourceAttributes },
+      ]
+    );
+  }
+
   return variables.reduce((acc: Record<string, string>, curr) => {
     acc[curr.key] = curr.value;
     return acc;

apps/webapp/app/routes/engine.v1.worker-actions.runs.$runFriendlyId.logs.debug.ts

@@ -3,7 +3,7 @@ import { RunId } from "@trigger.dev/core/v3/isomorphic";
 import { WorkerApiDebugLogBody } from "@trigger.dev/core/v3/runEngineWorker";
 import { z } from "zod";
 import { createActionWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
-import { recordRunDebugLog } from "~/v3/eventRepository.server";
+import { recordRunDebugLog } from "~/v3/eventRepository/eventRepository.server";
 
 export const action = createActionWorkerApiRoute(
   {

apps/webapp/app/routes/otel.v1.logs.ts

@@ -9,7 +9,7 @@ export async function action({ request }: ActionFunctionArgs) {
     if (contentType === "application/json") {
       const body = await request.json();
 
-      const exportResponse = await otlpExporter.exportLogs(body as ExportLogsServiceRequest, false);
+      const exportResponse = await otlpExporter.exportLogs(body as ExportLogsServiceRequest);
 
       return json(exportResponse, { status: 200 });
     } else if (contentType === "application/x-protobuf") {

apps/webapp/app/routes/otel.v1.traces.ts

@@ -9,10 +9,7 @@ export async function action({ request }: ActionFunctionArgs) {
     if (contentType === "application/json") {
       const body = await request.json();
 
-      const exportResponse = await otlpExporter.exportTraces(
-        body as ExportTraceServiceRequest,
-        false
-      );
+      const exportResponse = await otlpExporter.exportTraces(body as ExportTraceServiceRequest);
 
       return json(exportResponse, { status: 200 });
     } else if (contentType === "application/x-protobuf") {