diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index b7d6b0d6dc5..7be554390e0 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -297,7 +297,8 @@ jobs: bunx vitest run --mode integration lib/knowledge/access/group-membership.integration.ts bunx vitest run \ lib/knowledge/access/predicate.postgres.test.ts \ - lib/knowledge/connectors/external-directory.postgres.test.ts + lib/knowledge/connectors/external-directory.postgres.test.ts \ + lib/knowledge/connectors/sync-persistence.postgres.test.ts test-build: name: Lint and Test diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts new file mode 100644 index 00000000000..0b73882e0db --- /dev/null +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts @@ -0,0 +1,235 @@ +/** + * @vitest-environment node + */ +import { installProjectionSourceAcl } from '@sim/db/script-migrations/0021_embedding_search_connector' +import { generateId } from '@sim/utils/id' +import postgres, { type Sql } from 'postgres' +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' + +vi.unmock('drizzle-orm') +vi.unmock('@sim/db/schema') +vi.mock('@/lib/knowledge/documents/service', () => ({ hardDeleteDocuments: vi.fn() })) +vi.mock('@/lib/uploads', () => ({ StorageService: {} })) +vi.mock('@/lib/uploads/core/storage-service', () => ({ deleteFile: vi.fn() })) +vi.mock('@/lib/uploads/server/metadata', () => ({ deleteFileMetadata: vi.fn() })) +vi.mock('@/connectors/registry.server', () => ({ CONNECTOR_REGISTRY: {} })) + +const { drizzle } = await import('drizzle-orm/postgres-js') +const schema = await import('@sim/db/schema') +const { persistDocumentAcls } = await import('@/lib/knowledge/connectors/sync-persistence') + +const databaseUrl = process.env.KNOWLEDGE_ACL_TEST_DATABASE_URL + +const ALICE = 'u:alice@corp.com' +const BOB = 'u:bob@corp.com' + +/** + * The document and projections carry only the columns the ACL write and the projection trigger + * touch. A projection row whose `acl` is NULL is one the backfill has not filled yet: the trigger + * rewrites it on any ACL assignment, because NULL is distinct from every ACL. + */ +describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => { + let admin: Sql + let sql: Sql + const schemaName = `acl_write_${generateId().replaceAll('-', '')}` + + const projected = () => + sql<{ id: string; acl: string[] | null }[]>` + SELECT id, acl FROM embedding_search + UNION ALL SELECT id, acl FROM embedding_keyword_tin ORDER BY id` + + const persist = (acls: Map) => + persistDocumentAcls('admin', acls, drizzle(sql, { schema })) + + 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('ACL write 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, external_id text, connector_id text, + acl text[] NOT NULL DEFAULT '{ws}', acl_requirements jsonb NOT NULL DEFAULT '[]', + acl_verified_at timestamp + )` + 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 installProjectionSourceAcl(sql) + 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` + }, 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`INSERT INTO document (id, external_id, connector_id, acl, acl_verified_at) VALUES + ('doc-same', 'file-same', 'admin', ARRAY[${ALICE}], now() - interval '1 day'), + ('doc-moved', 'file-moved', 'admin', ARRAY[${ALICE}], now() - interval '1 day')` + for (const projection of ['embedding_search', 'embedding_keyword_tin']) { + const prefix = projection === 'embedding_search' ? 'vec' : 'kw' + await sql`INSERT INTO ${sql(projection)} (id, document_id, connector_id, acl) VALUES + (${`${prefix}-same-unfilled`}, 'doc-same', NULL, NULL), + (${`${prefix}-same-filled`}, 'doc-same', 'admin', ARRAY[${ALICE}]), + (${`${prefix}-moved-unfilled`}, 'doc-moved', NULL, NULL), + (${`${prefix}-moved-filled`}, 'doc-moved', 'admin', ARRAY[${ALICE}])` + } + }) + + it('refreshes the evidence of an unchanged ACL without rewriting any chunk projection row', async () => { + await expect(persist(new Map([['file-same', [ALICE]]]))).resolves.toEqual({ + updated: 1, + rejected: 0, + }) + + const [stored] = await sql<{ acl: string[]; fresh: boolean }[]>` + SELECT acl, acl_verified_at > now() AT TIME ZONE 'UTC' - interval '1 minute' AS fresh + FROM document WHERE id = 'doc-same'` + expect(stored).toEqual({ acl: [ALICE], fresh: true }) + expect((await projected()).filter((row) => row.id.includes('-same-'))).toEqual([ + { id: 'kw-same-filled', acl: [ALICE] }, + { id: 'kw-same-unfilled', acl: null }, + { id: 'vec-same-filled', acl: [ALICE] }, + { id: 'vec-same-unfilled', acl: null }, + ]) + }) + + it('propagates a changed ACL to every chunk projection row, filled or not', async () => { + await expect(persist(new Map([['file-moved', [BOB]]]))).resolves.toEqual({ + updated: 1, + rejected: 0, + }) + + const [stored] = await sql<{ acl: string[] }[]>`SELECT acl FROM document WHERE id = 'doc-moved'` + expect(stored.acl).toEqual([BOB]) + expect((await projected()).filter((row) => row.id.includes('-moved-'))).toEqual([ + { id: 'kw-moved-filled', acl: [BOB] }, + { id: 'kw-moved-unfilled', acl: [BOB] }, + { id: 'vec-moved-filled', acl: [BOB] }, + { id: 'vec-moved-unfilled', acl: [BOB] }, + ]) + }) + + it('treats a changed restriction under the same primary ACL as a change', async () => { + await persistDocumentAcls( + 'admin', + new Map([['file-same', { acl: [ALICE], requirements: [['g:confluence:tenant:space']] }]]), + drizzle(sql, { schema }) + ) + + const [stored] = await sql<{ requirements: string[][] }[]>` + SELECT acl_requirements AS requirements FROM document WHERE id = 'doc-same'` + expect(stored.requirements).toEqual([[ALICE], ['g:confluence:tenant:space']]) + }) + + it('writes each document once when one page mixes unchanged and changed ACLs', async () => { + await expect( + persist( + new Map([ + ['file-same', [ALICE]], + ['file-moved', [ALICE, BOB]], + ]) + ) + ).resolves.toEqual({ updated: 2, rejected: 0 }) + + const rows = await sql< + { id: string; acl: string[] }[] + >`SELECT id, acl FROM document ORDER BY id` + expect(rows).toEqual([ + { id: 'doc-moved', acl: [ALICE, BOB] }, + { id: 'doc-same', acl: [ALICE] }, + ]) + expect( + (await projected()).filter((row) => row.id.endsWith('-unfilled')).map((row) => row.acl) + ).toEqual([[ALICE, BOB], null, [ALICE, BOB], null]) + }) + it('writes a changed ACL group larger than one change batch completely', async () => { + const ids = Array.from( + { length: 60 }, + (_unused, index) => `bulk-${String(index).padStart(2, '0')}` + ) + await sql`INSERT INTO document ${sql( + ids.map((id) => ({ id: `doc-${id}`, external_id: id, connector_id: 'admin', acl: [ALICE] })) + )}` + await sql`INSERT INTO embedding_search ${sql( + ids.map((id) => ({ + id: `vec-${id}`, + document_id: `doc-${id}`, + connector_id: 'admin', + acl: [ALICE], + })) + )}` + + await expect(persist(new Map(ids.map((id) => [id, [BOB]])))).resolves.toEqual({ + updated: ids.length, + rejected: 0, + }) + + const documents = await sql<{ acl: string[] }[]>` + SELECT acl FROM document WHERE id LIKE 'doc-bulk-%'` + expect(documents).toHaveLength(ids.length) + expect(documents.every((row) => row.acl.join() === BOB)).toBe(true) + const chunks = await sql<{ acl: string[] }[]>` + SELECT acl FROM embedding_search WHERE id LIKE 'vec-bulk-%'` + expect(chunks).toHaveLength(ids.length) + expect(chunks.every((row) => row.acl.join() === BOB)).toBe(true) + }) + + it.each([ + { name: 'verified within this generation', verifiedOffsetMs: 1_000, preserved: true }, + { name: 'verified before this generation', verifiedOffsetMs: -1_000, preserved: false }, + { name: 'never verified', verifiedOffsetMs: null, preserved: false }, + ])( + 'applies the unresolved-evidence guard to an ACL the source could not answer: $name', + async ({ verifiedOffsetMs, preserved }) => { + const [clock] = await sql<{ now: string }[]>` + SELECT (now() AT TIME ZONE 'UTC')::text AS now` + const generationStartedAt = new Date(`${clock.now}Z`) + generationStartedAt.setTime(generationStartedAt.getTime() - 60_000) + const verifiedAt = + verifiedOffsetMs === null + ? null + : new Date(generationStartedAt.getTime() + verifiedOffsetMs).toISOString() + await sql`UPDATE document SET acl_verified_at = ${verifiedAt}::timestamptz AT TIME ZONE 'UTC' + WHERE id = 'doc-same'` + + const result = await persistDocumentAcls( + 'admin', + new Map([['file-same', []]]), + drizzle(sql, { schema }), + { + unresolvedExternalIds: new Set(['file-same']), + generationStartedAt, + } + ) + + expect(result).toEqual({ updated: preserved ? 0 : 1, rejected: 0 }) + const [stored] = await sql<{ acl: string[]; verified: boolean }[]>` + SELECT acl, acl_verified_at IS NOT NULL AS verified FROM document WHERE id = 'doc-same'` + expect(stored).toEqual( + preserved ? { acl: [ALICE], verified: true } : { acl: [], verified: false } + ) + const filled = (await projected()) + .filter((row) => row.id.endsWith('-same-filled')) + .map((row) => row.acl) + expect(filled).toEqual(preserved ? [[ALICE], [ALICE]] : [[], []]) + } + ) +}) diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts index fd90b59c1d4..ac84bae837a 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts @@ -71,13 +71,13 @@ describe('persistDocumentAcls', () => { * corpus every time somebody joined a group. */ it('refreshes only access fields, so no document is re-embedded', async () => { - queueUpdatedCounts(1) + queueUpdatedCounts(0, 1) await persistDocumentAcls(CONNECTOR, new Map([['file-1', ['u:alice@corp.com']]])) expect(dbChainMockFns.update).toHaveBeenCalledWith(schemaMock.document) - expect(dbChainMockFns.set).toHaveBeenCalledTimes(1) - expect(dbChainMockFns.set).toHaveBeenCalledWith({ + expect(dbChainMockFns.set).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.set).toHaveBeenNthCalledWith(2, { acl: ['u:alice@corp.com'], aclRequirements: [], aclVerifiedAt: expect.objectContaining({ @@ -87,6 +87,25 @@ describe('persistDocumentAcls', () => { }) }) + /** + * Assigning `acl` fires the projection trigger, which rewrites every chunk row whose copy + * differs, so a document whose ACL did not change must only have its evidence refreshed. + */ + it('refreshes the evidence of an unchanged ACL without assigning it', async () => { + queueUpdatedCounts(1, 0) + + await expect( + persistDocumentAcls(CONNECTOR, new Map([['file-1', ['u:alice@corp.com']]])) + ).resolves.toEqual({ updated: 1, rejected: 0 }) + + expect(dbChainMockFns.set).toHaveBeenNthCalledWith(1, { + aclVerifiedAt: expect.objectContaining({ + strings: ["statement_timestamp() AT TIME ZONE 'UTC'"], + values: [], + }), + }) + }) + it('reports how many documents received current permission evidence', async () => { queueUpdatedCounts(2) @@ -105,8 +124,8 @@ describe('persistDocumentAcls', () => { * Files under one folder overwhelmingly share an ACL, so grouping is what * keeps a crawl of thousands to a handful of statements. */ - it('writes one statement per distinct ACL, not per document', async () => { - queueUpdatedCounts(2, 1) + it('writes one refresh and one change statement per distinct ACL, not per document', async () => { + queueUpdatedCounts(0, 2, 0, 1) await persistDocumentAcls( CONNECTOR, @@ -117,8 +136,8 @@ describe('persistDocumentAcls', () => { ]) ) - expect(dbChainMockFns.set).toHaveBeenCalledTimes(2) - expect(dbChainMockFns.set).toHaveBeenNthCalledWith(1, { + expect(dbChainMockFns.set).toHaveBeenCalledTimes(4) + expect(dbChainMockFns.set).toHaveBeenNthCalledWith(2, { acl: ['u:alice@corp.com'], aclRequirements: [], aclVerifiedAt: expect.objectContaining({ @@ -126,7 +145,7 @@ describe('persistDocumentAcls', () => { values: [], }), }) - expect(dbChainMockFns.set).toHaveBeenNthCalledWith(2, { + expect(dbChainMockFns.set).toHaveBeenNthCalledWith(4, { acl: ['u:bob@corp.com'], aclRequirements: [], aclVerifiedAt: expect.objectContaining({ @@ -147,8 +166,8 @@ describe('persistDocumentAcls', () => { ]) ) - expect(dbChainMockFns.set).toHaveBeenCalledTimes(1) - expect(dbChainMockFns.set).toHaveBeenCalledWith({ + expect(dbChainMockFns.set).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.set).toHaveBeenNthCalledWith(2, { acl: ['u:alice@corp.com', 'u:bob@corp.com'], aclRequirements: [], aclVerifiedAt: expect.objectContaining({ @@ -174,7 +193,7 @@ describe('persistDocumentAcls', () => { }) it('retains an empty restriction and separately persists different clauses', async () => { - queueUpdatedCounts(1, 1) + queueUpdatedCounts(0, 1, 0, 1) await persistDocumentAcls( CONNECTOR, new Map([ @@ -182,8 +201,8 @@ describe('persistDocumentAcls', () => { ['file-2', { acl: ['u:alice@corp.com'], requirements: [['g:confluence:site:team']] }], ]) ) - expect(dbChainMockFns.set).toHaveBeenCalledTimes(2) - expect(dbChainMockFns.set).toHaveBeenNthCalledWith(1, { + expect(dbChainMockFns.set).toHaveBeenCalledTimes(4) + expect(dbChainMockFns.set).toHaveBeenNthCalledWith(2, { acl: ['u:alice@corp.com'], aclRequirements: [['u:alice@corp.com'], []], aclVerifiedAt: expect.objectContaining({ @@ -191,7 +210,7 @@ describe('persistDocumentAcls', () => { values: [], }), }) - expect(dbChainMockFns.set).toHaveBeenNthCalledWith(2, { + expect(dbChainMockFns.set).toHaveBeenNthCalledWith(4, { acl: ['u:alice@corp.com'], aclRequirements: [['u:alice@corp.com'], ['g:confluence:site:team']], aclVerifiedAt: expect.objectContaining({ diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.ts index 48946ea7067..b675700125d 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.ts @@ -4,7 +4,7 @@ import { createLogger } from '@sim/logger' import { chunkArray } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' import { truncateAtCodePoint } from '@sim/utils/string' -import { and, eq, exists, inArray, isNull, lt, or, sql } from 'drizzle-orm' +import { and, eq, exists, inArray, isNull, lt, not, or, sql } from 'drizzle-orm' import { getInternalApiBaseUrl } from '@/lib/core/utils/urls' import type { DbOrTx } from '@/lib/db/types' import { textArrayLiteral } from '@/lib/knowledge/access/predicate' @@ -79,12 +79,22 @@ export async function restoreWorkspaceDocumentAcls( } /** - * Documents whose ACL is rewritten per statement. Documents are grouped by - * identical ACL first — files under one folder overwhelmingly share theirs — so - * a crawl of thousands usually resolves to a handful of statements. + * Documents whose permission evidence is refreshed per statement. Documents are + * grouped by identical ACL first — files under one folder overwhelmingly share + * theirs — so a crawl of thousands usually resolves to a handful of statements. + * A refresh never assigns `acl`, so it fires no projection fan-out. */ const ACL_WRITE_BATCH_SIZE = 500 +/** + * Documents whose ACL actually changes, per statement. Assigning `acl` fires the + * document trigger that copies it onto every chunk's search projection rows, and + * each of those rows is re-inserted into the vector index, so one statement costs + * the chunks of every document in it rather than the documents. Kept small so a + * page of changed documents cannot outrun the statement timeout. + */ +const ACL_CHANGE_BATCH_SIZE = 25 + export interface DocumentAclWriteResult { /** Documents whose ACL or authoritative evidence timestamp was refreshed. */ updated: number @@ -94,8 +104,11 @@ export interface DocumentAclWriteResult { /** * Permission-only changes must not trigger re-embedding. Unchanged ACLs still refresh - * their evidence timestamp; failed fetches cannot extend it. Malformed or oversized - * ACLs are stored as unreadable so the previous grant cannot survive failed verification. + * their evidence timestamp, without assigning `acl`: every assignment fires the + * projection trigger, which rewrites any chunk row whose copy differs, including one + * the projection backfill has not filled yet. Failed fetches cannot extend it. + * Malformed or oversized ACLs are stored as unreadable so the previous grant cannot + * survive failed verification. * An unresolved duplicate may retain evidence verified during this durable crawl, * without refreshing its timestamp; explicit empty ACLs always revoke access. */ @@ -145,26 +158,36 @@ export async function persistDocumentAcls( let updated = 0 for (const { acl, requirements, externalIds, unresolved } of byAcl.values()) { + const aclVerifiedAt = acl.length > 0 ? sql`statement_timestamp() AT TIME ZONE 'UTC'` : null + const evidenceGuard = + unresolved && evidence + ? or( + isNull(document.aclVerifiedAt), + lt(document.aclVerifiedAt, evidence.generationStartedAt) + ) + : undefined + const stored = sql`(${document.acl} IS NOT DISTINCT FROM ${textArrayLiteral(acl)} AND ${document.aclRequirements} IS NOT DISTINCT FROM ${JSON.stringify(requirements)}::jsonb)` + const target = (batch: string[], unchanged: boolean) => + and( + eq(document.connectorId, connectorId), + inArray(document.externalId, batch), + unchanged ? stored : not(stored), + evidenceGuard + ) + /** Refreshed first, so a row the change write below has just rewritten is not counted twice. */ for (const batch of chunkArray(externalIds, ACL_WRITE_BATCH_SIZE)) { const rows = await executor .update(document) - .set({ - acl, - aclRequirements: requirements, - aclVerifiedAt: acl.length > 0 ? sql`statement_timestamp() AT TIME ZONE 'UTC'` : null, - }) - .where( - and( - eq(document.connectorId, connectorId), - inArray(document.externalId, batch), - unresolved && evidence - ? or( - isNull(document.aclVerifiedAt), - lt(document.aclVerifiedAt, evidence.generationStartedAt) - ) - : undefined - ) - ) + .set({ aclVerifiedAt }) + .where(target(batch, true)) + .returning({ id: document.id }) + updated += rows.length + } + for (const batch of chunkArray(externalIds, ACL_CHANGE_BATCH_SIZE)) { + const rows = await executor + .update(document) + .set({ acl, aclRequirements: requirements, aclVerifiedAt }) + .where(target(batch, false)) .returning({ id: document.id }) updated += rows.length }