From 7b0ee6f87da2d673b593c0ac5a1132c16dc72066 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 15:33:27 -0700 Subject: [PATCH 1/3] fix(knowledge): stop document ACL changes from rewriting unfilled projection chunks The document trigger copied a changed ACL onto every enabled chunk whose projected ACL differed, and a chunk the projection backfill has not filled yet (acl IS NULL) always differs. Each such write re-inserts the row into the vector index, so one document ACL change inside a sync or member rewrite could rewrite every unfilled chunk of that document in the writer's own statement. Unfilled chunks now keep a NULL ACL: search already decides them on their document, and the backfill copies the document's current values under a share lock. A document that moves to another source still carries the source onto them, since source filters read it from the row. The new body ships as script migration 0023, which replaces the function in place; 0022 installs the same body on a fresh database. --- .github/workflows/test-build.yml | 6 + ...rations-paused-billing-attribution.test.ts | 1 + ...mbedding_search_connector.postgres.test.ts | 171 +++++++++++++++++- .../0021_embedding_search_connector.ts | 45 +++-- .../0023_projection_acl_skip_unfilled.ts | 16 ++ packages/db/script-migrations/index.ts | 3 + 6 files changed, 227 insertions(+), 15 deletions(-) create mode 100644 packages/db/script-migrations/0023_projection_acl_skip_unfilled.ts diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index b7d6b0d6dc5..beceabc40ff 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -299,6 +299,12 @@ jobs: lib/knowledge/access/predicate.postgres.test.ts \ lib/knowledge/connectors/external-directory.postgres.test.ts + - name: Verify the projection source and ACL trigger and backfill in PostgreSQL + working-directory: packages/db + env: + KNOWLEDGE_ACL_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_acl_test + run: bunx vitest run script-migrations/0021_embedding_search_connector.postgres.test.ts + test-build: name: Lint and Test runs-on: ${{ (vars.CI_PROVIDER == '' || vars.CI_PROVIDER == 'blacksmith') && 'blacksmith-8vcpu-ubuntu-2404' || 'ubuntu-latest' }} diff --git a/packages/db/script-migrations-paused-billing-attribution.test.ts b/packages/db/script-migrations-paused-billing-attribution.test.ts index cea6d576863..33e7588dfae 100644 --- a/packages/db/script-migrations-paused-billing-attribution.test.ts +++ b/packages/db/script-migrations-paused-billing-attribution.test.ts @@ -454,6 +454,7 @@ describe('script migration registry', () => { '0018_repair_workspace_file_content_revision', '0019_tin_keyword_projection', '0022_projection_source_acl_backfill', + '0023_projection_acl_skip_unfilled', ]) }) }) 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 index b0e3ad37326..626c8ad4d18 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.postgres.test.ts +++ b/packages/db/script-migrations/0021_embedding_search_connector.postgres.test.ts @@ -1,7 +1,13 @@ -import { backfillProjectionSourceAcl } from '@sim/db/script-migrations/0021_embedding_search_connector' +import { + backfillProjectionSourceAcl, + PROJECTION_SOURCE_ACL_TABLES, + replaceProjectionSourceAclSync, +} from '@sim/db/script-migrations/0021_embedding_search_connector' import { projectionSourceAclBackfillMigration as embeddingSearchConnectorMigration } from '@sim/db/script-migrations/0022_projection_source_acl_backfill' +import { projectionAclSkipUnfilledMigration } from '@sim/db/script-migrations/0023_projection_acl_skip_unfilled' +import { sleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' -import postgres, { type Sql } from 'postgres' +import postgres, { type Sql, type TransactionSql } from 'postgres' import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest' const databaseUrl = process.env.KNOWLEDGE_ACL_TEST_DATABASE_URL @@ -125,4 +131,165 @@ describe.runIf(Boolean(databaseUrl))('projection source and ACL backfill in Post const again = await backfillProjectionSourceAcl(sql, 'embedding_keyword_tin', { pauseMs: 0 }) expect(again).toMatchObject({ scanned: 0, written: 0, afterId: '', done: true }) }) + describe('a document change on chunks the backfill has not filled', () => { + /** A promise the test resolves by hand, to hold a transaction open at a chosen point. */ + const gate = () => { + let resolve = () => {} + const promise = new Promise((done) => { + resolve = done + }) + return { promise, resolve } + } + + /** Waits until `pid` is blocked on a lock, so the interleaving under test really happened. */ + const blockedOnLock = async (pid: number) => { + for (let attempt = 0; attempt < 100; attempt++) { + const [row] = await admin<{ waiting: boolean }[]>` + SELECT wait_event_type = 'Lock' AS waiting FROM pg_stat_activity WHERE pid = ${pid}` + if (row?.waiting) return true + await sleep(20) + } + return false + } + + /** The backfill's page statement, run in a transaction the test holds open. */ + const backfillPage = (tx: TransactionSql) => + tx.unsafe(`WITH page AS ( + SELECT s.id, s.document_id, d.connector_id, d.acl + FROM embedding_search s JOIN document d ON d.id = s.document_id + WHERE s.acl IS NULL ORDER BY s.id LIMIT 100 + FOR SHARE OF d + ) + UPDATE embedding_search 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`) + + let other: Sql + beforeAll(() => { + other = postgres(databaseUrl!, { + max: 1, + onnotice: () => undefined, + connection: { search_path: schemaName }, + }) + }) + afterAll(async () => { + await other?.end() + }) + + beforeEach(async () => { + /** A test below installs an older body; each starts from the current one. */ + await replaceProjectionSourceAclSync(sql) + await sql`INSERT INTO document (id, connector_id, acl) VALUES ('doc', 'src', ARRAY['u:alice'])` + for (const projection of PROJECTION_SOURCE_ACL_TABLES) { + await sql`INSERT INTO ${sql(projection)} (id, document_id, connector_id, acl) VALUES + ('filled', 'doc', 'src', ARRAY['u:alice']), ('unfilled', 'doc', NULL, NULL), + ('unfilled-sourced', 'doc', 'src', NULL)` + } + }) + + it('replaces the body a database already has when its own migration runs', async () => { + /** The body `0022` installed before this change. */ + await sql.unsafe(`CREATE OR REPLACE FUNCTION sync_projection_source_acl() + RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN + UPDATE embedding_search SET connector_id = NEW.connector_id, acl = NEW.acl + WHERE document_id = NEW.id AND enabled + AND (connector_id IS DISTINCT FROM NEW.connector_id OR acl IS DISTINCT FROM NEW.acl); + UPDATE embedding_keyword_tin SET connector_id = NEW.connector_id, acl = NEW.acl + WHERE document_id = NEW.id AND enabled + AND (connector_id IS DISTINCT FROM NEW.connector_id OR acl IS DISTINCT FROM NEW.acl); + RETURN NEW; + END; + $$`) + await sql`UPDATE document SET acl = ARRAY['u:carol'] WHERE id = 'doc'` + expect((await projected('embedding_search')).map((row) => row.acl)).toEqual([ + ['u:carol'], + ['u:carol'], + ['u:carol'], + ]) + await sql`UPDATE embedding_search SET acl = NULL WHERE id LIKE 'unfilled%'` + + await projectionAclSkipUnfilledMigration.up(sql) + await sql`UPDATE document SET acl = ARRAY['u:bob'] WHERE id = 'doc'` + expect((await projected('embedding_search')).map((row) => row.acl)).toEqual([ + ['u:bob'], + null, + null, + ]) + }) + + it('writes a changed ACL onto filled chunks only, leaving unfilled ones to their document', async () => { + await sql`UPDATE document SET acl = ARRAY['u:bob'] WHERE id = 'doc'` + for (const projection of PROJECTION_SOURCE_ACL_TABLES) { + expect(await projected(projection)).toEqual([ + { id: 'filled', connector_id: 'src', acl: ['u:bob'] }, + { id: 'unfilled', connector_id: null, acl: null }, + { id: 'unfilled-sourced', connector_id: 'src', acl: null }, + ]) + } + }) + + it('still carries a changed source onto unfilled chunks, whose source filters read the row', async () => { + await sql`UPDATE document SET connector_id = 'moved' WHERE id = 'doc'` + for (const projection of PROJECTION_SOURCE_ACL_TABLES) { + expect(await projected(projection)).toEqual([ + { id: 'filled', connector_id: 'moved', acl: ['u:alice'] }, + { id: 'unfilled', connector_id: 'moved', acl: null }, + { id: 'unfilled-sourced', connector_id: 'moved', acl: null }, + ]) + } + }) + + it('fills the current ACL when the change commits before the backfill reads the document', async () => { + await sql`UPDATE document SET acl = ARRAY['u:bob'] WHERE id = 'doc'` + await backfillProjectionSourceAcl(sql, 'embedding_search', { pauseMs: 0 }) + expect((await projected('embedding_search')).map((row) => row.acl)).toEqual([ + ['u:bob'], + ['u:bob'], + ['u:bob'], + ]) + }) + + it('fans the change out after a backfill page that read the old ACL commits', async () => { + const [pageRead, release] = [gate(), gate()] + const page = sql.begin(async (tx) => { + await backfillPage(tx) + pageRead.resolve() + await release.promise + }) + await pageRead.promise + const [{ pid }] = await other<{ pid: number }[]>`SELECT pg_backend_pid() AS pid` + /** Blocks on the page's share lock on the document until the page commits. */ + const change = other`UPDATE document SET acl = ARRAY['u:bob'] WHERE id = 'doc'`.execute() + expect(await blockedOnLock(pid)).toBe(true) + release.resolve() + await page + await change + expect((await projected('embedding_search')).map((row) => row.acl)).toEqual([ + ['u:bob'], + ['u:bob'], + ['u:bob'], + ]) + }) + + it('fills the new ACL when the backfill waits on a change that has not committed yet', async () => { + const [changed, commit] = [gate(), gate()] + const change = other.begin(async (tx) => { + await tx`UPDATE document SET acl = ARRAY['u:bob'] WHERE id = 'doc'` + changed.resolve() + await commit.promise + }) + await changed.promise + const [{ pid }] = await sql<{ pid: number }[]>`SELECT pg_backend_pid() AS pid` + const fill = backfillProjectionSourceAcl(sql, 'embedding_search', { pauseMs: 0 }) + expect(await blockedOnLock(pid)).toBe(true) + commit.resolve() + await change + await fill + expect((await projected('embedding_search')).map((row) => row.acl)).toEqual([ + ['u:bob'], + ['u:bob'], + ['u:bob'], + ]) + }) + }) }) diff --git a/packages/db/script-migrations/0021_embedding_search_connector.ts b/packages/db/script-migrations/0021_embedding_search_connector.ts index d615b3235da..b5b6e3214ef 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.ts +++ b/packages/db/script-migrations/0021_embedding_search_connector.ts @@ -1,7 +1,7 @@ import { createLogger } from '@sim/logger' import { sleep } from '@sim/utils/helpers' import { backoffWithJitter } from '@sim/utils/retry' -import postgres, { type Sql } from 'postgres' +import postgres, { type Sql, type TransactionSql } from 'postgres' const logger = createLogger('ProjectionSourceAcl') @@ -69,6 +69,36 @@ const PROGRESS_EVERY_PAGES = 100 export const PROJECTION_SOURCE_ACL_TABLES = ['embedding_search', 'embedding_keyword_tin'] as const export type ProjectionSourceAclTable = (typeof PROJECTION_SOURCE_ACL_TABLES)[number] +/** + * The document trigger's body: fans a document's source and ACL out to its enabled chunks. + * + * A chunk the backfill has not filled yet (`acl IS NULL`) keeps a NULL ACL. Search decides such a + * row on its document, so writing the ACL there changes no answer, while every write to + * `embedding_search` re-inserts the row into its vector index: a document whose ACL changed would + * otherwise rewrite each of its unfilled chunks inside the writer's statement. The backfill fills + * the row later from the document under a share lock, so it copies whichever ACL is current. A + * document that moves to another source still carries the source onto its unfilled chunks, because + * source filters read it from the row; an ACL change alone leaves them untouched. + */ +export async function replaceProjectionSourceAclSync(sql: Sql | TransactionSql): Promise { + const fanOut = (projection: ProjectionSourceAclTable) => ` + UPDATE ${projection} + SET connector_id = NEW.connector_id, acl = CASE WHEN acl IS NULL THEN NULL ELSE NEW.acl END + WHERE document_id = NEW.id AND enabled + AND CASE WHEN acl IS NULL + THEN moved AND connector_id IS DISTINCT FROM NEW.connector_id + ELSE connector_id IS DISTINCT FROM NEW.connector_id OR acl IS DISTINCT FROM NEW.acl + END;` + await sql.unsafe(`CREATE OR REPLACE FUNCTION sync_projection_source_acl() + RETURNS trigger LANGUAGE plpgsql AS $$ + DECLARE + moved boolean := TG_OP = 'UPDATE' AND OLD.connector_id IS DISTINCT FROM NEW.connector_id; + BEGIN${PROJECTION_SOURCE_ACL_TABLES.map(fanOut).join('')} + RETURN NEW; + END; + $$`) +} + /** * Carries a chunk's source and ACL onto the ranking projections and keeps them there. * @@ -85,18 +115,7 @@ export type ProjectionSourceAclTable = (typeof PROJECTION_SOURCE_ACL_TABLES)[num export async function installProjectionSourceAcl(sql: Sql): Promise { await sql.begin(async (tx) => { await tx.unsafe("SET LOCAL lock_timeout = '5s'") - await tx.unsafe(`CREATE OR REPLACE FUNCTION sync_projection_source_acl() - RETURNS trigger LANGUAGE plpgsql AS $$ - BEGIN - UPDATE embedding_search SET connector_id = NEW.connector_id, acl = NEW.acl - WHERE document_id = NEW.id AND enabled - AND (connector_id IS DISTINCT FROM NEW.connector_id OR acl IS DISTINCT FROM NEW.acl); - UPDATE embedding_keyword_tin SET connector_id = NEW.connector_id, acl = NEW.acl - WHERE document_id = NEW.id AND enabled - AND (connector_id IS DISTINCT FROM NEW.connector_id OR acl IS DISTINCT FROM NEW.acl); - RETURN NEW; - END; - $$`) + await replaceProjectionSourceAclSync(tx) await tx.unsafe(`CREATE OR REPLACE TRIGGER projection_source_acl_sync AFTER INSERT OR UPDATE OF connector_id, acl ON document FOR EACH ROW EXECUTE FUNCTION sync_projection_source_acl()`) diff --git a/packages/db/script-migrations/0023_projection_acl_skip_unfilled.ts b/packages/db/script-migrations/0023_projection_acl_skip_unfilled.ts new file mode 100644 index 00000000000..00a20a35eda --- /dev/null +++ b/packages/db/script-migrations/0023_projection_acl_skip_unfilled.ts @@ -0,0 +1,16 @@ +import { replaceProjectionSourceAclSync } from '@sim/db/script-migrations/0021_embedding_search_connector' +import type { ScriptMigration } from '@sim/db/script-migrations/types' + +/** + * Replaces the document trigger's body so a document's ACL change no longer writes the ACL onto + * chunks the projection backfill has not filled yet; see {@link replaceProjectionSourceAclSync}. + * The function is replaced in place, so the trigger that calls it and every other object from + * `0022_projection_source_acl_backfill` stay as they are. A database that runs `0022` now installs + * the same body, so this is a no-op there. + */ +export const projectionAclSkipUnfilledMigration: ScriptMigration = { + name: '0023_projection_acl_skip_unfilled', + async up(sql) { + await replaceProjectionSourceAclSync(sql) + }, +} diff --git a/packages/db/script-migrations/index.ts b/packages/db/script-migrations/index.ts index 85a3b2f099c..4304a9336d4 100644 --- a/packages/db/script-migrations/index.ts +++ b/packages/db/script-migrations/index.ts @@ -6,6 +6,7 @@ import { indexSearchDocumentsMigration } from '@sim/db/script-migrations/0017_in 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 { projectionSourceAclBackfillMigration } from '@sim/db/script-migrations/0022_projection_source_acl_backfill' +import { projectionAclSkipUnfilledMigration } from '@sim/db/script-migrations/0023_projection_acl_skip_unfilled' import type { Sql } from 'postgres' import { backfillTableOrderKeys } from './0001_backfill_table_order_keys' import { backfillPausedBillingAttribution } from './0002_backfill_paused_billing_attribution' @@ -46,6 +47,8 @@ export const scriptMigrations: readonly ScriptMigration[] = [ tinKeywordProjectionMigration, /** 0022 supersedes 0021, whose synchronous backfill could not finish inside a deploy. */ projectionSourceAclBackfillMigration, + /** 0023 stops document ACL changes from writing chunks the 0022 backfill has not filled. */ + projectionAclSkipUnfilledMigration, ] /** From 20bb9feafad6d19cf2766c71f20a617ca6f856eb Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 15:39:37 -0700 Subject: [PATCH 2/3] test(knowledge): expect the projection ACL trigger migration in the upgrade run --- .../0016_backfill_search_vectors.postgres.test.ts | 1 + 1 file changed, 1 insertion(+) 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 6345cab6c3f..4565d84e4c2 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 @@ -393,6 +393,7 @@ describe.runIf(Boolean(databaseUrl))('search projection upgrade in PostgreSQL', { name: '0019_tin_keyword_projection' }, { name: '0021_embedding_search_connector' }, { name: '0022_projection_source_acl_backfill' }, + { name: '0023_projection_acl_skip_unfilled' }, ]) 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 From 8906c421051091600cf4d430f2a89d0be4abfd92 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 15:45:50 -0700 Subject: [PATCH 3/3] chore(ci): create the ACL test database in the projection trigger step itself --- .github/workflows/test-build.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index beceabc40ff..b9a844e42e2 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -303,7 +303,9 @@ jobs: working-directory: packages/db env: KNOWLEDGE_ACL_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_acl_test - run: bunx vitest run script-migrations/0021_embedding_search_connector.postgres.test.ts + run: | + bun -e 'import postgres from "postgres"; const sql = postgres(process.env.DATABASE_URL); const [row] = await sql`SELECT 1 FROM pg_database WHERE datname = ${"sim_acl_test"}`; if (!row) await sql`CREATE DATABASE sim_acl_test`; await sql.end()' + bunx vitest run script-migrations/0021_embedding_search_connector.postgres.test.ts test-build: name: Lint and Test