1
1
import { container , errors , logger } from '@powersync/lib-services-framework' ;
2
- import { getUuidReplicaIdentityBson , Metrics , SourceEntityDescriptor , storage } from '@powersync/service-core' ;
2
+ import {
3
+ BucketStorageBatch ,
4
+ getUuidReplicaIdentityBson ,
5
+ Metrics ,
6
+ SaveUpdate ,
7
+ SourceEntityDescriptor ,
8
+ storage
9
+ } from '@powersync/service-core' ;
3
10
import * as pgwire from '@powersync/service-jpgwire' ;
4
11
import { DatabaseInputRow , SqliteRow , SqlSyncRules , TablePattern , toSyncRulesRow } from '@powersync/service-sync-rules' ;
5
12
import * as pg_utils from '../utils/pgwire_utils.js' ;
6
13
import { PgManager } from './PgManager.js' ;
7
14
import { getPgOutputRelation , getRelId } from './PgRelation.js' ;
8
15
import { checkSourceConfiguration , getReplicationIdentityColumns } from './replication-utils.js' ;
9
- import { ChunkedSnapshotQuery , SimpleSnapshotQuery , SnapshotQuery } from './SnapshotQuery.js' ;
16
+ import {
17
+ ChunkedSnapshotQuery ,
18
+ IdSnapshotQuery ,
19
+ MissingRow ,
20
+ PrimaryKeyValue ,
21
+ SimpleSnapshotQuery ,
22
+ SnapshotQuery
23
+ } from './SnapshotQuery.js' ;
10
24
11
25
export const ZERO_LSN = '00000000/00000000' ;
12
26
export const PUBLICATION_NAME = 'powersync' ;
@@ -359,19 +373,8 @@ WHERE oid = $1::regclass`,
359
373
logger . info ( `${ this . slot_name } Skipping ${ table . qualifiedName } - snapshot already done` ) ;
360
374
continue ;
361
375
}
362
- let tableLsnNotBefore : string ;
363
- await db . query ( 'BEGIN' ) ;
364
- try {
365
- await this . snapshotTable ( batch , db , table ) ;
366
-
367
- const rs = await db . query ( `select pg_current_wal_lsn() as lsn` ) ;
368
- tableLsnNotBefore = rs . rows [ 0 ] [ 0 ] ;
369
- } finally {
370
- // Read-only transaction, commit does not actually do anything.
371
- await db . query ( 'COMMIT' ) ;
372
- }
373
376
374
- await batch . markSnapshotDone ( [ table ] , tableLsnNotBefore ) ;
377
+ await this . snapshotTableInTx ( batch , db , table ) ;
375
378
await touch ( ) ;
376
379
}
377
380
}
@@ -391,7 +394,38 @@ WHERE oid = $1::regclass`,
391
394
}
392
395
}
393
396
394
- private async snapshotTable ( batch : storage . BucketStorageBatch , db : pgwire . PgConnection , table : storage . SourceTable ) {
397
+ private async snapshotTableInTx (
398
+ batch : storage . BucketStorageBatch ,
399
+ db : pgwire . PgConnection ,
400
+ table : storage . SourceTable ,
401
+ limited ?: PrimaryKeyValue [ ]
402
+ ) : Promise < storage . SourceTable > {
403
+ await db . query ( 'BEGIN' ) ;
404
+ try {
405
+ let tableLsnNotBefore : string ;
406
+ await this . snapshotTable ( batch , db , table , limited ) ;
407
+
408
+ // Get the current LSN.
409
+ // The data will only be consistent once incremental replication
410
+ // has passed that point.
411
+ // We have to get this LSN _after_ we have started the snapshot query.
412
+ const rs = await db . query ( `select pg_current_wal_lsn() as lsn` ) ;
413
+ tableLsnNotBefore = rs . rows [ 0 ] [ 0 ] ;
414
+ await db . query ( 'COMMIT' ) ;
415
+ const [ resultTable ] = await batch . markSnapshotDone ( [ table ] , tableLsnNotBefore ) ;
416
+ return resultTable ;
417
+ } catch ( e ) {
418
+ await db . query ( 'ROLLBACK' ) ;
419
+ throw e ;
420
+ }
421
+ }
422
+
423
+ private async snapshotTable (
424
+ batch : storage . BucketStorageBatch ,
425
+ db : pgwire . PgConnection ,
426
+ table : storage . SourceTable ,
427
+ limited ?: PrimaryKeyValue [ ]
428
+ ) {
395
429
logger . info ( `${ this . slot_name } Replicating ${ table . qualifiedName } ` ) ;
396
430
const estimatedCount = await this . estimatedCount ( db , table ) ;
397
431
let at = 0 ;
@@ -401,13 +435,16 @@ WHERE oid = $1::regclass`,
401
435
// We do streaming on two levels:
402
436
// 1. Coarse level: DELCARE CURSOR, FETCH 10000 at a time.
403
437
// 2. Fine level: Stream chunks from each fetch call.
404
- if ( ChunkedSnapshotQuery . supports ( table ) ) {
438
+ if ( limited ) {
439
+ q = new IdSnapshotQuery ( db , table , limited ) ;
440
+ } else if ( ChunkedSnapshotQuery . supports ( table ) ) {
405
441
// Single primary key - we can use the primary key for chunking
406
442
const orderByKey = table . replicaIdColumns [ 0 ] ;
407
443
logger . info ( `Chunking ${ table . qualifiedName } by ${ orderByKey . name } ` ) ;
408
- q = new ChunkedSnapshotQuery ( db , table , 10_000 ) ;
444
+ q = new ChunkedSnapshotQuery ( db , table , 1000 ) ;
409
445
} else {
410
446
// Fallback case - query the entire table
447
+ logger . info ( `Snapshot ${ table . qualifiedName } without chunking` ) ;
411
448
q = new SimpleSnapshotQuery ( db , table , 10_000 ) ;
412
449
}
413
450
await q . initialize ( ) ;
@@ -501,37 +538,52 @@ WHERE oid = $1::regclass`,
501
538
// Truncate this table, in case a previous snapshot was interrupted.
502
539
await batch . truncate ( [ result . table ] ) ;
503
540
504
- let lsn : string = ZERO_LSN ;
505
541
// Start the snapshot inside a transaction.
506
542
// We use a dedicated connection for this.
507
543
const db = await this . connections . snapshotConnection ( ) ;
508
544
try {
509
- await db . query ( 'BEGIN' ) ;
510
- try {
511
- await this . snapshotTable ( batch , db , result . table ) ;
512
-
513
- // Get the current LSN.
514
- // The data will only be consistent once incremental replication
515
- // has passed that point.
516
- // We have to get this LSN _after_ we have started the snapshot query.
517
- const rs = await db . query ( `select pg_current_wal_lsn() as lsn` ) ;
518
- lsn = rs . rows [ 0 ] [ 0 ] ;
519
-
520
- await db . query ( 'COMMIT' ) ;
521
- } catch ( e ) {
522
- await db . query ( 'ROLLBACK' ) ;
523
- throw e ;
524
- }
545
+ const table = await this . snapshotTableInTx ( batch , db , result . table ) ;
546
+ return table ;
525
547
} finally {
526
548
await db . end ( ) ;
527
549
}
528
- const [ table ] = await batch . markSnapshotDone ( [ result . table ] , lsn ) ;
529
- return table ;
530
550
}
531
551
532
552
return result . table ;
533
553
}
534
554
555
+ /**
556
+ * Process rows that have missing TOAST values.
557
+ *
558
+ * This can happen during edge cases in the chunked intial snapshot process.
559
+ *
560
+ * We handle this similar to an inline table snapshot, but limited to the specific
561
+ * set of rows.
562
+ */
563
+ private async resnapshot ( batch : BucketStorageBatch , rows : MissingRow [ ] ) {
564
+ const byTable = new Map < string | number , MissingRow [ ] > ( ) ;
565
+ for ( let row of rows ) {
566
+ if ( ! byTable . has ( row . table . objectId ) ) {
567
+ byTable . set ( row . table . objectId , [ ] ) ;
568
+ }
569
+ byTable . get ( row . table . objectId ) ! . push ( row ) ;
570
+ }
571
+ const db = await this . connections . snapshotConnection ( ) ;
572
+ try {
573
+ for ( let rows of byTable . values ( ) ) {
574
+ const table = rows [ 0 ] . table ;
575
+ await this . snapshotTableInTx (
576
+ batch ,
577
+ db ,
578
+ table ,
579
+ rows . map ( ( r ) => r . key )
580
+ ) ;
581
+ }
582
+ } finally {
583
+ await db . end ( ) ;
584
+ }
585
+ }
586
+
535
587
private getTable ( relationId : number ) : storage . SourceTable {
536
588
const table = this . relation_cache . get ( relationId ) ;
537
589
if ( table == null ) {
@@ -640,8 +692,38 @@ WHERE oid = $1::regclass`,
640
692
// Auto-activate as soon as initial replication is done
641
693
await this . storage . autoActivate ( ) ;
642
694
695
+ let resnapshot : { table : storage . SourceTable ; key : PrimaryKeyValue } [ ] = [ ] ;
696
+
697
+ const markRecordUnavailable = ( record : SaveUpdate ) => {
698
+ if ( ! IdSnapshotQuery . supports ( record . sourceTable ) ) {
699
+ // If it's not supported, it's also safe to ignore
700
+ return ;
701
+ }
702
+ let key : PrimaryKeyValue = { } ;
703
+ for ( let column of record . sourceTable . replicaIdColumns ) {
704
+ const name = column . name ;
705
+ const value = record . after [ name ] ;
706
+ if ( value == null ) {
707
+ // We don't expect this to actually happen.
708
+ // The key should always be present in the "after" record.
709
+ return ;
710
+ }
711
+ key [ name ] = value ;
712
+ }
713
+ resnapshot . push ( {
714
+ table : record . sourceTable ,
715
+ key : key
716
+ } ) ;
717
+ } ;
718
+
643
719
await this . storage . startBatch (
644
- { zeroLSN : ZERO_LSN , defaultSchema : POSTGRES_DEFAULT_SCHEMA , storeCurrentData : true , skipExistingRows : false } ,
720
+ {
721
+ zeroLSN : ZERO_LSN ,
722
+ defaultSchema : POSTGRES_DEFAULT_SCHEMA ,
723
+ storeCurrentData : true ,
724
+ skipExistingRows : false ,
725
+ markRecordUnavailable
726
+ } ,
645
727
async ( batch ) => {
646
728
// Replication never starts in the middle of a transaction
647
729
let inTx = false ;
@@ -665,6 +747,16 @@ WHERE oid = $1::regclass`,
665
747
} else if ( msg . tag == 'commit' ) {
666
748
Metrics . getInstance ( ) . transactions_replicated_total . add ( 1 ) ;
667
749
inTx = false ;
750
+ // flush() must be before the resnapshot check - that is
751
+ // typically what reports the resnapshot records.
752
+ await batch . flush ( ) ;
753
+ // This _must_ be checked after the flush(), and before
754
+ // commit() or ack(). We never persist the resnapshot list,
755
+ // so we have to process it before marking our progress.
756
+ if ( resnapshot . length > 0 ) {
757
+ await this . resnapshot ( batch , resnapshot ) ;
758
+ resnapshot = [ ] ;
759
+ }
668
760
await batch . commit ( msg . lsn ! ) ;
669
761
await this . ack ( msg . lsn ! , replicationStream ) ;
670
762
} else {
0 commit comments