Add support for subscribing to run tags

Author: ericallam

apps/webapp/app/routes/realtime.v1.tags.$tagName.ts

@@ -0,0 +1,64 @@
+import { type ActionFunctionArgs, json } from "@remix-run/server-runtime";
+import { z } from "zod";
+import { $replica } from "~/db.server";
+import { permittedToReadRun, permittedToReadRunTag } from "~/services/accessControl.server";
+import { authenticateApiRequest } from "~/services/apiAuth.server";
+import { logger } from "~/services/logger.server";
+import { realtimeClient } from "~/services/realtimeClientGlobal.server";
+import { makeApiCors } from "~/utils/apiCors";
+
+const ParamsSchema = z.object({
+  tagName: z.string(),
+});
+
+export async function loader({ request, params }: ActionFunctionArgs) {
+  const apiCors = makeApiCors(request);
+
+  if (request.method.toUpperCase() === "OPTIONS") {
+    return apiCors(json({}));
+  }
+
+  // Authenticate the request
+  const authenticationResult = await authenticateApiRequest(request, { allowJWT: true });
+
+  if (!authenticationResult) {
+    return apiCors(json({ error: "Invalid or Missing API Key" }, { status: 401 }));
+  }
+
+  const parsedParams = ParamsSchema.safeParse(params);
+
+  if (!parsedParams.success) {
+    return apiCors(
+      json(
+        { error: "Invalid request parameters", issues: parsedParams.error.issues },
+        { status: 400 }
+      )
+    );
+  }
+
+  if (!permittedToReadRunTag(authenticationResult, parsedParams.data.tagName)) {
+    return apiCors(json({ error: "Unauthorized" }, { status: 403 }));
+  }
+
+  try {
+    return realtimeClient.streamRunsWhere(
+      request.url,
+      authenticationResult.environment,
+      `"runTags" @> ARRAY['${parsedParams.data.tagName}']`,
+      apiCors
+    );
+  } catch (error) {
+    if (error instanceof Response) {
+      // Error responses from longPollingFetch
+      return apiCors(error);
+    } else if (error instanceof TypeError) {
+      // Unexpected errors
+      logger.error("Unexpected error in loader:", { error: error.message });
+      return apiCors(new Response("An unexpected error occurred", { status: 500 }));
+    } else {
+      // Unknown errors
+      logger.error("Unknown error occurred in loader, not Error", { error: JSON.stringify(error) });
+      return apiCors(new Response("An unknown error occurred", { status: 500 }));
+    }
+  }
+}

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

@@ -42,6 +42,39 @@ export function permittedToReadRun(
   return false;
 }
 
+export function permittedToReadRunTag(
+  authenticationResult: ApiAuthenticationResult,
+  tagName: string
+): boolean {
+  if (authenticationResult.type === "PRIVATE") {
+    return true;
+  }
+
+  if (authenticationResult.type === "PUBLIC") {
+    return true;
+  }
+
+  if (!authenticationResult.claims) {
+    return false;
+  }
+
+  const parsedClaims = ClaimsSchema.safeParse(authenticationResult.claims);
+
+  if (!parsedClaims.success) {
+    return false;
+  }
+
+  if (parsedClaims.data.permissions?.includes("read:runs")) {
+    return true;
+  }
+
+  if (parsedClaims.data.permissions?.includes(`read:tags:${tagName}`)) {
+    return true;
+  }
+
+  return false;
+}
+
 export function permittedToReadBatch(
   authenticationResult: ApiAuthenticationResult,
   batchId: string

apps/webapp/test/realtimeClient.test.ts

@@ -134,4 +134,82 @@ describe("RealtimeClient", () => {
       expect(liveResponse2.status).toBe(429);
     }
   );
+
+  containerWithElectricTest(
+    "Should support subscribing to a run tag",
+    { timeout: 30_000 },
+    async ({ redis, electricOrigin, prisma }) => {
+      const client = new RealtimeClient({
+        electricOrigin,
+        keyPrefix: "test:realtime",
+        redis: redis.options,
+        expiryTimeInSeconds: 5,
+        cachedLimitProvider: {
+          async getCachedLimit() {
+            return 1;
+          },
+        },
+      });
+
+      const organization = await prisma.organization.create({
+        data: {
+          title: "test-org",
+          slug: "test-org",
+        },
+      });
+
+      const project = await prisma.project.create({
+        data: {
+          name: "test-project",
+          slug: "test-project",
+          organizationId: organization.id,
+          externalRef: "test-project",
+        },
+      });
+
+      const environment = await prisma.runtimeEnvironment.create({
+        data: {
+          projectId: project.id,
+          organizationId: organization.id,
+          slug: "test",
+          type: "DEVELOPMENT",
+          shortcode: "1234",
+          apiKey: "tr_dev_1234",
+          pkApiKey: "pk_test_1234",
+        },
+      });
+
+      const run = await prisma.taskRun.create({
+        data: {
+          taskIdentifier: "test-task",
+          friendlyId: "run_1234",
+          payload: "{}",
+          payloadType: "application/json",
+          traceId: "trace_1234",
+          spanId: "span_1234",
+          queue: "test-queue",
+          projectId: project.id,
+          runtimeEnvironmentId: environment.id,
+          runTags: ["test:tag:1234", "test:tag:5678"],
+        },
+      });
+
+      const response = await client.streamRunsWhere(
+        "http://localhost:3000?offset=-1",
+        environment,
+        `"runTags" @> ARRAY['test:tag:1234']`
+      );
+
+      const responseBody = await response.json();
+
+      const headers = Object.fromEntries(response.headers.entries());
+
+      const shapeId = headers["electric-shape-id"];
+      const chunkOffset = headers["electric-chunk-last-offset"];
+
+      expect(response.status).toBe(200);
+      expect(shapeId).toBeDefined();
+      expect(chunkOffset).toBe("0_0");
+    }
+  );
 });

docker/docker-compose.yml

@@ -60,7 +60,7 @@ services:
       - 6379:6379
 
   electric:
-    image: electricsql/electric:0.7.3
+    image: electricsql/electric:0.7.5
     restart: always
     environment:
       DATABASE_URL: postgresql://postgres:postgres@database:5432/postgres?sslmode=disable

internal-packages/testcontainers/src/utils.ts

@@ -55,7 +55,7 @@ export async function createElectricContainer(
     network.getName()
   )}:5432/${postgresContainer.getDatabase()}?sslmode=disable`;
 
-  const container = await new GenericContainer("electricsql/electric:0.7.3")
+  const container = await new GenericContainer("electricsql/electric:0.7.5")
     .withExposedPorts(3000)
     .withNetwork(network)
     .withEnvironment({

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

@@ -54,7 +54,7 @@ import {
   UpdateEnvironmentVariableParams,
 } from "./types.js";
 import { generateJWT } from "../jwt.js";
-import { TriggerJwtOptions } from "../types/tasks.js";
+import { AnyRunTypes, TriggerJwtOptions } from "../types/tasks.js";
 
 export type {
   CreateEnvironmentVariableParams,
@@ -590,15 +590,22 @@ export class ApiClient {
     );
   }
 
-  subscribeToRunChanges<TPayload = any, TOutput = any>(runId: string) {
-    return runShapeStream<TPayload, TOutput>(`${this.baseUrl}/realtime/v1/runs/${runId}`, {
+  subscribeToRunChanges<TRunTypes extends AnyRunTypes>(runId: string) {
+    return runShapeStream<TRunTypes>(`${this.baseUrl}/realtime/v1/runs/${runId}`, {
       closeOnComplete: true,
       headers: this.#getRealtimeHeaders(),
     });
   }
 
-  subscribeToBatchChanges<TPayload = any, TOutput = any>(batchId: string) {
-    return runShapeStream<TPayload, TOutput>(`${this.baseUrl}/realtime/v1/batches/${batchId}`, {
+  subscribeToRunTag<TRunTypes extends AnyRunTypes>(tag: string) {
+    return runShapeStream<TRunTypes>(`${this.baseUrl}/realtime/v1/tags/${tag}`, {
+      closeOnComplete: false,
+      headers: this.#getRealtimeHeaders(),
+    });
+  }
+
+  subscribeToBatchChanges<TRunTypes extends AnyRunTypes>(batchId: string) {
+    return runShapeStream<TRunTypes>(`${this.baseUrl}/realtime/v1/batches/${batchId}`, {
       closeOnComplete: false,
       headers: this.#getRealtimeHeaders(),
     });

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

@@ -1,45 +1,47 @@
 import { DeserializedJson } from "../../schemas/json.js";
 import { RunStatus, SubscribeRunRawShape } from "../schemas/api.js";
 import { SerializedError } from "../schemas/common.js";
-import { AnyTask, TaskOutput, TaskPayload } from "../types/tasks.js";
+import { AnyRunTypes, AnyTask, InferRunTypes } from "../types/tasks.js";
 import {
   conditionallyImportAndParsePacket,
   IOPacket,
   parsePacket,
 } from "../utils/ioSerialization.js";
 import { AsyncIterableStream, createAsyncIterableStream, zodShapeStream } from "./stream.js";
 
-export type RunShape<TPayload = any, TOutput = any> = {
-  id: string;
-  createdAt: Date;
-  updatedAt: Date;
-  taskIdentifier: string;
-  number: number;
-  status: RunStatus;
-  durationMs: number;
-  costInCents: number;
-  baseCostInCents: number;
-  payload: TPayload;
-  tags: string[];
-  idempotencyKey?: string;
-  expiredAt?: Date;
-  ttl?: string;
-  finishedAt?: Date;
-  startedAt?: Date;
-  delayedUntil?: Date;
-  queuedAt?: Date;
-  metadata?: Record<string, DeserializedJson>;
-  error?: SerializedError;
-  output?: TOutput;
-  isTest: boolean;
-};
+export type RunShape<TRunTypes extends AnyRunTypes> = TRunTypes extends AnyRunTypes
+  ? {
+      id: string;
+      taskIdentifier: TRunTypes["taskIdentifier"];
+      payload: TRunTypes["payload"];
+      output?: TRunTypes["output"];
+      createdAt: Date;
+      updatedAt: Date;
+      number: number;
+      status: RunStatus;
+      durationMs: number;
+      costInCents: number;
+      baseCostInCents: number;
+      tags: string[];
+      idempotencyKey?: string;
+      expiredAt?: Date;
+      ttl?: string;
+      finishedAt?: Date;
+      startedAt?: Date;
+      delayedUntil?: Date;
+      queuedAt?: Date;
+      metadata?: Record<string, DeserializedJson>;
+      error?: SerializedError;
+      isTest: boolean;
+    }
+  : never;
 
-export type AnyRunShape = RunShape<any, any>;
+export type AnyRunShape = RunShape<AnyRunTypes>;
 
-export type TaskRunShape<TTask extends AnyTask> = RunShape<TaskPayload<TTask>, TaskOutput<TTask>>;
+export type TaskRunShape<TTask extends AnyTask> = RunShape<InferRunTypes<TTask>>;
 
-export type RunStreamCallback<TPayload = any, TOutput = any> = (
-  run: RunShape<TPayload, TOutput>
+export type RunStreamCallback<TRunTypes extends AnyRunTypes> = (
+  run: RunShape<TRunTypes>
 ) => void | Promise<void>;
 
 export type RunShapeStreamOptions = {
@@ -48,19 +50,19 @@ export type RunShapeStreamOptions = {
   closeOnComplete?: boolean;
 };
 
-export function runShapeStream<TPayload = any, TOutput = any>(
+export function runShapeStream<TRunTypes extends AnyRunTypes>(
   url: string,
   options?: RunShapeStreamOptions
-): RunSubscription<TPayload, TOutput> {
-  const subscription = new RunSubscription<TPayload, TOutput>(url, options);
+): RunSubscription<TRunTypes> {
+  const subscription = new RunSubscription<TRunTypes>(url, options);
 
   return subscription.init();
 }
 
-export class RunSubscription<TPayload = any, TOutput = any> {
+export class RunSubscription<TRunTypes extends AnyRunTypes> {
   private abortController: AbortController;
   private unsubscribeShape: () => void;
-  private stream: AsyncIterableStream<RunShape<TPayload, TOutput>>;
+  private stream: AsyncIterableStream<RunShape<TRunTypes>>;
   private packetCache = new Map<string, any>();
 
   constructor(
@@ -117,15 +119,15 @@ export class RunSubscription<TPayload = any, TOutput = any> {
     this.unsubscribeShape?.();
   }
 
-  [Symbol.asyncIterator](): AsyncIterator<RunShape<TPayload, TOutput>> {
+  [Symbol.asyncIterator](): AsyncIterator<RunShape<TRunTypes>> {
     return this.stream[Symbol.asyncIterator]();
   }
 
-  getReader(): ReadableStreamDefaultReader<RunShape<TPayload, TOutput>> {
+  getReader(): ReadableStreamDefaultReader<RunShape<TRunTypes>> {
     return this.stream.getReader();
   }
 
-  private async transformRunShape(row: SubscribeRunRawShape): Promise<RunShape> {
+  private async transformRunShape(row: SubscribeRunRawShape): Promise<RunShape<TRunTypes>> {
     const payloadPacket = row.payloadType
       ? ({ data: row.payload ?? undefined, dataType: row.payloadType } satisfies IOPacket)
       : undefined;
@@ -183,7 +185,7 @@ export class RunSubscription<TPayload = any, TOutput = any> {
       error: row.error ?? undefined,
       isTest: row.isTest,
       metadata,
-    };
+    } as RunShape<TRunTypes>;
   }
 }