diff --git a/apps/sim/lib/knowledge/__integration__/listing-continuation.integration.ts b/apps/sim/lib/knowledge/__integration__/listing-continuation.integration.ts index 6a00d37dfd5..fdf3389ee8f 100644 --- a/apps/sim/lib/knowledge/__integration__/listing-continuation.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/listing-continuation.integration.ts @@ -966,8 +966,8 @@ describe('durable source and member cycles in PostgreSQL', () => { await db.execute(sql`ANALYZE document`) } const isWindowStatement = (query: string) => /^\s*with "document" as materialized/i.test(query) - const isWindowCount = (query: string) => query.includes('json_build_array') - /** Runs the pass, returning the window scans and counts it issued outside transactions. */ + const isWindowTally = (query: string) => query.includes('json_build_array') + /** Runs the pass, returning the window walks and tallies it issued outside transactions. */ const reconcile = async () => { const hardDelete = vi.spyOn(documentService, 'hardDeleteDocuments') const statements = vi.spyOn(db.$client, 'unsafe') @@ -1002,10 +1002,10 @@ describe('durable source and member cycles in PostgreSQL', () => { hardDeleted: hardDelete.mock.calls.flatMap(([batch]) => batch), walked: statements.mock.calls .map(([query, params]) => ({ query, params })) - .filter(({ query }) => isWindowStatement(query) && !isWindowCount(query)), - counted: statements.mock.calls + .filter(({ query }) => isWindowStatement(query) && !isWindowTally(query)), + tallied: statements.mock.calls .map(([query, params]) => ({ query, params })) - .filter(({ query }) => isWindowStatement(query) && isWindowCount(query)), + .filter(({ query }) => isWindowStatement(query) && isWindowTally(query)), } } finally { statements.mockRestore() @@ -1042,10 +1042,10 @@ describe('durable source and member cycles in PostgreSQL', () => { * must never happen. */ const expectBounded = async ( - walked: { query: string; params: Parameters[1] }[] + statements: { query: string; params: Parameters[1] }[] ) => { - expect(walked.length).toBeGreaterThan(0) - for (const { query, params } of walked) { + expect(statements.length).toBeGreaterThan(0) + for (const { query, params } of statements) { const { read, scans } = await explainWalk(query, params) expect(read, `${query}\n${JSON.stringify(scans, null, 2)}`).toBeLessThanOrEqual( RECONCILIATION_WINDOW_SIZE @@ -1069,7 +1069,7 @@ describe('durable source and member cycles in PostgreSQL', () => { deletedAt: new Date(startedAt.getTime() - 60_000), })) await seed([...present, ...absent, ...tombstoned], present.length) - const { pass, stats, hardDeleted, walked, counted } = await reconcile() + const { pass, stats, hardDeleted, walked, tallied } = await reconcile() expect(pass).toMatchObject({ complete: true, holdNotice: null }) expect(stats.docsDeleted).toBe(absent.length) expect(hardDeleted.sort()).toEqual(tombstoned.map((item) => item.id).sort()) @@ -1082,14 +1082,17 @@ describe('durable source and member cycles in PostgreSQL', () => { expect( stored.filter((item) => item.id.includes('-z-')).every((item) => item.acl.length === 0) ).toBe(true) - await expectBounded(walked) - /** The listing's counts read the connector window by window too, never in one statement. */ - expect(counted).toHaveLength( + /** + * The tally reads the connector window by window, never in one statement, and keeps the + * rows every removal acts on, so no removal reads the connector again. + */ + expect(walked).toEqual([]) + expect(tallied).toHaveLength( Math.floor( (present.length + absent.length + tombstoned.length) / RECONCILIATION_WINDOW_SIZE ) + 1 ) - await expectBounded(counted) + await expectBounded(tallied) }, 120_000) it('scans a dense window once and never reads every tombstone of the connector', async () => { @@ -1109,7 +1112,7 @@ describe('durable source and member cycles in PostgreSQL', () => { deletedAt: new Date(startedAt.getTime() - 60_000), })) await seed([...dense, ...seenTombstones], seenTombstones.length) - const { pass, hardDeleted, walked, counted } = await reconcile() + const { pass, hardDeleted, walked, tallied } = await reconcile() expect(pass).toMatchObject({ complete: true, holdNotice: null }) expect(hardDeleted.sort()).toEqual(dense.map((item) => item.id).sort()) expect( @@ -1119,15 +1122,15 @@ describe('durable source and member cycles in PostgreSQL', () => { .where(eq(document.connectorId, connectorId)) ).toHaveLength(seenTombstones.length) /** - * The only walk is the hard one: one statement per window, the dense first one included, - * then the tail, however many pages of matches a window holds. + * One tally statement per window, the dense first one included, then the tail, however + * many pages of matches a window holds; the purge pages the rows it kept. */ - expect(walked).toHaveLength( + expect(walked).toEqual([]) + expect(tallied).toHaveLength( Math.floor((dense.length + seenTombstones.length) / RECONCILIATION_WINDOW_SIZE) + 1 ) /** `deleted_at < $1` implies the tombstone index, which would read every connector tombstone. */ - await expectBounded(walked) - await expectBounded(counted) + await expectBounded(tallied) }, 120_000) }) diff --git a/apps/sim/lib/knowledge/connectors/member-sync-engine-content-credential.test.ts b/apps/sim/lib/knowledge/connectors/member-sync-engine-content-credential.test.ts index 0bbf0b30d2f..532eb48961d 100644 --- a/apps/sim/lib/knowledge/connectors/member-sync-engine-content-credential.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-sync-engine-content-credential.test.ts @@ -131,9 +131,9 @@ import { } from '@/lib/knowledge/connectors/listing-checkpoint' import { executeMemberSync } from '@/lib/knowledge/connectors/member-sync-engine' import { - routeWindowScans, - windowCount, - windowCounts, + routeWindowTallies, + windowTallies, + windowTally, } from '@/lib/knowledge/connectors/reconciliation-window.test-helpers' import { MEMBER_SCOPE_RENEWAL_PREFIX_BATCH, @@ -278,9 +278,7 @@ function arrange( : [] ) if (!options.contentIncomplete) { - windowCounts.push( - windowCount({ owned: options.existingDocument ? 2 : 1, listed: 1, soft: 0, hard: 0 }) - ) + windowTallies.push(windowTally({ owned: options.existingDocument ? 2 : 1, listed: 1 })) } if (!options.finishedPass || !options.contentIncomplete) queueTableRows(schemaMock.document, [{ count: 1 }]) @@ -367,7 +365,7 @@ describe('member engine with a dedicated content credential', () => { beforeEach(() => { resetEnvFlagsMock() resetDbChainMock() - routeWindowScans() + routeWindowTallies() }) it('refuses an already queued live Search crawl before resolving credentials or taking a lock', async () => { @@ -855,7 +853,7 @@ describe('member engine with a dedicated content credential', () => { vi.clearAllMocks() resetDbChainMock() - routeWindowScans() + routeWindowTallies() const retry = arrange({ members: true, memberContent: true, diff --git a/apps/sim/lib/knowledge/connectors/reconciliation-window.test-helpers.ts b/apps/sim/lib/knowledge/connectors/reconciliation-window.test-helpers.ts index 5bdbc4b6030..a80e838f5f1 100644 --- a/apps/sim/lib/knowledge/connectors/reconciliation-window.test-helpers.ts +++ b/apps/sim/lib/knowledge/connectors/reconciliation-window.test-helpers.ts @@ -1,25 +1,33 @@ import { dbChainMockFns } from '@sim/testing/mocks/database.mock' import { toRecord } from '@sim/utils/object' -/** Reconciliation window scans, which read through `db.execute`, queued in call order. */ -export const windowScans: unknown[][] = [] +/** Reconciliation window tallies, which read through `db.execute`, queued in call order. */ +export const windowTallies: unknown[][] = [] -/** Reconciliation window counts, which read through `db.execute`, queued in call order. */ -export const windowCounts: unknown[][] = [] +interface WindowRow { + id: string + tombstoned: boolean +} -/** The one-row result of a completed listing's counts over a connector shorter than a window. */ -export function windowCount(counts: { +/** + * The one-row result of a completed listing's tally over a connector shorter than a window: the + * listed count, then the absent documents with ACLs, to soft-delete and to hard-delete. + */ +export function windowTally(tally: { owned: number listed: number - acl?: number - soft: number - hard: number + acl?: WindowRow[] + soft?: WindowRow[] + hard?: WindowRow[] }) { + const matching: WindowRow[][] = [[], tally.acl ?? [], tally.soft ?? [], tally.hard ?? []] return [ { - size: counts.owned, - last: counts.owned > 0 ? 'last-owned' : null, - matching: [counts.listed, counts.acl ?? 0, counts.soft, counts.hard], + size: tally.owned, + last: tally.owned > 0 ? 'last-owned' : null, + counts: [tally.listed, ...matching.slice(1).map((rows) => rows.length)], + ids: matching.map((rows) => rows.map((row) => row.id)), + tombstoned: matching.map((rows) => rows.map((row) => row.tombstoned)), }, ] } @@ -43,24 +51,18 @@ function statementText(query: unknown): string { return [...strings, ...(Array.isArray(values) ? values.map(statementText) : [])].join(' ') } -function windowStatement(query: unknown): 'scan' | 'count' | null { - const text = statementText(query) - if (!text.includes('MATERIALIZED')) return null - return text.includes('json_build_array') ? 'count' : 'scan' -} - /** - * Empties {@link windowScans} and {@link windowCounts} and routes `db.execute`: a window scan or - * count takes the next queued result of its kind, and any other statement reads the database - * clock, `databaseTime` when given and the worker's clock otherwise. + * Empties {@link windowTallies} and routes `db.execute`: a window tally takes the next queued + * result, a window scan finds nothing, so a walk no test expects fails on its empty result, and + * any other statement reads the database clock, `databaseTime` when given and the worker's clock + * otherwise. */ -export function routeWindowScans(databaseTime?: Date) { - windowScans.length = 0 - windowCounts.length = 0 +export function routeWindowTallies(databaseTime?: Date) { + windowTallies.length = 0 dbChainMockFns.execute.mockImplementation(async (...args: unknown[]) => { - const statement = windowStatement(args[0]) - if (statement === 'scan') return windowScans.shift() ?? [] - if (statement === 'count') return windowCounts.shift() ?? [] + const text = statementText(args[0]) + if (text.includes('MATERIALIZED')) + return text.includes('json_build_array') ? (windowTallies.shift() ?? []) : [] return [{ startedAt: (databaseTime ?? new Date()).toISOString() }] }) } diff --git a/apps/sim/lib/knowledge/connectors/reconciliation-window.test.ts b/apps/sim/lib/knowledge/connectors/reconciliation-window.test.ts index 8e9ef650588..7938a551bfa 100644 --- a/apps/sim/lib/knowledge/connectors/reconciliation-window.test.ts +++ b/apps/sim/lib/knowledge/connectors/reconciliation-window.test.ts @@ -2,55 +2,104 @@ import { dbChainMockFns, resetDbChainMock } from '@sim/testing' import { sql } from 'drizzle-orm' import { beforeEach, describe, expect, it, vi } from 'vitest' import { - countReconciliationWindows, + RECONCILIATION_RETAINED_MATCHES, RECONCILIATION_WINDOW_SIZE, + tallyReconciliationWindows, } from '@/lib/knowledge/connectors/reconciliation-window' import { - routeWindowScans, - windowCounts, + routeWindowTallies, + windowTallies, } from '@/lib/knowledge/connectors/reconciliation-window.test-helpers' -describe('countReconciliationWindows', () => { - const conditions = [sql`listed`, sql`absent`] +describe('tallyReconciliationWindows', () => { + const conditions = [ + { where: sql`listed`, retain: false }, + { where: sql`absent`, retain: true }, + ] + /** One window's tally row: its size and listed count, then the live absent documents it read. */ + const window = (size: number, listed: number, absent: string[]) => [ + { + size, + last: size > 0 ? `last-${windowTallies.length}` : null, + counts: [listed, absent.length], + ids: [[], absent], + tombstoned: [[], absent.map(() => false)], + }, + ] + const absentIds = (prefix: string, count: number) => + Array.from({ length: count }, (_, index) => `${prefix}-${index}`) + const tally = (beforePage = async () => undefined) => + tallyReconciliationWindows({ + connectorId: 'connector', + conditions, + deadlineAt: Date.now() + 60_000, + beforePage, + }) beforeEach(() => { resetDbChainMock() - routeWindowScans() + routeWindowTallies() }) - it('sums a connector larger than one window across window statements', async () => { - windowCounts.push( - [{ size: RECONCILIATION_WINDOW_SIZE, last: 'doc-5000', matching: [4000, 7] }], - [{ size: 12, last: 'doc-5012', matching: [10, 2] }] + it('sums a connector larger than one window and keeps the retained rows of every window', async () => { + windowTallies.push( + window(RECONCILIATION_WINDOW_SIZE, 4000, ['doc-7', 'doc-9']), + window(12, 10, ['doc-5011']) ) const beforePage = vi.fn(async () => undefined) - const tally = await countReconciliationWindows({ - connectorId: 'connector', - conditions, - deadlineAt: Date.now() + 60_000, - beforePage, - }) - - expect(tally).toEqual({ + expect(await tally(beforePage)).toEqual({ finished: true, owned: RECONCILIATION_WINDOW_SIZE + 12, - matching: [4010, 9], + matching: [ + { count: 4010, rows: undefined }, + { + count: 3, + rows: ['doc-7', 'doc-9', 'doc-5011'].map((id) => ({ id, tombstoned: false })), + }, + ], }) expect(dbChainMockFns.execute).toHaveBeenCalledTimes(2) expect(beforePage).toHaveBeenCalledTimes(2) - expect(windowCounts).toEqual([]) + expect(windowTallies).toEqual([]) + }) + + it('keeps every row up to the retention bound and only the count past it', async () => { + const full = RECONCILIATION_WINDOW_SIZE + const windows = Math.ceil(RECONCILIATION_RETAINED_MATCHES / full) + const fullWindows = () => + Array.from({ length: windows }, (_, index) => + window(full, 0, absentIds(`w${index}`, RECONCILIATION_RETAINED_MATCHES / windows)) + ) + + windowTallies.push(...fullWindows(), window(0, 0, [])) + const atBound = await tally() + expect(atBound.finished && atBound.matching[1]).toEqual({ + count: RECONCILIATION_RETAINED_MATCHES, + rows: Array.from({ length: windows }, (_, index) => + absentIds(`w${index}`, RECONCILIATION_RETAINED_MATCHES / windows) + ) + .flat() + .map((id) => ({ id, tombstoned: false })), + }) + + windowTallies.push(...fullWindows(), window(1, 0, ['one-more'])) + const pastBound = await tally() + expect(pastBound.finished && pastBound.matching[1]).toEqual({ + count: RECONCILIATION_RETAINED_MATCHES + 1, + rows: undefined, + }) }) - it('reports an unfinished count instead of reading past its deadline', async () => { - const tally = await countReconciliationWindows({ + it('reports an unfinished tally instead of reading past its deadline', async () => { + const result = await tallyReconciliationWindows({ connectorId: 'connector', conditions, deadlineAt: Date.now() - 1, beforePage: async () => undefined, }) - expect(tally).toEqual({ finished: false }) + expect(result).toEqual({ finished: false }) expect(dbChainMockFns.execute).not.toHaveBeenCalled() }) }) diff --git a/apps/sim/lib/knowledge/connectors/reconciliation-window.ts b/apps/sim/lib/knowledge/connectors/reconciliation-window.ts index 01b8cb20ace..b0a08b62415 100644 --- a/apps/sim/lib/knowledge/connectors/reconciliation-window.ts +++ b/apps/sim/lib/knowledge/connectors/reconciliation-window.ts @@ -113,72 +113,174 @@ async function scanWindow(walk: ReconciliationWalk, afterId: string | undefined) return scan } -export interface ReconciliationCount { +/** + * Matching rows one tally keeps per condition. Past it a condition keeps only its count, and acting + * on its matches walks the connector again, so a tally's memory and the rows each window statement + * returns stay bounded however many documents a source drops. Far above the removals an ordinary + * sync reconciles, and at least the deletion cap of every connector of up to 40,000 documents. + */ +export const RECONCILIATION_RETAINED_MATCHES = 10_000 + +interface ReconciliationCondition { + /** Reads the window's columns, as {@link ReconciliationWalk.condition} does. */ + where: SQL + /** Whether the tally keeps the matching rows, for a caller that may act on them. */ + retain: boolean +} + +export interface ReconciliationTally { connectorId: string - /** Each condition reads the window's columns, as {@link ReconciliationWalk.condition} does. */ - conditions: readonly SQL[] + conditions: readonly ReconciliationCondition[] deadlineAt: number beforePage: () => Promise } -export type ReconciliationCountResult = +export interface ReconciliationMatches { + /** The owned documents matching the condition. */ + count: number + /** + * Every one of them in id order, as its window read it; undefined when the condition is not + * retained or matched more than {@link RECONCILIATION_RETAINED_MATCHES}. + */ + rows: ReconciliationRow[] | undefined +} + +export type ReconciliationTallyResult = | { finished: false } | { finished: true /** Every owned document of the connector. */ owned: number - /** The owned documents matching each condition, in the conditions' order. */ - matching: number[] + /** The matches of each condition, in the conditions' order. */ + matching: ReconciliationMatches[] } /** A type alias, not an interface, so it satisfies `db.execute`'s row-record constraint. */ -type WindowCount = { +type WindowTally = { size: number last: string | null - matching: number[] + counts: number[] + /** Per condition, its matching ids in id order; empty for a condition no longer retained. */ + ids: string[][] + tombstoned: boolean[][] } /** - * Counts a connector's owned documents, and those matching each condition, one window per - * statement. A connector's documents can number far past what one statement reads inside its - * timeout, so the counts are summed over the same windows a walk scans, each bounded by - * {@link RECONCILIATION_WINDOW_SIZE}. Ids never change, so each document is counted in exactly - * one window, in the state it had when that window's statement ran: the sums are not one - * snapshot, and differ from one only by the rows writers outside the sync lease changed during - * the count. They only decide whether a walk runs; each delete re-checks absence in its own - * statement, as it had to after the single count statement this replaced. Unfinished when the - * deadline arrives first. + * Counts a connector's owned documents and those matching each condition, keeping the matching + * rows of each retained condition, one window per statement. A connector's documents can number + * far past what one statement reads inside its timeout, so the tally is summed over the windows a + * walk scans, each bounded by {@link RECONCILIATION_WINDOW_SIZE}, and a caller acting on the + * retained rows ({@link visitReconciliationMatches}) never reads the connector a second time. + * + * Ids never change, so each document is counted in exactly one window, in the state it had when + * that window's statement ran: the sums are not one snapshot, and differ from one only by the rows + * writers outside the sync lease changed during the tally. The retained rows are the ones counted, + * so what a caller acts on is what its decision weighed; each removal still re-checks its target + * in its own statement. Unfinished when the deadline arrives first. */ -export async function countReconciliationWindows( - count: ReconciliationCount -): Promise { - const filters = sql.join( - count.conditions.map((condition) => sql`count(*) FILTER (WHERE ${condition})::int`), +export async function tallyReconciliationWindows( + tally: ReconciliationTally +): Promise { + const tallies: (ReconciliationMatches & { where: SQL })[] = tally.conditions.map( + ({ where, retain }) => ({ where, count: 0, rows: retain ? [] : undefined }) + ) + /** Each condition is evaluated once per row, into `matched.matches` at its 1-based position. */ + const evaluated = sql.join( + tallies.map(({ where }) => sql`(${where})`), sql`, ` ) - const matching = count.conditions.map(() => 0) + const perCondition = (aggregate: (match: SQL, rows: ReconciliationRow[] | undefined) => SQL) => + sql.join( + tallies.map(({ rows }, index) => aggregate(sql.raw(`matches[${index + 1}]`), rows)), + sql`, ` + ) + /** A condition no longer retained aggregates nothing, so its window returns only the count. */ + const retained = (column: string) => + perCondition((match, rows) => + rows + ? sql`coalesce(json_agg(${sql.raw(column)} ORDER BY id) FILTER (WHERE ${match}), '[]')` + : sql`'[]'::json` + ) let owned = 0 let covered: string | undefined for (;;) { - if (Date.now() >= count.deadlineAt) return { finished: false } - await count.beforePage() - if (Date.now() >= count.deadlineAt) return { finished: false } - const [window] = await db.execute(sql` - ${ownedWindow(count.connectorId, covered)} - SELECT count(*)::int AS "size", max(${document.id}) AS "last", - json_build_array(${filters}) AS "matching" - FROM ${document}`) + if (Date.now() >= tally.deadlineAt) return { finished: false } + await tally.beforePage() + if (Date.now() >= tally.deadlineAt) return { finished: false } + const [window] = await db.execute(sql` + ${ownedWindow(tally.connectorId, covered)}, + matched AS ( + SELECT ${document.id} AS id, ${document.deletedAt} IS NOT NULL AS tombstoned, + ARRAY[${evaluated}]::boolean[] AS matches + FROM ${document} + ) + SELECT count(*)::int AS "size", max(id) AS "last", + json_build_array(${perCondition((match) => sql`count(*) FILTER (WHERE ${match})::int`)}) AS "counts", + json_build_array(${retained('id')}) AS "ids", + json_build_array(${retained('tombstoned')}) AS "tombstoned" + FROM matched`) owned += window.size - window.matching.forEach((value, index) => { - matching[index] += value + tallies.forEach((matches, index) => { + matches.count += window.counts[index] ?? 0 + if (!matches.rows) return + if (matches.count > RECONCILIATION_RETAINED_MATCHES) { + matches.rows = undefined + return + } + const tombstoned = window.tombstoned[index] ?? [] + for (const [row, id] of (window.ids[index] ?? []).entries()) + matches.rows.push({ id, tombstoned: tombstoned[row] ?? false }) }) if (window.size < RECONCILIATION_WINDOW_SIZE || !window.last) { - return { finished: true, owned, matching } + return { + finished: true, + owned, + matching: tallies.map(({ count, rows }) => ({ count, rows })), + } } covered = window.last } } +export interface ReconciliationVisit + extends Pick { + connectorId: string + /** The tallied condition, walked again when its rows were not retained. */ + condition: SQL + matches: ReconciliationMatches +} + +/** + * Hands a tallied condition's matches to `onPage` in pages of at most `pageSize`: the rows the + * tally retained, or, when it kept only their count, a fresh walk of the connector for the + * condition. False when the deadline stopped it. + */ +export async function visitReconciliationMatches(visit: ReconciliationVisit): Promise { + if (!visit.matches.rows) return (await walkReconciliationWindows(visit)).finished + return (await handPages(visit.matches.rows, visit, false)).finished +} + +/** + * Hands `rows` to `onPage` in pages, beating before each one the caller has not already beaten + * for. Materialized rows are acted on only inside the budget, so a late page is left for the next + * run; `lastId` is the last row handed over. + */ +async function handPages( + rows: readonly ReconciliationRow[], + pages: Pick, + firstPageBeaten: boolean +): Promise<{ finished: boolean; lastId: string | undefined }> { + let lastId: string | undefined + for (let offset = 0; offset < rows.length; offset += pages.pageSize) { + if (offset > 0 || !firstPageBeaten) await pages.beforePage() + if (Date.now() >= pages.deadlineAt) return { finished: false, lastId } + const page = rows.slice(offset, offset + pages.pageSize) + await pages.onPage(page) + lastId = page.at(-1)?.id + } + return { finished: true, lastId } +} + /** * Walks a connector's owned documents in id order from `startAfterId`, one window per statement, * handing the rows matching `condition` to `onPage` in pages of at most `pageSize`. A window @@ -194,14 +296,13 @@ export async function walkReconciliationWindows( await walk.beforePage() if (Date.now() >= walk.deadlineAt) return { finished: false, lastId: covered } const { size, last, ids, tombstoned } = await scanWindow(walk, covered) - for (let offset = 0; offset < ids.length; offset += walk.pageSize) { - if (offset > 0) await walk.beforePage() - /** Materialized ids are acted on only inside the budget, so a late page is left for the next run. */ - if (Date.now() >= walk.deadlineAt) return { finished: false, lastId: covered } - const page = ids.slice(offset, offset + walk.pageSize) - await walk.onPage(page.map((id, index) => ({ id, tombstoned: tombstoned[offset + index] }))) - covered = page.at(-1) - } + const handed = await handPages( + ids.map((id, index) => ({ id, tombstoned: tombstoned[index] })), + walk, + true + ) + covered = handed.lastId ?? covered + if (!handed.finished) return { finished: false, lastId: covered } if (size < RECONCILIATION_WINDOW_SIZE || !last) { return { finished: true, lastId: last ?? covered } } diff --git a/apps/sim/lib/knowledge/connectors/sync-content-pass.integration.ts b/apps/sim/lib/knowledge/connectors/sync-content-pass.integration.ts index 3c33d773d94..78c644836be 100644 --- a/apps/sim/lib/knowledge/connectors/sync-content-pass.integration.ts +++ b/apps/sim/lib/knowledge/connectors/sync-content-pass.integration.ts @@ -26,6 +26,9 @@ const { drizzle } = await import('drizzle-orm/postgres-js') const { eq, inArray } = await import('drizzle-orm') const schema = await import('@sim/db/schema') const { beginListingCheckpoint } = await import('@/lib/knowledge/connectors/listing-checkpoint') +const { RECONCILIATION_RETAINED_MATCHES, RECONCILIATION_WINDOW_SIZE } = await import( + '@/lib/knowledge/connectors/reconciliation-window' +) const { runConnectorContentPass } = await import('@/lib/knowledge/connectors/sync-content-pass') const { revokeDocumentAcls } = await import('@/lib/knowledge/connectors/sync-persistence') const { leaseTransaction } = await import('@/lib/knowledge/connectors/sync-lock') @@ -38,6 +41,8 @@ const CONNECTOR = 'admin' const STARTED_AT = new Date('2026-09-08T11:00:00Z') const LISTED_AT = '2026-09-08 12:00:00' const ABSENT_SINCE = '2026-09-01 12:00:00' +/** Tombstoned by an earlier sync, before this listing started. */ +const TOMBSTONED_AT = '2026-09-02 12:00:00' /** * Absence reconciliation of an admin-mode connector against the real projection trigger. The @@ -54,19 +59,33 @@ describe('completed listing reconciliation in PostgreSQL', () => { await sql<{ document_id: string }[]>`SELECT document_id FROM fan_out ORDER BY document_id` ).map((row) => row.document_id) - const insertDocuments = ( - rows: { id: string; acl: string[]; seenAt: string; verified?: boolean }[] - ) => - sql`INSERT INTO document ${sql( - rows.map((row) => ({ - id: row.id, - external_id: row.id, - connector_id: CONNECTOR, - acl: row.acl, - acl_verified_at: row.verified ? '2026-09-01 12:00:00' : null, - source_seen_at: row.seenAt, - })) - )}` + interface DocumentRow { + id: string + acl: string[] + seenAt: string + verified?: boolean + deletedAt?: string + } + + /** + * Inserted in bounded statements, so a connector of several windows stays under the bind limit, + * then analyzed, so each step of a window's walk probes the primary key rather than scanning. + */ + const insertDocuments = async (rows: DocumentRow[]) => { + for (let offset = 0; offset < rows.length; offset += 5_000) + await sql`INSERT INTO document ${sql( + rows.slice(offset, offset + 5_000).map((row) => ({ + id: row.id, + external_id: row.id, + connector_id: CONNECTOR, + acl: row.acl, + acl_verified_at: row.verified ? '2026-09-01 12:00:00' : null, + source_seen_at: row.seenAt, + deleted_at: row.deletedAt ?? null, + })) + )}` + await sql`ANALYZE document` + } /** Listed documents keep both deletion guards open: the listing is neither empty nor collapsed. */ const insertListed = (count: number) => @@ -78,7 +97,10 @@ describe('completed listing reconciliation in PostgreSQL', () => { })) ) - async function reconcile(fullSync = false) { + async function reconcile( + fullSync = false, + options: { beatIfDue?: () => Promise; incomplete?: boolean } = {} + ) { const result = { docsAdded: 0, docsUpdated: 0, @@ -120,7 +142,7 @@ describe('completed listing reconciliation in PostgreSQL', () => { result, lease: { stillHeld: () => eq(schema.knowledgeConnector.id, CONNECTOR), - beatIfDue: async () => undefined, + beatIfDue: options.beatIfDue ?? (async () => undefined), beatLive: async () => undefined, }, leaseKind: 'content', @@ -133,7 +155,7 @@ describe('completed listing reconciliation in PostgreSQL', () => { fullSync, deadlineAt: Date.now() + 60_000, }) - expect(pass.complete).toBe(true) + expect(pass.complete).toBe(!options.incomplete) return { result, notice: pass.holdNotice } } @@ -274,4 +296,242 @@ describe('completed listing reconciliation in PostgreSQL', () => { { id: 'stale-evidence', acl: [], requirements: [], verified: null }, ]) }) + + describe('removal decisions over the tally', () => { + const pad = (index: number) => String(index).padStart(6, '0') + const listed = (count: number, prefix = 'listed'): DocumentRow[] => + Array.from({ length: count }, (_unused, index) => ({ + id: `${prefix}-${pad(index)}`, + acl: [ALICE], + seenAt: LISTED_AT, + })) + const absent = (count: number, prefix: string, deletedAt?: string): DocumentRow[] => + Array.from({ length: count }, (_unused, index) => ({ + id: `${prefix}-${pad(index)}`, + acl: [ALICE], + seenAt: ABSENT_SINCE, + deletedAt, + })) + const ids = (rows: DocumentRow[]) => rows.map((row) => row.id) + const tombstonedIds = async () => + ( + await sql<{ id: string }[]>`SELECT id FROM document WHERE deleted_at IS NOT NULL + AND deleted_at > ${STARTED_AT.toISOString()}::timestamptz ORDER BY id` + ).map((row) => row.id) + const revokedIds = async () => + ( + await sql<{ id: string }[]>`SELECT id FROM document WHERE cardinality(acl) = 0 ORDER BY id` + ).map((row) => row.id) + const purgedIds = () => hardDelete.mock.calls.flatMap(([batch]) => batch).sort() + /** The window statements a pass issues, outside transactions, by kind. */ + const windowStatements = () => { + const statements = vi.spyOn(sql, 'unsafe') + const windows = () => + statements.mock.calls + .map(([query]) => query) + .filter((query) => /^\s*with "document" as materialized/i.test(query)) + return { + tallies: () => windows().filter((query) => query.includes('json_build_array')).length, + walks: () => windows().filter((query) => !query.includes('json_build_array')).length, + restore: () => statements.mockRestore(), + } + } + + /** 100 owned documents allow 25 removals per generation. */ + it.each([ + { soft: 25, hard: 0, notice: null }, + { + soft: 26, + hard: 0, + notice: + /^Withheld 26 document removal\(s\) — more than the 25 allowed per generation in one sync of 100 documents\. Documents removed at the source are still indexed\. /, + }, + { soft: 0, hard: 25, notice: null }, + { + soft: 0, + hard: 26, + notice: + /^Withheld 26 document removal\(s\) — more than the 25 allowed per generation in one sync of 100 documents\. Documents already pending removal were not purged/, + }, + { + soft: 26, + hard: 26, + notice: + /^Withheld 52 document removal\(s\) — more than the 25 allowed per generation in one sync of 100 documents\. Documents removed at the source are still indexed, and documents already pending removal were not purged\. /, + }, + { + soft: 25, + hard: 26, + notice: + /^Withheld 26 document removal\(s\) — more than the 25 allowed per generation in one sync of 100 documents\. Documents already pending removal were not purged/, + }, + ])( + 'applies the per-generation cap to $soft soft and $hard hard removals of 100 documents', + async ({ soft, hard, notice }) => { + const live = absent(soft, 'gone') + const tombstones = absent(hard, 'pending', TOMBSTONED_AT) + await insertDocuments([...listed(100 - soft - hard), ...live, ...tombstones]) + + const { result, notice: holdNotice } = await reconcile() + + const softHeld = soft > 25 + const hardHeld = hard > 25 + expect(holdNotice ?? null).toEqual(notice ? expect.stringMatching(notice) : null) + expect(await tombstonedIds()).toEqual(softHeld ? [] : ids(live)) + expect(result.docsDeleted).toBe(softHeld ? 0 : soft) + expect(purgedIds()).toEqual(hardHeld ? [] : ids(tombstones)) + /** A hold withholds removals, never revocations: absent grants are revoked either way. */ + expect(await revokedIds()).toEqual([...ids(live), ...ids(tombstones)].sort()) + } + ) + + it('removes every absent document past the cap on a full sync, reporting only the newly hidden', async () => { + const live = absent(30, 'gone') + const tombstones = absent(30, 'pending', TOMBSTONED_AT) + await insertDocuments([...listed(40), ...live, ...tombstones]) + + const { result, notice } = await reconcile(true) + + expect(notice).toBeNull() + expect(purgedIds()).toEqual([...ids(live), ...ids(tombstones)].sort()) + expect(result.docsDeleted).toBe(live.length) + expect(await tombstonedIds()).toEqual([]) + }) + + it('tallies an empty connector in one statement and removes nothing', async () => { + const statements = windowStatements() + try { + const { result, notice } = await reconcile() + + expect(notice).toBeNull() + expect(result.docsDeleted).toBe(0) + expect(hardDelete).not.toHaveBeenCalled() + expect({ tallies: statements.tallies(), walks: statements.walks() }).toEqual({ + tallies: 1, + walks: 0, + }) + } finally { + statements.restore() + } + }) + + it('reads a connector of several windows once and acts on the rows it tallied', async () => { + /** Absent documents sort between the listed ones, so every window holds some of each. */ + const rows = Array.from( + { length: 2 * RECONCILIATION_WINDOW_SIZE + 2_000 }, + (_unused, index) => { + const id = `doc-${pad(index)}` + if (index % 1_000 === 7) return { id, acl: [ALICE], seenAt: ABSENT_SINCE } + if (index % 1_000 === 500) + return { id, acl: [ALICE], seenAt: ABSENT_SINCE, deletedAt: TOMBSTONED_AT } + return { id, acl: [ALICE], seenAt: LISTED_AT } + } + ) + const live = rows.filter((row) => row.seenAt === ABSENT_SINCE && !row.deletedAt) + const tombstones = rows.filter((row) => row.deletedAt) + await insertDocuments(rows) + const statements = windowStatements() + try { + const { result, notice } = await reconcile() + + expect(notice).toBeNull() + expect(await tombstonedIds()).toEqual(ids(live)) + expect(result.docsDeleted).toBe(live.length) + expect(purgedIds()).toEqual(ids(tombstones)) + expect(await revokedIds()).toEqual([...ids(live), ...ids(tombstones)].sort()) + expect(await fannedOut()).toEqual([...ids(live), ...ids(tombstones)].sort()) + /** One statement per window decides and finds every removal; no removal reads it again. */ + expect({ tallies: statements.tallies(), walks: statements.walks() }).toEqual({ + tallies: 3, + walks: 0, + }) + } finally { + statements.restore() + } + }, 60_000) + + it('walks again for removals too many to keep, removing each exactly once', async () => { + const gone = Array.from( + { length: RECONCILIATION_RETAINED_MATCHES + 1 }, + (_unused, index) => ({ + id: `gone-${pad(index)}`, + acl: [], + seenAt: ABSENT_SINCE, + }) + ) + const present = listed(1_200) + await insertDocuments([...present, ...gone]) + const statements = windowStatements() + try { + const { result, notice } = await reconcile(true) + + expect(notice).toBeNull() + expect(purgedIds()).toEqual(ids(gone)) + expect(hardDelete.mock.calls.every(([batch]) => batch.length <= 25)).toBe(true) + expect(result.docsDeleted).toBe(gone.length) + const windows = Math.ceil((present.length + gone.length) / RECONCILIATION_WINDOW_SIZE) + expect({ tallies: statements.tallies(), walks: statements.walks() }).toEqual({ + tallies: windows, + walks: windows, + }) + } finally { + statements.restore() + } + }, 120_000) + + it('skips a tallied removal the user excluded before it ran, and never tallies a later one', async () => { + const early = absent(3, 'a-gone') + const late = absent(3, 'z-gone') + await insertDocuments([...early, ...listed(RECONCILIATION_WINDOW_SIZE, 'm-listed'), ...late]) + const statements = windowStatements() + let excluded = false + try { + /** Once the first window is tallied, the user excludes one document from each window. */ + const { result, notice } = await reconcile(false, { + beatIfDue: async () => { + if (excluded || statements.tallies() !== 1) return + excluded = true + await sql`UPDATE document SET user_excluded = true + WHERE id IN (${early[0].id}, ${late[0].id})` + }, + }) + + expect(excluded).toBe(true) + expect(notice).toBeNull() + const removed = [...ids(early.slice(1)), ...ids(late.slice(1))] + expect(await tombstonedIds()).toEqual(removed) + expect(result.docsDeleted).toBe(removed.length) + expect(await revokedIds()).toEqual(removed) + } finally { + statements.restore() + } + }, 60_000) + + it('resumes removals a deadline cut short on the next pass, removing each exactly once', async () => { + /** 2,400 owned documents allow 600 removals: exactly this many, in two soft-delete pages. */ + const gone = absent(600, 'gone') + await insertDocuments([...listed(1_800), ...gone]) + vi.useFakeTimers({ toFake: ['Date'] }) + try { + let cut = false + const first = await reconcile(false, { + incomplete: true, + beatIfDue: async () => { + if (cut || (await tombstonedIds()).length === 0) return + cut = true + vi.setSystemTime(Date.now() + 120_000) + }, + }) + expect(first.result.docsDeleted).toBe(500) + + const second = await reconcile() + + expect(second.result.docsDeleted).toBe(100) + expect(second.notice).toBeNull() + expect(await tombstonedIds()).toEqual(ids(gone)) + } finally { + vi.useRealTimers() + } + }) + }) }) diff --git a/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts b/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts index d141075a0c4..cdfb4984faf 100644 --- a/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts @@ -27,11 +27,9 @@ import { type ListingCheckpoint, } from '@/lib/knowledge/connectors/listing-checkpoint' import { - routeWindowScans, - windowCount, - windowCounts, - windowScan, - windowScans, + routeWindowTallies, + windowTallies, + windowTally, } from '@/lib/knowledge/connectors/reconciliation-window.test-helpers' import { removeDisabledImages, @@ -58,7 +56,7 @@ uploadsMetadataMockFns.mockInsertImmutableFileMetadata.mockImplementation( function resetDbChainMock() { resetDatabaseMock() - routeWindowScans() + routeWindowTallies() } const hoisted = vi.hoisted(() => ({ @@ -226,18 +224,9 @@ describe('completed listing removal counts', () => { complete: true, listedCount: 8, } - windowCounts.push( - windowCount({ - owned: 10, - listed: 8, - acl: options.revoked?.length ?? 0, - soft: soft.length, - hard: hard.length, - }) - ) - /** A walk whose count is zero never scans. Each walk here fits in one partial window. */ + /** One partial window tallies every count and keeps each removal's rows, so nothing walks. */ + windowTallies.push(windowTally({ owned: 10, listed: 8, acl: options.revoked, soft, hard })) if (options.revoked?.length) { - windowScans.push(windowScan(options.revoked)) /** Each window reads what still grants someone; pages are bounded by their chunks' rows. */ const granting = options.revoked.map(({ id }) => ({ id, chunkCount: 10 })) queueTableRows(schemaMock.document, granting) @@ -252,15 +241,11 @@ describe('completed listing removal counts', () => { queueTableRows(schemaMock.knowledgeConnector, [{ id: 'connector' }]) } if (!options.fullSync && soft.length) { - windowScans.push(windowScan(soft)) queueTableRows(schemaMock.knowledgeConnector, [{ id: 'connector' }]) dbChainMockFns.returning.mockResolvedValueOnce( options.updated ?? soft.map(({ id }) => ({ id })) ) } - if (hard.length) { - windowScans.push(windowScan(hard)) - } const result: SyncResult = { docsAdded: 0, docsUpdated: 0, @@ -349,7 +334,7 @@ async function runPass( vi.clearAllMocks() resetDbChainMock() vi.setSystemTime(new Date(Date.now() + 60_000)) - if (options.databaseTime) routeWindowScans(options.databaseTime) + if (options.databaseTime) routeWindowTallies(options.databaseTime) for (let index = 0; index < 16; index++) { queueTableRows(schemaMock.knowledgeConnector, [ { @@ -386,7 +371,7 @@ async function runPass( schemaMock.document, options.permissionStoredAfter ? [options.permissionStoredAfter] : [] ) - windowCounts.push(windowCount({ owned: 1, listed: 1, soft: 0, hard: 0 })) + windowTallies.push(windowTally({ owned: 1, listed: 1 })) for (let index = 0; index < 3; index++) queueTableRows(schemaMock.document, []) dbChainMockFns.returning.mockResolvedValue([{ id: 'document' }]) diff --git a/apps/sim/lib/knowledge/connectors/sync-content-pass.ts b/apps/sim/lib/knowledge/connectors/sync-content-pass.ts index 64636bcecaa..622437e8790 100644 --- a/apps/sim/lib/knowledge/connectors/sync-content-pass.ts +++ b/apps/sim/lib/knowledge/connectors/sync-content-pass.ts @@ -27,9 +27,10 @@ import { connectorPartitionWorkStore, } from '@/lib/knowledge/connectors/partition-store' import { - countReconciliationWindows, + type ReconciliationMatches, type ReconciliationWalk, - walkReconciliationWindows, + tallyReconciliationWindows, + visitReconciliationMatches, } from '@/lib/knowledge/connectors/reconciliation-window' import { IMAGES_DEFERRED_NOTICE, @@ -480,34 +481,44 @@ async function reconcileCompletedListing( const soft = and(absent, isNull(document.deletedAt)) const hard = checkpoint.fullSync ? absent : and(absent, lt(document.deletedAt, startedAt)) /** - * Walks by id through `doc_connector_reconciliation_v2_idx`, filtering absence per row, so no - * walk depends on `source_seen_at` order and the listing's seen stamps stay HOT. + * Reads by id through `doc_connector_reconciliation_v2_idx`, filtering absence per row, so no + * read depends on `source_seen_at` order and the listing's seen stamps stay HOT. One tally + * counts every decision's inputs and keeps the rows each removal may act on, paged like a walk + * because one statement over a large connector outlasts its timeout. */ - const walk = ( + const tally = await tallyReconciliationWindows({ + connectorId: input.connectorId, + conditions: [ + { where: gte(document.sourceSeenAt, startedAt), retain: false }, + { where: aclAbsent, retain: input.documentAccess === 'admin' }, + { where: soft, retain: !checkpoint.fullSync }, + { where: hard, retain: true }, + ].map(({ where, retain }) => ({ where: where ?? sql`true`, retain })), + deadlineAt: input.deadlineAt, + beforePage: input.lease.beatIfDue, + }) + if (!tally.finished) return { finished: false, notice: null } + const ownedCount = tally.owned + const [listed, aclMatches, softMatches, hardMatches] = tally.matching + const listedCount = listed.count + const aclCount = aclMatches.count + const softCount = softMatches.count + const hardCount = hardMatches.count + const visit = ( condition: SQL | undefined, + matches: ReconciliationMatches, pageSize: number, onPage: ReconciliationWalk['onPage'] ) => - walkReconciliationWindows({ + visitReconciliationMatches({ connectorId: input.connectorId, - condition, + condition: condition ?? sql`true`, + matches, pageSize, deadlineAt: input.deadlineAt, beforePage: input.lease.beatIfDue, onPage, }) - /** Paged like the walks: one statement over a large connector outlasts its timeout. */ - const tally = await countReconciliationWindows({ - connectorId: input.connectorId, - conditions: [gte(document.sourceSeenAt, startedAt), aclAbsent, soft, hard].map( - (condition) => condition ?? sql`true` - ), - deadlineAt: input.deadlineAt, - beforePage: input.lease.beatIfDue, - }) - if (!tally.finished) return { finished: false, notice: null } - const ownedCount = tally.owned - const [listedCount, aclCount, softCount, hardCount] = tally.matching let notice: string | null = null let allowDeletion = true if (classifySuspectListing(listedCount, ownedCount)) { @@ -531,7 +542,7 @@ async function reconcileCompletedListing( hardHeld ) if (input.documentAccess === 'admin' && aclCount > 0) { - const { finished } = await walk(aclAbsent, 500, async (rows) => { + const finished = await visit(aclAbsent, aclMatches, 500, async (rows) => { await revokeDocumentAcls( withAclPage, rows.map((row) => row.id), @@ -543,7 +554,7 @@ async function reconcileCompletedListing( } if (!allowDeletion) return { finished: true, notice } if (!checkpoint.fullSync && !softHeld && softCount > 0) { - const { finished } = await walk(soft, 500, async (rows) => { + const finished = await visit(soft, softMatches, 500, async (rows) => { const removed = await withLease((tx) => tx .update(document) @@ -564,7 +575,7 @@ async function reconcileCompletedListing( if (!finished) return { finished: false, notice } } if (!hardHeld && hardCount > 0) { - const { finished } = await walk(hard, 25, async (rows) => { + const finished = await visit(hard, hardMatches, 25, async (rows) => { /** Report newly removed documents once; purging existing tombstones is storage cleanup. */ for (const tombstoned of [false, true]) { const ids = rows.filter((row) => row.tombstoned === tombstoned).map((row) => row.id) diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts index f0ae97cbe15..cb218d47599 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -27,11 +27,9 @@ import { DrizzleQueryError } from 'drizzle-orm/errors' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import * as connectorTokens from '@/lib/knowledge/connectors/access-token' import { - routeWindowScans, - windowCount, - windowCounts, - windowScan, - windowScans, + routeWindowTallies, + windowTallies, + windowTally, } from '@/lib/knowledge/connectors/reconciliation-window.test-helpers' import { buildSyncDatabaseRetryUpdate, @@ -75,7 +73,7 @@ beforeEach(resetEnvFlagsMock) function resetDbChainMock() { resetDatabaseMock() - routeWindowScans() + routeWindowTallies() } vi.mock('@/lib/knowledge/documents/service', () => knowledgeDocumentsServiceMock) @@ -2250,11 +2248,16 @@ describe('executeSync hard-delete reconciliation', () => { queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-1' }]) queueTableRows(schemaMock.document, []) queueTableRows(schemaMock.document, ownedDocs.slice(0, LISTED_DOC_COUNT)) - windowCounts.push( - windowCount({ owned: OWNED_DOC_COUNT, listed: LISTED_DOC_COUNT, soft: 40, hard: 40 }) + /** Less than a window: one tally counts the absent documents and keeps them for the purge. */ + const missing = missingIds.map((id) => ({ id, tombstoned: false })) + windowTallies.push( + windowTally({ + owned: OWNED_DOC_COUNT, + listed: LISTED_DOC_COUNT, + soft: missing, + hard: missing, + }) ) - /** Less than a window remains, so the hard walk scans it once, then pages its matches. */ - windowScans.push(windowScan(missingIds.map((id) => ({ id, tombstoned: false })))) queueTableRows(schemaMock.document, []) queueTableRows(schemaMock.document, [{ count: LISTED_DOC_COUNT }]) dbChainMockFns.returning.mockResolvedValueOnce([CONNECTOR]) @@ -2516,7 +2519,7 @@ describe('executeSync terminal exits under a lost lock', () => { }, ]) queueTableRows(schemaMock.document, []) - windowCounts.push(windowCount({ owned: 0, listed: 0, soft: 0, hard: 0 })) + windowTallies.push(windowTally({ owned: 0, listed: 0 })) queueTableRows(schemaMock.document, []) queueTableRows(schemaMock.document, []) queueTableRows(schemaMock.document, [])