diff --git a/fluxer_api/src/api/database/PostgresKvKeysetScan.test.ts b/fluxer_api/src/api/database/PostgresKvKeysetScan.test.ts new file mode 100644 index 000000000..89ba1fb45 --- /dev/null +++ b/fluxer_api/src/api/database/PostgresKvKeysetScan.test.ts @@ -0,0 +1,334 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +import {defineTable} from '@app/api/database/CassandraTableDsl'; +import type {PreparedQuery} from '@app/api/database/CassandraTypes'; +import {keyFromColumns, PostgresKvQueryExecutor} from '@app/api/database/PostgresKvQueryExecutor'; +import {dockerAvailable, startPostgresKv, stopPostgresKv} from '@app/api/test/DockerTestContainer'; +import type {IPostgresClient} from '@pkgs/postgres/src/Client'; +import {afterAll, beforeAll, beforeEach, describe, expect, it} from 'vitest'; + +const KV_TABLE = 'kv_keyset_scan'; + +type Row = Record; + +interface UserRow { + user_id: bigint; + name: string; + kind: number; +} + +interface BucketRow { + channel_id: bigint; + bucket: number; + note: string; +} + +interface NameRow { + name: string; + value: number; +} + +const KeysetUsers = defineTable({ + name: 'keyset_users', + columns: ['user_id', 'name', 'kind'], + primaryKey: ['user_id'], +}); + +const KeysetBuckets = defineTable({ + name: 'keyset_buckets', + columns: ['channel_id', 'bucket', 'note'], + primaryKey: ['channel_id', 'bucket'], + partitionKey: ['channel_id'], +}); + +const KeysetNames = defineTable({ + name: 'keyset_names', + columns: ['name', 'value'], + primaryKey: ['name'], +}); + +interface Statement { + text: string; + rows: number; +} + +function recordingClient(client: IPostgresClient, statements: Array): IPostgresClient { + return { + query: (async (text: string, values?: Array, name?: string) => { + const result = await client.query(text, values, name); + statements.push({text, rows: result.rows.length}); + return result; + }) as IPostgresClient['query'], + connect: () => client.connect(), + shutdown: () => client.shutdown(), + transaction: (fn) => client.transaction(fn), + kvTable: () => client.kvTable(), + }; +} + +function pageState(state: Row): string { + return Buffer.from(JSON.stringify(state)).toString('base64url'); +} + +function readPageState(state: string): Row { + return JSON.parse(Buffer.from(state, 'base64url').toString('utf8')) as Row; +} + +function userId(index: number): bigint { + return 1_000_000_000_000_000_000n + BigInt(index) * 7_919n; +} + +describe.skipIf(!dockerAvailable)('Postgres KV keyset scans', () => { + let raw: IPostgresClient; + let statements: Array; + let executor: PostgresKvQueryExecutor; + let postgresContainer: string; + + beforeAll(async () => { + const started = await startPostgresKv({namePrefix: 'fluxer-kvkeyset', kvTable: KV_TABLE, maxConnections: 4}); + postgresContainer = started.container; + raw = started.client; + }, 900_000); + + beforeEach(async () => { + await raw.query(`DELETE FROM ${KV_TABLE}`); + statements = []; + executor = new PostgresKvQueryExecutor(recordingClient(raw, statements)); + }); + + afterAll(async () => { + await stopPostgresKv(postgresContainer); + }); + + async function seed(queries: ReadonlyArray): Promise { + await executor.executeBatch(queries.map((query) => ({query: query.cql, params: query.params, meta: query.kvMeta}))); + } + + async function seedUsers(count: number): Promise> { + const rows = Array.from({length: count}, (_, index) => ({ + user_id: userId((index * 37) % count), + name: `user-${index}`, + kind: index % 7, + })); + await seed(rows.map((row) => KeysetUsers.insert(row))); + return rows; + } + + async function seedBuckets(channels: number, buckets: number): Promise { + const queries: Array = []; + for (let channel = 0; channel < channels; channel += 1) { + for (let bucket = 0; bucket < buckets; bucket += 1) { + queries.push( + KeysetBuckets.insert({ + channel_id: userId((channel * 13) % channels), + bucket: bucket * 11 - 20, + note: `note-${(channel + bucket) % 5}`, + }), + ); + } + } + await seed(queries); + } + + async function oracle(query: PreparedQuery): Promise> { + const rows = await executor.executeQuery(query); + statements.length = 0; + return rows; + } + + async function walk( + query: PreparedQuery, + pageSize: number, + start: string | null = null, + ): Promise<{pages: Array>; states: Array}> { + const pages: Array> = []; + const states: Array = []; + let state = start; + do { + const page = await executor.executePagedQuery(query, {pageSize, pageState: state}); + pages.push(page.rows); + state = page.pageState; + if (state !== null) states.push(state); + expect(pages.length).toBeLessThan(10_000); + } while (state !== null); + return {pages, states}; + } + + function expectBounded(batch: number): void { + expect(statements.length).toBeGreaterThan(0); + for (const statement of statements) { + expect(statement.rows).toBeLessThanOrEqual(batch); + if (statement.text.includes('row_data FROM')) expect(statement.text).toMatch(/LIMIT \$2$/); + } + } + + function byKey(columns: ReadonlyArray, rows: ReadonlyArray): Array { + return [...rows].sort((left, right) => { + const l = keyFromColumns(columns, left); + const r = keyFromColumns(columns, right); + return l < r ? -1 : l > r ? 1 : 0; + }); + } + + it('walks a numeric single key table in the same order without a full load', async () => { + await seedUsers(250); + const query = KeysetUsers.select({}).bind({}); + const expected = await oracle(query); + expect(expected).toHaveLength(250); + const {pages, states} = await walk(query, 40); + expect(pages.flat()).toEqual(expected); + expect(pages.slice(0, -1).every((page) => page.length === 40)).toBe(true); + expect(readPageState(states[0]!)).toEqual({ + offset: 40, + after: keyFromColumns(['user_id'], expected[39]!), + keyed: true, + }); + expectBounded(41); + }); + + it('walks a composite key table by row key without reloading it per page', async () => { + await seedBuckets(30, 7); + const query = KeysetBuckets.select({}).bind({}); + const expected = await oracle(query); + expect(expected).toHaveLength(210); + const {pages} = await walk(query, 25); + const rows = pages.flat(); + expect(byKey(['channel_id', 'bucket'], rows)).toEqual(byKey(['channel_id', 'bucket'], expected)); + expect(new Set(rows.map((row) => keyFromColumns(['channel_id', 'bucket'], row))).size).toBe(210); + const channelRuns = rows.filter((row, index) => index === 0 || rows[index - 1]!.channel_id !== row.channel_id); + expect(channelRuns).toHaveLength(30); + expect(statements.filter((statement) => statement.text.includes('row_data FROM'))).toHaveLength(pages.length); + expectBounded(26); + }); + + it('walks a text key table in the same order', async () => { + await seed( + Array.from({length: 90}, (_, index) => KeysetNames.insert({name: `name-${(index * 31) % 90}-x`, value: index})), + ); + const query = KeysetNames.select({}).bind({}); + const expected = await oracle(query); + const {pages} = await walk(query, 16); + expect(pages.flat()).toEqual(expected); + expectBounded(17); + }); + + it('fills filtered pages from bounded batches on numeric and composite tables', async () => { + await seedUsers(300); + await seedBuckets(20, 9); + const users = KeysetUsers.select({where: KeysetUsers.where.eq('kind')}).bind({kind: 3}); + const expectedUsers = await oracle(users); + const userWalk = await walk(users, 10); + expect(userWalk.pages.flat()).toEqual(expectedUsers); + expect(userWalk.pages.slice(0, -1).every((page) => page.length === 10)).toBe(true); + expectBounded(11); + const buckets = KeysetBuckets.select({where: KeysetBuckets.where.eq('note')}).bind({note: 'note-2'}); + const expectedBuckets = await oracle(buckets); + const bucketWalk = await walk(buckets, 6); + expect(byKey(['channel_id', 'bucket'], bucketWalk.pages.flat())).toEqual( + byKey(['channel_id', 'bucket'], expectedBuckets), + ); + expect(bucketWalk.pages.slice(0, -1).every((page) => page.length === 6)).toBe(true); + expectBounded(7); + }); + + it('skips expired rows', async () => { + await seedUsers(60); + await seedBuckets(6, 5); + await raw.query( + `UPDATE ${KV_TABLE} SET expires_at = now() - interval '1 minute' WHERE (row_data ->> 'kind')::int = 2 OR row_data ->> 'note' = 'note-1'`, + ); + const users = KeysetUsers.select({}).bind({}); + const expectedUsers = await oracle(users); + expect(expectedUsers.length).toBeLessThan(60); + expect((await walk(users, 7)).pages.flat()).toEqual(expectedUsers); + const buckets = KeysetBuckets.select({}).bind({}); + const expectedBuckets = await oracle(buckets); + expect(expectedBuckets.length).toBeLessThan(30); + expect(byKey(['channel_id', 'bucket'], (await walk(buckets, 7)).pages.flat())).toEqual( + byKey(['channel_id', 'bucket'], expectedBuckets), + ); + expectBounded(8); + }); + + it('returns one empty page for an empty table', async () => { + for (const query of [KeysetUsers.select({}).bind({}), KeysetBuckets.select({}).bind({})]) { + expect(await executor.executePagedQuery(query, {pageSize: 10})).toEqual({rows: [], pageState: null}); + } + expectBounded(1); + }); + + it('never drops rows with non-numeric keys from a single key table', async () => { + await seedUsers(40); + await executor.executeQuery(KeysetUsers.insert({user_id: 'stray' as unknown as bigint, name: 'stray', kind: 0})); + const {pages} = await walk(KeysetUsers.select({}).bind({}), 9); + const names = pages.flat().map((row) => row.name); + expect(names).toHaveLength(41); + expect(new Set(names).size).toBe(41); + expect(names).toContain('stray'); + expectBounded(10); + }); + + it('resumes every page state shape the previous release produced', async () => { + await seedUsers(120); + await seedBuckets(12, 6); + const users = KeysetUsers.select({}).bind({}); + const expectedUsers = await oracle(users); + const userKey = keyFromColumns(['user_id'], expectedUsers[49]!); + for (const state of [{offset: 50, after: userKey, keyed: true}, {offset: 50, after: userKey}, {offset: 50}]) { + expect((await walk(users, 15, pageState(state))).pages.flat()).toEqual(expectedUsers.slice(50)); + } + const buckets = KeysetBuckets.select({}).bind({}); + const expectedBuckets = await oracle(buckets); + const bucketKey = keyFromColumns(['channel_id', 'bucket'], expectedBuckets[29]!); + statements.length = 0; + const resumed = await walk(buckets, 11, pageState({offset: 30, after: bucketKey})); + expect(byKey(['channel_id', 'bucket'], resumed.pages.flat())).toEqual( + byKey(['channel_id', 'bucket'], expectedBuckets.slice(30)), + ); + expect(readPageState(resumed.states[0]!)).toMatchObject({byRowKey: true, since: bucketKey}); + expectBounded(12); + }); + + it('streams limited scans in key order instead of loading the table', async () => { + const seeded = await seedUsers(200); + const sorted = [...seeded].sort((left, right) => (left.user_id < right.user_id ? -1 : 1)); + const first = KeysetUsers.select({limit: 20}).bind({}); + const next = KeysetUsers.select({where: KeysetUsers.where.tokenGt('user_id', 'last_user_id'), limit: 20}); + const expectedFirst = await oracle(KeysetUsers.select({}).bind({})); + expect(await executor.executeQuery(first)).toEqual(expectedFirst.slice(0, 20)); + expectBounded(21); + statements.length = 0; + const rows = await executor.executeQuery(next.bind({last_user_id: sorted[99]!.user_id})); + expect(rows).toEqual(expectedFirst.slice(100, 120)); + expect(statements.filter((statement) => statement.text.includes('row_data FROM'))).toHaveLength(1); + expectBounded(21); + statements.length = 0; + expect(await executor.executeQuery(next.bind({last_user_id: sorted[195]!.user_id}))).toEqual( + expectedFirst.slice(196), + ); + expectBounded(21); + }); + + it('probes for non-numeric keys through the primary key while other tables hold string keys', async () => { + await raw.query( + `INSERT INTO ${KV_TABLE} (table_name, partition_key, row_key, row_data) SELECT 'keyset_users', k, k, '{}'::jsonb FROM generate_series(1, 20000) g, LATERAL (SELECT '{"__fluxer_type":"bigint","value":"' || g || '"}' AS k) s`, + ); + await raw.query( + `INSERT INTO ${KV_TABLE} (table_name, partition_key, row_key, row_data) SELECT 'keyset_strings', g::text, '"' || lpad(g::text, 8, '0') || '"', '{}'::jsonb FROM generate_series(1, 200000) g`, + ); + await raw.query(`ANALYZE ${KV_TABLE}`); + await executor.executePagedQuery(KeysetUsers.select({}).bind({}), {pageSize: 10}); + const probe = statements.find((statement) => statement.text.endsWith('AS found')); + expect(probe).toBeDefined(); + const plan = await raw.query<{'QUERY PLAN': string}>(`EXPLAIN (COSTS OFF) ${probe!.text}`, ['keyset_users']); + expect(plan.rows.map((row) => row['QUERY PLAN']).join('\n')).not.toContain('Seq Scan'); + }); + + it('keeps the sorted load when a limited scan asks for an order the key walk cannot give', async () => { + await seedUsers(30); + const descending = KeysetUsers.select({orderBy: {col: 'user_id', direction: 'DESC'}, limit: 5}).bind({}); + const rows = await executor.executeQuery(descending); + const all = await oracle(KeysetUsers.select({}).bind({})); + expect(rows).toEqual(all.slice(-5).reverse()); + }); +}); diff --git a/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts b/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts index 41f976183..271de7516 100644 --- a/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts +++ b/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts @@ -28,6 +28,19 @@ interface PageState { offset: number; after?: string; keyed?: boolean; + byRowKey?: boolean; + since?: string; +} + +interface ScanWalk { + numeric: boolean; + after: string | null; + since: string | null; +} + +interface NumericFloor { + value: string; + inclusive: boolean; } interface PageEntry { @@ -90,6 +103,12 @@ const EXPIRED_STORED_ROW = 'kv.expires_at IS NOT NULL AND kv.expires_at <= now() const MERGED_ROW_DATA = `CASE WHEN ${EXPIRED_STORED_ROW} THEN EXCLUDED.row_data ELSE kv.row_data || EXCLUDED.row_data END`; const KEPT_EXPIRES_AT = `CASE WHEN ${EXPIRED_STORED_ROW} THEN NULL ELSE kv.expires_at END`; const NO_EXPIRY = 'infinity'; +const BIGINT_ROW_KEY_PREFIX = `{"${ENCODED_TYPE_KEY}":"bigint","value":"`; +const NON_NUMERIC_ROW_KEY_RANGES = [ + `kv.row_key COLLATE "C" < '-'`, + `kv.row_key COLLATE "C" >= ':' AND kv.row_key COLLATE "C" < '${BIGINT_ROW_KEY_PREFIX}'`, + `kv.row_key COLLATE "C" >= '${BIGINT_ROW_KEY_PREFIX}:'`, +]; export async function postgresKvPassIsFresh( client: IPostgresClient, @@ -676,15 +695,18 @@ function logError(details: Record, message: string): void { const fullScanLoggedAt = new Map(); -function logFullScan(meta: KvQueryMeta): void { +function logFullScan(meta: KvQueryMeta, walk?: 'numeric' | 'row_key'): void { const shape = queryShape(meta); - const key = `${meta.table.name}|${meta.action}|${shape.summary}`; + const key = `${meta.table.name}|${meta.action}|${shape.summary}|${walk ?? 'load'}`; const now = Date.now(); const last = fullScanLoggedAt.get(key); if (last !== undefined && now - last < FULL_SCAN_LOG_INTERVAL_MS) return; if (fullScanLoggedAt.size >= FULL_SCAN_LOG_KEY_LIMIT) fullScanLoggedAt.clear(); fullScanLoggedAt.set(key, now); - logWarn({table: meta.table.name, action: meta.action, where: shape.summary || 'none'}, 'Postgres KV full table scan'); + logWarn( + {table: meta.table.name, action: meta.action, where: shape.summary || 'none', walk: walk ?? 'load'}, + 'Postgres KV full table scan', + ); } function ttlExpiresAt(meta: KvQueryMeta, params: CassandraParams): Date | typeof NO_EXPIRY | null | undefined { @@ -715,7 +737,13 @@ function decodePageState(pageState: string | null | undefined): PageState { throw new Error('Invalid Postgres KV page state'); } if (typeof decoded.after !== 'string') return {offset: decoded.offset}; - return {offset: decoded.offset, after: decoded.after, keyed: decoded.keyed === true}; + return { + offset: decoded.offset, + after: decoded.after, + keyed: decoded.keyed === true, + byRowKey: decoded.byRowKey === true, + ...(typeof decoded.since === 'string' ? {since: decoded.since} : {}), + }; } function pageableSelect(meta: KvQueryMeta, pageSize: number): boolean { @@ -734,22 +762,30 @@ function numericKeyValue(value: unknown): string | null { return null; } -function numericScanPlan(meta: KvQueryMeta, plan: QueryPlan): boolean { - return plan.exact && plan.candidates.kind === 'scan' && (meta.table.primaryKey as ReadonlyArray).length === 1; +function numericRowKeyValue(key: string): string | null { + const values = decodeRowKey(key, 1); + return values === null ? null : numericKeyValue(values[0]); } -function numericScanKeyed(meta: KvQueryMeta, plan: QueryPlan, entries: ReadonlyArray): boolean { - if (!numericScanPlan(meta, plan)) return false; - const column = (meta.table.primaryKey as ReadonlyArray)[0]!; - return entries.every((entry) => numericKeyValue(entry.row[column]) !== null); +function singlePrimaryKey(meta: KvQueryMeta): boolean { + return (meta.table.primaryKey as ReadonlyArray).length === 1; } -function numericScanCursor(state: PageState): {rowKey: string; value: string} | null { - if (state.keyed !== true || state.after === undefined) return null; - const values = decodeRowKey(state.after, 1); - if (values === null) return null; - const value = numericKeyValue(values[0]); - return value === null ? null : {rowKey: state.after, value}; +function walksInPrimaryKeyOrder(meta: KvQueryMeta): boolean { + if (!singlePrimaryKey(meta)) return false; + if (meta.orderBy === undefined) return true; + return meta.orderBy.col === meta.table.primaryKey[0] && meta.orderBy.direction !== 'DESC'; +} + +function numericFloor(meta: KvQueryMeta, params: CassandraParams): NumericFloor | null { + const column = meta.table.primaryKey[0]; + for (const clause of whereClauses(meta)) { + if (clause.col !== column) continue; + if (clause.kind !== 'gt' && clause.kind !== 'tokenGt' && clause.kind !== 'gte') continue; + const value = numericKeyValue(getParam(params, clause.param)); + if (value !== null) return {value, inclusive: clause.kind === 'gte'}; + } + return null; } function pageStart(meta: KvQueryMeta, entries: ReadonlyArray, state: PageState): number { @@ -1018,10 +1054,10 @@ export class PostgresKvQueryExecutor { const meta = this.meta(query); if (pageableSelect(meta, options.pageSize) && (state.after !== undefined || state.offset === 0)) { const plan = buildCandidatePlan(meta, query.params); - const cursor = numericScanPlan(meta, plan) ? numericScanCursor(state) : null; - const page = cursor - ? await this.numericScanPage(meta, query.params, cursor, state.offset, options.pageSize) - : await this.sortedPage(meta, query.params, plan, state, options.pageSize); + const page = + plan.candidates.kind === 'scan' + ? await this.scanPage(meta, query.params, await this.scanWalk(meta, state), state.offset, options.pageSize) + : await this.sortedPage(meta, query.params, plan, state, options.pageSize); return {rows: page.rows as Array, pageState: page.pageState}; } const rows = await this.executeQuery(query); @@ -1033,29 +1069,111 @@ export class PostgresKvQueryExecutor { }; } - private async numericScanPage( + private async scanWalk(meta: KvQueryMeta, state: PageState): Promise { + if (state.byRowKey) return {numeric: false, after: state.after ?? null, since: state.since ?? null}; + const after = state.after ?? null; + const numericAfter = after === null || numericRowKeyValue(after) !== null; + if (state.keyed && after !== null && numericAfter && singlePrimaryKey(meta)) { + return {numeric: true, after, since: null}; + } + const numeric = singlePrimaryKey(meta) && !(await this.hasNonNumericRows(meta, this.client)); + if (numeric && numericAfter) return {numeric: true, after, since: null}; + return {numeric: false, after: null, since: after}; + } + + private async hasNonNumericRows(meta: KvQueryMeta, db: PostgresQueryable): Promise { + const probes = NON_NUMERIC_ROW_KEY_RANGES.map( + (range) => + `(SELECT kv.row_key FROM ${this.table} kv WHERE kv.table_name = $1 AND ${range} AND (kv.expires_at IS NULL OR kv.expires_at > now()) AND ${NUMERIC_ROW_KEY} IS NULL ORDER BY kv.row_key COLLATE "C" LIMIT 1) IS NOT NULL`, + ); + const result = await db.query<{found: boolean}>(`SELECT ${probes.join(' OR ')} AS found`, [meta.table.name]); + return result.rows[0]?.found === true; + } + + private async scanPage( meta: KvQueryMeta, params: CassandraParams, - cursor: {rowKey: string; value: string}, + walk: ScanWalk, offset: number, pageSize: number, ): Promise<{rows: Array; pageState: string | null}> { - logFullScan(meta); - const result = await this.client.query( - `SELECT kv.row_key, kv.row_data FROM ${this.table} kv WHERE kv.table_name = $1 AND ${NUMERIC_ROW_KEY} IS NOT NULL AND (${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C") > ($2::numeric, $3) AND (kv.expires_at IS NULL OR kv.expires_at > now()) ORDER BY ${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C" LIMIT $4`, - [meta.table.name, cursor.value, cursor.rowKey, pageSize + 1], - ); - const entries = this.matchingEntries(meta, result.rows.slice(0, pageSize), params); - const last = entries[entries.length - 1]; + const page = await this.walkScan(meta, params, walk, pageSize, this.client); + if (!page.more || page.after === null) return {rows: this.projected(meta, page.entries), pageState: null}; + const position = {offset: offset + page.entries.length, after: page.after}; return { - rows: this.projected(meta, entries), - pageState: - result.rows.length > pageSize && last - ? encodePageState({offset: offset + entries.length, after: last.key, keyed: true}) - : null, + rows: this.projected(meta, page.entries), + pageState: encodePageState( + walk.numeric + ? {...position, keyed: true} + : {...position, byRowKey: true, ...(walk.since === null ? {} : {since: walk.since})}, + ), }; } + private async walkScan( + meta: KvQueryMeta, + params: CassandraParams, + walk: ScanWalk, + want: number, + db: PostgresQueryable, + ): Promise<{entries: Array; after: string | null; more: boolean}> { + logFullScan(meta, walk.numeric ? 'numeric' : 'row_key'); + const floor = walk.numeric ? numericFloor(meta, params) : null; + const since = walk.since === null ? null : decodeRowKey(walk.since, meta.table.primaryKey.length); + if (walk.since !== null && since === null) throw new Error('Invalid Postgres KV page state'); + const limit = want + 1; + const entries: Array = []; + let after = walk.after; + while (entries.length < want) { + const batch = await this.scanBatch(meta, walk.numeric, floor, after, limit, db); + for (let index = 0; index < batch.length; index += 1) { + const stored = batch[index]!; + after = stored.row_key; + const entry = this.matchingEntry(meta, stored, params); + if (entry && (since === null || compareRowToKeyValues(meta, entry.row, since) > 0)) entries.push(entry); + if (entries.length === want) { + return {entries, after, more: index < batch.length - 1 || batch.length === limit}; + } + } + if (batch.length < limit) return {entries, after, more: false}; + } + return {entries, after, more: false}; + } + + private async scanBatch( + meta: KvQueryMeta, + numeric: boolean, + floor: NumericFloor | null, + after: string | null, + limit: number, + db: PostgresQueryable, + ): Promise> { + const params: Array = [meta.table.name, limit]; + let predicate = ''; + if (numeric) { + predicate = ` AND ${NUMERIC_ROW_KEY} IS NOT NULL`; + if (floor) { + params.push(floor.value); + predicate += ` AND ${NUMERIC_ROW_KEY} ${floor.inclusive ? '>=' : '>'} $${params.length}::numeric`; + } + if (after !== null) { + const value = numericRowKeyValue(after); + if (value === null) throw new Error('Invalid Postgres KV page state'); + params.push(value, after); + predicate += ` AND (${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C") > ($${params.length - 1}::numeric, $${params.length})`; + } + } else if (after !== null) { + params.push(after); + predicate = ` AND kv.row_key COLLATE "C" > $${params.length}`; + } + const order = numeric ? `${NUMERIC_ROW_KEY}, kv.row_key COLLATE "C"` : 'kv.row_key COLLATE "C"'; + const result = await db.query( + `SELECT kv.row_key, kv.row_data FROM ${this.table} kv WHERE kv.table_name = $1${predicate} AND (kv.expires_at IS NULL OR kv.expires_at > now()) ORDER BY ${order} LIMIT $2`, + params, + ); + return result.rows; + } + private async sortedPage( meta: KvQueryMeta, params: CassandraParams, @@ -1071,12 +1189,9 @@ export class PostgresKvQueryExecutor { const last = page[page.length - 1]; const nextOffset = start + page.length; if (nextOffset >= entries.length || !last) return {rows: this.projected(meta, page), pageState: null}; - const keyed = numericScanKeyed(meta, plan, entries); return { rows: this.projected(meta, page), - pageState: encodePageState( - keyed ? {offset: nextOffset, after: last.key, keyed: true} : {offset: nextOffset, after: last.key}, - ), + pageState: encodePageState({offset: nextOffset, after: last.key}), }; } @@ -1146,17 +1261,22 @@ export class PostgresKvQueryExecutor { stored: ReadonlyArray, params: CassandraParams, ): Array { - const required = queryShape(meta).requiredColumns; const entries: Array = []; for (const entry of stored) { - const row = required ? decodeRowColumns(entry.row_data, required) : decodeRow(entry.row_data); - if (matchesWhere(row, meta.where as ReadonlyArray> | undefined, params)) { - entries.push({key: entry.row_key, row}); - } + const matched = this.matchingEntry(meta, entry, params); + if (matched) entries.push(matched); } return entries; } + private matchingEntry(meta: KvQueryMeta, stored: StoredRow, params: CassandraParams): PageEntry | null { + const required = queryShape(meta).requiredColumns; + const row = required ? decodeRowColumns(stored.row_data, required) : decodeRow(stored.row_data); + return matchesWhere(row, meta.where as ReadonlyArray> | undefined, params) + ? {key: stored.row_key, row} + : null; + } + private matchingRows(meta: KvQueryMeta, stored: ReadonlyArray, params: CassandraParams): Array { return this.matchingEntries(meta, stored, params).map((entry) => entry.row); } @@ -1185,6 +1305,16 @@ export class PostgresKvQueryExecutor { ) { return this.unorderedScan(meta, meta.limit, db); } + if ( + typeof meta.limit === 'number' && + Number.isInteger(meta.limit) && + plan.candidates.kind === 'scan' && + walksInPrimaryKeyOrder(meta) && + !(await this.hasNonNumericRows(meta, db)) + ) { + const walk = {numeric: true, after: null, since: null}; + return this.projected(meta, (await this.walkScan(meta, params, walk, meta.limit, db)).entries); + } let rows = this.matchingRows(meta, await this.candidates(meta, plan, db), params); rows = sortRows(meta, rows); if (typeof meta.limit === 'number') rows = rows.slice(0, meta.limit); diff --git a/fluxer_api/src/api/database/PostgresKvUnorderedScan.test.ts b/fluxer_api/src/api/database/PostgresKvUnorderedScan.test.ts index 558060c46..b870364de 100644 --- a/fluxer_api/src/api/database/PostgresKvUnorderedScan.test.ts +++ b/fluxer_api/src/api/database/PostgresKvUnorderedScan.test.ts @@ -55,15 +55,15 @@ describe('Postgres KV unordered scan', () => { expect(rows).toEqual([{user_id: '2'}, {user_id: '1'}]); }); - it('keeps the sorted full scan when order is asked for', async () => { + it('keeps the sorted full scan when the order cannot be walked by key', async () => { const statements: Array<{text: string; values: Array | undefined}> = []; const executor = new PostgresKvQueryExecutor(recordingClient(statements)); const rows = await executor.executeQuery({ cql: '__ordered_probe', params: {}, - kvMeta: selectMeta({limit: 1}), + kvMeta: selectMeta({limit: 1, orderBy: {col: 'user_id', direction: 'DESC'}}), }); expect(statements[0]!.text).not.toContain('LIMIT'); - expect(rows).toEqual([{user_id: '1'}]); + expect(rows).toEqual([{user_id: '2'}]); }); });