diff --git a/apps/docs/content/docs/search/slack.mdx b/apps/docs/content/docs/search/slack.mdx index be46d622a9b..4040bf7d0d9 100644 --- a/apps/docs/content/docs/search/slack.mdx +++ b/apps/docs/content/docs/search/slack.mdx @@ -115,7 +115,11 @@ To switch apps, first remove Slack connections under **Settings → Sources → ## Permissions reference -New Search member connections request these read-only **User Token Scopes**. DM scopes are requested even when DM indexing is off; the source settings determine what is indexed. Bot scopes are separate and allow Sim to receive and answer questions in Slack. +New Search member connections request these read-only **User Token Scopes**. DM scopes are requested even when DM indexing is off; the source settings determine what is indexed. Bot scopes are separate and allow Sim to receive and answer questions in Slack and list channels during source setup. + +The custom app's **Bot Token Scopes** include `channels:read` and `groups:read` for the channel picker. Private channels appear only when the bot has access. For an existing app, add these scopes under **OAuth & Permissions**, reinstall the app in Slack to approve the changes, then reconnect the bot in Sim. + +Custom and official app manifests declare the same full bot and user scope sets, including permissions reserved for additional capabilities. The table below lists the scopes requested by member indexing; declaring additional user scopes in the manifest does not automatically grant them to each member connection. | Purpose | User scopes | |---|---| @@ -146,5 +150,6 @@ See Slack's [app manifest reference](https://docs.slack.dev/reference/app-manife | Redirect mismatch | Check all three redirect URLs above against your Sim origin. | | App or workspace mismatch | Use the App ID and client credentials from the same app, and the ID of the workspace being authorized. | | Missing scopes | Compare User Token Scopes with the table above, update the Slack app, reinstall as Slack requires, and reconnect. In workspace setup, select **Search documents** in both setup screens. | +| Channel picker says Options unavailable | Check that the selected custom bot has `channels:read` and `groups:read` under Bot Token Scopes. After adding them, reinstall the app in Slack and reconnect the bot in Sim. | | Missing private-channel results | Confirm the member is in the channel and it is within the source filters. With an indexing account, confirm that account can read it too. | | Slow initial indexing | Check sync status and Slack rate limits. A large history can take multiple background runs. | diff --git a/apps/sim/background/projection-source-acl-backfill.ts b/apps/sim/background/projection-source-acl-backfill.ts index e6b75c0794a..71c80c86369 100644 --- a/apps/sim/background/projection-source-acl-backfill.ts +++ b/apps/sim/background/projection-source-acl-backfill.ts @@ -1,8 +1,10 @@ import { task, tasks } from '@trigger.dev/sdk' import { resolveTriggerRegion } from '@/lib/core/async-jobs/region' import { + PROJECTION_SOURCE_ACL_BACKFILL_SHARDS, PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, type ProjectionSourceAclBackfillPayload, + projectionSourceAclChainTag, runProjectionSourceAclBackfill, } from '@/lib/knowledge/search/projection-source-acl-backfill' @@ -13,16 +15,21 @@ const RUN_BUDGET_MS = 60 * 60 * 1000 * Trigger.dev wrapper around `runProjectionSourceAclBackfill`. A run fills unset rows for up to * {@link RUN_BUDGET_MS}, then triggers its continuation from the cursor it reached, so the whole * projection is filled across as many bounded runs as it takes. Retry-safe: every run writes only - * rows still unset, so a retried or restarted run repeats no write. The queue admits one run at a - * time, so two starts never fill the same pages against each other. + * rows still unset, so a retried or restarted run repeats no write. A shard's continuation keeps + * its shard, so a sliced fill stays sliced until every slice is done. */ export const projectionSourceAclBackfillTask = task({ id: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, machine: 'small-1x', retry: { maxAttempts: 3 }, + /** + * One run per shard the id space may be sliced into. Shards fill disjoint ranges, so runs never + * fill the same page against each other; an unsliced chain still runs one at a time because each + * run triggers its continuation only as it ends. + */ queue: { name: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, - concurrencyLimit: 1, + concurrencyLimit: PROJECTION_SOURCE_ACL_BACKFILL_SHARDS, }, run: async (payload: ProjectionSourceAclBackfillPayload) => { const cursor = await runProjectionSourceAclBackfill(payload, { budgetMs: RUN_BUDGET_MS }) @@ -30,6 +37,8 @@ export const projectionSourceAclBackfillTask = task({ const continuation: ProjectionSourceAclBackfillPayload = { ...payload, cursor } await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, continuation, { region: await resolveTriggerRegion(), + /** The chain's tag rides on every continuation, so a start finds the chain wherever it is. */ + tags: [projectionSourceAclChainTag(payload.shard)], }) }, }) diff --git a/apps/sim/lib/internal/slack/oauth.test.ts b/apps/sim/lib/internal/slack/oauth.test.ts index a8d5d6cca6f..9c518785976 100644 --- a/apps/sim/lib/internal/slack/oauth.test.ts +++ b/apps/sim/lib/internal/slack/oauth.test.ts @@ -76,6 +76,17 @@ describe('Slack bot grant policy and cleanup', () => { it('accepts the existing indexing bot scope policy', () => { expect(() => validateSlackBotAuthorization(grant)).not.toThrow() }) + it.each(['channels:read', 'groups:read'] as const)( + 'rejects a bot grant missing channel picker scope %s', + (missingScope) => { + expect(() => + validateSlackBotAuthorization({ + ...grant, + scope: SLACK_SEARCH_SCOPES.filter((scope) => scope !== missingScope).join(','), + }) + ).toThrow(`Reinstall the app with these scopes: ${missingScope}`) + } + ) it('requires the additional command scope for shared installs', () => { expect(() => validateSlackBotAuthorization(grant, [...SLACK_SEARCH_SCOPES, 'commands']) diff --git a/apps/sim/lib/internal/slack/search-client.test.ts b/apps/sim/lib/internal/slack/search-client.test.ts index 864a82df65c..ecbd3175517 100644 --- a/apps/sim/lib/internal/slack/search-client.test.ts +++ b/apps/sim/lib/internal/slack/search-client.test.ts @@ -45,6 +45,21 @@ describe('Slack Search provider verification', () => { fetchMock.mockResolvedValue(new Response(JSON.stringify(auth))) await expect(verifySlackSearchBot('token')).rejects.toThrow('Reinstall') }) + it.each(['channels:read', 'groups:read'] as const)( + 'rejects an installed bot missing channel picker scope %s', + async (missingScope) => { + fetchMock.mockResolvedValue( + reply( + auth, + SLACK_SEARCH_SCOPES.filter((scope) => scope !== missingScope) + ) + ) + await expect(verifySlackSearchBot('token')).rejects.toThrow( + `Reinstall the Slack bot with these scopes: ${missingScope}` + ) + expect(fetchMock).toHaveBeenCalledOnce() + } + ) it.each([ { deleted: true }, { is_bot: true }, diff --git a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts index 4a357e3e9a3..f5d9434e2a1 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts @@ -3,12 +3,24 @@ */ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { mockBackfill, mockEnd, mockPostgres, mockPrewarm, mockTasksTrigger } = vi.hoisted(() => ({ +const { + mockBackfill, + mockEnd, + mockPostgres, + mockPrewarm, + mockRunsList, + mockTasksTrigger, + mockUnsafe, +} = vi.hoisted(() => ({ mockBackfill: vi.fn(), mockEnd: vi.fn(async () => undefined), mockPostgres: vi.fn(), mockPrewarm: vi.fn(async () => []), + mockRunsList: vi.fn( + (_query: unknown): AsyncIterable<{ id: string; status: string }> => (async function* () {})() + ), mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })), + mockUnsafe: vi.fn(async () => [{ unfilled: false }]), })) vi.mock('@sim/db', () => ({ resolveDbUrl: () => 'postgres://localhost:5432/sim' })) @@ -18,7 +30,10 @@ vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({ })) vi.mock('postgres', () => ({ default: mockPostgres })) vi.mock('@/lib/knowledge/search/prewarm', () => ({ prewarmSearchProjection: mockPrewarm })) -vi.mock('@trigger.dev/sdk', () => ({ tasks: { trigger: mockTasksTrigger } })) +vi.mock('@trigger.dev/sdk', () => ({ + runs: { list: mockRunsList }, + tasks: { trigger: mockTasksTrigger }, +})) vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' })) vi.mock('@/lib/core/utils/background', () => ({ runDetached: (_label: string, work: () => Promise) => { @@ -29,10 +44,11 @@ vi.mock('@/lib/core/utils/background', () => ({ import { enqueueProjectionSourceAclBackfill, PROJECTION_PREWARM_BUDGET_MS, + projectionSourceAclShardRange, runProjectionSourceAclBackfill, } from '@/lib/knowledge/search/projection-source-acl-backfill' -const connection = { end: mockEnd } +const connection = { end: mockEnd, unsafe: mockUnsafe } describe('runProjectionSourceAclBackfill', () => { beforeEach(() => { @@ -60,8 +76,17 @@ describe('runProjectionSourceAclBackfill', () => { expect(mockEnd).toHaveBeenCalledTimes(1) }) - it('warms the projections on the same connection once both are filled, before closing it', async () => { + it('analyzes and warms the projections on the same connection once both are filled, before closing it', async () => { await runProjectionSourceAclBackfill({}) + /** A row whose document is gone is not the fill's to finish; the probe joins the document. */ + expect( + mockUnsafe.mock.calls.some(([query]) => + String(query).includes('JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL') + ) + ).toBe(true) + expect(mockUnsafe.mock.calls.map(([query]) => query)).toEqual( + expect.arrayContaining(['ANALYZE embedding_search', 'ANALYZE embedding_keyword_tin']) + ) expect(mockPrewarm).toHaveBeenCalledTimes(1) expect(mockPrewarm).toHaveBeenCalledWith(connection, { budgetMs: PROJECTION_PREWARM_BUDGET_MS }) expect(mockPrewarm.mock.invocationCallOrder[0]).toBeLessThan( @@ -101,11 +126,65 @@ describe('runProjectionSourceAclBackfill', () => { await expect(runProjectionSourceAclBackfill({})).rejects.toThrow('statement timeout') expect(mockEnd).toHaveBeenCalledTimes(1) }) + + it('fills only its shard of the id space in both projections', async () => { + await runProjectionSourceAclBackfill({ shard: { index: 1, count: 4 } }) + for (const [, , options] of mockBackfill.mock.calls) { + expect(options).toMatchObject({ afterId: '4', beforeId: '8' }) + } + }) + + it('resumes a shard after its cursor and keeps its upper bound', async () => { + await runProjectionSourceAclBackfill({ + shard: { index: 1, count: 4 }, + cursor: { projection: 'embedding_search', afterId: '5a' }, + }) + expect(mockBackfill.mock.calls[0][2]).toMatchObject({ afterId: '5a', beforeId: '8' }) + expect(mockBackfill.mock.calls[1][2]).toMatchObject({ afterId: '4', beforeId: '8' }) + }) + + it('leaves the analysis and the warm to whoever fills the rows another shard still holds', async () => { + mockUnsafe.mockResolvedValueOnce([{ unfilled: true }]) + await expect( + runProjectionSourceAclBackfill({ shard: { index: 0, count: 4 } }) + ).resolves.toBeNull() + expect(mockUnsafe.mock.calls.map(([query]) => query)).not.toContain('ANALYZE embedding_search') + expect(mockPrewarm).not.toHaveBeenCalled() + expect(mockEnd).toHaveBeenCalledTimes(1) + }) +}) + +describe('projectionSourceAclShardRange', () => { + it('slices the hex id space into contiguous ranges', () => { + expect(projectionSourceAclShardRange({ index: 0, count: 4 })).toEqual({ + afterId: '', + beforeId: '4', + }) + expect(projectionSourceAclShardRange({ index: 3, count: 4 })).toEqual({ + afterId: 'c', + beforeId: undefined, + }) + expect(projectionSourceAclShardRange({ index: 0, count: 1 })).toEqual({ + afterId: '', + beforeId: undefined, + }) + }) + + it.each([ + [{ index: 0, count: 3 }, 'shard count must divide 16'], + [{ index: 0, count: 8 }, 'shard count must be at most 4'], + [{ index: 4, count: 4 }, 'shard index must be within 0..3'], + [{ index: 0.5, count: 2 }, 'shard index must be within 0..1'], + ])('refuses %j', (shard, message) => { + expect(() => projectionSourceAclShardRange(shard)).toThrow(message) + }) }) describe('enqueueProjectionSourceAclBackfill', () => { beforeEach(() => { vi.clearAllMocks() + /** No chain in flight unless a case says so. */ + mockRunsList.mockImplementation(() => (async function* () {})()) mockPostgres.mockReturnValue(connection) mockBackfill.mockResolvedValue({ projection: 'embedding_search', @@ -118,13 +197,91 @@ describe('enqueueProjectionSourceAclBackfill', () => { it('hands the backfill to the Trigger.dev worker when one is configured', async () => { await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 })).resolves.toEqual({ - runId: 'run-1', + runIds: ['run-1'], + inFlight: [], }) + expect(mockRunsList).toHaveBeenCalledWith( + expect.objectContaining({ tag: 'projection-source-acl-backfill:shard:0/1' }) + ) expect(mockTasksTrigger).toHaveBeenCalledWith( 'projection-source-acl-backfill', { pageSize: 25 }, - { region: 'us-east-1' } + { + region: 'us-east-1', + tags: ['projection-source-acl-backfill:shard:0/1'], + idempotencyKey: 'projection-source-acl-backfill:shard:0/1:after:none', + idempotencyKeyTTL: '2m', + } ) expect(mockBackfill).not.toHaveBeenCalled() }) + + it('keys a start after a chain that ended on that chain, so a restart is its own start', async () => { + mockRunsList.mockImplementation(() => + (async function* () { + yield { id: 'run-done', status: 'COMPLETED' } + })() + ) + await expect(enqueueProjectionSourceAclBackfill({})).resolves.toEqual({ + runIds: ['run-1'], + inFlight: [], + }) + expect(mockTasksTrigger.mock.calls[0][2].idempotencyKey).toBe( + 'projection-source-acl-backfill:shard:0/1:after:run-done' + ) + }) + + it('refuses a shard the id space cannot be sliced into before starting anything', async () => { + await expect( + enqueueProjectionSourceAclBackfill({ shard: { index: 5, count: 4 } }) + ).rejects.toThrow('shard index must be within 0..3') + expect(mockTasksTrigger).not.toHaveBeenCalled() + }) + + it('leaves a range whose chain is still in flight to that chain', async () => { + mockRunsList.mockImplementation((query: unknown) => + (async function* () { + if ((query as { tag: string }).tag.endsWith(':shard:1/4')) + yield { id: 'run-live', status: 'EXECUTING' } + })() + ) + await expect(enqueueProjectionSourceAclBackfill({}, 4)).resolves.toEqual({ + runIds: ['run-1', 'run-1', 'run-1'], + inFlight: ['run-live'], + }) + expect(mockTasksTrigger.mock.calls.map(([, payload]) => payload.shard?.index)).toEqual([ + 0, 2, 3, + ]) + }) + + it('starts one run per shard, each on its own slice under its own chain tag', async () => { + await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 }, 4)).resolves.toEqual({ + runIds: ['run-1', 'run-1', 'run-1', 'run-1'], + inFlight: [], + }) + expect(mockTasksTrigger.mock.calls.map(([, payload]) => payload)).toEqual( + [0, 1, 2, 3].map((index) => ({ pageSize: 25, shard: { index, count: 4 } })) + ) + expect(mockTasksTrigger.mock.calls.map(([, , options]) => options.tags)).toEqual( + [0, 1, 2, 3].map((index) => [`projection-source-acl-backfill:shard:${index}/4`]) + ) + }) + + it.each([ + [3, 'must divide 16'], + [8, 'must be at most 4'], + ])('refuses %s shards before starting anything', async (shards, message) => { + await expect(enqueueProjectionSourceAclBackfill({}, shards)).rejects.toThrow(message) + expect(mockTasksTrigger).not.toHaveBeenCalled() + }) + + it('refuses to slice a start that carries a cursor, which belongs to one chain', async () => { + await expect( + enqueueProjectionSourceAclBackfill( + { cursor: { projection: 'embedding_search', afterId: '5a' } }, + 4 + ) + ).rejects.toThrow('cannot start from a cursor') + expect(mockTasksTrigger).not.toHaveBeenCalled() + }) }) diff --git a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts index 6ac3947d5d3..4e9b0f0e813 100644 --- a/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts +++ b/apps/sim/lib/knowledge/search/projection-source-acl-backfill.ts @@ -26,13 +26,61 @@ export interface ProjectionSourceAclBackfillCursor { afterId: string } +/** + * One of `count` equal slices of the chunk id space. Chunk ids are lowercase hex UUIDs, so the + * space is sliced on the first hex digit and `count` must divide sixteen; every slice is a + * contiguous id range, so workers on different slices never fill the same page. + */ +export interface ProjectionSourceAclBackfillShard { + index: number + count: number +} + +/** + * Shards the id space may be sliced into at most: the runs the task's queue admits at once, so a + * sliced fill runs its slices together rather than in waves. + */ +export const PROJECTION_SOURCE_ACL_BACKFILL_SHARDS = 4 + export interface ProjectionSourceAclBackfillPayload { /** The first projection's first page when absent. */ cursor?: ProjectionSourceAclBackfillCursor + /** The whole id space when absent. */ + shard?: ProjectionSourceAclBackfillShard pageSize?: number pauseMs?: number } +/** Refuses a shard the id space cannot be sliced into. */ +function assertProjectionSourceAclShard({ index, count }: ProjectionSourceAclBackfillShard): void { + if (!Number.isInteger(count) || count < 1 || 16 % count !== 0) { + throw new Error(`Projection backfill shard count must divide 16, got ${count}`) + } + if (count > PROJECTION_SOURCE_ACL_BACKFILL_SHARDS) { + throw new Error( + `Projection backfill shard count must be at most ${PROJECTION_SOURCE_ACL_BACKFILL_SHARDS}, got ${count}` + ) + } + if (!Number.isInteger(index) || index < 0 || index >= count) { + throw new Error(`Projection backfill shard index must be within 0..${count - 1}, got ${index}`) + } +} + +/** The id range a shard covers: the id its first page follows, and the first id past it. */ +export function projectionSourceAclShardRange(shard: ProjectionSourceAclBackfillShard): { + afterId: string + beforeId: string | undefined +} { + assertProjectionSourceAclShard(shard) + const { index, count } = shard + const width = 16 / count + const digit = (value: number) => value.toString(16) + return { + afterId: index === 0 ? '' : digit(index * width), + beforeId: index + 1 < count ? digit((index + 1) * width) : undefined, + } +} + export interface ProjectionSourceAclBackfillRunOptions { /** Stop once this much time has passed and return where to resume; unbounded otherwise. */ budgetMs?: number @@ -42,8 +90,9 @@ export interface ProjectionSourceAclBackfillRunOptions { * Fills the ranking projections' source and ACL columns from the cursor onwards, one projection * after the other, on a connection of its own: the page statement binds the keyset cursor as a * scalar and needs no array parameter, so the pool's options would serve, but a run this long - * should not hold one of the worker's pooled connections. Returns the cursor to continue from when - * the budget ran out, `null` once both projections are filled. + * should not hold one of the worker's pooled connections. A shard fills its own slice of the id + * space in every projection. Returns the cursor to continue from when the budget ran out, `null` + * once the run's range is filled in both projections. */ export async function runProjectionSourceAclBackfill( payload: ProjectionSourceAclBackfillPayload, @@ -58,13 +107,16 @@ export async function runProjectionSourceAclBackfill( ? PROJECTION_SOURCE_ACL_TABLES.indexOf(payload.cursor.projection) : 0 if (start < 0) throw new Error(`Unknown projection ${payload.cursor?.projection}`) + const range = payload.shard ? projectionSourceAclShardRange(payload.shard) : undefined for (const projection of PROJECTION_SOURCE_ACL_TABLES.slice(start)) { const budgetMs = options.budgetMs === undefined ? undefined : Math.max(0, options.budgetMs - (Date.now() - startedAt)) const progress = await backfillProjectionSourceAcl(sql, projection, { - afterId: payload.cursor?.projection === projection ? payload.cursor.afterId : undefined, + afterId: + payload.cursor?.projection === projection ? payload.cursor.afterId : range?.afterId, + beforeId: range?.beforeId, pageSize: payload.pageSize, pauseMs: payload.pauseMs, budgetMs, @@ -72,10 +124,21 @@ export async function runProjectionSourceAclBackfill( if (!progress.done) return { projection, afterId: progress.afterId } } logger.info('Projection source and ACL backfill complete', { + shard: payload.shard, elapsedMs: Date.now() - startedAt, }) - /** The fill just streamed through both projections; put the ranking pages back before anyone searches. */ - await prewarmSearchProjection(sql, { budgetMs: PROJECTION_PREWARM_BUDGET_MS }) + /** + * The fill just streamed through both projections: the planner last saw every row unfilled and + * should see the finished projections, and the ranking pages should be back before anyone + * searches. With shards, the one that finishes last does both: a shard that still finds + * unfilled rows anywhere leaves them to whoever fills those. Two shards ending in the same + * moment can both read none left and both do this, which repeats reads and nothing else. + */ + if (await projectionsFilled(sql)) { + for (const projection of PROJECTION_SOURCE_ACL_TABLES) + await sql.unsafe(`ANALYZE ${projection}`) + await prewarmSearchProjection(sql, { budgetMs: PROJECTION_PREWARM_BUDGET_MS }) + } return null } finally { await sql.end() @@ -83,16 +146,96 @@ export async function runProjectionSourceAclBackfill( } /** - * Starts the backfill on the deployment's Trigger.dev worker, where bounded runs chain until both - * projections are filled. Safe to call again at any time: a run only fills rows still unset. + * Whether no projection still holds a row the fill could give its source and ACL: a row without + * them whose document exists. A row whose document is gone is not the fill's to finish and never + * counts as left. Each read is one index probe while any such row remains. + */ +async function projectionsFilled(sql: postgres.Sql): Promise { + for (const projection of PROJECTION_SOURCE_ACL_TABLES) { + const [row] = await sql.unsafe>( + `SELECT EXISTS ( + SELECT 1 FROM ${projection} s JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL + ) AS unfilled` + ) + if (row?.unfilled) return false + } + return true +} + +/** The tag every run of one chain carries, so a chain in flight is found before another is started. */ +export function projectionSourceAclChainTag(shard?: ProjectionSourceAclBackfillShard): string { + const { index, count } = shard ?? { index: 0, count: 1 } + return `${PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID}:shard:${index}/${count}` +} + +/** + * How long a start's trigger stays idempotent. The lookup and the trigger are two calls, so two + * starts in the same instant could both find no chain in flight; the key is what they both saw, + * the chain's latest run, so they collapse into one start, while a start after another chain has + * ended sees a different latest run and is a new key. + */ +const START_IDEMPOTENCY_TTL = '2m' + +/** A run that has not ended: it, or the continuation it triggers, still owns its range. */ +const IN_FLIGHT_RUN_STATUSES: ReadonlySet = new Set([ + 'PENDING_VERSION', + 'QUEUED', + 'DEQUEUED', + 'EXECUTING', + 'WAITING', + 'DELAYED', +]) + +/** + * Starts the backfill on the deployment's Trigger.dev worker, where bounded runs chain until the + * projections are filled: one chain over the whole id space, or one per shard, each filling its + * own slice at the same time. Safe to call again at any time: a run only fills rows still unset, + * and a range whose chain is still in flight is left to that chain rather than given a second. + * A cursor belongs to one chain, so a sliced start takes none: each slice begins at its own bound. */ export async function enqueueProjectionSourceAclBackfill( - payload: ProjectionSourceAclBackfillPayload = {} -): Promise<{ runId: string }> { - const { tasks } = await import('@trigger.dev/sdk') - const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, payload, { - region: await resolveTriggerRegion(), - }) - logger.info('Projection source and ACL backfill enqueued', { runId: handle.id }) - return { runId: handle.id } + payload: ProjectionSourceAclBackfillPayload = {}, + shards = 1 +): Promise<{ runIds: string[]; inFlight: string[] }> { + if (payload.shard) assertProjectionSourceAclShard(payload.shard) + if (shards !== 1) { + assertProjectionSourceAclShard({ index: 0, count: shards }) + if (payload.cursor) throw new Error('A sliced projection backfill cannot start from a cursor') + } + const { runs, tasks } = await import('@trigger.dev/sdk') + const region = await resolveTriggerRegion() + const payloads: ProjectionSourceAclBackfillPayload[] = + shards === 1 + ? [payload] + : Array.from({ length: shards }, (_, index) => ({ + ...payload, + shard: { index, count: shards }, + })) + const runIds: string[] = [] + const inFlight: string[] = [] + for (const shardPayload of payloads) { + const tag = projectionSourceAclChainTag(shardPayload.shard) + /** The chain's latest run, newest first, whatever its state. */ + let latest: { id: string; status: string } | undefined + for await (const run of runs.list({ + taskIdentifier: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, + tag, + limit: 1, + })) { + latest = run + } + if (latest && IN_FLIGHT_RUN_STATUSES.has(latest.status)) { + inFlight.push(latest.id) + continue + } + const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, shardPayload, { + region, + tags: [tag], + idempotencyKey: `${tag}:after:${latest?.id ?? 'none'}`, + idempotencyKeyTTL: START_IDEMPOTENCY_TTL, + }) + runIds.push(handle.id) + } + logger.info('Projection source and ACL backfill enqueued', { runIds, inFlight, shards }) + return { runIds, inFlight } } diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 6e76b50d677..219f323d72e 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -2051,6 +2051,89 @@ describe('permitted-document planner', () => { expect(reachCounts()).toHaveLength(2) }) + it('reports a caller who reaches nothing as a bounded set of nothing, counted every time', async () => { + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + if (statement.includes('EXPLAIN')) + return [{ 'QUERY PLAN': [{ Plan: { 'Plan Rows': 1_000_000 } }] }] + if (statement.includes(') reached')) return [{ n: 0 }] + return [] + }) + const reachCounts = () => statements().filter((query) => query.sql.includes(') reached')) + const plan = { + connectors: { workspace: [], admin: [], members: [], liveProofRequired: [] }, + observers: { confirmed: [], observed: [] }, + memberSources: [], + connectorTypes: new Map(), + uploads: true, + } + const budget = () => new SearchBudget('vector', performance.now() + 10_000) + await expect( + resolveReach(['org-index'], scope('reaches-nothing'), budget(), plan) + ).resolves.toEqual({ kind: 'bounded', documents: [] }) + /** Emptiness decides completeness, so it is never remembered: the next search counts again. */ + await expect( + resolveReach(['org-index'], scope('reaches-nothing'), budget(), plan) + ).resolves.toEqual({ kind: 'bounded', documents: [] }) + expect(reachCounts()).toHaveLength(2) + }) + + it('reports a saturated probe whose count then finds nothing as a bounded set of nothing', async () => { + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + if (isProbeStatement(statement)) return [{ id: null, connectorId: null, saturated: true }] + if (statement.includes('EXPLAIN')) + return [{ 'QUERY PLAN': [{ Plan: { 'Plan Rows': 1_000_000 } }] }] + if (statement.includes(') reached')) return [{ n: 0 }] + return [] + }) + await expect( + resolvePermittedDocuments({ + knowledgeBaseIds: ['org-index'], + access: scope('saturated-then-nothing'), + budget: new SearchBudget('vector', performance.now() + 10_000), + }) + ).resolves.toEqual({ kind: 'bounded', documents: [] }) + }) + + it('does not read an unanalyzed index as a reach of nothing', async () => { + /** The planner knows no rows yet, so the bound is zero and the count looked at nothing. */ + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + if (statement.includes('EXPLAIN')) return [{ 'QUERY PLAN': [{ Plan: { 'Plan Rows': 0 } }] }] + if (statement.includes(') reached')) return [{ n: 0 }] + return [] + }) + await expect( + resolveReach( + ['org-index'], + scope('unanalyzed'), + new SearchBudget('vector', performance.now() + 10_000), + { + connectors: { workspace: [], admin: [], members: [], liveProofRequired: [] }, + observers: { confirmed: [], observed: [] }, + memberSources: [], + connectorTypes: new Map(), + uploads: true, + } + ) + ).resolves.toEqual({ kind: 'unbounded', broad: true }) + }) + + it('reports a leg whose own deadline passed during the count as short, not failed', async () => { + const budget = new SearchBudget('vector', performance.now() - 1) + await expect( + resolveReach(['org-index'], scope('spent-leg'), budget, { + connectors: { workspace: [], admin: [], members: [], liveProofRequired: [] }, + observers: { confirmed: [], observed: [] }, + memberSources: [], + connectorTypes: new Map(), + uploads: true, + }) + ).resolves.toEqual({ kind: 'unbounded', broad: true }) + expect(budget.timedOut).toBe(true) + }) + it('counts a resolved reach against a small index instead of assuming it broad', async () => { /** A bound inside the probe limit proves nothing without a saturated probe. */ dbChainMockFns.execute.mockImplementation(async (query) => { @@ -2075,9 +2158,10 @@ describe('permitted-document planner', () => { ) expect(reach).toEqual({ kind: 'unbounded', broad: false }) expect(reachCounts()).toHaveLength(1) - /** The count is the search's own read: it runs inside the leg's deadline statement. */ + /** The count is the search's own read, under the probe's share of the deadline, not the leg's. */ const countAt = statements().findIndex((query) => query.sql.includes(') reached')) expect(statements()[countAt - 1].sql).toContain('statement_timeout') + expect(Number(statements()[countAt - 1].params[0])).toBeLessThanOrEqual(600) }) it('does not remember a reach whose count ran out of time', async () => { @@ -2102,6 +2186,8 @@ describe('permitted-document planner', () => { }) expect(plan).toEqual({ kind: 'unbounded', broad: true }) expect(reachCounts()).toHaveLength(1) + /** Only the count's share of the deadline was spent; the leg is not the one that timed out. */ + expect(budget.timedOut).toBe(false) /** The next search counts again rather than trusting an answer that never came. */ await resolveReach( ['org-index'], diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index af01100b26e..f2ff5bf066c 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1184,7 +1184,19 @@ export const BROAD_REACH_SHARE = 0.25 */ const SATURATED_REACH_TTL_MS = 5 * 60 * 1000 -/** A saturated reach, and whether it is broad enough to walk the whole graph for. */ +/** + * A counted reach: whether it is broad enough to walk the whole graph for, or empty, in which + * case the caller reads nothing in these bases and no leg has anything to rank. + */ +interface CountedReach { + broad: boolean + empty: boolean +} + +/** + * Only breadth is remembered. Emptiness decides completeness, not strategy, so it is counted on + * every search: the count of a reach of nothing finds nothing and costs almost nothing. + */ const saturatedReach = new LRUCache({ max: 10_000, ttl: SATURATED_REACH_TTL_MS, @@ -1249,39 +1261,58 @@ async function estimateFilteredDocuments( } /** - * Whether a reach is broad: the caller reaches at least {@link BROAD_REACH_SHARE} of the bases' - * documents. Counted once against that bound and remembered, so the first search after the - * window pays for it and the rest do not. A caller whose probe already saturated is known to - * reach past the probe's limit, so a bound inside that limit is met without counting. + * How far a caller reaches: broad when they reach at least {@link BROAD_REACH_SHARE} of the + * bases' documents, empty when they reach none. A reach of nothing is a bounded set of nothing: a + * caller who reads no document in these bases, such as a member with no source of their own yet, + * has nothing for any leg to rank, where an unbounded set would have each leg scan to its + * deadline for rows it cannot find. Breadth is counted once against the bound and remembered, so + * the first search after the window pays for it and the rest do not. A caller whose probe already + * saturated is known to reach past the probe's limit, so a bound inside that limit is met without + * counting. + * + * The count reads as many index entries as the caller reaches, so on a large index it can cost + * more than the leg it serves; it gets the probe's share of the deadline, never the whole leg's. + * A count that runs out of that share answers `null`: the leg keeps its time and its deadline + * intact, and the caller decides this search alone without remembering anything. */ -async function reachIsBroad( +async function countReach( knowledgeBaseIds: string[], access: KnowledgeAccessScope, budget: SearchBudget | undefined, plan: SearchAccessPlan | undefined, saturated: boolean -): Promise { - if (access.kind !== 'user') return true - const total = - (await indexDocumentCounts.fetch([...knowledgeBaseIds].sort().join(','), { - context: budget, - })) ?? 0 - const bound = Math.ceil(total * BROAD_REACH_SHARE) - if (saturated && bound <= VECTOR_PROBE_DOCUMENT_LIMIT) return true - const [row] = await runSearchQuery(budget, 'permitted_documents', (executor) => - executor.execute<{ n: number }>(sql` - SELECT count(*) AS n FROM ( - SELECT 1 FROM ${document} - WHERE ${and( - isNull(document.deletedAt), - knowledgeAclOverlapCondition(access), - inArray(document.knowledgeBaseId, knowledgeBaseIds), - planSourceCondition(plan) - )} - LIMIT ${bound} - ) reached`) - ) - return Number(row?.n ?? 0) >= bound +): Promise { + if (access.kind !== 'user') return { broad: true, empty: false } + const countBudget = budget?.capped(VECTOR_PROBE_BUDGET_MS) + try { + const total = + (await indexDocumentCounts.fetch([...knowledgeBaseIds].sort().join(','), { + context: countBudget, + })) ?? 0 + const bound = Math.ceil(total * BROAD_REACH_SHARE) + if (saturated && bound <= VECTOR_PROBE_DOCUMENT_LIMIT) return { broad: true, empty: false } + const [row] = await runSearchQuery(countBudget, 'permitted_documents', (executor) => + executor.execute<{ n: number }>(sql` + SELECT count(*) AS n FROM ( + SELECT 1 FROM ${document} + WHERE ${and( + isNull(document.deletedAt), + knowledgeAclOverlapCondition(access), + inArray(document.knowledgeBaseId, knowledgeBaseIds), + planSourceCondition(plan) + )} + LIMIT ${bound} + ) reached`) + ) + const reached = Number(row?.n ?? 0) + /** A count that looked and found nothing: only a bound of zero looks at nothing. */ + return { broad: reached >= bound, empty: bound > 0 && reached === 0 } + } catch (error) { + if (!budget || !countBudget?.isTimeout(error)) throw error + /** Only the count's share was spent; the leg's own deadline still governs. */ + budget.remaining() + return null + } } /** @@ -1332,12 +1363,15 @@ export async function resolveReach( const remembered = key ? saturatedReach.get(key) : undefined if (remembered) return { kind: 'unbounded', broad: remembered.broad } try { - const broad = await reachIsBroad(knowledgeBaseIds, access, budget, plan, false) - if (key) saturatedReach.set(key, { broad }) - return { kind: 'unbounded', broad } + const reach = await countReach(knowledgeBaseIds, access, budget, plan, false) + /** A count that ran out of time decides this search only; the next one counts again. */ + if (reach === null) return { kind: 'unbounded', broad: true } + if (reach.empty) return { kind: 'bounded', documents: [] } + if (key) saturatedReach.set(key, { broad: reach.broad }) + return { kind: 'unbounded', broad: reach.broad } } catch (error) { + /** The leg's own deadline passed during the count: the leg is short, the search is not failed. */ if (!budget?.isTimeout(error)) throw error - /** A count that ran out of time decides this search only; the next one counts again. */ return { kind: 'unbounded', broad: true } } } @@ -1393,17 +1427,22 @@ export async function resolvePermittedDocuments(params: { } if (probe.kind === 'saturated') { try { - broad = await reachIsBroad( + const reach = await countReach( params.knowledgeBaseIds, params.access, params.budget, params.accessPlan, true ) - if (key) saturatedReach.set(key, { broad }) + /** A count that ran out of time decides this search only; the next one counts again. */ + if (reach?.empty) probe = { kind: 'documents', documents: [] } + else if (reach !== null) { + broad = reach.broad + if (key) saturatedReach.set(key, { broad }) + } } catch (error) { + /** The leg's own deadline passed during the count: the leg is short, the search is not failed. */ if (!params.budget?.isTimeout(error)) throw error - /** A count that ran out of time decides this search only; the next one counts again. */ } } } @@ -2658,9 +2697,7 @@ export async function retrieveKnowledgeSearch( * readable documents enumerated ahead of ranking: its reach alone chooses between one * walk over the whole graph and a search of each source. */ - await measureSearchStage('permitted_documents', () => - resolveReach(knowledgeBaseIds, access, budgets.vector, accessPlan) - ) + await resolveReach(knowledgeBaseIds, access, budgets.vector, accessPlan) : await resolvePermittedDocuments({ knowledgeBaseIds, access, diff --git a/apps/sim/lib/slack-search/constants.ts b/apps/sim/lib/slack-search/constants.ts index def905139d5..8010af77cbc 100644 --- a/apps/sim/lib/slack-search/constants.ts +++ b/apps/sim/lib/slack-search/constants.ts @@ -1,6 +1,8 @@ export const SLACK_SEARCH_SCOPES = [ 'assistant:write', 'chat:write', + 'channels:read', + 'groups:read', 'im:history', 'im:write', 'app_mentions:read', diff --git a/apps/sim/lib/slack-search/manifest.test.ts b/apps/sim/lib/slack-search/manifest.test.ts index 68bf1a3804c..733357be584 100644 --- a/apps/sim/lib/slack-search/manifest.test.ts +++ b/apps/sim/lib/slack-search/manifest.test.ts @@ -15,9 +15,10 @@ describe('Search app manifest', () => { 'app_mentions:read', 'im:write', 'im:history', + 'channels:read', + 'groups:read', ]) ) - expect(manifest.oauth_config.scopes.bot).not.toContain('groups:history') expect(manifest.oauth_config.scopes.user).toEqual( expect.arrayContaining([ 'users:read', @@ -47,26 +48,21 @@ describe('Search app manifest', () => { ).toBe(true) }) it('preserves existing member grants when updating a bot manifest', () => { - expect( - createSlackSearchManifest('Sim Search', 'Search', 'https://sim.test', ['files:read']) - .oauth_config.scopes.user - ).toContain('files:read') + const manifest = createSlackSearchManifest('Sim Search', 'Search', 'https://sim.test', [ + 'files:read', + 'files:write', + ]) + expect(manifest.oauth_config.scopes.user).toContain('files:write') + expect(manifest.oauth_config.scopes.user.filter((scope) => scope === 'files:read')).toEqual([ + 'files:read', + ]) }) - it('uses one origin for unified ingress and OAuth with only the required bot permissions', () => { + it('uses one origin for unified ingress and OAuth', () => { const manifest = createSlackSearchManifest( 'Sim Search', 'Search with sources', 'https://search-test.ngrok.app' ) - expect(manifest.oauth_config.scopes.bot).toEqual([ - 'assistant:write', - 'chat:write', - 'im:history', - 'im:write', - 'app_mentions:read', - 'users:read', - 'users:read.email', - ]) expect(manifest.settings.event_subscriptions.request_url).toBe( 'https://search-test.ngrok.app/api/webhooks/slack' ) @@ -92,23 +88,26 @@ describe('Search app manifest', () => { }) }) -it('official app declares expanded permissions without subscribing to member message events', () => { - const manifest = createSharedSlackSearchManifest('https://www.sim.ai') - expect(manifest.oauth_config.scopes.user).toEqual([ +it.each([ + { + name: 'custom', + manifest: createSlackSearchManifest('Sim Search', 'Search', 'https://sim.test'), + }, + { name: 'shared', manifest: createSharedSlackSearchManifest('https://sim.test') }, +])('$name app declares the complete bot and user scope sets without duplicates', ({ manifest }) => { + expect([...manifest.oauth_config.scopes.user].sort()).toEqual([ + 'canvases:read', + 'canvases:write', 'channels:history', 'channels:read', + 'chat:write', + 'files:read', 'groups:history', 'groups:read', 'im:history', 'im:read', 'mpim:history', 'mpim:read', - 'users:read', - 'users:read.email', - 'canvases:read', - 'canvases:write', - 'chat:write', - 'files:read', 'search:read.files', 'search:read.im', 'search:read.mpim', @@ -117,32 +116,38 @@ it('official app declares expanded permissions without subscribing to member mes 'search:read.users', 'team:read', 'usergroups:read', - ]) - expect(manifest.oauth_config.scopes.bot).toEqual([ - 'assistant:write', - 'chat:write', - 'im:history', - 'im:write', - 'app_mentions:read', 'users:read', 'users:read.email', - 'commands', + ]) + expect([...manifest.oauth_config.scopes.bot].sort()).toEqual([ + 'app_mentions:read', + 'assistant:write', 'channels:history', 'channels:manage', 'channels:read', 'channels:write.invites', + 'chat:write', 'chat:write.public', + 'commands', 'groups:history', 'groups:read', 'groups:write', 'groups:write.invites', + 'im:history', + 'im:write', 'links:read', 'links:write', 'mpim:history', 'mpim:read', 'mpim:write', 'reactions:write', + 'users:read', + 'users:read.email', ]) +}) + +it('official app declares commands and lifecycle events without member message events', () => { + const manifest = createSharedSlackSearchManifest('https://www.sim.ai') expect(manifest.features.slash_commands.map((command) => command.command)).toEqual([ '/query', '/connect', diff --git a/apps/sim/lib/slack-search/manifest.ts b/apps/sim/lib/slack-search/manifest.ts index c8ce9b5bd2e..a8d3a27f153 100644 --- a/apps/sim/lib/slack-search/manifest.ts +++ b/apps/sim/lib/slack-search/manifest.ts @@ -3,7 +3,7 @@ import { SLACK_MANAGED_USER_ENROLLMENT_CALLBACK_PATH, SLACK_SEARCH_USER_SCOPES, } from '@/lib/credential-groups/slack-managed-user-scopes' -import { SLACK_SEARCH_SCOPES, SLACK_SHARED_SEARCH_BOT_SCOPES } from '@/lib/slack-search/constants' +import { SLACK_SHARED_SEARCH_BOT_SCOPES } from '@/lib/slack-search/constants' export const SLACK_SEARCH_CALLBACK_PATH = '/api/knowledge/slack/oauth/callback' export const SLACK_SEARCH_WEBHOOK_PATH = '/api/webhooks/slack' @@ -11,7 +11,10 @@ export const SLACK_SEARCH_DEFAULT_NAME = 'Sim Search' export const SLACK_SEARCH_DEFAULT_DESCRIPTION = 'Ask questions about your organization’s knowledge and get answers with sources.' -/** Bot conversations and member indexing share one manifest and app identity. */ +/** + * Custom and shared apps declare the same permissions, including planned capabilities. + * Runtime OAuth validation requires only scopes used by implemented features. + */ export function createSlackSearchManifest( name: string, description: string, @@ -38,8 +41,40 @@ export function createSlackSearchManifest( SLACK_MANAGED_USER_ENROLLMENT_CALLBACK_PATH, ].map((path) => new URL(path, url).href), scopes: { - bot: [...SLACK_SEARCH_SCOPES], - user: [...new Set([...SLACK_SEARCH_USER_SCOPES, ...existingUserScopes])], + bot: [ + ...SLACK_SHARED_SEARCH_BOT_SCOPES, + 'channels:history', + 'channels:manage', + 'channels:write.invites', + 'chat:write.public', + 'groups:history', + 'groups:write', + 'groups:write.invites', + 'links:read', + 'links:write', + 'mpim:history', + 'mpim:read', + 'mpim:write', + 'reactions:write', + ], + user: [ + ...new Set([ + ...SLACK_SEARCH_USER_SCOPES, + 'canvases:read', + 'canvases:write', + 'chat:write', + 'files:read', + 'search:read.files', + 'search:read.im', + 'search:read.mpim', + 'search:read.private', + 'search:read.public', + 'search:read.users', + 'team:read', + 'usergroups:read', + ...existingUserScopes, + ]), + ], }, }, settings: { @@ -55,10 +90,7 @@ export function createSlackSearchManifest( } } -/** - * Declares the company app's permissions, including planned capabilities. - * Runtime OAuth validation continues to require only scopes used by implemented features. - */ +/** Adds the official app's commands and lifecycle events to the common manifest. */ export function createSharedSlackSearchManifest(origin: string) { const manifest = createSlackSearchManifest( SLACK_SEARCH_DEFAULT_NAME, @@ -87,44 +119,6 @@ export function createSharedSlackSearchManifest(origin: string) { }, ], }, - oauth_config: { - ...manifest.oauth_config, - scopes: { - bot: [ - ...SLACK_SHARED_SEARCH_BOT_SCOPES, - 'channels:history', - 'channels:manage', - 'channels:read', - 'channels:write.invites', - 'chat:write.public', - 'groups:history', - 'groups:read', - 'groups:write', - 'groups:write.invites', - 'links:read', - 'links:write', - 'mpim:history', - 'mpim:read', - 'mpim:write', - 'reactions:write', - ], - user: [ - ...SLACK_SEARCH_USER_SCOPES, - 'canvases:read', - 'canvases:write', - 'chat:write', - 'files:read', - 'search:read.files', - 'search:read.im', - 'search:read.mpim', - 'search:read.private', - 'search:read.public', - 'search:read.users', - 'team:read', - 'usergroups:read', - ], - }, - }, settings: { ...manifest.settings, event_subscriptions: { diff --git a/apps/sim/scripts/backfill-projection-source-acl.ts b/apps/sim/scripts/backfill-projection-source-acl.ts index be982412df6..ebc4a1ec390 100644 --- a/apps/sim/scripts/backfill-projection-source-acl.ts +++ b/apps/sim/scripts/backfill-projection-source-acl.ts @@ -22,13 +22,31 @@ import { const logger = createLogger('BackfillProjectionSourceAcl') -/** A script has no long-lived process to detach into, so without a worker it fills inline. */ +/** `--shards ` as given, or one; a value that is not a whole number is refused here. */ +function shardsFlag(): number { + const flag = process.argv.indexOf('--shards') + if (flag === -1) return 1 + const shards = Number(process.argv[flag + 1]) + if (!Number.isInteger(shards) || shards < 1) { + throw new Error(`--shards must be a positive whole number, got ${process.argv[flag + 1]}`) + } + return shards +} + +/** + * A script has no long-lived process to detach into, so without a worker it fills inline, one + * pass over the whole id space. With a worker, `--shards ` slices the id space so that many + * runs fill at once; inline there is nothing to slice across, so the flag is refused there. + */ async function main(): Promise { + const shards = shardsFlag() if (isTriggerDevEnabled && env.TRIGGER_SECRET_KEY) { - const handle = await enqueueProjectionSourceAclBackfill() - logger.info('Backfill enqueued on the Trigger.dev worker', handle) + const started = await enqueueProjectionSourceAclBackfill({}, shards) + logger.info('Backfill enqueued on the Trigger.dev worker', started) return } + if (shards !== 1) + throw new Error('--shards needs the Trigger.dev worker; the inline fill is one pass') await runProjectionSourceAclBackfill({}) logger.info('Backfill complete') } diff --git a/packages/db/script-migrations/0021_embedding_search_connector.test.ts b/packages/db/script-migrations/0021_embedding_search_connector.test.ts index d4e030b449a..94c4246100e 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.test.ts +++ b/packages/db/script-migrations/0021_embedding_search_connector.test.ts @@ -25,4 +25,24 @@ describe('backfillProjectionSourceAcl', () => { ).rejects.toThrow('pause must be a non-negative number') expect(untouched.begin).not.toHaveBeenCalled() }) + + it('binds the range it was given to every page, so shards never meet', async () => { + const statements: Array<{ query: string; params?: unknown[] }> = [] + const tx = { + unsafe: vi.fn(async (query: string, params?: unknown[]) => { + statements.push({ query, params }) + return query.includes('WITH page') ? [{ scanned: 0, filled: 0, last_id: null }] : [] + }), + } + const session = { + begin: vi.fn(async (work: (tx: unknown) => Promise) => work(tx)), + unsafe: vi.fn(async () => []), + } as unknown as Sql + await expect( + backfillProjectionSourceAcl(session, 'embedding_search', { afterId: '4', beforeId: '8' }) + ).resolves.toMatchObject({ done: true }) + const page = statements.find((statement) => statement.query.includes('WITH page'))! + expect(page.query).toContain('($2::text IS NULL OR s.id < $2)') + expect(page.params).toEqual(['4', '8']) + }) }) diff --git a/packages/db/script-migrations/0021_embedding_search_connector.ts b/packages/db/script-migrations/0021_embedding_search_connector.ts index 4fc87e10fe3..3a05f4e76d3 100644 --- a/packages/db/script-migrations/0021_embedding_search_connector.ts +++ b/packages/db/script-migrations/0021_embedding_search_connector.ts @@ -83,6 +83,8 @@ export async function installProjectionSourceAcl(sql: Sql): Promise { export interface ProjectionSourceAclBackfillOptions { /** Resume after this chunk id; the projection's first page otherwise. */ afterId?: string + /** Stop before this chunk id; the projection's end otherwise. Lets workers fill disjoint ranges. */ + beforeId?: string pageSize?: number pauseMs?: number /** Stop once this much time has passed and report where to resume; unbounded otherwise. */ @@ -112,7 +114,9 @@ export interface ProjectionSourceAclBackfillProgress { * cost of a page and far more than a deploy can wait for; the run paces itself with a pause between * pages and stops at its budget so a background task can chain runs until the projection is filled. * Search does not wait: an unfilled row is decided on its document by the on-row candidate - * predicate, the join per candidate every row paid before the columns existed. + * predicate, the join per candidate every row paid before the columns existed. A run fills the + * range it was given and reports that range done; whether the projection as a whole is done, and + * the analysis the planner then needs, is the caller's, since several runs may share a projection. */ export async function backfillProjectionSourceAcl( sql: Sql, @@ -135,6 +139,7 @@ export async function backfillProjectionSourceAcl( const deadline = options.budgetMs === undefined ? Number.POSITIVE_INFINITY : startedAt + options.budgetMs let afterId = options.afterId ?? '' + const beforeId = options.beforeId ?? null let scanned = 0 let written = 0 let pages = 0 @@ -149,7 +154,7 @@ export async function backfillProjectionSourceAcl( `WITH page AS ( SELECT s.id, s.document_id, d.connector_id, d.acl FROM ${projection} s JOIN document d ON d.id = s.document_id - WHERE s.id > $1 AND s.acl IS NULL + WHERE s.id > $1 AND ($2::text IS NULL OR s.id < $2) AND s.acl IS NULL ORDER BY s.id LIMIT ${pageSize} FOR SHARE OF d ), updated AS ( @@ -161,14 +166,12 @@ export async function backfillProjectionSourceAcl( SELECT (SELECT count(*)::int FROM page) AS scanned, (SELECT count(*)::int FROM updated) AS filled, (SELECT max(id) FROM page) AS last_id`, - [afterId] + [afterId, beforeId] ) return row }) if (page.last_id === null) { done = true - /** The planner last saw every row unfilled; it should see the finished projection. */ - await sql.unsafe(`ANALYZE ${projection}`) break } afterId = page.last_id @@ -189,9 +192,12 @@ export async function backfillProjectionSourceAcl( if (Date.now() >= deadline) break } logger.info( - done ? 'Projection source and ACL backfilled' : 'Projection source and ACL backfill paused', + done + ? 'Projection source and ACL range backfilled' + : 'Projection source and ACL backfill paused', { projection, + beforeId, scanned, written, afterId,