diff --git a/apps/sim/lib/uploads/contexts/workspace/creator-handoff.ts b/apps/sim/lib/uploads/contexts/workspace/creator-handoff.ts index 947b59e9b48..e1b3f530b6c 100644 --- a/apps/sim/lib/uploads/contexts/workspace/creator-handoff.ts +++ b/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()) ) ) } diff --git a/apps/sim/lib/uploads/contexts/workspace/file-version-store.ts b/apps/sim/lib/uploads/contexts/workspace/file-version-store.ts index a5844e4d9e7..a6c3208b4fd 100644 --- a/apps/sim/lib/uploads/contexts/workspace/file-version-store.ts +++ b/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> { - 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`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`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, diff --git a/apps/sim/lib/workspaces/access/shared-resource-handoff.integration.ts b/apps/sim/lib/workspaces/access/shared-resource-handoff.integration.ts index 83d23a3dffc..26f117eaaa6 100644 --- a/apps/sim/lib/workspaces/access/shared-resource-handoff.integration.ts +++ b/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`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 })