diff --git a/apps/sim/lib/table/rows/row-writes.integration.ts b/apps/sim/lib/table/rows/row-writes.integration.ts index 84c3e7bb17b..78ca4e30eda 100644 --- a/apps/sim/lib/table/rows/row-writes.integration.ts +++ b/apps/sim/lib/table/rows/row-writes.integration.ts @@ -1989,4 +1989,202 @@ describe('table row writes against real PostgreSQL', () => { } }) }) + + describe('bounded row reads', () => { + /** + * A custom sort is read once as ids and its rows fetched by id in byte-sized batches, so the + * page must still be the view's order — ties broken by `order_key` then `id`, NULL and missing + * cells last — across every batch boundary and from a mid-view offset. A 120-row limit drains + * in at least two batches (the first is capped at 51 rows); an unbounded drain from the offset + * outgrows its first id pass (205 ids) and has to read the order again from where it stopped. + */ + it('drains a custom sort across batches in the order of the whole view', async () => { + const table = await createTable([ + { id: 'score', name: 'score', type: 'number' }, + ...textColumns('name'), + ]) + const fixtures = Array.from({ length: 260 }, (_, index) => { + const data: RowData = { name: `n${index}` } + if (index % 9 === 0) data.score = null + else if (index % 4 !== 0) data.score = index % 5 + return { + id: `${table.id}-${String(index).padStart(3, '0')}`, + data, + orderKey: index % 7 === 0 ? null : `a${index % 3}`, + } + }) + await seedRows(table.id, fixtures) + const score = (row: (typeof fixtures)[number]) => + typeof row.data.score === 'number' ? row.data.score : null + const expected = [...fixtures] + .sort((left, right) => { + const [l, r] = [score(left), score(right)] + if (l !== r) return l === null ? 1 : r === null ? -1 : r - l + if (left.orderKey !== right.orderKey) { + if (left.orderKey === null) return 1 + if (right.orderKey === null) return -1 + return left.orderKey < right.orderKey ? -1 : 1 + } + return left.id < right.id ? -1 : 1 + }) + .map((row) => row.id) + + for (const offset of [0, 37]) { + const page = await queryRows( + table, + { + sort: { score: 'desc' }, + limit: 120, + offset, + includeTotal: false, + withExecutions: false, + }, + 'sorted-drain' + ) + expect(page.rows.map((row) => row.id)).toEqual(expected.slice(offset, offset + 120)) + expect(page.nextCursor).not.toBeNull() + } + + const rest = await queryRows( + table, + { sort: { score: 'desc' }, offset: 37, includeTotal: false, withExecutions: false }, + 'sorted-drain' + ) + expect(rest.rows.map((row) => row.id)).toEqual(expected.slice(37)) + expect(rest.nextCursor).toBeNull() + }) + + /** + * The selection is applied in SQL; the response must still keep exactly the selected keys a + * row wrote, JSON nulls included, in the caller's order rather than JSONB's key order. + */ + it('narrows row data to the selected columns exactly as returned without a selection', async () => { + const table = await createTable(textColumns('b', 'c', 'other')) + const nested = { a: 'ü "q"', z: [1, { k: null }] } + await seedRows(table.id, [ + { id: `${table.id}-0`, data: { b: null, c: nested, other: 'x' }, orderKey: 'a0' }, + { id: `${table.id}-1`, data: { other: 'only' }, orderKey: 'a1' }, + ]) + + const page = await queryRows( + table, + { + columnIds: new Set(['c', 'b', 'missing']), + includeTotal: false, + withExecutions: false, + }, + 'projection' + ) + + expect(page.rows.map((row) => JSON.stringify(row.data))).toEqual([ + JSON.stringify({ c: nested, b: null }), + JSON.stringify({}), + ]) + }) + + /** + * A cursor past keyed rows reads through the keyed/unkeyed `UNION ALL` seek, so both branches + * must project alike: every page — keyed rows, then the NULL-`order_key` tail — keeps exactly + * the selected keys. + */ + it('projects every page read through a keyset cursor, keyed rows and unkeyed tail alike', async () => { + const table = await createTable(textColumns('b', 'c', 'other')) + const id = (suffix: string) => `${table.id}-${suffix}` + await seedRows(table.id, [ + { id: id('k0'), data: { b: 'k0', c: [0], other: 'x' }, orderKey: 'a0' }, + { id: id('k1'), data: { c: { k: 1 }, other: 'x' }, orderKey: 'a1' }, + { id: id('k2'), data: { b: null, other: 'x' }, orderKey: 'a2' }, + { id: id('u0'), data: { b: 'u0', c: 'u', other: 'x' }, orderKey: null }, + { id: id('u1'), data: { other: 'only' }, orderKey: null }, + ]) + const expected = [ + [id('k0'), { c: [0], b: 'k0' }], + [id('k1'), { c: { k: 1 } }], + [id('k2'), { b: null }], + [id('u0'), { c: 'u', b: 'u0' }], + [id('u1'), {}], + ].map(([rowId, data]) => [rowId, JSON.stringify(data)]) + + for (const limit of [1, 2]) { + const paged: string[][] = [] + let cursor: ReturnType | undefined + do { + const page = await queryRows( + table, + { + ...cursor, + limit, + columnIds: new Set(['c', 'b']), + includeTotal: false, + withExecutions: false, + }, + 'projected-keyset-paging' + ) + paged.push(...page.rows.map((row) => [row.id, JSON.stringify(row.data)])) + cursor = page.nextCursor ? decodeCursor(page.nextCursor) : undefined + } while (cursor) + expect(paged).toEqual(expected) + } + }) + + /** + * A sorted drain reads its order as ids and then fetches those rows in separate statements. A + * writer that commits between the two must not leak into the page: a deleted row still counts + * as the witness that more rows follow, and a re-scored row keeps the value it was ordered by. + * Under READ COMMITTED the page lost its witness (a null cursor with rows left) and returned + * rows out of order. + */ + it('drains a custom sort in one snapshot while a writer commits between its statements', async () => { + const table = await createTable([{ id: 'score', name: 'score', type: 'number' }]) + const id = (suffix: string) => `${table.id}-${suffix}` + await seedRows( + table.id, + ['a', 'b', 'c', 'd'].map((suffix, index) => ({ + id: id(suffix), + data: { score: 40 - index * 10 }, + orderKey: `a${index}`, + })) + ) + + const transaction = db.transaction.bind(db) + const spy = vi.spyOn(db, 'transaction').mockImplementation((callback, config) => + transaction((tx) => { + const select = tx.select.bind(tx) + let selects = 0 + tx.select = ((...args: Parameters) => { + const builder = select(...args) + if (++selects !== 1) return builder + const from = builder.from.bind(builder) + builder.from = ((source: Parameters[0]) => { + const query = from(source) + const execute = query.execute.bind(query) + query.execute = (async (...executeArgs: Parameters) => { + const result = await execute(...executeArgs) + await control`DELETE FROM user_table_rows WHERE id = ${id('c')}` + await control`UPDATE user_table_rows SET data = '{"score": 0}'::jsonb WHERE id = ${id('a')}` + return result + }) as typeof execute + return query + }) as typeof from + return builder + }) as typeof select + return callback(tx) + }, config) + ) + try { + const page = await queryRows( + table, + { sort: { score: 'desc' }, limit: 2, includeTotal: false, withExecutions: false }, + 'sorted-snapshot' + ) + expect(page.rows.map((row) => [row.id, row.data.score])).toEqual([ + [id('a'), 40], + [id('b'), 30], + ]) + expect(page.nextCursor).not.toBeNull() + } finally { + spy.mockRestore() + } + }) + }) }) diff --git a/apps/sim/lib/table/rows/service.ts b/apps/sim/lib/table/rows/service.ts index 4874906e43e..62d70d4e839 100644 --- a/apps/sim/lib/table/rows/service.ts +++ b/apps/sim/lib/table/rows/service.ts @@ -15,7 +15,7 @@ import { userTableRows } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { generateId } from '@sim/utils/id' import { escapeLikePattern } from '@sim/utils/string' -import { and, asc, count, eq, inArray, isNull, type SQL, sql } from 'drizzle-orm' +import { and, asc, count, eq, getTableColumns, inArray, isNull, type SQL, sql } from 'drizzle-orm' import { OrchestrationError } from '@/lib/core/orchestration/types' import { assertRowCapacity, @@ -1299,6 +1299,9 @@ export async function queryRows( budgetBytes: TABLE_LIMITS.MAX_QUERY_RESULT_BYTES, pageCutBytes: getMaxPageBytes(), columnIds, + projectInSql: Boolean( + columnIds && columns.some((column) => !columnIds.has(getColumnId(column))) + ), readProvenance, }) @@ -1385,8 +1388,14 @@ export interface BoundedFetchParams { budgetBytes: number /** Byte cut for a **bounded** page; defaults to 5MB and is environment-overridable. */ pageCutBytes: number - /** Stable column ids to keep in `data`; projected before a row is measured. */ + /** Stable column ids to keep in `data`; a row is measured as returned. */ columnIds?: ReadonlySet + /** + * Narrow `data` in SQL, so the columns left out are never sent. Worth it only when the + * selection drops some of the table's columns: rebuilding nearly the whole value costs the + * server more than sending it. Without it `columnIds` is applied to the fetched rows. + */ + projectInSql?: boolean } export interface BoundedFetchResult { @@ -1409,6 +1418,30 @@ function projectRowData(data: RowData, columnIds: ReadonlySet): RowData return projected } +/** + * The SELECT list of a drained row, with `data` narrowed to `sqlColumnIds` when given: a drain's + * wall time tracks the bytes it pulls off the socket, not the server's work on them. + * + * Every function that reads a TOASTed `data` detoasts it again, so the value is read once: the + * `|| '{}'` copy (kept from inlining by `OFFSET 0`) is the one detoast, and the object check, the + * key scan and the key removal all read that copy. Removing the unselected keys, rather than + * rebuilding the selected ones, returns the value untouched when a row wrote nothing else. A + * value that is not an object projects to `{}`. + */ +function boundedRowFields(sqlColumnIds: ReadonlySet | undefined) { + if (!sqlColumnIds) return getTableColumns(userTableRows) + const keys = sql.join( + [...sqlColumnIds].map((columnId) => sql`${columnId}`), + sql`, ` + ) + return { + ...getTableColumns(userTableRows), + data: sql`(SELECT o.d - ARRAY(SELECT k FROM jsonb_object_keys(o.d) AS k WHERE k <> ALL(ARRAY[${keys}]::text[])) FROM (SELECT CASE WHEN jsonb_typeof(s.d) = 'object' THEN s.d ELSE '{}'::jsonb END AS d FROM (SELECT ${userTableRows.data} || '{}'::jsonb AS d OFFSET 0) AS s OFFSET 0) AS o)`.mapWith( + userTableRows.data + ), + } +} + /** * Drains rows in adaptively-sized bounded batches until the caller's `limit` * or the byte ceiling ends the page. Never issues an unbounded SELECT: the @@ -1423,9 +1456,10 @@ function projectRowData(data: RowData, columnIds: ReadonlySet): RowData * * Advance strategy: when `keysetValid`, the loop re-anchors on each consumed * keyed row and seeks `(order_key, id) > (anchor)` — delete-tolerant and an - * index seek. Otherwise (custom sort, flag off) it advances by OFFSET from the - * inbound position; rows deleted mid-drain can skip/duplicate exactly as - * cross-request offset paging already does. + * index seek. A custom sort reads its order as ids from the inbound offset + * and fetches the rows by id, all in one snapshot (see `drainSortedById`); any + * other unanchored drain advances by OFFSET, where rows deleted mid-drain can + * skip/duplicate exactly as cross-request offset paging already does. * * Always returns at least one row when any match exists, even if that row * alone exceeds the budget. @@ -1443,11 +1477,12 @@ export async function fetchRowsBounded(params: BoundedFetchParams): Promise = [] let bytes = 0 let maxRowBytes = 0 - // What each consumed row cost to FETCH, as stored. With a projection the - // response bytes above can be tiny while the SELECT still returns whole rows, - // so batch sizing must be bounded by this, not by the projected average. - let storedBytes = 0 - let maxStoredRowBytes = 0 + // What each consumed row cost to FETCH. A projection applied to the fetched + // rows leaves the response bytes above tiny while the SELECT still returned + // whole rows, so batch sizing must be bounded by this, not by the projected + // average. Projected in SQL, the SELECT returns only what is measured. + let fetchedBytes = 0 + let maxFetchedRowBytes = 0 let hasMore = false let anchor = params.seek let anchorOffset = params.startOffset @@ -1463,16 +1498,16 @@ export async function fetchRowsBounded(params: BoundedFetchParams): Promise + ): boolean => { + // Project before measuring — the budget is a promise about the response, so columns the + // caller will never receive must not count — and in the caller's column order, which a + // projection built in SQL does not keep. + const row = columnIds + ? { ...fetchedRow, data: projectRowData(fetchedRow.data as RowData, columnIds) } + : fetchedRow + const rowBytes = Buffer.byteLength(JSON.stringify(row.data)) + const rowFetchedBytes = + columnIds && !params.projectInSql + ? Buffer.byteLength(JSON.stringify(fetchedRow.data)) + : rowBytes + if (cutBytes !== undefined && rows.length > 0 && bytes + rowBytes > cutBytes) { + // Unbounded queries promise the ENTIRE result — a partial page would be + // silent truncation, so fail fast instead (the drain has only fetched + // ~budget bytes at this point, never the whole table). + if (limit === undefined) { + throw new TableQueryValidationError( + `Query result exceeds the ${Math.floor(cutBytes / (1024 * 1024))}MB limit. Add a filter or a limit to narrow the result.`, + 'TABLE_QUERY_RESULT_TOO_LARGE' + ) + } + // Bounded page, byte cut opted in: `row` is the witness. Requires a + // non-empty page so a single over-budget row is still returned alone. + hasMore = true + return true + } + // Limit cut: `row` is the +1 peek witness. + if (rows.length === limit) { + hasMore = true + return true + } + rows.push(row) + returnedRows.push(row) + bytes += rowBytes + fetchedBytes += rowFetchedBytes + consumedSinceAnchor++ + if (rowBytes > maxRowBytes) maxRowBytes = rowBytes + if (rowFetchedBytes > maxFetchedRowBytes) maxFetchedRowBytes = rowFetchedBytes + if (keysetValid && row.orderKey) { + anchor = { orderKey: row.orderKey, id: row.id } + anchorOffset = 0 + consumedSinceAnchor = 0 + } + return false + } + + const runBatch = async ( trx: DbTransaction, batchSeek: TableRowsCursor | undefined, batchOffset: number, @@ -1494,7 +1582,7 @@ export async function fetchRowsBounded(params: BoundedFetchParams): Promise { if (!batchSeek) { const query = executor - .select() + .select(rowFields) .from(userTableRows) .where(baseWhere) .orderBy(orderBy) @@ -1512,7 +1600,7 @@ export async function fetchRowsBounded(params: BoundedFetchParams): Promise 0 ? query.offset(batchOffset) : query } - return (async () => { - const batch = await buildQuery(trx) - let cut = false + const batch = await buildQuery(trx) + let cut = false + const returnedRows: Array = [] + for (const fetchedRow of batch) { + if (consumeRow(fetchedRow, returnedRows)) { + cut = true + break + } + } + await params.readProvenance?.capture(trx, returnedRows) + return { cut, batchLength: batch.length } + } + + /** + * A custom sort orders by a key computed from every matching row's `data` — no index holds it, + * so each statement that sorts re-reads and detoasts the whole view, however few rows it + * returns. Advancing by OFFSET paid that pass once per batch. Instead the order is read as ids + * only, in passes that each reach at least as far as a batch would have, and the rows are then + * fetched by primary key in the same byte-sized batches, consumed in that order. + * The id pass carries the same OFFSET, LIMIT reach, and tiebreak the batches did, and every + * statement of the drain reads one REPEATABLE READ snapshot, so the passes and fetches return + * exactly what one statement would have: every id read is fetched, in an order that still matches + * its values, and the witness row and the next pass's offset count the same rows. + */ + const drainSortedById = async (trx: DbTransaction) => { + let idOffset = params.startOffset + let idsExhausted = false + let pendingIds: string[] = [] + while (true) { + if (pendingIds.length === 0) { + if (idsExhausted) return + // A bounded page needs its remaining rows plus the witness. An unbounded drain starts + // with a few first batches' worth, since most either end there or fail fast at the + // budget; any later pass reaches as far as the largest batch could. + const idsWanted = + limit !== undefined + ? limit - rows.length + 1 + : idOffset === params.startOffset + ? 4 * firstBatchCap + 1 + : Number.POSITIVE_INFINITY + const idAsk = Math.min(TABLE_LIMITS.QUERY_BATCH_MAX_ROWS, idsWanted) + const idPage = await trx + .select({ id: userTableRows.id }) + .from(userTableRows) + .where(baseWhere) + .orderBy(orderBy) + .limit(idAsk) + .offset(idOffset) + pendingIds = idPage.map((idRow) => idRow.id) + idOffset += pendingIds.length + if (pendingIds.length < idAsk) idsExhausted = true + if (pendingIds.length === 0) return + } + const limitRemaining = limit === undefined ? Number.POSITIVE_INFINITY : limit - rows.length + const batchIds = pendingIds.splice(0, Math.min(nextBatchRows(), limitRemaining) + 1) + const fetched = await trx + .select(rowFields) + .from(userTableRows) + .where(and(baseWhere, inArray(userTableRows.id, batchIds))) + const fetchedById = new Map(fetched.map((fetchedRow) => [fetchedRow.id, fetchedRow])) const returnedRows: Array = [] - for (const fetchedRow of batch) { - // Project before measuring: the budget is a promise about the response, - // so columns the caller will never receive must not count against it. - const row = columnIds - ? { ...fetchedRow, data: projectRowData(fetchedRow.data as RowData, columnIds) } - : fetchedRow - const rowBytes = Buffer.byteLength(JSON.stringify(row.data)) - const rowStoredBytes = columnIds - ? Buffer.byteLength(JSON.stringify(fetchedRow.data)) - : rowBytes - if (cutBytes !== undefined && rows.length > 0 && bytes + rowBytes > cutBytes) { - // Unbounded queries promise the ENTIRE result — a partial page would be - // silent truncation, so fail fast instead (the drain has only fetched - // ~budget bytes at this point, never the whole table). - if (limit === undefined) { - throw new TableQueryValidationError( - `Query result exceeds the ${Math.floor(cutBytes / (1024 * 1024))}MB limit. Add a filter or a limit to narrow the result.`, - 'TABLE_QUERY_RESULT_TOO_LARGE' - ) - } - // Bounded page, byte cut opted in: `row` is the witness. Requires a - // non-empty page so a single over-budget row is still returned alone. - hasMore = true - cut = true - break - } - // Limit cut: `row` is the +1 peek witness. - if (rows.length === limit) { - hasMore = true + let cut = false + for (const id of batchIds) { + const fetchedRow = fetchedById.get(id) + if (fetchedRow && consumeRow(fetchedRow, returnedRows)) { cut = true break } - rows.push(row) - returnedRows.push(row) - bytes += rowBytes - storedBytes += rowStoredBytes - consumedSinceAnchor++ - if (rowBytes > maxRowBytes) maxRowBytes = rowBytes - if (rowStoredBytes > maxStoredRowBytes) maxStoredRowBytes = rowStoredBytes - if (keysetValid && row.orderKey) { - anchor = { orderKey: row.orderKey, id: row.id } - anchorOffset = 0 - consumedSinceAnchor = 0 - } } await params.readProvenance?.capture(trx, returnedRows) - return { cut, batchLength: batch.length } - })() + if (cut) return + } } /** @@ -1600,10 +1708,17 @@ export async function fetchRowsBounded(params: BoundedFetchParams): Promise { + if (drainsSortedById) { + await drainSortedById(trx) + return + } while (true) { const limitRemaining = limit === undefined ? Number.POSITIVE_INFINITY : limit - rows.length const target = Math.min(nextBatchRows(), limitRemaining) @@ -1619,7 +1734,7 @@ export async function fetchRowsBounded(params: BoundedFetchParams): Promise { expect(secondBatchAsk).toBeGreaterThan(firstBatchAsk) }) - it('sort path advances by offset and emits an offset cursor with inbound offset applied', async () => { + it('sort path reads its order from the inbound offset and emits an offset cursor past it', async () => { const perRow = Math.floor(TABLE_LIMITS.MAX_QUERY_RESULT_BYTES * 0.6) dbChainMockFns.limit.mockResolvedValueOnce([]) - // Sorted batch with inbound offset 40 terminates at .offset(40). - dbChainMockFns.offset.mockResolvedValueOnce([row(1, perRow), row(2, perRow)]) + // The sorted id pass with inbound offset 40 terminates at .offset(40); its rows are then + // fetched by id. + dbChainMockFns.offset.mockResolvedValueOnce([{ id: 'row_1' }, { id: 'row_2' }]) + queueTableRows(schemaMock.userTableRows, [row(1, perRow), row(2, perRow)]) const result = await queryRows( TABLE,