Queue an admin suspension's token rotations in one statement

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: Brooooooklyn

platform/packages/api/src/db/admin.ts

@@ -16,7 +16,10 @@ import {
 import { alias } from 'drizzle-orm/pg-core';
 import { lockAuthSettings } from '../auth/configuration.ts';
 import { generateId } from '../id.ts';
-import { requestProxyTokenRotation } from '../proxy-token-rotation-request.ts';
+import {
+  requestProxyTokenRotation,
+  requestProxyTokenRotationsForMember,
+} from '../proxy-token-rotation-request.ts';
 import {
   users,
   projects,
@@ -178,16 +181,7 @@ async function withAdminRemovalGuard<T>(
 
 /** An admin suspension ends deploy access, like account deletion. */
 async function requestRotationsForAdminSuspension(tx: Db, id: string): Promise<void> {
-  const deployMemberships = await tx
-    .select({ projectId: projectMembers.project_id })
-    .from(projectMembers)
-    .where(
-      and(eq(projectMembers.user_id, id), inArray(projectMembers.role, ['collaborator', 'admin'])),
-    )
-    .orderBy(asc(projectMembers.project_id));
-  for (const { projectId } of deployMemberships) {
-    await requestProxyTokenRotation(tx, projectId, 'user-suspended', 'member');
-  }
+  await requestProxyTokenRotationsForMember(tx, id, 'user-suspended');
 }
 
 /**

platform/packages/api/src/proxy-token-rotation-request.ts

@@ -1,6 +1,6 @@
-import { sql } from 'drizzle-orm';
+import { and, eq, inArray, sql } from 'drizzle-orm';
 import type { DbOrTx } from './db/client.ts';
-import { projectCleanupJobs } from './schema.ts';
+import { projectCleanupJobs, projectMembers } from './schema.ts';
 
 // Rotation requests, apart from the runner in `proxy-token-rotation.ts` so the
 // `db/` modules can write one inside their transactions without importing
@@ -26,6 +26,13 @@ export function rotationLog(event: string, projectId: string, detail = ''): void
   console.log(`[proxy-token-rotation] ${event} ${projectId}${detail ? ` ${detail}` : ''}`);
 }
 
+// What a request does to an existing job. Both writers below share these, so
+// one statement for many projects leaves the rows one request per project would.
+const jobStarted = sql`${projectCleanupJobs.payload}->>'phase' <> 'requested'`;
+const flagAgain = sql`CASE WHEN ${projectCleanupJobs.payload}->>'phase' <> 'requested'
+    THEN ${projectCleanupJobs.payload} || '{"again":"1"}'::jsonb
+    ELSE ${projectCleanupJobs.payload} END`;
+
 /**
  * Asks for a new token generation. A job that has not started yet already
  * covers the request; one past START may have handed out tok(T), so it
@@ -38,9 +45,6 @@ export async function requestProxyTokenRotation(
   reason: ProxyTokenRotationReason,
   mode: 'member' | 'operator',
 ): Promise<void> {
-  const flagAgain = sql`CASE WHEN ${projectCleanupJobs.payload}->>'phase' <> 'requested'
-    THEN ${projectCleanupJobs.payload} || '{"again":"1"}'::jsonb
-    ELSE ${projectCleanupJobs.payload} END`;
   const rows = await tx
     .insert(projectCleanupJobs)
     .values({
@@ -56,12 +60,39 @@ export async function requestProxyTokenRotation(
         mode === 'operator'
           ? { payload: flagAgain, next_attempt_at: sql`now()` }
           : { payload: flagAgain },
-      ...(mode === 'member'
-        ? { setWhere: sql`${projectCleanupJobs.payload}->>'phase' <> 'requested'` }
-        : {}),
+      ...(mode === 'member' ? { setWhere: jobStarted } : {}),
     })
     .returning({ phase: sql<string>`${projectCleanupJobs.payload}->>'phase'` });
   if (rows.length > 0) {
     rotationLog('requested', projectId, `reason=${reason}`);
   }
 }
+
+/**
+ * `requestProxyTokenRotation` in `member` mode for every project where the
+ * user is a collaborator or admin, in one statement whatever the number of
+ * projects. Call inside the caller's transaction.
+ */
+export async function requestProxyTokenRotationsForMember(
+  tx: DbOrTx,
+  userId: string,
+  reason: ProxyTokenRotationReason,
+): Promise<void> {
+  const deployMember = and(
+    eq(projectMembers.user_id, userId),
+    inArray(projectMembers.role, ['collaborator', 'admin']),
+  );
+  const result = await tx.execute<{ project_id: string }>(sql`
+    INSERT INTO ${projectCleanupJobs} (target, project_id, kind, payload, next_attempt_at)
+    SELECT ${PROXY_TOKEN_ROTATION_JOB_KIND}::text || ':' || ${projectMembers.project_id},
+      ${projectMembers.project_id}, ${PROXY_TOKEN_ROTATION_JOB_KIND}::text,
+      jsonb_build_object('phase', 'requested', 'reason', ${reason}::text), now()
+    FROM ${projectMembers}
+    WHERE ${deployMember}
+    ORDER BY ${projectMembers.project_id}
+    ON CONFLICT (target) DO UPDATE SET payload = ${flagAgain} WHERE ${jobStarted}
+    RETURNING project_id`);
+  for (const { project_id: projectId } of result.rows) {
+    rotationLog('requested', projectId, `reason=${reason}`);
+  }
+}

platform/packages/api/test/integration/project-team-lifecycle.test.ts

@@ -1,4 +1,4 @@
-import { eq } from 'drizzle-orm';
+import { eq, sql } from 'drizzle-orm';
 import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vite-plus/test';
 import {
   createGitHubConnection,
@@ -14,6 +14,7 @@ import {
   respondToProjectInvitation,
   suspendUser,
 } from '../../src/db.ts';
+import { proxyTokenRotationTarget } from '../../src/proxy-token-rotation-request.ts';
 import {
   projectAccessEvents,
   projectCleanupJobs,
@@ -196,6 +197,75 @@ describe('project team lifecycle', () => {
     }
   });
 
+  it('queues the rotations of an admin suspension in one statement, keeping queued jobs', async () => {
+    const owner = await createTestUser(env.db, 'many-owner');
+    const member = await createTestUser(env.db, 'many-member');
+    const collaborated = await createProject(env.db, owner.id, 'many-collaborator');
+    const administered = await createProject(env.db, owner.id, 'many-admin');
+    const read = await createProject(env.db, owner.id, 'many-reader');
+    const queued = await createProject(env.db, owner.id, 'many-queued');
+    const started = await createProject(env.db, owner.id, 'many-started');
+    await env.db.insert(projectMembers).values([
+      { project_id: collaborated.id, user_id: member.id, role: 'collaborator' },
+      { project_id: administered.id, user_id: member.id, role: 'admin' },
+      { project_id: read.id, user_id: member.id, role: 'reader' },
+      { project_id: queued.id, user_id: member.id, role: 'collaborator' },
+      { project_id: started.id, user_id: member.id, role: 'admin' },
+    ]);
+    const later = sql`now() + interval '1 hour'`;
+    await env.db.insert(projectCleanupJobs).values([
+      {
+        target: proxyTokenRotationTarget(queued.id),
+        project_id: queued.id,
+        kind: 'proxy-token-rotation',
+        payload: { phase: 'requested', reason: 'member-role-changed' },
+        next_attempt_at: later,
+      },
+      {
+        target: proxyTokenRotationTarget(started.id),
+        project_id: started.id,
+        kind: 'proxy-token-rotation',
+        payload: { phase: 'patch', reason: 'operator', gen: '1', cursor: '' },
+        attempts: 2,
+        next_attempt_at: later,
+      },
+    ]);
+    const before = new Map(
+      (await env.db.select().from(projectCleanupJobs)).map((job) => [job.project_id, job]),
+    );
+
+    const { queries } = await captureQueries(env.db, () =>
+      suspendUser(env.db, member.id, 'admin_suspended'),
+    );
+
+    const jobs = new Map(
+      (await env.db.select().from(projectCleanupJobs)).map((job) => [job.project_id, job]),
+    );
+    expect([...jobs.keys()].sort()).toEqual(
+      [collaborated.id, administered.id, queued.id, started.id].sort(),
+    );
+    for (const projectId of [collaborated.id, administered.id]) {
+      expect(jobs.get(projectId)).toMatchObject({
+        target: proxyTokenRotationTarget(projectId),
+        kind: 'proxy-token-rotation',
+        payload: { phase: 'requested', reason: 'user-suspended' },
+        attempts: 0,
+        last_error: null,
+      });
+    }
+    // A job that has not started covers the request as it is.
+    expect(jobs.get(queued.id)).toEqual(before.get(queued.id));
+    // A started job finishes and runs once more, on its own schedule.
+    const startedBefore = before.get(started.id)!;
+    expect(jobs.get(started.id)).toEqual({
+      ...startedBefore,
+      payload: { ...startedBefore.payload, again: '1' },
+    });
+    expect(
+      queries.filter((query) => /^insert into "project_cleanup_jobs"/i.test(query.trim())),
+    ).toHaveLength(1);
+  });
+
   it('atomically discovers and locks pending-invitation projects before user deletion', async () => {
     const owner = await createTestUser(env.db, 'owner');
     const invitee = await createTestUser(env.db, 'invitee');