fix(workspaces): hand off a departing member's files without binding every file id

apps/sim/lib/uploads/contexts/workspace/creator-handoff.ts

@@ -1,31 +1,49 @@
 import { workspaceFiles } from '@sim/db/schema'
-import { asc, inArray, type SQL } from 'drizzle-orm'
+import { and, asc, count, inArray, not, or, type SQL } from 'drizzle-orm'
 import type { DbTransaction } from '@/lib/db/types'
 import { snapshotWorkspaceFileSecretProvenanceInTx } from '@/lib/uploads/contexts/workspace/file-secret-provenance'
 import {
   isVersionHeadCurrent,
   loadWorkspaceFileVersionHeads,
   materializeWorkspaceFileVersionInTx,
+  workspaceFileVersionHeadCurrentSql,
 } from '@/lib/uploads/contexts/workspace/file-version-store'
 
-/** Freezes implicit history under the writer lock before replacing a shared file's live creator. */
+/** Contexts whose files keep version history, frozen before their creator changes. */
+const VERSIONED_FILE_CONTEXTS = ['workspace', 'project']
+
+/**
+ * Freezes implicit history under the writer lock before replacing a shared file's live creator.
+ * Every statement is keyed by `condition`, never by a list of file ids: a departing member can
+ * own more files than Postgres accepts bind parameters.
+ */
 export async function handoffFileCreatorsInTx(
   tx: DbTransaction,
   condition: SQL | undefined,
   successorId: string
 ) {
   if (!condition) throw new Error('File creator handoff requires a scope')
-  const files = await tx
-    .select()
+  // Takes every row lock in id order without loading the rows into memory.
+  const locked = tx
+    .select({ id: workspaceFiles.id })
     .from(workspaceFiles)
     .where(condition)
     .orderBy(asc(workspaceFiles.id))
     .for('update')
-  const versioned = files.filter(
-    (file) => file.context === 'workspace' || file.context === 'project'
-  )
+    .as('locked')
+  await tx.select({ count: count() }).from(locked)
+
+  const isVersioned = inArray(workspaceFiles.context, VERSIONED_FILE_CONTEXTS)
+  const versionedCondition = and(condition, isVersioned)
+  // Also locks a versioned file committed after the lock pass before its history is frozen.
+  const versioned = await tx
+    .select()
+    .from(workspaceFiles)
+    .where(versionedCondition)
+    .orderBy(asc(workspaceFiles.id))
+    .for('update')
   const heads = await loadWorkspaceFileVersionHeads(
-    versioned.map((file) => file.id),
+    tx.select({ id: workspaceFiles.id }).from(workspaceFiles).where(versionedCondition),
     tx
   )
   for (const file of versioned) {
@@ -42,14 +60,15 @@ export async function handoffFileCreatorsInTx(
           )
     await materializeWorkspaceFileVersionInTx(tx, file, head, provenance, new Date())
   }
-  if (!files.length) return
+
   await tx
     .update(workspaceFiles)
     .set({ userId: successorId, updatedAt: new Date() })
     .where(
-      inArray(
-        workspaceFiles.id,
-        files.map((file) => file.id)
+      and(
+        condition,
+        // A versioned file committed after the history pass keeps its creator until frozen.
+        or(not(isVersioned), workspaceFileVersionHeadCurrentSql())
       )
     )
 }

apps/sim/lib/uploads/contexts/workspace/file-version-store.ts

@@ -7,7 +7,17 @@ import {
   workspaceFileVersion,
 } from '@sim/db/schema'
 import { generateId } from '@sim/utils/id'
-import { and, desc, eq, inArray, isNotNull, lt, notInArray, sql } from 'drizzle-orm'
+import {
+  and,
+  desc,
+  eq,
+  inArray,
+  isNotNull,
+  lt,
+  notInArray,
+  type SQLWrapper,
+  sql,
+} from 'drizzle-orm'
 import {
   type CursorKey,
   type KeysetKey,
@@ -105,14 +115,14 @@ export async function loadWorkspaceFileVersionHead(
 
 /**
  * Batched {@link loadWorkspaceFileVersionHead}: each file's highest version, keyed by file id.
- * The max is resolved from the `(file_id, version)` index first, so only head rows are fetched
- * rather than every retained version.
+ * `fileIds` is a subquery, so the file set never becomes a bind-parameter list. The max is
+ * resolved from the `(file_id, version)` index first, so only head rows are fetched rather than
+ * every retained version.
  */
 export async function loadWorkspaceFileVersionHeads(
-  fileIds: string[],
+  fileIds: SQLWrapper,
   executor: DbOrTx = db
 ): Promise<Map<string, WorkspaceFileVersionSummaryRow>> {
-  if (!fileIds.length) return new Map()
   const heads = await executor
     .select(versionSummaryColumns)
     .from(workspaceFileVersion)
@@ -713,27 +723,36 @@ export function getCurrentWorkspaceFileVersion(
   return withVersionSnapshot(file, currentVersionInSnapshot, executor)
 }
 
+/**
+ * Correlated columns for subqueries over the enclosing query's `workspace_files` row. Both sides
+ * are table-qualified because Drizzle renders single-table columns bare, which would bind the outer
+ * columns to the subquery's own table.
+ */
+function correlatedVersionColumns() {
+  const qualified = (table: typeof workspaceFileVersion | typeof workspaceFiles, name: string) =>
+    sql`${table}.${sql.identifier(name)}`
+  return {
+    head: {
+      fileId: qualified(workspaceFileVersion, workspaceFileVersion.fileId.name),
+      version: qualified(workspaceFileVersion, workspaceFileVersion.version.name),
+      key: qualified(workspaceFileVersion, workspaceFileVersion.key.name),
+      supersededAt: qualified(workspaceFileVersion, workspaceFileVersion.supersededAt.name),
+    },
+    file: {
+      id: qualified(workspaceFiles, workspaceFiles.id.name),
+      key: qualified(workspaceFiles, workspaceFiles.key.name),
+    },
+  }
+}
+
 /**
  * The current version number of the enclosing query's `workspace_files` row, as a correlated
  * subquery so the row and its number come from one statement's snapshot. The newest row numbers
  * the file's bytes while it still holds them; otherwise the bytes are on the implicit version after
  * it, and a file with no rows is on version 1 — the numbering {@link implicitCurrentVersion} gives.
- * Both sides of the correlation are table-qualified because Drizzle renders single-table columns
- * bare, which would bind the outer columns to this subquery's own table.
  */
 export function currentWorkspaceFileVersionNumberSql() {
-  const qualified = (table: typeof workspaceFileVersion | typeof workspaceFiles, name: string) =>
-    sql`${table}.${sql.identifier(name)}`
-  const head = {
-    fileId: qualified(workspaceFileVersion, workspaceFileVersion.fileId.name),
-    version: qualified(workspaceFileVersion, workspaceFileVersion.version.name),
-    key: qualified(workspaceFileVersion, workspaceFileVersion.key.name),
-    supersededAt: qualified(workspaceFileVersion, workspaceFileVersion.supersededAt.name),
-  }
-  const file = {
-    id: qualified(workspaceFiles, workspaceFiles.id.name),
-    key: qualified(workspaceFiles, workspaceFiles.key.name),
-  }
+  const { head, file } = correlatedVersionColumns()
   const number = sql`case when ${head.supersededAt} is null and ${head.key} = ${file.key} then ${head.version} else ${head.version} + 1 end`
   return sql<number>`coalesce((
     select ${number} from ${workspaceFileVersion}
@@ -743,6 +762,17 @@ export function currentWorkspaceFileVersionNumberSql() {
   ), ${INITIAL_WORKSPACE_FILE_VERSION})`.mapWith(Number)
 }
 
+/** {@link isVersionHeadCurrent} for the enclosing query's `workspace_files` row. */
+export function workspaceFileVersionHeadCurrentSql() {
+  const { head, file } = correlatedVersionColumns()
+  return sql<boolean>`coalesce((
+    select ${head.supersededAt} is null and ${head.key} = ${file.key} from ${workspaceFileVersion}
+    where ${head.fileId} = ${file.id}
+    order by ${head.version} desc
+    limit 1
+  ), false)`
+}
+
 /** One version of a file, or null when it never existed or retention removed it. */
 export function getWorkspaceFileVersion(
   file: WorkspaceFileVersionSubject,

apps/sim/lib/workspaces/access/shared-resource-handoff.integration.ts

@@ -492,6 +492,33 @@ describe('shared resource retention across workspace departure and account erasu
     ).toMatchObject([{ version: 1, authorUserIds: [fixture.departingId] }])
   })
 
+  it('hands off more files than Postgres accepts bind parameters', async () => {
+    const fixture = await seedResources(false)
+    /** More ids than Postgres's 65,535 bind parameters per statement. */
+    const bulkFileCount = 66_000
+    await db.execute(sql`
+      INSERT INTO ${workspaceFiles} (id, key, user_id, workspace_id, context, original_name, content_type)
+      SELECT ${fixture.workspaceId} || '-kb-' || n, ${fixture.workspaceId} || '/kb/' || n,
+        ${fixture.departingId}, ${fixture.workspaceId}, 'knowledge-base', 'bulk.txt', 'text/plain'
+      FROM generate_series(1, ${bulkFileCount}) AS n
+    `)
+    const result = await db.transaction((tx) =>
+      revokeWorkspaceAccessTx(tx, {
+        workspaceId: fixture.workspaceId,
+        userId: fixture.departingId,
+      })
+    )
+    expect(result.revoked).toBe(true)
+    expect(
+      await db
+        .select({ userId: workspaceFiles.userId, files: sql<number>`count(*)::int` })
+        .from(workspaceFiles)
+        .where(eq(workspaceFiles.workspaceId, fixture.workspaceId))
+        .groupBy(workspaceFiles.userId)
+    ).toEqual([{ userId: fixture.ownerId, files: bulkFileCount + 1 }])
+    await assertTransferred(fixture)
+  })
+
   it('repairs references left by departures before the handoff fix', async () => {
     const fixture = await seedResources(true, false)
     await deleteUserAccount(fixture.departingId, { recordIdentity: false })