diff --git a/apps/sim/background/projection-source-acl-backfill.ts b/apps/sim/background/projection-source-acl-backfill.ts new file mode 100644 index 00000000000..e6b75c0794a --- /dev/null +++ b/apps/sim/background/projection-source-acl-backfill.ts @@ -0,0 +1,35 @@ +import { task, tasks } from '@trigger.dev/sdk' +import { resolveTriggerRegion } from '@/lib/core/async-jobs/region' +import { + PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, + type ProjectionSourceAclBackfillPayload, + runProjectionSourceAclBackfill, +} from '@/lib/knowledge/search/projection-source-acl-backfill' + +/** One run's share of the backfill, inside the worker's run ceiling with room to end its page. */ +const RUN_BUDGET_MS = 60 * 60 * 1000 + +/** + * Trigger.dev wrapper around `runProjectionSourceAclBackfill`. A run fills unset rows for up to + * {@link RUN_BUDGET_MS}, then triggers its continuation from the cursor it reached, so the whole + * projection is filled across as many bounded runs as it takes. Retry-safe: every run writes only + * rows still unset, so a retried or restarted run repeats no write. The queue admits one run at a + * time, so two starts never fill the same pages against each other. + */ +export const projectionSourceAclBackfillTask = task({ + id: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, + machine: 'small-1x', + retry: { maxAttempts: 3 }, + queue: { + name: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, + concurrencyLimit: 1, + }, + run: async (payload: ProjectionSourceAclBackfillPayload) => { + const cursor = await runProjectionSourceAclBackfill(payload, { budgetMs: RUN_BUDGET_MS }) + if (!cursor) return + const continuation: ProjectionSourceAclBackfillPayload = { ...payload, cursor } + await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, continuation, { + region: await resolveTriggerRegion(), + }) + }, +}) diff --git a/apps/sim/lib/knowledge/access/predicate.postgres.test.ts b/apps/sim/lib/knowledge/access/predicate.postgres.test.ts index 028bb285eae..11a1d8559df 100644 --- a/apps/sim/lib/knowledge/access/predicate.postgres.test.ts +++ b/apps/sim/lib/knowledge/access/predicate.postgres.test.ts @@ -622,6 +622,24 @@ describe.runIf(Boolean(databaseUrl))('knowledge ACLs in PostgreSQL', () => { ) expect(await admits(perRow, 'members-other')).toBe(false) expect(await onRowAdmits('members-other')).toBe(false) + /** + * A chunk the backfill has not reached carries no source or ACL yet and is decided on its + * document, as every candidate was before the columns existed: search must not depend on the + * backfill, must not lose a document to it, and must not rank an unreadable one because of it. + */ + await connection.unsafe( + "INSERT INTO document(id, connector_id, acl, acl_verified_at) VALUES ('unfilled-other', 'members', ARRAY['s:slack:-:bob'], statement_timestamp())" + ) + await connection.unsafe(`INSERT INTO embedding_search(id, document_id) VALUES + ('unfilled-other-chunk', 'unfilled-other'), ('unfilled-admin-chunk', 'admin-current'), + ('unfilled-members-chunk', 'members-current'), ('unfilled-gone-chunk', 'deleted-connector')`) + await connection.unsafe( + "DELETE FROM embedding_search WHERE id IN ('admin-current-chunk', 'members-current-chunk', 'deleted-connector-chunk')" + ) + expect(await onRowAdmits('unfilled-other')).toBe(false) + expect(await onRowAdmits('admin-current')).toBe(true) + expect(await onRowAdmits('members-current')).toBe(true) + expect(await onRowAdmits('deleted-connector')).toBe(false) /** * Candidate ranking defers the live source proof, as the per-row candidate predicate does: a * caller holds those grants only after authorization, so applying the clause during ranking diff --git a/apps/sim/lib/knowledge/access/predicate.test.ts b/apps/sim/lib/knowledge/access/predicate.test.ts index 11582cc360e..4e2e7327e71 100644 --- a/apps/sim/lib/knowledge/access/predicate.test.ts +++ b/apps/sim/lib/knowledge/access/predicate.test.ts @@ -17,13 +17,56 @@ vi.unmock('@sim/db/schema') process.env.DATABASE_URL ??= 'postgresql://user:pass@localhost:5432/test' const { PgDialect } = await import('drizzle-orm/pg-core') -const { knowledgeAccessCondition } = await import('@/lib/knowledge/access/predicate') +const { embeddingSearch } = await import('@sim/db/schema') +const { knowledgeAccessCondition, projectionCandidateAccessCondition } = await import( + '@/lib/knowledge/access/predicate' +) const { SYSTEM_ACCESS_SCOPE } = await import('@/lib/knowledge/access/types') function render(condition: ReturnType) { return new PgDialect().sqlToQuery(condition) } +describe('projectionCandidateAccessCondition', () => { + const plan = { + connectors: { workspace: ['ws-src'], admin: [], members: [], liveProofRequired: [] }, + observers: { confirmed: [], observed: [] }, + memberSources: [], + } + + it('decides a filled row on its mirrored columns and an unfilled row on its document', () => { + const { sql, params } = render( + projectionCandidateAccessCondition( + embeddingSearch, + { kind: 'user', userId: 'user-1', tokens: ['ws', 'u:alice'] }, + plan + ) + ) + expect(sql).toContain( + '("embedding_search"."acl" IS NULL AND EXISTS (\n SELECT 1 FROM "document"\n WHERE "document"."id" = "embedding_search"."document_id"\n AND (' + ) + expect(sql).toContain('"document"."acl" && ARRAY[$1, $2]::text[]') + expect(sql).toMatch( + /OR \("embedding_search"\."acl" && ARRAY\[\$\d+, \$\d+\]::text\[\]\n {4}AND \("embedding_search"\."connector_id" IS NULL OR "embedding_search"\."connector_id" = ANY\(ARRAY\[\$\d+\]::text\[\]\)\)\)\)$/ + ) + expect(params.slice(0, 2)).toEqual(['ws', 'u:alice']) + expect(params.slice(-3)).toEqual(['ws', 'u:alice', 'ws-src']) + for (const param of params) expect(Array.isArray(param)).toBe(false) + }) + + it('still denies everything for an empty token set', () => { + expect( + render( + projectionCandidateAccessCondition( + embeddingSearch, + { kind: 'user', userId: 'user-1', tokens: [] }, + plan + ) + ).sql + ).toBe('false') + }) +}) + describe('knowledgeAccessCondition', () => { it('overlaps the ACL with the tokens as a literal array of scalar binds', () => { const { sql, params } = render( diff --git a/apps/sim/lib/knowledge/access/predicate.ts b/apps/sim/lib/knowledge/access/predicate.ts index a6a988c723e..9c8cd680c32 100644 --- a/apps/sim/lib/knowledge/access/predicate.ts +++ b/apps/sim/lib/knowledge/access/predicate.ts @@ -312,9 +312,19 @@ export function knowledgeCandidateAccessConditionForConnectors( * observation's freshness, and requirement clauses live on the document. Both are refused there, * under the full predicate, before content is returned — this predicate only decides what is worth * ranking. + * + * A row the backfill has not reached yet carries no ACL (`acl IS NULL`) and is decided on its + * document instead, under {@link knowledgeCandidateAccessConditionForConnectors} — the join per + * candidate that every row paid before the columns existed. The backfill runs in the background, + * so search never waits on it, never loses a document to it, and never ranks an unreadable one + * into a bounded candidate pool because of it. */ export function projectionCandidateAccessCondition( - projection: { connectorId: AnyPgColumn | SQL; acl: AnyPgColumn | SQL }, + projection: { + connectorId: AnyPgColumn | SQL + acl: AnyPgColumn | SQL + documentId: AnyPgColumn | SQL + }, scope: KnowledgeAccessScope | SystemAccessScope, plan: SearchAccessPlan ): SQL { @@ -330,8 +340,13 @@ export function projectionCandidateAccessCondition( ...plan.connectors.admin, ...plan.connectors.members, ] - return sql`(${projection.acl} && ${tokens} - AND (${projection.connectorId} IS NULL OR ${inSources(mirrored)}))` + const unfilled = sql`(${projection.acl} IS NULL AND EXISTS ( + SELECT 1 FROM ${document} + WHERE ${document.id} = ${projection.documentId} + AND ${knowledgeCandidateAccessConditionForConnectors(scope, plan)} + ))` + return sql`(${unfilled} OR (${projection.acl} && ${tokens} + AND (${projection.connectorId} IS NULL OR ${inSources(mirrored)})))` } /** diff --git a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts new file mode 100644 index 00000000000..a314999a4bf --- /dev/null +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts @@ -0,0 +1,123 @@ +/** + * @vitest-environment node + */ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockBackfill, mockEnd, mockPostgres, mockTasksTrigger } = vi.hoisted(() => ({ + mockBackfill: vi.fn(), + mockEnd: vi.fn(async () => undefined), + mockPostgres: vi.fn(), + mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })), +})) + +vi.mock('@sim/db', () => ({ resolveDbUrl: () => 'postgres://localhost:5432/sim' })) +vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({ + PROJECTION_SOURCE_ACL_TABLES: ['embedding_search', 'embedding_keyword_tin'], + backfillProjectionSourceAcl: mockBackfill, +})) +vi.mock('postgres', () => ({ default: mockPostgres })) +vi.mock('@trigger.dev/sdk', () => ({ tasks: { trigger: mockTasksTrigger } })) +vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' })) +vi.mock('@/lib/core/utils/background', () => ({ + runDetached: (_label: string, work: () => Promise) => { + void work() + }, +})) + +import { + enqueueProjectionSourceAclBackfill, + runProjectionSourceAclBackfill, +} from '@/lib/knowledge/search/projection-source-acl-backfill' + +const connection = { end: mockEnd } + +describe('runProjectionSourceAclBackfill', () => { + beforeEach(() => { + vi.clearAllMocks() + mockPostgres.mockReturnValue(connection) + mockBackfill.mockImplementation(async (_sql, projection, options) => ({ + projection, + scanned: 0, + written: 0, + afterId: options?.afterId ?? '', + done: true, + })) + }) + + it('fills both projections in order on its own connection and closes it', async () => { + await expect(runProjectionSourceAclBackfill({ pageSize: 50, pauseMs: 10 })).resolves.toBeNull() + expect(mockBackfill.mock.calls.map(([, projection]) => projection)).toEqual([ + 'embedding_search', + 'embedding_keyword_tin', + ]) + for (const [sql, , options] of mockBackfill.mock.calls) { + expect(sql).toBe(connection) + expect(options).toMatchObject({ afterId: undefined, pageSize: 50, pauseMs: 10 }) + } + expect(mockEnd).toHaveBeenCalledTimes(1) + }) + + it('resumes after the cursor in its projection and from the start of the next', async () => { + await runProjectionSourceAclBackfill({ + cursor: { projection: 'embedding_keyword_tin', afterId: 'chunk-9' }, + }) + expect(mockBackfill).toHaveBeenCalledTimes(1) + expect(mockBackfill.mock.calls[0][1]).toBe('embedding_keyword_tin') + expect(mockBackfill.mock.calls[0][2]).toMatchObject({ afterId: 'chunk-9' }) + }) + + it('returns where a budgeted run stopped so the next run can carry on', async () => { + mockBackfill.mockResolvedValueOnce({ + projection: 'embedding_search', + scanned: 100, + written: 100, + afterId: 'chunk-100', + done: false, + }) + await expect(runProjectionSourceAclBackfill({}, { budgetMs: 1000 })).resolves.toEqual({ + projection: 'embedding_search', + afterId: 'chunk-100', + }) + expect(mockBackfill).toHaveBeenCalledTimes(1) + expect(mockBackfill.mock.calls[0][2].budgetMs).toBeLessThanOrEqual(1000) + expect(mockEnd).toHaveBeenCalledTimes(1) + }) + + it('closes the connection when a page fails', async () => { + mockBackfill.mockRejectedValueOnce(new Error('canceling statement due to statement timeout')) + await expect(runProjectionSourceAclBackfill({})).rejects.toThrow('statement timeout') + expect(mockEnd).toHaveBeenCalledTimes(1) + }) +}) + +describe('enqueueProjectionSourceAclBackfill', () => { + beforeEach(() => { + vi.clearAllMocks() + mockPostgres.mockReturnValue(connection) + mockBackfill.mockResolvedValue({ + projection: 'embedding_search', + scanned: 0, + written: 0, + afterId: '', + done: true, + }) + }) + + it('hands the backfill to the Trigger.dev worker when one is configured', async () => { + await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 }, true)).resolves.toEqual({ + runId: 'run-1', + }) + expect(mockTasksTrigger).toHaveBeenCalledWith( + 'projection-source-acl-backfill', + { pageSize: 25 }, + { region: 'us-east-1' } + ) + expect(mockBackfill).not.toHaveBeenCalled() + }) + + it('fills the projections detached in this process without one', async () => { + await expect(enqueueProjectionSourceAclBackfill({}, false)).resolves.toBeNull() + expect(mockTasksTrigger).not.toHaveBeenCalled() + await vi.waitFor(() => expect(mockBackfill).toHaveBeenCalledTimes(2)) + }) +}) diff --git a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts new file mode 100644 index 00000000000..fec8a6f489a --- /dev/null +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts @@ -0,0 +1,98 @@ +import { resolveDbUrl } from '@sim/db' +import { + backfillProjectionSourceAcl, + PROJECTION_SOURCE_ACL_TABLES, + type ProjectionSourceAclTable, +} from '@sim/db/script-migrations/0021_embedding_search_connector' +import { createLogger } from '@sim/logger' +import postgres from 'postgres' +import { resolveTriggerRegion } from '@/lib/core/async-jobs/region' +import { env } from '@/lib/core/config/env' +import { isTriggerDevEnabled } from '@/lib/core/config/env-flags' +import { runDetached } from '@/lib/core/utils/background' + +const logger = createLogger('ProjectionSourceAclBackfill') + +export const PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID = 'projection-source-acl-backfill' + +/** Where a run stopped, so the next one carries on from there instead of rescanning. */ +export interface ProjectionSourceAclBackfillCursor { + projection: ProjectionSourceAclTable + afterId: string +} + +export interface ProjectionSourceAclBackfillPayload { + /** The first projection's first page when absent. */ + cursor?: ProjectionSourceAclBackfillCursor + pageSize?: number + pauseMs?: number +} + +export interface ProjectionSourceAclBackfillRunOptions { + /** Stop once this much time has passed and return where to resume; unbounded otherwise. */ + budgetMs?: number +} + +/** + * Fills the ranking projections' source and ACL columns from the cursor onwards, one projection + * after the other, on a connection of its own: the page statement binds the keyset cursor as a + * scalar and needs no array parameter, so the pool's options would serve, but a run this long + * should not hold one of the worker's pooled connections. Returns the cursor to continue from when + * the budget ran out, `null` once both projections are filled. + */ +export async function runProjectionSourceAclBackfill( + payload: ProjectionSourceAclBackfillPayload, + options: ProjectionSourceAclBackfillRunOptions = {} +): Promise { + const url = resolveDbUrl('DATABASE_URL', process.env.SIM_DB_ROLE?.trim() || 'web') + if (!url) throw new Error('DATABASE_URL is required to backfill the projection source and ACL') + const sql = postgres(url, { max: 1, max_lifetime: null, onnotice: () => undefined }) + const startedAt = Date.now() + try { + const start = payload.cursor + ? PROJECTION_SOURCE_ACL_TABLES.indexOf(payload.cursor.projection) + : 0 + if (start < 0) throw new Error(`Unknown projection ${payload.cursor?.projection}`) + for (const projection of PROJECTION_SOURCE_ACL_TABLES.slice(start)) { + const budgetMs = + options.budgetMs === undefined + ? undefined + : Math.max(0, options.budgetMs - (Date.now() - startedAt)) + const progress = await backfillProjectionSourceAcl(sql, projection, { + afterId: payload.cursor?.projection === projection ? payload.cursor.afterId : undefined, + pageSize: payload.pageSize, + pauseMs: payload.pauseMs, + budgetMs, + }) + if (!progress.done) return { projection, afterId: progress.afterId } + } + logger.info('Projection source and ACL backfill complete', { + elapsedMs: Date.now() - startedAt, + }) + return null + } finally { + await sql.end() + } +} + +/** + * Starts the backfill the way the table backfill is started: on the deployment's Trigger.dev + * worker when one is configured, where bounded runs chain until both projections are filled, and + * detached in this process otherwise. Safe to call again at any time — a run only fills rows still + * unset. + */ +export async function enqueueProjectionSourceAclBackfill( + payload: ProjectionSourceAclBackfillPayload = {}, + useTrigger = Boolean(isTriggerDevEnabled && env.TRIGGER_SECRET_KEY) +): Promise<{ runId: string } | null> { + if (useTrigger) { + const { tasks } = await import('@trigger.dev/sdk') + const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, payload, { + region: await resolveTriggerRegion(), + }) + logger.info('Projection source and ACL backfill enqueued', { runId: handle.id }) + return { runId: handle.id } + } + runDetached(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, () => runProjectionSourceAclBackfill(payload)) + return null +} diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 8137a007d31..76ca0005a6a 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -1527,6 +1527,9 @@ describe('permitted-document planner', () => { /** The ranked CTE carries the mirrored source and ACL the predicate tests. */ expect(statement).toContain('AS connector_id') expect(statement).toContain('ranked_tin_chunks.acl') + /** A row the backfill has not filled (`acl IS NULL`) is decided on its document instead. */ + expect(statement).toContain('IS NULL AND EXISTS (') + expect(statement).toContain('ranked_tin_chunks.document_id') }) it('widens the window for a broad resolved scope whose first page came back short', async () => { diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index ac7174e6082..abb9aa76113 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1819,7 +1819,11 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise const onRowKeywordVisibility = (excludedSources: readonly string[]) => and( projectionCandidateAccessCondition( - { connectorId: sql`ranked_tin_chunks.connector_id`, acl: sql`ranked_tin_chunks.acl` }, + { + connectorId: sql`ranked_tin_chunks.connector_id`, + acl: sql`ranked_tin_chunks.acl`, + documentId: sql`ranked_tin_chunks.document_id`, + }, access, accessPlan! ), diff --git a/apps/sim/scripts/backfill-projection-source-acl.ts b/apps/sim/scripts/backfill-projection-source-acl.ts new file mode 100644 index 00000000000..def5181c49c --- /dev/null +++ b/apps/sim/scripts/backfill-projection-source-acl.ts @@ -0,0 +1,44 @@ +#!/usr/bin/env bun + +/** + * Starts the projection source and ACL backfill that script migration + * `0022_projection_source_acl_backfill` leaves to the background. On a deployment with Trigger.dev + * it enqueues the `projection-source-acl-backfill` task, which chains bounded runs until both + * ranking projections are filled; without one it fills them here, paced the same way. Safe to run + * again at any time — a run only fills rows still unset. + * + * Usage: + * bun apps/sim/scripts/backfill-projection-source-acl.ts + */ + +import { createLogger } from '@sim/logger' +import { toError } from '@sim/utils/errors' +import { env } from '@/lib/core/config/env' +import { isTriggerDevEnabled } from '@/lib/core/config/env-flags' +import { + enqueueProjectionSourceAclBackfill, + runProjectionSourceAclBackfill, +} from '@/lib/knowledge/search/projection-source-acl-backfill' + +const logger = createLogger('BackfillProjectionSourceAcl') + +/** A script has no long-lived process to detach into, so without a worker it fills inline. */ +async function main(): Promise { + if (isTriggerDevEnabled && env.TRIGGER_SECRET_KEY) { + const handle = await enqueueProjectionSourceAclBackfill({}, true) + logger.info('Backfill enqueued on the Trigger.dev worker', handle ?? {}) + return + } + await runProjectionSourceAclBackfill({}) + logger.info('Backfill complete') +} + +if (import.meta.main) { + main().then( + () => process.exit(0), + (error) => { + logger.error('Backfill failed', toError(error)) + process.exit(1) + } + ) +} diff --git a/packages/db/script-migrations-paused-billing-attribution.test.ts b/packages/db/script-migrations-paused-billing-attribution.test.ts index bc9b2eb474f..cea6d576863 100644 --- a/packages/db/script-migrations-paused-billing-attribution.test.ts +++ b/packages/db/script-migrations-paused-billing-attribution.test.ts @@ -453,7 +453,7 @@ describe('script migration registry', () => { '0017_index_search_documents', '0018_repair_workspace_file_content_revision', '0019_tin_keyword_projection', - '0021_embedding_search_connector', + '0022_projection_source_acl_backfill', ]) }) }) diff --git a/packages/db/script-migrations/0016_backfill_search_vectors.postgres.test.ts b/packages/db/script-migrations/0016_backfill_search_vectors.postgres.test.ts index 3781437a856..6345cab6c3f 100644 --- a/packages/db/script-migrations/0016_backfill_search_vectors.postgres.test.ts +++ b/packages/db/script-migrations/0016_backfill_search_vectors.postgres.test.ts @@ -392,6 +392,7 @@ describe.runIf(Boolean(databaseUrl))('search projection upgrade in PostgreSQL', { name: '0018_repair_workspace_file_content_revision' }, { name: '0019_tin_keyword_projection' }, { name: '0021_embedding_search_connector' }, + { name: '0022_projection_source_acl_backfill' }, ]) const [{ complete }] = await sql`SELECT count(*)::int AS complete FROM embedding e JOIN embedding_search s ON s.id = e.id JOIN embedding_keyword_search k ON k.id = e.id diff --git a/packages/db/script-migrations/0021_embedding_search_connector.postgres.test.ts b/packages/db/script-migrations/0021_embedding_search_connector.postgres.test.ts new file mode 100644 index 00000000000..b0e3ad37326 --- /dev/null +++ b/packages/db/script-migrations/0021_embedding_search_connector.postgres.test.ts @@ -0,0 +1,128 @@ +import { backfillProjectionSourceAcl } from '@sim/db/script-migrations/0021_embedding_search_connector' +import { projectionSourceAclBackfillMigration as embeddingSearchConnectorMigration } from '@sim/db/script-migrations/0022_projection_source_acl_backfill' +import { generateId } from '@sim/utils/id' +import postgres, { type Sql } from 'postgres' +import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest' + +const databaseUrl = process.env.KNOWLEDGE_ACL_TEST_DATABASE_URL + +/** + * The projections here carry only the columns the source and ACL triggers and backfill touch; the + * vector and lexeme columns, and their indexes, are what make the backfill slow, not what decides + * which rows it writes. + */ +describe.runIf(Boolean(databaseUrl))('projection source and ACL backfill in PostgreSQL', () => { + let admin: Sql + let sql: Sql + const schemaName = `projection_acl_${generateId().replaceAll('-', '')}` + + const projected = (projection: 'embedding_search' | 'embedding_keyword_tin') => + sql<{ id: string; connector_id: string | null; acl: string[] | null }[]>` + SELECT id, connector_id, acl FROM ${sql(projection)} ORDER BY id` + + beforeAll(async () => { + const url = new URL(databaseUrl!) + if ( + !['localhost', '127.0.0.1'].includes(url.hostname) || + !url.pathname.startsWith('/sim_acl_test') + ) { + throw new Error('Projection tests require a disposable local integration database') + } + admin = postgres(url.toString(), { max: 1, onnotice: () => undefined }) + await admin.unsafe(`CREATE SCHEMA "${schemaName}"`) + sql = postgres(url.toString(), { + max: 1, + onnotice: () => undefined, + connection: { search_path: schemaName }, + }) + await sql`CREATE TABLE document ( + id text PRIMARY KEY, connector_id text, acl text[] NOT NULL DEFAULT '{ws}' + )` + for (const projection of ['embedding_search', 'embedding_keyword_tin']) { + await sql`CREATE TABLE ${sql(projection)} ( + id text PRIMARY KEY, document_id text NOT NULL, enabled boolean NOT NULL DEFAULT true, + connector_id text, acl text[] + )` + } + await embeddingSearchConnectorMigration.up(sql) + }, 60_000) + + afterAll(async () => { + await sql?.end() + await admin?.unsafe(`DROP SCHEMA IF EXISTS "${schemaName}" CASCADE`) + await admin?.end() + }) + + beforeEach(async () => { + await sql`TRUNCATE embedding_search, embedding_keyword_tin, document` + await sql`ALTER TABLE embedding_search DISABLE TRIGGER embedding_search_source_acl_set` + await sql`ALTER TABLE embedding_keyword_tin DISABLE TRIGGER embedding_keyword_tin_source_acl_set` + }) + + it('installs its triggers and indexes again without failing, so a cut-short deploy completes', async () => { + await expect(embeddingSearchConnectorMigration.up(sql)).resolves.toBeUndefined() + const indexes = await sql<{ indexname: string }[]>` + SELECT indexname FROM pg_indexes WHERE schemaname = ${schemaName} ORDER BY indexname` + expect(indexes.map((row) => row.indexname)).toEqual( + expect.arrayContaining([ + 'embedding_keyword_tin_acl_gin_idx', + 'embedding_keyword_tin_acl_unfilled_idx', + 'embedding_search_acl_gin_idx', + 'embedding_search_acl_unfilled_idx', + 'embedding_search_source_idx', + ]) + ) + }) + + it('fills only the rows still unset, in pages, and leaves a chunk that changed documents to its trigger', async () => { + await sql`INSERT INTO document (id, connector_id, acl) VALUES + ('doc-a', 'src-a', ARRAY['u:alice']), ('doc-b', NULL, ARRAY['ws']), ('doc-c', 'src-c', ARRAY['u:carol'])` + await sql`INSERT INTO embedding_search (id, document_id, connector_id, acl) VALUES + ('c1', 'doc-a', NULL, NULL), ('c2', 'doc-b', NULL, NULL), + ('c3', 'doc-c', 'src-old', ARRAY['u:stale']), ('c4', 'doc-a', NULL, NULL), ('c5', 'missing', NULL, NULL)` + const progress = await backfillProjectionSourceAcl(sql, 'embedding_search', { + pageSize: 2, + pauseMs: 0, + }) + expect(progress).toEqual({ + projection: 'embedding_search', + scanned: 3, + written: 3, + afterId: 'c4', + done: true, + }) + expect(await projected('embedding_search')).toEqual([ + { id: 'c1', connector_id: 'src-a', acl: ['u:alice'] }, + { id: 'c2', connector_id: null, acl: ['ws'] }, + { id: 'c3', connector_id: 'src-old', acl: ['u:stale'] }, + { id: 'c4', connector_id: 'src-a', acl: ['u:alice'] }, + { id: 'c5', connector_id: null, acl: null }, + ]) + expect(await projected('embedding_keyword_tin')).toEqual([]) + }) + + it('stops at its budget with the cursor to resume from, and resumes after it', async () => { + await sql`INSERT INTO document (id, connector_id, acl) VALUES ('doc', 'src', ARRAY['u:alice'])` + await sql`INSERT INTO embedding_keyword_tin (id, document_id) VALUES + ('k1', 'doc'), ('k2', 'doc'), ('k3', 'doc')` + const paused = await backfillProjectionSourceAcl(sql, 'embedding_keyword_tin', { + pageSize: 1, + pauseMs: 0, + budgetMs: 0, + }) + expect(paused).toMatchObject({ scanned: 1, written: 1, afterId: 'k1', done: false }) + const resumed = await backfillProjectionSourceAcl(sql, 'embedding_keyword_tin', { + afterId: paused.afterId, + pageSize: 1, + pauseMs: 0, + }) + expect(resumed).toMatchObject({ scanned: 2, written: 2, afterId: 'k3', done: true }) + expect((await projected('embedding_keyword_tin')).map((row) => row.acl)).toEqual([ + ['u:alice'], + ['u:alice'], + ['u:alice'], + ]) + const again = await backfillProjectionSourceAcl(sql, 'embedding_keyword_tin', { pauseMs: 0 }) + expect(again).toMatchObject({ scanned: 0, written: 0, afterId: '', done: true }) + }) +}) diff --git a/packages/db/script-migrations/0021_embedding_search_connector.test.ts b/packages/db/script-migrations/0021_embedding_search_connector.test.ts new file mode 100644 index 00000000000..d4e030b449a --- /dev/null +++ b/packages/db/script-migrations/0021_embedding_search_connector.test.ts @@ -0,0 +1,28 @@ +/** + * @vitest-environment node + */ +import { backfillProjectionSourceAcl } from '@sim/db/script-migrations/0021_embedding_search_connector' +import type { Sql } from 'postgres' +import { describe, expect, it, vi } from 'vitest' + +/** A session that must never be reached: every case below is refused before the first page. */ +const untouched = { begin: vi.fn() } as unknown as Sql + +describe('backfillProjectionSourceAcl', () => { + it.each([0, -1, 1.5, Number.NaN, Number.POSITIVE_INFINITY])( + 'refuses a page size of %s instead of reporting the projection filled', + async (pageSize) => { + await expect( + backfillProjectionSourceAcl(untouched, 'embedding_search', { pageSize }) + ).rejects.toThrow('page size must be a positive integer') + expect(untouched.begin).not.toHaveBeenCalled() + } + ) + + it.each([-1, Number.NaN])('refuses a pause of %s', async (pauseMs) => { + await expect( + backfillProjectionSourceAcl(untouched, 'embedding_search', { pauseMs }) + ).rejects.toThrow('pause must be a non-negative number') + expect(untouched.begin).not.toHaveBeenCalled() + }) +}) diff --git a/packages/db/script-migrations/0021_embedding_search_connector.ts b/packages/db/script-migrations/0021_embedding_search_connector.ts index 3ddb1ad9430..4fc87e10fe3 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.ts +++ b/packages/db/script-migrations/0021_embedding_search_connector.ts @@ -1,15 +1,32 @@ -import type { ScriptMigration } from '@sim/db/script-migrations/types' import { createLogger } from '@sim/logger' +import { sleep } from '@sim/utils/helpers' import postgres, { type Sql } from 'postgres' const logger = createLogger('ProjectionSourceAcl') -/** Chunks read per page; each page commits on its own, as the other projection backfills do. */ -const BATCH_SIZE = 500 +/** + * Chunks filled per page. Every write to `embedding_search` re-inserts the row into each of its + * HNSW indexes, so a page's cost is index maintenance rather than the plan, and a small page keeps + * each transaction short and its locks brief. + */ +export const PROJECTION_SOURCE_ACL_PAGE_SIZE = 100 + +/** Pause between pages, so the backfill shares the database with the search it serves. */ +export const PROJECTION_SOURCE_ACL_PAGE_PAUSE_MS = 250 + +/** + * Longest a page may run before the database cancels it; the run then fails and resumes. A caller + * that bounds a run leaves at least this much headroom after its budget, since the budget is + * checked between pages and the page in flight runs to this limit. + */ +export const PROJECTION_SOURCE_ACL_PAGE_TIMEOUT_MS = 60_000 + +/** Pages between progress log lines. */ +const PROGRESS_EVERY_PAGES = 100 /** The projections that carry their document's source and ACL. */ -const PROJECTIONS = ['embedding_search', 'embedding_keyword_tin'] as const -type Projection = (typeof PROJECTIONS)[number] +export const PROJECTION_SOURCE_ACL_TABLES = ['embedding_search', 'embedding_keyword_tin'] as const +export type ProjectionSourceAclTable = (typeof PROJECTION_SOURCE_ACL_TABLES)[number] /** * Carries a chunk's source and ACL onto the ranking projections and keeps them there. @@ -50,7 +67,7 @@ export async function installProjectionSourceAcl(sql: Sql): Promise { RETURN NEW; END; $$`) - for (const projection of PROJECTIONS) { + for (const projection of PROJECTION_SOURCE_ACL_TABLES) { await tx.unsafe(`CREATE OR REPLACE TRIGGER ${projection}_source_acl_set BEFORE INSERT OR UPDATE OF document_id, enabled ON ${projection} FOR EACH ROW EXECUTE FUNCTION set_projection_source_acl()`) @@ -63,87 +80,153 @@ export async function installProjectionSourceAcl(sql: Sql): Promise { }) } +export interface ProjectionSourceAclBackfillOptions { + /** Resume after this chunk id; the projection's first page otherwise. */ + afterId?: string + pageSize?: number + pauseMs?: number + /** Stop once this much time has passed and report where to resume; unbounded otherwise. */ + budgetMs?: number +} + +export interface ProjectionSourceAclBackfillProgress { + projection: ProjectionSourceAclTable + /** Unfilled chunks this run read, including any a concurrent write filled first. */ + scanned: number + written: number + /** The last chunk id this run reached; the next run resumes after it while `done` is false. */ + afterId: string + done: boolean +} + /** * Fills a projection's source and ACL for chunks written before the trigger existed, in - * independently committed keyset pages. Each page bounds its own locking and runtime and writes only - * rows still unset, so an interrupted run resumes by rerunning and a row the trigger has since - * written is left alone. The documents are share-locked before their values are copied, so a change - * in flight waits for the page and then fans its own values out; and a chunk that moved to another - * document meanwhile is left to that document's trigger, since the write requires the document the - * values were read from. + * independently committed keyset pages of unfilled rows. Each page is one statement that bounds + * its own locking and runtime and writes only rows still unset, so an interrupted run resumes by + * rerunning and a row the trigger has since written is left alone. The documents are share-locked + * before their values are copied, so a change in flight waits for the page and then fans its own + * values out; and a chunk that moved to another document meanwhile is left to that document's + * trigger, since the write requires the document the values were read from. + * + * On `embedding_search` every filled row is re-inserted into each HNSW index, which is the whole + * cost of a page and far more than a deploy can wait for; the run paces itself with a pause between + * pages and stops at its budget so a background task can chain runs until the projection is filled. + * Search does not wait: an unfilled row is decided on its document by the on-row candidate + * predicate, the join per candidate every row paid before the columns existed. */ export async function backfillProjectionSourceAcl( sql: Sql, - projection: Projection -): Promise { + projection: ProjectionSourceAclTable, + options: ProjectionSourceAclBackfillOptions = {} +): Promise { + const pageSize = options.pageSize ?? PROJECTION_SOURCE_ACL_PAGE_SIZE + const pauseMs = options.pauseMs ?? PROJECTION_SOURCE_ACL_PAGE_PAUSE_MS + /** + * The page size is interpolated into the statement and a page of nothing would report the + * projection filled; a payload that asks for either is refused rather than quietly reshaped. + */ + if (!Number.isSafeInteger(pageSize) || pageSize < 1) { + throw new Error(`Projection backfill page size must be a positive integer, got ${pageSize}`) + } + if (!Number.isFinite(pauseMs) || pauseMs < 0) { + throw new Error(`Projection backfill pause must be a non-negative number, got ${pauseMs}`) + } const startedAt = Date.now() - let afterId = '' + const deadline = + options.budgetMs === undefined ? Number.POSITIVE_INFINITY : startedAt + options.budgetMs + let afterId = options.afterId ?? '' let scanned = 0 let written = 0 + let pages = 0 + let done = false for (;;) { const page = await sql.begin(async (tx) => { await tx.unsafe("SET LOCAL lock_timeout = '5s'") - await tx.unsafe("SET LOCAL statement_timeout = '60s'") - const rows = await tx.unsafe>( - `SELECT id FROM ${projection} WHERE id > $1 ORDER BY id LIMIT ${BATCH_SIZE}`, - [afterId] - ) - if (rows.length === 0) return null - const ids = rows.map((row) => row.id) - const [{ filled }] = await tx.unsafe>( + await tx.unsafe(`SET LOCAL statement_timeout = ${PROJECTION_SOURCE_ACL_PAGE_TIMEOUT_MS}`) + const [row] = await tx.unsafe< + Array<{ scanned: number; filled: number; last_id: string | null }> + >( `WITH page AS ( SELECT s.id, s.document_id, d.connector_id, d.acl FROM ${projection} s JOIN document d ON d.id = s.document_id - WHERE s.id = ANY($1::text[]) AND s.acl IS NULL + WHERE s.id > $1 AND s.acl IS NULL + ORDER BY s.id LIMIT ${pageSize} FOR SHARE OF d ), updated AS ( UPDATE ${projection} s SET connector_id = page.connector_id, acl = page.acl FROM page WHERE s.id = page.id AND s.document_id = page.document_id AND s.acl IS NULL RETURNING s.id - ) SELECT count(*)::int AS filled FROM updated`, - [ids] + ) + SELECT (SELECT count(*)::int FROM page) AS scanned, + (SELECT count(*)::int FROM updated) AS filled, + (SELECT max(id) FROM page) AS last_id`, + [afterId] ) - return { afterId: ids[ids.length - 1], scanned: ids.length, filled } + return row }) - if (!page) break - afterId = page.afterId + if (page.last_id === null) { + done = true + /** The planner last saw every row unfilled; it should see the finished projection. */ + await sql.unsafe(`ANALYZE ${projection}`) + break + } + afterId = page.last_id scanned += page.scanned written += page.filled - if (scanned % (BATCH_SIZE * 100) === 0) { + pages += 1 + if (pages % PROGRESS_EVERY_PAGES === 0) { logger.info('Projection source and ACL backfill progress', { projection, scanned, written, + afterId, elapsedMs: Date.now() - startedAt, }) } + if (pauseMs > 0) await sleep(pauseMs) + /** Checked after the pause, so the pause cannot carry a run past its budget into another page. */ + if (Date.now() >= deadline) break } - logger.info('Projection source and ACL backfilled', { - projection, - scanned, - written, - elapsedMs: Date.now() - startedAt, - }) - return written + logger.info( + done ? 'Projection source and ACL backfilled' : 'Projection source and ACL backfill paused', + { + projection, + scanned, + written, + afterId, + elapsedMs: Date.now() - startedAt, + } + ) + return { projection, scanned, written, afterId, done } } /** * The indexes exact ranking of a readable set needs: the ACL index on each projection, and on the * vector projection the source index that lets the planner lead with a few sources when the - * caller's tokens alone would match most of the index. Built after the bulk load and concurrently, - * so the triggers keep writing. `CONCURRENTLY` cannot run in a transaction, and the pool's lock - * timeout would cancel a build that merely waits for a long transaction to finish. + * caller's tokens alone would match most of the index. Built concurrently, so the triggers and the + * backfill keep writing. `CONCURRENTLY` cannot run in a transaction, and the pool's lock timeout + * would cancel a build that merely waits for a long transaction to finish. The timeout is a + * session setting, so one connection is reserved for it, the builds, and the reset — a pool + * would otherwise hand the builds to connections that never saw the setting. + * + * The unfilled index on each projection lists the rows the backfill has not reached: each page + * reads its rows from it instead of walking past every filled one, and the on-row predicate's + * unfilled branch, an `OR` beside the ACL overlap, stays an index probe for the planner — once the + * projection is filled, a probe of an empty index. */ -export async function indexProjectionAcl(sql: Sql): Promise { +export async function indexProjectionAcl(pool: Sql): Promise { + const sql = await pool.reserve() const [{ timeout }] = await sql`SELECT current_setting('lock_timeout') AS timeout` await sql.unsafe('SET lock_timeout = 0') try { const builds: Array<[name: string, definition: string]> = [ - ...PROJECTIONS.map((projection): [string, string] => [ - `${projection}_acl_gin_idx`, - `ON ${projection} USING gin (acl) WHERE enabled`, - ]), + ...PROJECTION_SOURCE_ACL_TABLES.flatMap( + (projection): Array<[string, string]> => [ + [`${projection}_acl_gin_idx`, `ON ${projection} USING gin (acl) WHERE enabled`], + [`${projection}_acl_unfilled_idx`, `ON ${projection} (id) WHERE acl IS NULL`], + ] + ), ['embedding_search_source_idx', 'ON embedding_search (connector_id) WHERE enabled'], ] for (const [name, definition] of builds) { @@ -159,27 +242,28 @@ export async function indexProjectionAcl(sql: Sql): Promise { await sql.unsafe(`CREATE INDEX CONCURRENTLY IF NOT EXISTS ${name} ${definition}`) } /** The new columns carry no statistics until analyzed; the reach and source predicates plan on them. */ - for (const projection of PROJECTIONS) await sql.unsafe(`ANALYZE ${projection}`) + for (const projection of PROJECTION_SOURCE_ACL_TABLES) await sql.unsafe(`ANALYZE ${projection}`) } finally { await sql`SELECT set_config('lock_timeout', ${timeout}, false)` + sql.release() } } -export const embeddingSearchConnectorMigration: ScriptMigration = { - name: '0021_embedding_search_connector', - async up(sql) { - await installProjectionSourceAcl(sql) - for (const projection of PROJECTIONS) await backfillProjectionSourceAcl(sql, projection) - await indexProjectionAcl(sql) - }, -} - +/** + * Run directly — `db:push`, or an operator filling a database by hand — the projections are filled + * here, paced the same way, rather than left to the app. The registered migration is + * `0022_projection_source_acl_backfill`, which supersedes this file's earlier, synchronous shape. + */ if (import.meta.main) { const url = process.env.MIGRATION_DATABASE_URL ?? process.env.DATABASE_URL if (!url) throw new Error('DATABASE_URL is required to backfill the projection source and ACL') const sql = postgres(url, { max: 1, max_lifetime: null, onnotice: () => undefined }) try { - await embeddingSearchConnectorMigration.up(sql) + await installProjectionSourceAcl(sql) + await indexProjectionAcl(sql) + for (const projection of PROJECTION_SOURCE_ACL_TABLES) { + await backfillProjectionSourceAcl(sql, projection) + } } finally { await sql.end() } diff --git a/packages/db/script-migrations/0022_projection_source_acl_backfill.ts b/packages/db/script-migrations/0022_projection_source_acl_backfill.ts new file mode 100644 index 00000000000..9d3718ec84f --- /dev/null +++ b/packages/db/script-migrations/0022_projection_source_acl_backfill.ts @@ -0,0 +1,30 @@ +import { + indexProjectionAcl, + installProjectionSourceAcl, +} from '@sim/db/script-migrations/0021_embedding_search_connector' +import type { ScriptMigration } from '@sim/db/script-migrations/types' + +/** + * Installs the projection source and ACL triggers and builds their indexes; both are idempotent, + * so a run that was cut short completes on the next deploy. The columns are not filled here: on + * `embedding_search` every filled row is re-inserted into each HNSW index, which puts the full + * projection far beyond what a deploy job can wait for. The backfill is started after the deploy + * with `bun apps/sim/scripts/backfill-projection-source-acl.ts`, which enqueues the + * `projection-source-acl-backfill` Trigger.dev task the way the table backfill is enqueued (or runs + * the same backfill detached in the app where no worker is configured); the task chains bounded + * runs until both projections are filled and is safe to start again at any time. Until it + * completes, an unfilled row is decided on its document, the join per candidate every row paid + * before the columns existed. + * + * Supersedes `0021_embedding_search_connector`, whose synchronous backfill this replaces: a + * database that recorded it still gains the unfilled indexes, and one where it was cut short + * records both names once this completes. + */ +export const projectionSourceAclBackfillMigration: ScriptMigration = { + name: '0022_projection_source_acl_backfill', + supersedes: ['0021_embedding_search_connector'], + async up(sql) { + await installProjectionSourceAcl(sql) + await indexProjectionAcl(sql) + }, +} diff --git a/packages/db/script-migrations/index.ts b/packages/db/script-migrations/index.ts index aab5835488a..85a3b2f099c 100644 --- a/packages/db/script-migrations/index.ts +++ b/packages/db/script-migrations/index.ts @@ -5,7 +5,7 @@ import { backfillSearchVectorsMigration } from '@sim/db/script-migrations/0016_b import { indexSearchDocumentsMigration } from '@sim/db/script-migrations/0017_index_search_documents' import { repairWorkspaceFileContentRevisionMigration } from '@sim/db/script-migrations/0018_repair_workspace_file_content_revision' import { tinKeywordProjectionMigration } from '@sim/db/script-migrations/0019_tin_keyword_projection' -import { embeddingSearchConnectorMigration } from '@sim/db/script-migrations/0021_embedding_search_connector' +import { projectionSourceAclBackfillMigration } from '@sim/db/script-migrations/0022_projection_source_acl_backfill' import type { Sql } from 'postgres' import { backfillTableOrderKeys } from './0001_backfill_table_order_keys' import { backfillPausedBillingAttribution } from './0002_backfill_paused_billing_attribution' @@ -44,7 +44,8 @@ export const scriptMigrations: readonly ScriptMigration[] = [ /** 0358 stops new sub-millisecond revisions; this retires the ones that predate it. */ repairWorkspaceFileContentRevisionMigration, tinKeywordProjectionMigration, - embeddingSearchConnectorMigration, + /** 0022 supersedes 0021, whose synchronous backfill could not finish inside a deploy. */ + projectionSourceAclBackfillMigration, ] /**