diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index b7d6b0d6dc5..b9a844e42e2 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -299,6 +299,14 @@ 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: | + 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 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/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 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, ] /**