perf(kv): merge row data in one statement on upsert and patch (#2161)

This commit is contained in:
Hampus
2026-08-30 22:37:18 +02:00
committed by GitHub
parent 1054962008
commit 4a2a29f154
2 changed files with 138 additions and 29 deletions
@@ -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(
@@ -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<Array<Row>> {
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<void> {
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<string>);
const next = {...base};
const incoming = paramsRow(params, (meta.pkColumns ?? meta.table.primaryKey) as ReadonlyArray<string>);
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<StoredRow & {expires_at: Date | null}>(
`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<Row | null> {
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<Row | null> {
return (await this.getStoredRow(meta, key, db))?.row ?? null;
return row ? decodeRow(row.row_data) : null;
}
}