Filter realtime colums and expose ability to skip some columns

Author: ericallam

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

@@ -13,6 +13,37 @@ export interface CachedLimitProvider {
 
 const MAXIMUM_CREATED_AT_FILTER_AGE_IN_MS = 7 * 24 * 60 * 60 * 1000;
 
+const DEFAULT_ELECTRIC_COLUMNS = [
+  "id",
+  "taskIdentifier",
+  "createdAt",
+  "updatedAt",
+  "startedAt",
+  "delayUntil",
+  "queuedAt",
+  "expiredAt",
+  "completedAt",
+  "friendlyId",
+  "number",
+  "isTest",
+  "status",
+  "usageDurationMs",
+  "costInCents",
+  "baseCostInCents",
+  "ttl",
+  "payload",
+  "payloadType",
+  "metadata",
+  "metadataType",
+  "output",
+  "outputType",
+  "runTags",
+  "error",
+];
+
+const RESERVED_COLUMNS = ["id", "taskIdentifier", "friendlyId", "status", "createdAt"];
+const RESERVED_SEARCH_PARAMS = ["createdAt", "tags", "skipColumns"];
+
 export type RealtimeClientOptions = {
   electricOrigin: string;
   redis: RedisWithClusterOptions;
@@ -202,6 +233,10 @@ export class RealtimeClient {
 
     // Copy over all the url search params to the electric url
     $url.searchParams.forEach((value, key) => {
+      if (RESERVED_SEARCH_PARAMS.includes(key)) {
+        return;
+      }
+
       electricUrl.searchParams.set(key, value);
     });
 
@@ -214,6 +249,27 @@ export class RealtimeClient {
       electricUrl.searchParams.set("handle", electricUrl.searchParams.get("shape_id") ?? "");
     }
 
+    const skipColumnsRaw = $url.searchParams.get("skipColumns");
+
+    if (skipColumnsRaw) {
+      const skipColumns = skipColumnsRaw
+        .split(",")
+        .map((c) => c.trim())
+        .filter((c) => c !== "" && !RESERVED_COLUMNS.includes(c));
+
+      electricUrl.searchParams.set(
+        "columns",
+        DEFAULT_ELECTRIC_COLUMNS.filter((c) => !skipColumns.includes(c))
+          .map((c) => `"${c}"`)
+          .join(",")
+      );
+    } else {
+      electricUrl.searchParams.set(
+        "columns",
+        DEFAULT_ELECTRIC_COLUMNS.map((c) => `"${c}"`).join(",")
+      );
+    }
+
     return electricUrl;
   }
 

packages/core/src/v3/apiClient/index.ts

@@ -66,6 +66,7 @@ import {
   SSEStreamSubscriptionFactory,
   TaskRunShape,
   runShapeStream,
+  RealtimeRunSkipColumns,
 } from "./runStream.js";
 import {
   CreateEnvironmentVariableParams,
@@ -88,6 +89,7 @@ export type {
   ImportEnvironmentVariablesParams,
   SubscribeToRunsQueryParams,
   UpdateEnvironmentVariableParams,
+  RealtimeRunSkipColumns,
 };
 
 export type ClientTriggerOptions = {
@@ -890,26 +892,37 @@ export class ApiClient {
       signal?: AbortSignal;
       closeOnComplete?: boolean;
       onFetchError?: (error: Error) => void;
+      skipColumns?: string[];
     }
   ) {
-    return runShapeStream<TRunTypes>(`${this.baseUrl}/realtime/v1/runs/${runId}`, {
-      closeOnComplete:
-        typeof options?.closeOnComplete === "boolean" ? options.closeOnComplete : true,
-      headers: this.#getRealtimeHeaders(),
-      client: this,
-      signal: options?.signal,
-      onFetchError: options?.onFetchError,
-    });
+    const queryParams = new URLSearchParams();
+
+    if (options?.skipColumns) {
+      queryParams.append("skipColumns", options.skipColumns.join(","));
+    }
+
+    return runShapeStream<TRunTypes>(
+      `${this.baseUrl}/realtime/v1/runs/${runId}${queryParams ? `?${queryParams}` : ""}`,
+      {
+        closeOnComplete:
+          typeof options?.closeOnComplete === "boolean" ? options.closeOnComplete : true,
+        headers: this.#getRealtimeHeaders(),
+        client: this,
+        signal: options?.signal,
+        onFetchError: options?.onFetchError,
+      }
+    );
   }
 
   subscribeToRunsWithTag<TRunTypes extends AnyRunTypes>(
     tag: string | string[],
-    filters?: { createdAt?: string },
+    filters?: { createdAt?: string; skipColumns?: string[] },
     options?: { signal?: AbortSignal; onFetchError?: (error: Error) => void }
   ) {
     const searchParams = createSearchQueryForSubscribeToRuns({
       tags: tag,
       ...(filters ? { createdAt: filters.createdAt } : {}),
+      ...(filters?.skipColumns ? { skipColumns: filters.skipColumns } : {}),
     });
 
     return runShapeStream<TRunTypes>(
@@ -926,15 +939,28 @@ export class ApiClient {
 
   subscribeToBatch<TRunTypes extends AnyRunTypes>(
     batchId: string,
-    options?: { signal?: AbortSignal; onFetchError?: (error: Error) => void }
+    options?: {
+      signal?: AbortSignal;
+      onFetchError?: (error: Error) => void;
+      skipColumns?: string[];
+    }
   ) {
-    return runShapeStream<TRunTypes>(`${this.baseUrl}/realtime/v1/batches/${batchId}`, {
-      closeOnComplete: false,
-      headers: this.#getRealtimeHeaders(),
-      client: this,
-      signal: options?.signal,
-      onFetchError: options?.onFetchError,
-    });
+    const queryParams = new URLSearchParams();
+
+    if (options?.skipColumns) {
+      queryParams.append("skipColumns", options.skipColumns.join(","));
+    }
+
+    return runShapeStream<TRunTypes>(
+      `${this.baseUrl}/realtime/v1/batches/${batchId}${queryParams ? `?${queryParams}` : ""}`,
+      {
+        closeOnComplete: false,
+        headers: this.#getRealtimeHeaders(),
+        client: this,
+        signal: options?.signal,
+        onFetchError: options?.onFetchError,
+      }
+    );
   }
 
   async fetchStream<T>(
@@ -1049,6 +1075,10 @@ function createSearchQueryForSubscribeToRuns(query?: SubscribeToRunsQueryParams)
     if (query.createdAt) {
       searchParams.append("createdAt", query.createdAt);
     }
+
+    if (query.skipColumns) {
+      searchParams.append("skipColumns", query.skipColumns.join(","));
+    }
   }
 
   return searchParams;

packages/core/src/v3/apiClient/runStream.ts

@@ -55,6 +55,27 @@ export type TaskRunShape<TTask extends AnyTask> = RunShape<InferRunTypes<TTask>>
 export type RealtimeRun<TTask extends AnyTask> = TaskRunShape<TTask>;
 export type AnyRealtimeRun = RealtimeRun<AnyTask>;
 
+export type RealtimeRunSkipColumns = Array<
+  | "startedAt"
+  | "delayUntil"
+  | "queuedAt"
+  | "expiredAt"
+  | "completedAt"
+  | "number"
+  | "isTest"
+  | "usageDurationMs"
+  | "costInCents"
+  | "baseCostInCents"
+  | "ttl"
+  | "payload"
+  | "payloadType"
+  | "metadata"
+  | "output"
+  | "outputType"
+  | "runTags"
+  | "error"
+>;
+
 export type RunStreamCallback<TRunTypes extends AnyRunTypes> = (
   run: RunShape<TRunTypes>
 ) => void | Promise<void>;
@@ -395,16 +416,16 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
 
     return {
       id: row.friendlyId,
-      payload,
-      output,
       createdAt: row.createdAt,
       updatedAt: row.updatedAt,
       taskIdentifier: row.taskIdentifier,
-      number: row.number,
       status: apiStatusFromRunStatus(row.status),
-      durationMs: row.usageDurationMs,
-      costInCents: row.costInCents,
-      baseCostInCents: row.baseCostInCents,
+      payload,
+      output,
+      number: row.number ?? 0,
+      durationMs: row.usageDurationMs ?? 0,
+      costInCents: row.costInCents ?? 0,
+      baseCostInCents: row.baseCostInCents ?? 0,
       tags: row.runTags ?? [],
       idempotencyKey: row.idempotencyKey ?? undefined,
       expiredAt: row.expiredAt ?? undefined,
@@ -413,7 +434,7 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
       delayedUntil: row.delayUntil ?? undefined,
       queuedAt: row.queuedAt ?? undefined,
       error: row.error ? createJsonErrorObject(row.error) : undefined,
-      isTest: row.isTest,
+      isTest: row.isTest ?? false,
       metadata,
     } as RunShape<TRunTypes>;
   }

packages/core/src/v3/apiClient/stream.ts

@@ -102,19 +102,16 @@ class ReadableShapeStream<T extends Row<unknown> = Row> {
     // Create the source stream that will receive messages
     const source = new ReadableStream<Message<T>[]>({
       start: (controller) => {
-        this.#unsubscribe = this.#stream.subscribe(
-          (messages) => {
-            if (!this.#isStreamClosed) {
-              controller.enqueue(messages);
-            }
-          },
-          this.#handleError.bind(this)
-        );
+        this.#unsubscribe = this.#stream.subscribe((messages) => {
+          if (!this.#isStreamClosed) {
+            controller.enqueue(messages);
+          }
+        }, this.#handleError.bind(this));
       },
       cancel: () => {
         this.#isStreamClosed = true;
         this.#unsubscribe?.();
-      }
+      },
     });
 
     // Create the transformed stream that processes messages and emits complete rows

packages/core/src/v3/apiClient/types.ts

@@ -42,6 +42,7 @@ export interface SubscribeToRunsQueryParams {
   tasks?: Array<string> | string;
   tags?: Array<string> | string;
   createdAt?: string;
+  skipColumns?: string[];
 }
 
 export interface ListWaitpointTokensQueryParams extends CursorPageParams {

packages/core/src/v3/schemas/api.ts

@@ -870,22 +870,22 @@ const RawOptionalShapeDate = z
 
 export const SubscribeRunRawShape = z.object({
   id: z.string(),
-  idempotencyKey: z.string().nullish(),
+  taskIdentifier: z.string(),
+  friendlyId: z.string(),
+  status: z.string(),
   createdAt: RawShapeDate,
   updatedAt: RawShapeDate,
   startedAt: RawOptionalShapeDate,
   delayUntil: RawOptionalShapeDate,
   queuedAt: RawOptionalShapeDate,
   expiredAt: RawOptionalShapeDate,
   completedAt: RawOptionalShapeDate,
-  taskIdentifier: z.string(),
-  friendlyId: z.string(),
-  number: z.number(),
-  isTest: z.boolean(),
-  status: z.string(),
-  usageDurationMs: z.number(),
-  costInCents: z.number(),
-  baseCostInCents: z.number(),
+  idempotencyKey: z.string().nullish(),
+  number: z.number().default(0),
+  isTest: z.boolean().default(false),
+  usageDurationMs: z.number().default(0),
+  costInCents: z.number().default(0),
+  baseCostInCents: z.number().default(0),
   ttl: z.string().nullish(),
   payload: z.string().nullish(),
   payloadType: z.string().nullish(),