From 4a2a29f15411903c77d14bcb2ca9215e87a24f10 Mon Sep 17 00:00:00 2001 From: Hampus Date: Sun, 30 Aug 2026 22:37:18 +0200 Subject: [PATCH] perf(kv): merge row data in one statement on upsert and patch (#2161) --- .../database/PostgresKvNonSelectDrift.test.ts | 118 ++++++++++++++++++ .../api/database/PostgresKvQueryExecutor.ts | 49 +++----- 2 files changed, 138 insertions(+), 29 deletions(-) diff --git a/fluxer_api/src/api/database/PostgresKvNonSelectDrift.test.ts b/fluxer_api/src/api/database/PostgresKvNonSelectDrift.test.ts index 1ed3f0f80..404bc10f1 100644 --- a/fluxer_api/src/api/database/PostgresKvNonSelectDrift.test.ts +++ b/fluxer_api/src/api/database/PostgresKvNonSelectDrift.test.ts @@ -208,6 +208,124 @@ describe.skipIf(!dockerAvailable)('postgres kv non-select drift', () => { expect(await dump(NEXT_TABLE, Parent.name)).toBe(await dump(LEGACY_TABLE, Parent.name)); }); + it('partial upsert merges with the stored row the way legacy does', async () => { + await reset(Parent, seedRows); + const meta = {action: 'upsert', table: Parent} as AnyMeta; + const partial = {user_id: 1n, channel_id: 10n, blob_: 'merged'} as CassandraParams; + await legacy.executeQuery({cql: '__mu__', params: partial, kvMeta: meta}); + await next.executeQuery({cql: '__mu__', params: partial, kvMeta: meta}); + expect(await dump(NEXT_TABLE, Parent.name)).toBe(await dump(LEGACY_TABLE, Parent.name)); + const merged = await raw.query<{note: string | null; blob_: string | null}>( + `SELECT row_data->>'note' AS note, row_data->>'blob_' AS blob_ FROM ${NEXT_TABLE} WHERE table_name = $1 AND row_data->>'blob_' IS NOT NULL`, + [Parent.name], + ); + expect(merged.rows).toEqual([{note: 'a', blob_: 'merged'}]); + }); + + it('upsert and patch replace an expired row instead of merging into it', async () => { + await reset(Parent, seedRows); + for (const kv of [LEGACY_TABLE, NEXT_TABLE]) { + await raw.query(`UPDATE ${kv} SET expires_at = timestamptz '2000-01-01T00:00:00Z' WHERE table_name = $1`, [ + Parent.name, + ]); + } + const upsertMeta = {action: 'upsert', table: Parent} as AnyMeta; + const patchMeta = { + action: 'patch', + table: Parent, + patchKeys: ['note'], + pkColumns: ['user_id', 'channel_id'], + } as unknown as AnyMeta; + const u = {user_id: 1n, channel_id: 10n, blob_: 'fresh'} as CassandraParams; + const p = {user_id: 2n, channel_id: 10n, note: 'fresh'} as CassandraParams; + for (const exec of [legacy, next]) { + await exec.executeQuery({cql: '__xu__', params: u, kvMeta: upsertMeta}); + await exec.executeQuery({cql: '__xp__', params: p, kvMeta: patchMeta}); + } + expect(await dump(NEXT_TABLE, Parent.name)).toBe(await dump(LEGACY_TABLE, Parent.name)); + const rows = await raw.query<{note: string | null; blob_: string | null; expires_at: Date | null}>( + `SELECT row_data->>'note' AS note, row_data->>'blob_' AS blob_, expires_at FROM ${NEXT_TABLE} WHERE table_name = $1 ORDER BY row_key COLLATE "C"`, + [Parent.name], + ); + expect(rows.rows.map((row) => [row.note, row.blob_, row.expires_at === null])).toEqual([ + [null, 'fresh', true], + ['b', null, false], + ['fresh', null, true], + ]); + }); + + it('patch without a ttl parameter keeps the stored expiry', async () => { + await reset(Parent, seedRows); + for (const kv of [LEGACY_TABLE, NEXT_TABLE]) { + await raw.query(`UPDATE ${kv} SET expires_at = timestamptz '2099-01-01T00:00:00Z' WHERE table_name = $1`, [ + Parent.name, + ]); + } + const patchMeta = { + action: 'patch', + table: Parent, + patchKeys: ['note'], + pkColumns: ['user_id', 'channel_id'], + } as unknown as AnyMeta; + const p = {user_id: 1n, channel_id: 10n, note: 'kept'} as CassandraParams; + await legacy.executeQuery({cql: '__pk__', params: p, kvMeta: patchMeta}); + await next.executeQuery({cql: '__pk__', params: p, kvMeta: patchMeta}); + expect(await dump(NEXT_TABLE, Parent.name)).toBe(await dump(LEGACY_TABLE, Parent.name)); + const kept = await raw.query<{expires_at: Date | null}>( + `SELECT expires_at FROM ${NEXT_TABLE} WHERE table_name = $1 AND row_data->>'note' = 'kept'`, + [Parent.name], + ); + expect(kept.rows.map((row) => row.expires_at?.toISOString() ?? null)).toEqual(['2099-01-01T00:00:00.000Z']); + }); + + it('patch with a ttl parameter replaces the stored expiry', async () => { + await reset(Parent, seedRows); + for (const kv of [LEGACY_TABLE, NEXT_TABLE]) { + await raw.query(`UPDATE ${kv} SET expires_at = timestamptz '2099-01-01T00:00:00Z' WHERE table_name = $1`, [ + Parent.name, + ]); + } + const ttlMeta = { + action: 'patch', + table: Parent, + patchKeys: ['note'], + pkColumns: ['user_id', 'channel_id'], + ttlParamName: 'ttl_', + } as unknown as AnyMeta; + const p = {user_id: 1n, channel_id: 10n, note: 'ttl', ttl_: 600} as CassandraParams; + const before = Date.now(); + await legacy.executeQuery({cql: '__pt2__', params: p, kvMeta: ttlMeta}); + await next.executeQuery({cql: '__pt2__', params: p, kvMeta: ttlMeta}); + const after = Date.now(); + expect(await dump(NEXT_TABLE, Parent.name)).toBe(await dump(LEGACY_TABLE, Parent.name)); + const set = await raw.query<{expires_at: Date}>( + `SELECT expires_at FROM ${NEXT_TABLE} WHERE table_name = $1 AND row_data->>'note' = 'ttl'`, + [Parent.name], + ); + expect(set.rows.length).toBe(1); + const expiry = set.rows[0]!.expires_at.getTime(); + expect(expiry).toBeGreaterThanOrEqual(before + 600_000); + expect(expiry).toBeLessThanOrEqual(after + 600_000); + }); + + it('insert if not exists over an expired row matches legacy', async () => { + await reset(Parent, seedRows); + for (const kv of [LEGACY_TABLE, NEXT_TABLE]) { + await raw.query(`UPDATE ${kv} SET expires_at = now() - interval '1 hour' WHERE table_name = $1`, [Parent.name]); + } + const insMeta = {action: 'upsert', table: Parent, ifNotExists: true} as AnyMeta; + const revived = {user_id: 1n, channel_id: 10n, note: 'revived', blob_: null} as CassandraParams; + const l1 = await legacy.executeQuery({cql: '__ie__', params: revived, kvMeta: insMeta}); + const n1 = await next.executeQuery({cql: '__ie__', params: revived, kvMeta: insMeta}); + expect(n1).toEqual(l1); + expect(n1).toEqual([{'[applied]': true}]); + const l2 = await legacy.executeQuery({cql: '__ie__', params: revived, kvMeta: insMeta}); + const n2 = await next.executeQuery({cql: '__ie__', params: revived, kvMeta: insMeta}); + expect(n2).toEqual(l2); + expect(n2).toEqual([{'[applied]': false}]); + expect(await dump(NEXT_TABLE, Parent.name)).toBe(await dump(LEGACY_TABLE, Parent.name)); + }); + it('count returns the same value and the same javascript type', async () => { await reset(Parent, seedRows); await raw.query( diff --git a/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts b/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts index 10648d915..726e77b39 100644 --- a/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts +++ b/fluxer_api/src/api/database/PostgresKvQueryExecutor.ts @@ -63,6 +63,9 @@ const FULL_SCAN_LOG_KEY_LIMIT = 1024; const ENCODED_TYPE_KEY = '__fluxer_type'; const POSTGRES_KV_SCHEMA_LOCK_NAMESPACE = 0x46584b56; const POSTGRES_KV_SCHEMA_LOCK_TIMEOUT = '120s'; +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`; function normalizeCql(cql: string): string { return cql.replace(/\s+/g, ' ').trim(); @@ -904,29 +907,27 @@ export class PostgresKvQueryExecutor { private async upsert(meta: KvQueryMeta, params: CassandraParams, db: PostgresQueryable): Promise> { const incoming = rowFromParams(meta, params); const key = rowKey(meta, incoming); - const existing = await this.getRow(meta, key, db); - if (meta.ifNotExists && existing) { - return [{'[applied]': false}]; - } if (meta.ifNotExists) { + if (await this.getRow(meta, key, db)) { + return [{'[applied]': false}]; + } await db.query( `DELETE FROM ${this.table} WHERE table_name = $1 AND row_key = $2 AND expires_at IS NOT NULL AND expires_at <= now()`, [meta.table.name, key], ); } - const next = {...(existing ?? {}), ...incoming}; const expiresAt = ttlExpiresAt(meta, params) ?? null; const result = await db.query( - `INSERT INTO ${this.table} (table_name, partition_key, row_key, row_data, expires_at, updated_at) + `INSERT INTO ${this.table} AS kv (table_name, partition_key, row_key, row_data, expires_at, updated_at) VALUES ($1, $2, $3, $4::jsonb, $5, now()) ON CONFLICT (table_name, row_key) -DO UPDATE SET partition_key = EXCLUDED.partition_key, row_data = EXCLUDED.row_data, expires_at = EXCLUDED.expires_at, updated_at = now() +DO UPDATE SET partition_key = EXCLUDED.partition_key, row_data = ${MERGED_ROW_DATA}, expires_at = EXCLUDED.expires_at, updated_at = now() WHERE NOT $6`, [ meta.table.name, - partitionKey(meta, next), + partitionKey(meta, incoming), key, - JSON.stringify(encodeRow(next)), + JSON.stringify(encodeRow(incoming)), expiresAt, meta.ifNotExists === true, ], @@ -939,20 +940,18 @@ WHERE NOT $6`, private async patch(meta: KvQueryMeta, params: CassandraParams, db: PostgresQueryable): Promise { const key = rowKeyFromParams(meta, params); - const stored = await this.getStoredRow(meta, key, db); - const base = stored?.row ?? paramsRow(params, (meta.pkColumns ?? meta.table.primaryKey) as ReadonlyArray); - const next = {...base}; + const incoming = paramsRow(params, (meta.pkColumns ?? meta.table.primaryKey) as ReadonlyArray); for (const column of meta.patchKeys ?? []) { - next[column] = column in params ? params[column] : null; + incoming[column] = column in params ? params[column] : null; } const ttl = ttlExpiresAt(meta, params); - const expiresAt = ttl === undefined ? (stored?.expiresAt ?? null) : ttl; + const expiresAtExpr = ttl === undefined ? KEPT_EXPIRES_AT : 'EXCLUDED.expires_at'; await db.query( - `INSERT INTO ${this.table} (table_name, partition_key, row_key, row_data, expires_at, updated_at) + `INSERT INTO ${this.table} AS kv (table_name, partition_key, row_key, row_data, expires_at, updated_at) VALUES ($1, $2, $3, $4::jsonb, $5, now()) ON CONFLICT (table_name, row_key) -DO UPDATE SET partition_key = EXCLUDED.partition_key, row_data = EXCLUDED.row_data, expires_at = EXCLUDED.expires_at, updated_at = now()`, - [meta.table.name, partitionKey(meta, next), key, JSON.stringify(encodeRow(next)), expiresAt ?? null], +DO UPDATE SET partition_key = EXCLUDED.partition_key, row_data = ${MERGED_ROW_DATA}, expires_at = ${expiresAtExpr}, updated_at = now()`, + [meta.table.name, partitionKey(meta, incoming), key, JSON.stringify(encodeRow(incoming)), ttl ?? null], ); } @@ -987,20 +986,12 @@ DO UPDATE SET partition_key = EXCLUDED.partition_key, row_data = EXCLUDED.row_da ]); } - private async getStoredRow( - meta: KvQueryMeta, - key: string, - db: PostgresQueryable, - ): Promise<{row: Row; expiresAt: Date | null} | null> { - const result = await db.query( - `SELECT row_key, row_data, expires_at FROM ${this.table} WHERE table_name = $1 AND row_key = $2 AND (expires_at IS NULL OR expires_at > now()) LIMIT 1`, + private async getRow(meta: KvQueryMeta, key: string, db: PostgresQueryable): Promise { + const result = await db.query<{row_data: unknown}>( + `SELECT row_data FROM ${this.table} WHERE table_name = $1 AND row_key = $2 AND (expires_at IS NULL OR expires_at > now()) LIMIT 1`, [meta.table.name, key], ); const row = result.rows[0]; - return row ? {row: decodeRow(row.row_data), expiresAt: row.expires_at ?? null} : null; - } - - private async getRow(meta: KvQueryMeta, key: string, db: PostgresQueryable): Promise { - return (await this.getStoredRow(meta, key, db))?.row ?? null; + return row ? decodeRow(row.row_data) : null; } }