fix(kv): mark the messages migration done instead of rescanning (#2221)

This commit is contained in:
Hampus authored and GitHub committed 2026-08-31 00:22:56 +02:00
1 parent c4b1471923
commit 0a920def2b
4 files changed
+70 -22

No files matched your search

@@ -1082,8 +1082,8 @@ suite('PostgresKvQueryExecutor differential', () => {
);
await raw.query(`CREATE TABLE ${SEED_TABLE} AS TABLE ${LEGACY_TABLE}`);
await raw.query(`CREATE INDEX ${SEED_TABLE}_table_idx ON ${SEED_TABLE} (table_name)`);
await raw.query(`TRUNCATE ${LEGACY_TABLE}`);
for (const kv of [LEGACY_TABLE, NEXT_TABLE, ICU_TABLE]) {
await raw.query(`TRUNCATE ${kv}`);
await raw.query(`INSERT INTO ${kv} SELECT * FROM ${SEED_TABLE} ORDER BY table_name, row_key COLLATE "C"`);
}
}, 900_000);
@@ -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 POSTGRES_KV_SCHEMA_MIGRATION_TIMEOUT = '30min';
export const POSTGRES_KV_MIGRATION_TABLE = '__fluxer_schema_migrations';
const POSTGRES_KV_MESSAGES_PARTITION_MIGRATION = 'messages_partition_key_v1';
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`;
@@ -723,7 +726,7 @@ export async function ensurePostgresKvSchema(client: IPostgresClient): Promise<v
await client.transaction(async (db) => {
await db.query("SELECT set_config('statement_timeout', $1, true)", [POSTGRES_KV_SCHEMA_LOCK_TIMEOUT]);
await db.query('SELECT pg_advisory_xact_lock($1, hashtext($2))', [POSTGRES_KV_SCHEMA_LOCK_NAMESPACE, kvTable]);
await db.query("SELECT set_config('statement_timeout', '0', true)");
await db.query("SELECT set_config('statement_timeout', $1, true)", [POSTGRES_KV_SCHEMA_MIGRATION_TIMEOUT]);
await db.query(`
CREATE TABLE IF NOT EXISTS ${table} (
table_name text NOT NULL,
@@ -752,20 +755,32 @@ CREATE TABLE IF NOT EXISTS ${table} (
`CREATE INDEX IF NOT EXISTS ${quoteIdentifier(`${kvTable}_message_reactions_message_idx`)} ON ${table} (partition_key, ((CASE WHEN row_data -> 'message_id' ->> 'value' ~ '^-?[0-9]+$' THEN (row_data -> 'message_id' ->> 'value')::bigint END))) WHERE table_name = 'message_reactions'`,
);
await repairInvalidPostgresKvIndexes(db, kvTable);
const pending = await db.query(`
const migrated = await db.query(`SELECT 1 FROM ${table} WHERE table_name = $1 AND row_key = $2 LIMIT 1`, [
POSTGRES_KV_MIGRATION_TABLE,
POSTGRES_KV_MESSAGES_PARTITION_MIGRATION,
]);
if (migrated.rows.length === 0) {
const pending = await db.query(`
SELECT 1
FROM ${table}
WHERE table_name = 'messages'
AND partition_key = row_key
AND split_part(row_key, chr(31), 3) <> ''
LIMIT 1`);
if (pending.rows.length > 0) {
await db.query(`
if (pending.rows.length > 0) {
await db.query(`
UPDATE ${table}
SET partition_key = split_part(row_key, chr(31), 1) || chr(31) || split_part(row_key, chr(31), 2)
WHERE table_name = 'messages'
AND partition_key = row_key
AND split_part(row_key, chr(31), 3) <> ''`);
}
await db.query(
`INSERT INTO ${table} (table_name, partition_key, row_key, row_data)
VALUES ($1, $2, $2, jsonb_build_object('applied_at', now()))
ON CONFLICT (table_name, row_key) DO NOTHING`,
[POSTGRES_KV_MIGRATION_TABLE, POSTGRES_KV_MESSAGES_PARTITION_MIGRATION],
);
}
await db.query(`DROP INDEX IF EXISTS ${quoteIdentifier(`${kvTable}_partition_idx`)}`);
});
@@ -2,9 +2,17 @@
import cassandra from 'cassandra-driver';
import {describe, expect, it} from 'vitest';
import '../Tables';
import {getTableMetadata} from './CassandraMetaRegistry';
import {defineTable} from './CassandraTableDsl';
import type {CassandraParams, PreparedQuery} from './CassandraTypes';
import {buildCandidatePlan, type CandidatePlan, keyFromColumns, matchesWhere} from './PostgresKvQueryExecutor';
import {
buildCandidatePlan,
type CandidatePlan,
keyFromColumns,
matchesWhere,
POSTGRES_KV_MIGRATION_TABLE,
} from './PostgresKvQueryExecutor';
type Row = Record<string, unknown>;
@@ -488,3 +496,10 @@ describe('PostgresKvQueryExecutor key equality', () => {
expect(matchesWhere({a: '2020-01-01'}, eqWhere, {a: localDate} as CassandraParams)).toBe(false);
});
});
describe('PostgresKvQueryExecutor migration bookkeeping', () => {
it('reserves a table name that no logical table can claim', () => {
expect(POSTGRES_KV_MIGRATION_TABLE.startsWith('__')).toBe(true);
expect(getTableMetadata(POSTGRES_KV_MIGRATION_TABLE)).toBeUndefined();
});
});
@@ -12,7 +12,7 @@ import {
import {afterAll, beforeAll, describe, expect, it} from 'vitest';
import {GuildMembers, ReadStates, Users} from '../Tables';
import {LegacyPostgresKvQueryExecutor, legacyEnsurePostgresKvSchema} from './__testref__/LegacyPostgresKvQueryExecutor';
import {ensurePostgresKvSchema, PostgresKvQueryExecutor} from './PostgresKvQueryExecutor';
import {ensurePostgresKvSchema, POSTGRES_KV_MIGRATION_TABLE, PostgresKvQueryExecutor} from './PostgresKvQueryExecutor';
type Row = Record<string, unknown>;
@@ -307,11 +307,11 @@ suite('postgres kv upgrade safety', () => {
}
const diff = await raw.query<{n: string}>(`
SELECT count(*) AS n FROM (
(SELECT table_name, partition_key, row_key FROM ${OLD}
EXCEPT SELECT table_name, partition_key, row_key FROM ${NEW})
(SELECT table_name, partition_key, row_key FROM ${OLD} WHERE table_name <> '${POSTGRES_KV_MIGRATION_TABLE}'
EXCEPT SELECT table_name, partition_key, row_key FROM ${NEW} WHERE table_name <> '${POSTGRES_KV_MIGRATION_TABLE}')
UNION ALL
(SELECT table_name, partition_key, row_key FROM ${NEW}
EXCEPT SELECT table_name, partition_key, row_key FROM ${OLD})
(SELECT table_name, partition_key, row_key FROM ${NEW} WHERE table_name <> '${POSTGRES_KV_MIGRATION_TABLE}'
EXCEPT SELECT table_name, partition_key, row_key FROM ${OLD} WHERE table_name <> '${POSTGRES_KV_MIGRATION_TABLE}')
) d`);
const schemaDiff = await raw.query<{n: string}>(`
SELECT count(*) AS n FROM (
@@ -377,30 +377,48 @@ suite('postgres kv upgrade safety', () => {
expect(failures.length).toBe(0);
}, 120_000);
it('backfills the messages partition key once and then skips the update', async () => {
it('backfills the messages partition key once and never scans for it again', async () => {
const BACKFILL = 'kv_backfill';
const legacyKey = `"c"${String.fromCharCode(31)}"b"${String.fromCharCode(31)}"m"`;
const SEP = String.fromCharCode(31);
const backfillClient = new TableClient(raw, BACKFILL);
const legacyKey = (id: string) => `"c"${SEP}"b"${SEP}"${id}"`;
const insertLegacy = async (id: string) => {
await raw.query(
`INSERT INTO ${BACKFILL} (table_name, partition_key, row_key, row_data) VALUES ('messages', $1, $1, '{}'::jsonb)`,
[legacyKey(id)],
);
};
const partitionOf = async (id: string) => {
const result = await raw.query<{partition_key: string}>(
`SELECT partition_key FROM ${BACKFILL} WHERE table_name = 'messages' AND row_key = $1`,
[legacyKey(id)],
);
return result.rows[0]?.partition_key;
};
await raw.query(`DROP TABLE IF EXISTS ${BACKFILL}`);
await ensurePostgresKvSchema(backfillClient);
await raw.query(
`INSERT INTO ${BACKFILL} (table_name, partition_key, row_key, row_data) VALUES ('messages', $1, $1, '{}'::jsonb)`,
[legacyKey],
);
await legacyEnsurePostgresKvSchema(backfillClient);
await insertLegacy('m1');
const pendingBefore = await raw.query(
`SELECT 1 FROM ${BACKFILL} WHERE table_name = 'messages' AND partition_key = row_key AND split_part(row_key, chr(31), 3) <> ''`,
);
expect(pendingBefore.rows.length).toBe(1);
await ensurePostgresKvSchema(backfillClient);
const after = await raw.query<{partition_key: string}>(`SELECT partition_key FROM ${BACKFILL}`);
expect(after.rows[0]?.partition_key).toBe(`"c"${String.fromCharCode(31)}"b"`);
const marker = await raw.query(
`SELECT 1 FROM ${BACKFILL} WHERE table_name = $1 AND row_key = 'messages_partition_key_v1'`,
[POSTGRES_KV_MIGRATION_TABLE],
);
const pendingAfter = await raw.query(
`SELECT 1 FROM ${BACKFILL} WHERE table_name = 'messages' AND partition_key = row_key AND split_part(row_key, chr(31), 3) <> ''`,
);
expect(await partitionOf('m1')).toBe(`"c"${SEP}"b"`);
expect(pendingAfter.rows.length).toBe(0);
expect(marker.rows.length).toBe(1);
await insertLegacy('m2');
await ensurePostgresKvSchema(backfillClient);
const stable = await raw.query<{partition_key: string}>(`SELECT partition_key FROM ${BACKFILL}`);
expect(stable.rows[0]?.partition_key).toBe(`"c"${String.fromCharCode(31)}"b"`);
const skipped = await partitionOf('m2');
console.log(`second boot left the new legacy row alone: partition_key=${skipped === legacyKey('m2')}`);
expect(skipped).toBe(legacyKey('m2'));
expect(await partitionOf('m1')).toBe(`"c"${SEP}"b"`);
await raw.query(`DROP TABLE IF EXISTS ${BACKFILL}`);
}, 120_000);