Skip to content

Commit 3cb3248

Browse files
committed
Override chunk size for testing.
1 parent 50ec4db commit 3cb3248

File tree

3 files changed

+68
-46
lines changed

3 files changed

+68
-46
lines changed

modules/module-postgres/src/replication/WalStream.ts

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,15 @@ export interface WalStreamOptions {
3030
connections: PgManager;
3131
storage: storage.SyncRulesBucketStorage;
3232
abort_signal: AbortSignal;
33+
34+
/**
35+
* Override snapshot chunk size, for testing.
36+
*
37+
* Defaults to 10_000.
38+
*
39+
* Note that queries are streamed, so we don't actually keep that much data in memory.
40+
*/
41+
snapshotChunkSize?: number;
3342
}
3443

3544
interface InitResult {
@@ -63,12 +72,15 @@ export class WalStream {
6372

6473
private startedStreaming = false;
6574

75+
private snapshotChunkSize: number;
76+
6677
constructor(options: WalStreamOptions) {
6778
this.storage = options.storage;
6879
this.sync_rules = options.storage.getParsedSyncRules({ defaultSchema: POSTGRES_DEFAULT_SCHEMA });
6980
this.group_id = options.storage.group_id;
7081
this.slot_name = options.storage.slot_name;
7182
this.connections = options.connections;
83+
this.snapshotChunkSize = options.snapshotChunkSize ?? 10_000;
7284

7385
this.abort_signal = options.abort_signal;
7486
this.abort_signal.addEventListener(
@@ -441,11 +453,11 @@ WHERE oid = $1::regclass`,
441453
// Single primary key - we can use the primary key for chunking
442454
const orderByKey = table.replicaIdColumns[0];
443455
logger.info(`Chunking ${table.qualifiedName} by ${orderByKey.name}`);
444-
q = new ChunkedSnapshotQuery(db, table, 1000);
456+
q = new ChunkedSnapshotQuery(db, table, this.snapshotChunkSize);
445457
} else {
446458
// Fallback case - query the entire table
447459
logger.info(`Snapshot ${table.qualifiedName} without chunking`);
448-
q = new SimpleSnapshotQuery(db, table, 10_000);
460+
q = new SimpleSnapshotQuery(db, table, this.snapshotChunkSize);
449461
}
450462
await q.initialize();
451463

modules/module-postgres/test/src/large_batch.test.ts

Lines changed: 48 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -369,66 +369,74 @@ function defineBatchTests(factory: StorageFactory) {
369369
expect(data.length).toEqual(11002 + deletedRowOps.length);
370370
}
371371

372-
test.only('chunked snapshot edge case', async () => {
372+
test('chunked snapshot edge case', async () => {
373373
// 1. Start with 10k rows, one row with id = 10000, and a large TOAST value in another column.
374374
// 2. Replicate one batch of rows (id < 10000).
375375
// 3. `UPDATE table SET id = 0 WHERE id = 10000`
376376
// 4. Replicate the rest of the table.
377377
// 5. Logical replication picks up the UPDATE above, but it is missing the TOAST column.
378378
// 6. We end up with a row that has a missing TOAST column.
379379

380-
try {
381-
await using context = await WalStreamTestContext.open(factory);
380+
await using context = await WalStreamTestContext.open(factory, {
381+
// We need to use a smaller chunk size here, so that we can run a query in between chunks
382+
walStreamOptions: { snapshotChunkSize: 100 }
383+
});
382384

383-
await context.updateSyncRules(`bucket_definitions:
385+
await context.updateSyncRules(`bucket_definitions:
384386
global:
385387
data:
386388
- SELECT * FROM test_data`);
387-
const { pool } = context;
389+
const { pool } = context;
388390

389-
await pool.query(`CREATE TABLE test_data(id integer primary key, description text)`);
391+
await pool.query(`CREATE TABLE test_data(id integer primary key, description text)`);
390392

391-
await pool.query({
392-
statement: `INSERT INTO test_data(id, description) SELECT i, 'foo' FROM generate_series(1, 10000) i`
393-
});
393+
// 1. Start with 10k rows, one row with id = 10000...
394+
await pool.query({
395+
statement: `INSERT INTO test_data(id, description) SELECT i, 'foo' FROM generate_series(1, 10000) i`
396+
});
394397

395-
// Toast value, must be > 8kb after compression
396-
const largeDescription = crypto.randomBytes(20_000).toString('hex');
397-
await pool.query({
398-
statement: 'UPDATE test_data SET description = $1 WHERE id = 10000',
399-
params: [{ type: 'varchar', value: largeDescription }]
400-
});
398+
// ...and a large TOAST value in another column.
399+
// Toast value, must be > 8kb after compression
400+
const largeDescription = crypto.randomBytes(20_000).toString('hex');
401+
await pool.query({
402+
statement: 'UPDATE test_data SET description = $1 WHERE id = 10000',
403+
params: [{ type: 'varchar', value: largeDescription }]
404+
});
401405

402-
const p = context.replicateSnapshot();
406+
// 2. Replicate one batch of rows (id < 10000).
407+
// Our "stopping point" here is not quite deterministic.
408+
const p = context.replicateSnapshot();
403409

404-
const stopAfter = 1_000;
405-
const startRowCount =
406-
(await Metrics.getInstance().getMetricValueForTests('powersync_rows_replicated_total')) ?? 0;
410+
const stopAfter = 1_000;
411+
const startRowCount = (await Metrics.getInstance().getMetricValueForTests('powersync_rows_replicated_total')) ?? 0;
407412

408-
while (true) {
409-
const count =
410-
((await Metrics.getInstance().getMetricValueForTests('powersync_rows_replicated_total')) ?? 0) -
411-
startRowCount;
413+
while (true) {
414+
const count =
415+
((await Metrics.getInstance().getMetricValueForTests('powersync_rows_replicated_total')) ?? 0) - startRowCount;
412416

413-
if (count >= stopAfter) {
414-
break;
415-
}
416-
await timers.setTimeout(1);
417+
if (count >= stopAfter) {
418+
break;
417419
}
418-
await pool.query('UPDATE test_data SET id = 0 WHERE id = 10000');
419-
await p;
420-
421-
context.startStreaming();
422-
const data = await context.getBucketData('global[]', undefined, {});
423-
const reduced = reduceBucket(data);
424-
expect(reduced.length).toEqual(10_001);
425-
426-
const movedRow = reduced.find((row) => row.object_id === '0');
427-
expect(movedRow?.data).toEqual(`{"id":0,"description":"${largeDescription}"}`);
428-
} catch (e) {
429-
console.error(e);
430-
throw e;
420+
await timers.setTimeout(1);
431421
}
422+
423+
// 3. `UPDATE table SET id = 0 WHERE id = 10000`
424+
await pool.query('UPDATE test_data SET id = 0 WHERE id = 10000');
425+
426+
// 4. Replicate the rest of the table.
427+
await p;
428+
429+
// 5. Logical replication picks up the UPDATE above, but it is missing the TOAST column.
430+
context.startStreaming();
431+
432+
// 6. If all went well, the "resnapshot" process would take care of this.
433+
const data = await context.getBucketData('global[]', undefined, {});
434+
const reduced = reduceBucket(data);
435+
436+
const movedRow = reduced.find((row) => row.object_id === '0');
437+
expect(movedRow?.data).toEqual(`{"id":0,"description":"${largeDescription}"}`);
438+
439+
expect(reduced.length).toEqual(10_001);
432440
});
433441

434442
function printMemoryUsage() {

modules/module-postgres/test/src/wal_stream_utils.ts

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ export class WalStreamTestContext implements AsyncDisposable {
2121
*/
2222
static async open(
2323
factory: (options: StorageOptions) => Promise<BucketStorageFactory>,
24-
options?: { doNotClear?: boolean }
24+
options?: { doNotClear?: boolean; walStreamOptions?: Partial<WalStreamOptions> }
2525
) {
2626
const f = await factory({ doNotClear: options?.doNotClear });
2727
const connectionManager = new PgManager(TEST_CONNECTION_OPTIONS, {});
@@ -30,12 +30,13 @@ export class WalStreamTestContext implements AsyncDisposable {
3030
await clearTestDb(connectionManager.pool);
3131
}
3232

33-
return new WalStreamTestContext(f, connectionManager);
33+
return new WalStreamTestContext(f, connectionManager, options?.walStreamOptions);
3434
}
3535

3636
constructor(
3737
public factory: BucketStorageFactory,
38-
public connectionManager: PgManager
38+
public connectionManager: PgManager,
39+
private walStreamOptions?: Partial<WalStreamOptions>
3940
) {}
4041

4142
async [Symbol.asyncDispose]() {
@@ -97,7 +98,8 @@ export class WalStreamTestContext implements AsyncDisposable {
9798
const options: WalStreamOptions = {
9899
storage: this.storage,
99100
connections: this.connectionManager,
100-
abort_signal: this.abortController.signal
101+
abort_signal: this.abortController.signal,
102+
...this.walStreamOptions
101103
};
102104
this._walStream = new WalStream(options);
103105
return this._walStream!;

0 commit comments

Comments
 (0)