Realtime: improve scope access to runs with tags and batches

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

@@ -15,7 +15,7 @@ import assertNever from "assert-never";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { generatePresignedUrl } from "~/v3/r2.server";
 import { BasePresenter } from "./basePresenter.server";
-import { prisma } from "~/db.server";
+import { $replica, prisma } from "~/db.server";
 
 // Build 'select' object
 const commonRunSelect = {
@@ -59,48 +59,46 @@ type CommonRelatedRun = Prisma.Result<
   "findFirstOrThrow"
 >;
 
+type FoundRun = NonNullable<Awaited<ReturnType<typeof ApiRetrieveRunPresenter.findRun>>>;
+
 export class ApiRetrieveRunPresenter extends BasePresenter {
-  public async call(
-    friendlyId: string,
-    env: AuthenticatedEnvironment
-  ): Promise<RetrieveRunResponse | undefined> {
-    return this.traceWithEnv("call", env, async (span) => {
-      const taskRun = await this._replica.taskRun.findFirst({
-        where: {
-          friendlyId,
-          runtimeEnvironmentId: env.id,
-        },
-        include: {
-          attempts: true,
-          lockedToVersion: true,
-          schedule: true,
-          tags: true,
-          batch: {
-            select: {
-              id: true,
-              friendlyId: true,
-            },
+  public static async findRun(friendlyId: string, env: AuthenticatedEnvironment) {
+    return $replica.taskRun.findFirst({
+      where: {
+        friendlyId,
+        runtimeEnvironmentId: env.id,
+      },
+      include: {
+        attempts: true,
+        lockedToVersion: true,
+        schedule: true,
+        tags: true,
+        batch: {
+          select: {
+            id: true,
+            friendlyId: true,
           },
-          parentTaskRun: {
-            select: commonRunSelect,
-          },
-          rootTaskRun: {
-            select: commonRunSelect,
-          },
-          childRuns: {
-            select: {
-              ...commonRunSelect,
-            },
+        },
+        parentTaskRun: {
+          select: commonRunSelect,
+        },
+        rootTaskRun: {
+          select: commonRunSelect,
+        },
+        childRuns: {
+          select: {
+            ...commonRunSelect,
           },
         },
-      });
-
-      if (!taskRun) {
-        logger.debug("Task run not found", { friendlyId, envId: env.id });
-
-        return undefined;
-      }
+      },
+    });
+  }
 
+  public async call(
+    taskRun: FoundRun,
+    env: AuthenticatedEnvironment
+  ): Promise<RetrieveRunResponse | undefined> {
+    return this.traceWithEnv("call", env, async (span) => {
       let $payload: any;
       let $payloadPresignedUrl: string | undefined;
       let $output: any;

apps/webapp/app/routes/api.v1.batches.$batchId.ts

@@ -1,6 +1,6 @@
 import { json } from "@remix-run/server-runtime";
 import { z } from "zod";
-import { ApiRetrieveBatchPresenter } from "~/presenters/v3/ApiRetrieveBatchPresenter.server";
+import { $replica } from "~/db.server";
 import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
 const ParamsSchema = z.object({
@@ -12,20 +12,28 @@ export const loader = createLoaderApiRoute(
     params: ParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    findResource: (params, auth) => {
+      return $replica.batchTaskRun.findFirst({
+        where: {
+          friendlyId: params.batchId,
+          runtimeEnvironmentId: auth.environment.id,
+        },
+      });
+    },
     authorization: {
       action: "read",
-      resource: (params) => ({ batch: params.batchId }),
+      resource: (batch) => ({ batch: batch.friendlyId }),
       superScopes: ["read:runs", "read:all", "admin"],
     },
   },
-  async ({ params, authentication }) => {
-    const presenter = new ApiRetrieveBatchPresenter();
-    const result = await presenter.call(params.batchId, authentication.environment);
-
-    if (!result) {
-      return json({ error: "Batch not found" }, { status: 404 });
-    }
-
-    return json(result);
+  async ({ resource: batch }) => {
+    return json({
+      id: batch.friendlyId,
+      status: batch.status,
+      idempotencyKey: batch.idempotencyKey ?? undefined,
+      createdAt: batch.createdAt,
+      updatedAt: batch.updatedAt,
+      runCount: batch.runCount,
+    });
   }
 );

apps/webapp/app/routes/api.v1.packets.$.ts

@@ -45,6 +45,7 @@ export const loader = createLoaderApiRoute(
     params: ParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    findResource: async () => 1, // This is a dummy function, we don't need to find a resource
   },
   async ({ params, authentication }) => {
     const filename = params["*"];

apps/webapp/app/routes/api.v1.runs.$runParam.reschedule.ts

@@ -61,8 +61,17 @@ export async function action({ request, params }: ActionFunctionArgs) {
       return json({ error: "An unknown error occurred" }, { status: 500 });
     }
 
+    const run = await ApiRetrieveRunPresenter.findRun(
+      updatedRun.friendlyId,
+      authenticationResult.environment
+    );
+
+    if (!run) {
+      return json({ error: "Run not found" }, { status: 404 });
+    }
+
     const presenter = new ApiRetrieveRunPresenter();
-    const result = await presenter.call(updatedRun.friendlyId, authenticationResult.environment);
+    const result = await presenter.call(run, authenticationResult.environment);
 
     if (!result) {
       return json({ error: "Run not found" }, { status: 404 });

apps/webapp/app/routes/api.v1.runs.ts

@@ -12,9 +12,10 @@ export const loader = createLoaderApiRoute(
     corsStrategy: "all",
     authorization: {
       action: "read",
-      resource: (_, searchParams) => ({ tasks: searchParams["filter[taskIdentifier]"] }),
+      resource: (_, __, searchParams) => ({ tasks: searchParams["filter[taskIdentifier]"] }),
       superScopes: ["read:runs", "read:all", "admin"],
     },
+    findResource: async () => 1, // This is a dummy function, we don't need to find a resource
   },
   async ({ searchParams, authentication }) => {
     const presenter = new ApiRunListPresenter();

apps/webapp/app/routes/api.v1.tasks.batch.ts

@@ -133,7 +133,7 @@ async function responseHeaders(
     const claims = {
       sub: environment.id,
       pub: true,
-      scopes: [`read:batch:${batch.id}`].concat(batch.runs.map((r) => `read:runs:${r.id}`)),
+      scopes: [`read:batch:${batch.id}`],
     };
 
     const jwt = await generateJWT({

apps/webapp/app/routes/api.v3.runs.$runId.ts

@@ -12,15 +12,22 @@ export const loader = createLoaderApiRoute(
     params: ParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    findResource: (params, auth) => {
+      return ApiRetrieveRunPresenter.findRun(params.runId, auth.environment);
+    },
     authorization: {
       action: "read",
-      resource: (params) => ({ runs: params.runId }),
+      resource: (run) => ({
+        runs: run.friendlyId,
+        tags: run.runTags,
+        batch: run.batch?.friendlyId,
+      }),
       superScopes: ["read:runs", "read:all", "admin"],
     },
   },
-  async ({ params, authentication }) => {
+  async ({ authentication, resource }) => {
     const presenter = new ApiRetrieveRunPresenter();
-    const result = await presenter.call(params.runId, authentication.environment);
+    const result = await presenter.call(resource, authentication.environment);
 
     if (!result) {
       return json(

apps/webapp/app/routes/realtime.v1.batches.$batchId.ts

@@ -1,4 +1,3 @@
-import { json } from "@remix-run/server-runtime";
 import { z } from "zod";
 import { $replica } from "~/db.server";
 import { realtimeClient } from "~/services/realtimeClientGlobal.server";
@@ -13,24 +12,21 @@ export const loader = createLoaderApiRoute(
     params: ParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    findResource: (params, auth) => {
+      return $replica.batchTaskRun.findFirst({
+        where: {
+          friendlyId: params.batchId,
+          runtimeEnvironmentId: auth.environment.id,
+        },
+      });
+    },
     authorization: {
       action: "read",
-      resource: (params) => ({ batch: params.batchId }),
+      resource: (batch) => ({ batch: batch.friendlyId }),
       superScopes: ["read:runs", "read:all", "admin"],
     },
   },
-  async ({ params, authentication, request }) => {
-    const batchRun = await $replica.batchTaskRun.findFirst({
-      where: {
-        friendlyId: params.batchId,
-        runtimeEnvironmentId: authentication.environment.id,
-      },
-    });
-
-    if (!batchRun) {
-      return json({ error: "Batch not found" }, { status: 404 });
-    }
-
+  async ({ authentication, request, resource: batchRun }) => {
     return realtimeClient.streamBatch(
       request.url,
       authentication.environment,

apps/webapp/app/routes/realtime.v1.runs.$runId.ts

@@ -13,24 +13,32 @@ export const loader = createLoaderApiRoute(
     params: ParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    findResource: async (params, authentication) => {
+      return $replica.taskRun.findFirst({
+        where: {
+          friendlyId: params.runId,
+          runtimeEnvironmentId: authentication.environment.id,
+        },
+        include: {
+          batch: {
+            select: {
+              friendlyId: true,
+            },
+          },
+        },
+      });
+    },
     authorization: {
       action: "read",
-      resource: (params) => ({ runs: params.runId }),
+      resource: (run) => ({
+        runs: run.friendlyId,
+        tags: run.runTags,
+        batch: run.batch?.friendlyId,
+      }),
       superScopes: ["read:runs", "read:all", "admin"],
     },
   },
-  async ({ params, authentication, request }) => {
-    const run = await $replica.taskRun.findFirst({
-      where: {
-        friendlyId: params.runId,
-        runtimeEnvironmentId: authentication.environment.id,
-      },
-    });
-
-    if (!run) {
-      return json({ error: "Run not found" }, { status: 404 });
-    }
-
+  async ({ authentication, request, resource: run }) => {
     return realtimeClient.streamRun(
       request.url,
       authentication.environment,

apps/webapp/app/routes/realtime.v1.runs.ts

@@ -16,9 +16,10 @@ export const loader = createLoaderApiRoute(
     searchParams: SearchParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    findResource: async () => 1, // This is a dummy value, it's not used
     authorization: {
       action: "read",
-      resource: (_, searchParams) => searchParams,
+      resource: (_, __, searchParams) => searchParams,
       superScopes: ["read:runs", "read:all", "admin"],
     },
   },

apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts

@@ -24,24 +24,32 @@ export const loader = createLoaderApiRoute(
     params: ParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    findResource: async (params, auth) => {
+      return $replica.taskRun.findFirst({
+        where: {
+          friendlyId: params.runId,
+          runtimeEnvironmentId: auth.environment.id,
+        },
+        include: {
+          batch: {
+            select: {
+              friendlyId: true,
+            },
+          },
+        },
+      });
+    },
     authorization: {
       action: "read",
-      resource: (params) => ({ runs: params.runId }),
+      resource: (run) => ({
+        runs: run.friendlyId,
+        tags: run.runTags,
+        batch: run.batch?.friendlyId,
+      }),
       superScopes: ["read:runs", "read:all", "admin"],
     },
   },
-  async ({ params, authentication, request }) => {
-    const run = await $replica.taskRun.findFirst({
-      where: {
-        friendlyId: params.runId,
-        runtimeEnvironmentId: authentication.environment.id,
-      },
-    });
-
-    if (!run) {
-      return new Response("Run not found", { status: 404 });
-    }
-
+  async ({ params, request, resource: run }) => {
     return realtimeStreams.streamResponse(run.friendlyId, params.streamId, request.signal);
   }
 );