|
| 1 | +/** |
| 2 | + * @vitest-environment node |
| 3 | + */ |
| 4 | +import { installProjectionSourceAcl } from '@sim/db/script-migrations/0021_embedding_search_connector' |
| 5 | +import { generateId } from '@sim/utils/id' |
| 6 | +import postgres, { type Sql } from 'postgres' |
| 7 | +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' |
| 8 | + |
| 9 | +vi.unmock('drizzle-orm') |
| 10 | +vi.unmock('@sim/db/schema') |
| 11 | +vi.mock('@/lib/knowledge/documents/service', () => ({ hardDeleteDocuments: vi.fn() })) |
| 12 | +vi.mock('@/lib/uploads', () => ({ StorageService: {} })) |
| 13 | +vi.mock('@/lib/uploads/core/storage-service', () => ({ deleteFile: vi.fn() })) |
| 14 | +vi.mock('@/lib/uploads/server/metadata', () => ({ deleteFileMetadata: vi.fn() })) |
| 15 | +vi.mock('@/connectors/registry.server', () => ({ CONNECTOR_REGISTRY: {} })) |
| 16 | + |
| 17 | +const { drizzle } = await import('drizzle-orm/postgres-js') |
| 18 | +const schema = await import('@sim/db/schema') |
| 19 | +const { persistDocumentAcls } = await import('@/lib/knowledge/connectors/sync-persistence') |
| 20 | + |
| 21 | +const databaseUrl = process.env.KNOWLEDGE_ACL_TEST_DATABASE_URL |
| 22 | + |
| 23 | +const ALICE = 'u:alice@corp.com' |
| 24 | +const BOB = 'u:bob@corp.com' |
| 25 | + |
| 26 | +/** |
| 27 | + * The document and projections carry only the columns the ACL write and the projection trigger |
| 28 | + * touch. A projection row whose `acl` is NULL is one the backfill has not filled yet: the trigger |
| 29 | + * rewrites it on any ACL assignment, because NULL is distinct from every ACL. |
| 30 | + */ |
| 31 | +describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => { |
| 32 | + let admin: Sql |
| 33 | + let sql: Sql |
| 34 | + const schemaName = `acl_write_${generateId().replaceAll('-', '')}` |
| 35 | + |
| 36 | + const projected = () => |
| 37 | + sql<{ id: string; acl: string[] | null }[]>` |
| 38 | + SELECT id, acl FROM embedding_search |
| 39 | + UNION ALL SELECT id, acl FROM embedding_keyword_tin ORDER BY id` |
| 40 | + |
| 41 | + const persist = (acls: Map<string, string[]>) => |
| 42 | + persistDocumentAcls('admin', acls, drizzle(sql, { schema })) |
| 43 | + |
| 44 | + beforeAll(async () => { |
| 45 | + const url = new URL(databaseUrl!) |
| 46 | + if ( |
| 47 | + !['localhost', '127.0.0.1'].includes(url.hostname) || |
| 48 | + !url.pathname.startsWith('/sim_acl_test') |
| 49 | + ) { |
| 50 | + throw new Error('ACL write tests require a disposable local integration database') |
| 51 | + } |
| 52 | + admin = postgres(url.toString(), { max: 1, onnotice: () => undefined }) |
| 53 | + await admin.unsafe(`CREATE SCHEMA "${schemaName}"`) |
| 54 | + sql = postgres(url.toString(), { |
| 55 | + max: 1, |
| 56 | + onnotice: () => undefined, |
| 57 | + connection: { search_path: schemaName }, |
| 58 | + }) |
| 59 | + await sql`CREATE TABLE document ( |
| 60 | + id text PRIMARY KEY, external_id text, connector_id text, |
| 61 | + acl text[] NOT NULL DEFAULT '{ws}', acl_requirements jsonb NOT NULL DEFAULT '[]', |
| 62 | + acl_verified_at timestamp |
| 63 | + )` |
| 64 | + for (const projection of ['embedding_search', 'embedding_keyword_tin']) { |
| 65 | + await sql`CREATE TABLE ${sql(projection)} ( |
| 66 | + id text PRIMARY KEY, document_id text NOT NULL, enabled boolean NOT NULL DEFAULT true, |
| 67 | + connector_id text, acl text[] |
| 68 | + )` |
| 69 | + } |
| 70 | + await installProjectionSourceAcl(sql) |
| 71 | + await sql`ALTER TABLE embedding_search DISABLE TRIGGER embedding_search_source_acl_set` |
| 72 | + await sql`ALTER TABLE embedding_keyword_tin DISABLE TRIGGER embedding_keyword_tin_source_acl_set` |
| 73 | + }, 60_000) |
| 74 | + |
| 75 | + afterAll(async () => { |
| 76 | + await sql?.end() |
| 77 | + await admin?.unsafe(`DROP SCHEMA IF EXISTS "${schemaName}" CASCADE`) |
| 78 | + await admin?.end() |
| 79 | + }) |
| 80 | + |
| 81 | + beforeEach(async () => { |
| 82 | + await sql`TRUNCATE embedding_search, embedding_keyword_tin, document` |
| 83 | + await sql`INSERT INTO document (id, external_id, connector_id, acl, acl_verified_at) VALUES |
| 84 | + ('doc-same', 'file-same', 'admin', ARRAY[${ALICE}], now() - interval '1 day'), |
| 85 | + ('doc-moved', 'file-moved', 'admin', ARRAY[${ALICE}], now() - interval '1 day')` |
| 86 | + for (const projection of ['embedding_search', 'embedding_keyword_tin']) { |
| 87 | + const prefix = projection === 'embedding_search' ? 'vec' : 'kw' |
| 88 | + await sql`INSERT INTO ${sql(projection)} (id, document_id, connector_id, acl) VALUES |
| 89 | + (${`${prefix}-same-unfilled`}, 'doc-same', NULL, NULL), |
| 90 | + (${`${prefix}-same-filled`}, 'doc-same', 'admin', ARRAY[${ALICE}]), |
| 91 | + (${`${prefix}-moved-unfilled`}, 'doc-moved', NULL, NULL), |
| 92 | + (${`${prefix}-moved-filled`}, 'doc-moved', 'admin', ARRAY[${ALICE}])` |
| 93 | + } |
| 94 | + }) |
| 95 | + |
| 96 | + it('refreshes the evidence of an unchanged ACL without rewriting any chunk projection row', async () => { |
| 97 | + await expect(persist(new Map([['file-same', [ALICE]]]))).resolves.toEqual({ |
| 98 | + updated: 1, |
| 99 | + rejected: 0, |
| 100 | + }) |
| 101 | + |
| 102 | + const [stored] = await sql<{ acl: string[]; fresh: boolean }[]>` |
| 103 | + SELECT acl, acl_verified_at > now() AT TIME ZONE 'UTC' - interval '1 minute' AS fresh |
| 104 | + FROM document WHERE id = 'doc-same'` |
| 105 | + expect(stored).toEqual({ acl: [ALICE], fresh: true }) |
| 106 | + expect((await projected()).filter((row) => row.id.includes('-same-'))).toEqual([ |
| 107 | + { id: 'kw-same-filled', acl: [ALICE] }, |
| 108 | + { id: 'kw-same-unfilled', acl: null }, |
| 109 | + { id: 'vec-same-filled', acl: [ALICE] }, |
| 110 | + { id: 'vec-same-unfilled', acl: null }, |
| 111 | + ]) |
| 112 | + }) |
| 113 | + |
| 114 | + it('propagates a changed ACL to every chunk projection row, filled or not', async () => { |
| 115 | + await expect(persist(new Map([['file-moved', [BOB]]]))).resolves.toEqual({ |
| 116 | + updated: 1, |
| 117 | + rejected: 0, |
| 118 | + }) |
| 119 | + |
| 120 | + const [stored] = await sql<{ acl: string[] }[]>`SELECT acl FROM document WHERE id = 'doc-moved'` |
| 121 | + expect(stored.acl).toEqual([BOB]) |
| 122 | + expect((await projected()).filter((row) => row.id.includes('-moved-'))).toEqual([ |
| 123 | + { id: 'kw-moved-filled', acl: [BOB] }, |
| 124 | + { id: 'kw-moved-unfilled', acl: [BOB] }, |
| 125 | + { id: 'vec-moved-filled', acl: [BOB] }, |
| 126 | + { id: 'vec-moved-unfilled', acl: [BOB] }, |
| 127 | + ]) |
| 128 | + }) |
| 129 | + |
| 130 | + it('treats a changed restriction under the same primary ACL as a change', async () => { |
| 131 | + await persistDocumentAcls( |
| 132 | + 'admin', |
| 133 | + new Map([['file-same', { acl: [ALICE], requirements: [['g:confluence:tenant:space']] }]]), |
| 134 | + drizzle(sql, { schema }) |
| 135 | + ) |
| 136 | + |
| 137 | + const [stored] = await sql<{ requirements: string[][] }[]>` |
| 138 | + SELECT acl_requirements AS requirements FROM document WHERE id = 'doc-same'` |
| 139 | + expect(stored.requirements).toEqual([[ALICE], ['g:confluence:tenant:space']]) |
| 140 | + }) |
| 141 | + |
| 142 | + it('writes each document once when one page mixes unchanged and changed ACLs', async () => { |
| 143 | + await expect( |
| 144 | + persist( |
| 145 | + new Map([ |
| 146 | + ['file-same', [ALICE]], |
| 147 | + ['file-moved', [ALICE, BOB]], |
| 148 | + ]) |
| 149 | + ) |
| 150 | + ).resolves.toEqual({ updated: 2, rejected: 0 }) |
| 151 | + |
| 152 | + const rows = await sql< |
| 153 | + { id: string; acl: string[] }[] |
| 154 | + >`SELECT id, acl FROM document ORDER BY id` |
| 155 | + expect(rows).toEqual([ |
| 156 | + { id: 'doc-moved', acl: [ALICE, BOB] }, |
| 157 | + { id: 'doc-same', acl: [ALICE] }, |
| 158 | + ]) |
| 159 | + expect( |
| 160 | + (await projected()).filter((row) => row.id.endsWith('-unfilled')).map((row) => row.acl) |
| 161 | + ).toEqual([[ALICE, BOB], null, [ALICE, BOB], null]) |
| 162 | + }) |
| 163 | + it('writes a changed ACL group larger than one change batch completely', async () => { |
| 164 | + const ids = Array.from( |
| 165 | + { length: 60 }, |
| 166 | + (_unused, index) => `bulk-${String(index).padStart(2, '0')}` |
| 167 | + ) |
| 168 | + await sql`INSERT INTO document ${sql( |
| 169 | + ids.map((id) => ({ id: `doc-${id}`, external_id: id, connector_id: 'admin', acl: [ALICE] })) |
| 170 | + )}` |
| 171 | + await sql`INSERT INTO embedding_search ${sql( |
| 172 | + ids.map((id) => ({ |
| 173 | + id: `vec-${id}`, |
| 174 | + document_id: `doc-${id}`, |
| 175 | + connector_id: 'admin', |
| 176 | + acl: [ALICE], |
| 177 | + })) |
| 178 | + )}` |
| 179 | + |
| 180 | + await expect(persist(new Map(ids.map((id) => [id, [BOB]])))).resolves.toEqual({ |
| 181 | + updated: ids.length, |
| 182 | + rejected: 0, |
| 183 | + }) |
| 184 | + |
| 185 | + const documents = await sql<{ acl: string[] }[]>` |
| 186 | + SELECT acl FROM document WHERE id LIKE 'doc-bulk-%'` |
| 187 | + expect(documents).toHaveLength(ids.length) |
| 188 | + expect(documents.every((row) => row.acl.join() === BOB)).toBe(true) |
| 189 | + const chunks = await sql<{ acl: string[] }[]>` |
| 190 | + SELECT acl FROM embedding_search WHERE id LIKE 'vec-bulk-%'` |
| 191 | + expect(chunks).toHaveLength(ids.length) |
| 192 | + expect(chunks.every((row) => row.acl.join() === BOB)).toBe(true) |
| 193 | + }) |
| 194 | + |
| 195 | + it.each([ |
| 196 | + { name: 'verified within this generation', verifiedOffsetMs: 1_000, preserved: true }, |
| 197 | + { name: 'verified before this generation', verifiedOffsetMs: -1_000, preserved: false }, |
| 198 | + { name: 'never verified', verifiedOffsetMs: null, preserved: false }, |
| 199 | + ])( |
| 200 | + 'applies the unresolved-evidence guard to an ACL the source could not answer: $name', |
| 201 | + async ({ verifiedOffsetMs, preserved }) => { |
| 202 | + const [clock] = await sql<{ now: string }[]>` |
| 203 | + SELECT (now() AT TIME ZONE 'UTC')::text AS now` |
| 204 | + const generationStartedAt = new Date(`${clock.now}Z`) |
| 205 | + generationStartedAt.setTime(generationStartedAt.getTime() - 60_000) |
| 206 | + const verifiedAt = |
| 207 | + verifiedOffsetMs === null |
| 208 | + ? null |
| 209 | + : new Date(generationStartedAt.getTime() + verifiedOffsetMs).toISOString() |
| 210 | + await sql`UPDATE document SET acl_verified_at = ${verifiedAt}::timestamptz AT TIME ZONE 'UTC' |
| 211 | + WHERE id = 'doc-same'` |
| 212 | + |
| 213 | + const result = await persistDocumentAcls( |
| 214 | + 'admin', |
| 215 | + new Map([['file-same', []]]), |
| 216 | + drizzle(sql, { schema }), |
| 217 | + { |
| 218 | + unresolvedExternalIds: new Set(['file-same']), |
| 219 | + generationStartedAt, |
| 220 | + } |
| 221 | + ) |
| 222 | + |
| 223 | + expect(result).toEqual({ updated: preserved ? 0 : 1, rejected: 0 }) |
| 224 | + const [stored] = await sql<{ acl: string[]; verified: boolean }[]>` |
| 225 | + SELECT acl, acl_verified_at IS NOT NULL AS verified FROM document WHERE id = 'doc-same'` |
| 226 | + expect(stored).toEqual( |
| 227 | + preserved ? { acl: [ALICE], verified: true } : { acl: [], verified: false } |
| 228 | + ) |
| 229 | + const filled = (await projected()) |
| 230 | + .filter((row) => row.id.endsWith('-same-filled')) |
| 231 | + .map((row) => row.acl) |
| 232 | + expect(filled).toEqual(preserved ? [[ALICE], [ALICE]] : [[], []]) |
| 233 | + } |
| 234 | + ) |
| 235 | +}) |
0 commit comments