|
| 1 | +/** Real PostgreSQL cancellation must roll back preparation and release its advisory lock. */ |
| 2 | +import { withUtcTimestamps } from '@sim/db/timestamps' |
| 3 | +import { getPostgresErrorCode } from '@sim/utils/errors' |
| 4 | +import { generateId } from '@sim/utils/id' |
| 5 | +import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js' |
| 6 | +import postgres from 'postgres' |
| 7 | +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' |
| 8 | + |
| 9 | +const database = vi.hoisted(() => ({ current: undefined as PostgresJsDatabase | undefined })) |
| 10 | +const mocks = vi.hoisted(() => ({ batchTrigger: vi.fn() })) |
| 11 | +vi.mock('@sim/db', () => ({ |
| 12 | + get db() { |
| 13 | + if (!database.current) throw new Error('Dispatcher test database is not initialized') |
| 14 | + return database.current |
| 15 | + }, |
| 16 | +})) |
| 17 | +vi.mock('@/lib/workspace-files/search/indexing', () => ({ |
| 18 | + indexWorkspaceFileForSearch: vi.fn(), |
| 19 | + markWorkspaceFileSearchIndexFailed: vi.fn(), |
| 20 | +})) |
| 21 | +vi.mock('@trigger.dev/sdk', () => ({ tasks: { batchTrigger: mocks.batchTrigger } })) |
| 22 | +vi.mock('@/lib/core/config/env-flags', () => ({ isTriggerDevEnabled: true })) |
| 23 | +vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' })) |
| 24 | + |
| 25 | +import { |
| 26 | + dispatchWorkspaceFileSearchIndexJobs, |
| 27 | + prepareWorkspaceFileSearchDispatch, |
| 28 | +} from '@/lib/workspace-files/search/dispatcher' |
| 29 | + |
| 30 | +describe('workspace file search dispatch PostgreSQL deadlines', () => { |
| 31 | + const schemaName = `dispatch_test_${generateId().replaceAll('-', '')}` |
| 32 | + const databaseUrl = process.env.KNOWLEDGE_ACL_TEST_DATABASE_URL |
| 33 | + if (!databaseUrl) throw new Error('Dispatcher tests require a disposable local database') |
| 34 | + const connection = postgres( |
| 35 | + databaseUrl, |
| 36 | + withUtcTimestamps({ |
| 37 | + max: 3, |
| 38 | + prepare: false, |
| 39 | + fetch_types: false, |
| 40 | + connection: { search_path: schemaName }, |
| 41 | + onnotice: () => {}, |
| 42 | + }) |
| 43 | + ) |
| 44 | + |
| 45 | + beforeAll(async () => { |
| 46 | + await connection`CREATE SCHEMA ${connection(schemaName)}` |
| 47 | + await connection`CREATE TABLE workspace_file_search_backfill ( |
| 48 | + id text PRIMARY KEY, after_workspace_id text, after_file_id text, |
| 49 | + completed_at timestamp, updated_at timestamp NOT NULL |
| 50 | + )` |
| 51 | + await connection`CREATE TABLE workspace_files ( |
| 52 | + id text PRIMARY KEY, workspace_id text NOT NULL, context text NOT NULL, |
| 53 | + deleted_at timestamp, content_updated_at timestamp NOT NULL |
| 54 | + )` |
| 55 | + await connection`CREATE TABLE workspace_file_search_index ( |
| 56 | + file_id text NOT NULL, workspace_id text NOT NULL, source_content_updated_at timestamp NOT NULL, |
| 57 | + status text NOT NULL, dispatched_at timestamp, updated_at timestamp NOT NULL, |
| 58 | + PRIMARY KEY (file_id, source_content_updated_at) |
| 59 | + )` |
| 60 | + await connection`CREATE TABLE workspace_file_search_dispatch_queue ( |
| 61 | + workspace_id text PRIMARY KEY, enqueued_at timestamp NOT NULL, |
| 62 | + updated_at timestamp NOT NULL, last_dispatched_at timestamp |
| 63 | + )` |
| 64 | + await connection`INSERT INTO workspace_file_search_backfill (id, updated_at) |
| 65 | + VALUES ('workspace-file-search-v1', '2026-09-16 00:00:00')` |
| 66 | + database.current = drizzle(connection) |
| 67 | + }) |
| 68 | + |
| 69 | + beforeEach(async () => { |
| 70 | + mocks.batchTrigger.mockReset() |
| 71 | + await connection`DROP TRIGGER IF EXISTS slow_backfill ON workspace_file_search_backfill` |
| 72 | + await connection`TRUNCATE workspace_files, workspace_file_search_index, workspace_file_search_dispatch_queue` |
| 73 | + await connection`UPDATE workspace_file_search_backfill |
| 74 | + SET updated_at = '2026-09-16 00:00:00', completed_at = NULL` |
| 75 | + }) |
| 76 | + |
| 77 | + afterAll(async () => { |
| 78 | + try { |
| 79 | + await connection`DROP SCHEMA ${connection(schemaName)} CASCADE` |
| 80 | + } finally { |
| 81 | + await connection.end() |
| 82 | + database.current = undefined |
| 83 | + } |
| 84 | + }) |
| 85 | + |
| 86 | + async function expectAdvisoryLockReleased() { |
| 87 | + await connection.begin(async (tx) => { |
| 88 | + const [row] = await tx`SELECT pg_try_advisory_xact_lock( |
| 89 | + hashtextextended('workspace-file-search-dispatch', 0) |
| 90 | + ) AS acquired` |
| 91 | + expect(row.acquired).toBe(true) |
| 92 | + }) |
| 93 | + } |
| 94 | + |
| 95 | + it('fails on a locked backfill row and releases the dispatcher lock', async () => { |
| 96 | + let release = () => {} |
| 97 | + let locked = () => {} |
| 98 | + const releaseLock = new Promise<void>((resolve) => { |
| 99 | + release = resolve |
| 100 | + }) |
| 101 | + const lockReady = new Promise<void>((resolve) => { |
| 102 | + locked = resolve |
| 103 | + }) |
| 104 | + const blocker = connection.begin(async (tx) => { |
| 105 | + await tx`SELECT id FROM workspace_file_search_backfill FOR UPDATE` |
| 106 | + locked() |
| 107 | + await releaseLock |
| 108 | + }) |
| 109 | + await lockReady |
| 110 | + try { |
| 111 | + const failure = await prepareWorkspaceFileSearchDispatch().catch((error: unknown) => error) |
| 112 | + expect(getPostgresErrorCode(failure)).toBe('55P03') |
| 113 | + await expectAdvisoryLockReleased() |
| 114 | + } finally { |
| 115 | + release() |
| 116 | + await blocker |
| 117 | + } |
| 118 | + }) |
| 119 | + |
| 120 | + it('cancels a slow statement and rolls back its earlier writes', async () => { |
| 121 | + await connection`CREATE FUNCTION slow_backfill() RETURNS trigger LANGUAGE plpgsql AS $$ |
| 122 | + BEGIN |
| 123 | + UPDATE workspace_file_search_backfill SET updated_at = '2099-01-01'; |
| 124 | + PERFORM pg_sleep(15); |
| 125 | + RETURN NEW; |
| 126 | + END |
| 127 | + $$` |
| 128 | + await connection`CREATE TRIGGER slow_backfill BEFORE INSERT ON workspace_file_search_backfill |
| 129 | + FOR EACH ROW EXECUTE FUNCTION slow_backfill()` |
| 130 | + |
| 131 | + const failure = await prepareWorkspaceFileSearchDispatch().catch((error: unknown) => error) |
| 132 | + |
| 133 | + expect(getPostgresErrorCode(failure)).toBe('57014') |
| 134 | + const [row] = |
| 135 | + await connection`SELECT updated_at::text AS updated_at FROM workspace_file_search_backfill` |
| 136 | + expect(row.updated_at).toBe('2026-09-16 00:00:00') |
| 137 | + await expectAdvisoryLockReleased() |
| 138 | + }, 20_000) |
| 139 | + |
| 140 | + it('releases committed claims without the preparation deadlines', async () => { |
| 141 | + const fileId = generateId() |
| 142 | + const workspaceId = generateId() |
| 143 | + await connection`UPDATE workspace_file_search_backfill SET completed_at = now()` |
| 144 | + await connection`INSERT INTO workspace_files (id, workspace_id, context, content_updated_at) |
| 145 | + VALUES (${fileId}, ${workspaceId}, 'workspace', '2026-09-16')` |
| 146 | + await connection`INSERT INTO workspace_file_search_index |
| 147 | + (file_id, workspace_id, source_content_updated_at, status, updated_at) |
| 148 | + VALUES (${fileId}, ${workspaceId}, '2026-09-16', 'pending', now())` |
| 149 | + await connection`INSERT INTO workspace_file_search_dispatch_queue |
| 150 | + (workspace_id, enqueued_at, updated_at) VALUES (${workspaceId}, now(), now())` |
| 151 | + await connection`CREATE TABLE cleanup_timeouts ( |
| 152 | + lock_timeout text, statement_timeout text, transaction_timeout text |
| 153 | + )` |
| 154 | + await connection`CREATE FUNCTION record_cleanup_timeouts() RETURNS trigger LANGUAGE plpgsql AS $$ |
| 155 | + BEGIN |
| 156 | + INSERT INTO cleanup_timeouts VALUES ( |
| 157 | + current_setting('lock_timeout'), |
| 158 | + current_setting('statement_timeout'), |
| 159 | + current_setting('transaction_timeout') |
| 160 | + ); |
| 161 | + RETURN NEW; |
| 162 | + END |
| 163 | + $$` |
| 164 | + await connection`CREATE TRIGGER record_cleanup_timeouts AFTER UPDATE OF dispatched_at |
| 165 | + ON workspace_file_search_index FOR EACH ROW |
| 166 | + WHEN (OLD.dispatched_at IS NOT NULL AND NEW.dispatched_at IS NULL) |
| 167 | + EXECUTE FUNCTION record_cleanup_timeouts()` |
| 168 | + |
| 169 | + const enqueueError = new Error('Queue unavailable') |
| 170 | + mocks.batchTrigger.mockRejectedValueOnce(enqueueError) |
| 171 | + |
| 172 | + await expect(dispatchWorkspaceFileSearchIndexJobs()).rejects.toBe(enqueueError) |
| 173 | + expect(mocks.batchTrigger).toHaveBeenCalledWith('workspace-file-search-index', [ |
| 174 | + expect.objectContaining({ |
| 175 | + payload: { |
| 176 | + fileId, |
| 177 | + workspaceId, |
| 178 | + sourceContentUpdatedAt: '2026-09-16T00:00:00.000Z', |
| 179 | + }, |
| 180 | + }), |
| 181 | + ]) |
| 182 | + const [index] = await connection`SELECT dispatched_at FROM workspace_file_search_index |
| 183 | + WHERE file_id = ${fileId}` |
| 184 | + expect(index.dispatched_at).toBeNull() |
| 185 | + const [queued] = await connection`SELECT workspace_id FROM workspace_file_search_dispatch_queue |
| 186 | + WHERE workspace_id = ${workspaceId}` |
| 187 | + expect(queued.workspace_id).toBe(workspaceId) |
| 188 | + const timeouts = await connection`SELECT * FROM cleanup_timeouts` |
| 189 | + expect([...timeouts]).toEqual([ |
| 190 | + { lock_timeout: '0', statement_timeout: '0', transaction_timeout: '0' }, |
| 191 | + ]) |
| 192 | + }) |
| 193 | +}) |
0 commit comments