fix(knowledge): stop document ACL changes from rewriting unfilled projection chunks

.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' }}

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',
     ])
   })
 })

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

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<void>((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'],
+      ])
+    })
+  })
 })

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<void> {
+  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<void> {
   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()`)

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)
+  },
+}

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,
 ]
 
 /**