WIP: Record replication progress for Postgres.

Author: rkistner

modules/module-mongodb-storage/src/storage/implementation/MongoBucketBatch.ts

@@ -859,6 +859,32 @@ export class MongoBucketBatch
     return last_op!;
   }
 
+  async updateTableProgress(
+    table: storage.SourceTable,
+    progress: Partial<storage.TableSnapshotStatus>
+  ): Promise<storage.SourceTable> {
+    const copy = table.clone();
+    copy.snapshotStatus = {
+      totalEstimatedCount: progress.totalEstimatedCount ?? copy.snapshotStatus?.totalEstimatedCount ?? 0,
+      replicatedCount: progress.replicatedCount ?? copy.snapshotStatus?.replicatedCount ?? 0,
+      lastKey: progress.lastKey ?? copy.snapshotStatus?.lastKey ?? null
+    };
+
+    await this.withTransaction(async () => {
+      await this.db.source_tables.updateOne(
+        { _id: table.id },
+        {
+          $set: {
+            snapshot_status: copy.snapshotStatus
+          }
+        },
+        { session: this.session }
+      );
+    });
+
+    return copy;
+  }
+
   async markSnapshotDone(tables: storage.SourceTable[], no_checkpoint_before_lsn: string) {
     const session = this.session;
     const ids = tables.map((table) => table.id);
@@ -869,6 +895,9 @@ export class MongoBucketBatch
         {
           $set: {
             snapshot_done: true
+          },
+          $unset: {
+            snapshot_status: 1
           }
         },
         { session }
@@ -892,17 +921,8 @@ export class MongoBucketBatch
       }
     });
     return tables.map((table) => {
-      const copy = new storage.SourceTable(
-        table.id,
-        table.connectionTag,
-        table.objectId,
-        table.schema,
-        table.table,
-        table.replicaIdColumns,
-        table.snapshotComplete
-      );
-      copy.syncData = table.syncData;
-      copy.syncParameters = table.syncParameters;
+      const copy = table.clone();
+      copy.snapshotComplete = true;
       return copy;
     });
   }

modules/module-mongodb-storage/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -194,7 +194,8 @@ export class MongoSyncBucketStorage
           table_name: table,
           replica_id_columns: null,
           replica_id_columns2: columns,
-          snapshot_done: false
+          snapshot_done: false,
+          snapshot_status: undefined
         };
 
         await col.insertOne(doc, { session });
@@ -211,6 +212,8 @@ export class MongoSyncBucketStorage
       sourceTable.syncEvent = options.sync_rules.tableTriggersEvent(sourceTable);
       sourceTable.syncData = options.sync_rules.tableSyncsData(sourceTable);
       sourceTable.syncParameters = options.sync_rules.tableSyncsParameters(sourceTable);
+      // TODO: standardize serialization format
+      sourceTable.snapshotStatus = doc.snapshot_status;
 
       let dropTables: storage.SourceTable[] = [];
       // Detect tables that are either renamed, or have different replica_id_columns

modules/module-mongodb-storage/src/storage/implementation/models.ts

@@ -1,4 +1,4 @@
-import { storage } from '@powersync/service-core';
+import { storage, TableSnapshotStatus } from '@powersync/service-core';
 import { SqliteJsonValue } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 
@@ -73,6 +73,7 @@ export interface SourceTableDocument {
   replica_id_columns: string[] | null;
   replica_id_columns2: { name: string; type_oid?: number; type?: string }[] | undefined;
   snapshot_done: boolean | undefined;
+  snapshot_status: TableSnapshotStatus | undefined;
 }
 
 /**

modules/module-mongodb/src/replication/ChangeStream.ts

@@ -8,7 +8,14 @@ import {
   ReplicationAssertionError,
   ServiceError
 } from '@powersync/lib-services-framework';
-import { MetricsEngine, SaveOperationTag, SourceEntityDescriptor, SourceTable, storage } from '@powersync/service-core';
+import {
+  MetricsEngine,
+  RelationCache,
+  SaveOperationTag,
+  SourceEntityDescriptor,
+  SourceTable,
+  storage
+} from '@powersync/service-core';
 import { DatabaseInputRow, SqliteRow, SqlSyncRules, TablePattern } from '@powersync/service-sync-rules';
 import { ReplicationMetric } from '@powersync/service-types';
 import { MongoLSN } from '../common/MongoLSN.js';
@@ -73,7 +80,7 @@ export class ChangeStream {
 
   private abort_signal: AbortSignal;
 
-  private relation_cache = new Map<string | number, storage.SourceTable>();
+  private relationCache = new RelationCache(getCacheIdentifier);
 
   private checkpointStreamId = new mongo.ObjectId();
 
@@ -178,11 +185,14 @@ export class ChangeStream {
   }
 
   async estimatedCount(table: storage.SourceTable): Promise<string> {
-    const db = this.client.db(table.schema);
-    const count = await db.collection(table.table).estimatedDocumentCount();
+    const count = await this.estimatedCountNumber(table);
     return `~${count}`;
   }
 
+  async estimatedCountNumber(table: storage.SourceTable): Promise<number> {
+    const db = this.client.db(table.schema);
+    return await db.collection(table.table).estimatedDocumentCount();
+  }
   /**
    * Start initial replication.
    *
@@ -228,7 +238,17 @@ export class ChangeStream {
           allSourceTables.push(...tables);
         }
 
+        let tablesWithStatus: SourceTable[] = [];
         for (let table of allSourceTables) {
+          let count = await this.estimatedCountNumber(table);
+          const updated = await batch.updateTableProgress(table, {
+            totalEstimatedCount: count
+          });
+          tablesWithStatus.push(updated);
+          this.relationCache.update(updated);
+        }
+
+        for (let table of tablesWithStatus) {
           await this.snapshotTable(batch, table);
           await batch.markSnapshotDone([table], MongoLSN.ZERO.comparable);
 
@@ -358,8 +378,7 @@ export class ChangeStream {
     descriptor: SourceEntityDescriptor,
     options: { snapshot: boolean }
   ): Promise<SourceTable> {
-    const cacheId = getCacheIdentifier(descriptor);
-    const existing = this.relation_cache.get(cacheId);
+    const existing = this.relationCache.get(descriptor);
     if (existing != null) {
       return existing;
     }
@@ -425,7 +444,7 @@ export class ChangeStream {
       entity_descriptor: descriptor,
       sync_rules: this.sync_rules
     });
-    this.relation_cache.set(getCacheIdentifier(descriptor), result.table);
+    this.relationCache.update(result.table);
 
     // Drop conflicting collections.
     // This is generally not expected for MongoDB source dbs, so we log an error.
@@ -800,7 +819,7 @@ export class ChangeStream {
             });
             if (table.syncAny) {
               await batch.drop([table]);
-              this.relation_cache.delete(getCacheIdentifier(rel));
+              this.relationCache.delete(table);
             }
           } else if (changeDocument.operationType == 'rename') {
             const relFrom = getMongoRelation(changeDocument.ns);
@@ -811,7 +830,7 @@ export class ChangeStream {
             });
             if (tableFrom.syncAny) {
               await batch.drop([tableFrom]);
-              this.relation_cache.delete(getCacheIdentifier(relFrom));
+              this.relationCache.delete(relFrom);
             }
             // Here we do need to snapshot the new table
             const collection = await this.getCollectionInfo(relTo.schema, relTo.name);

modules/module-mongodb/src/replication/MongoRelation.ts

@@ -20,7 +20,10 @@ export function getMongoRelation(source: mongo.ChangeStreamNameSpace): storage.S
 /**
  * For in-memory cache only.
  */
-export function getCacheIdentifier(source: storage.SourceEntityDescriptor): string {
+export function getCacheIdentifier(source: storage.SourceEntityDescriptor | storage.SourceTable): string {
+  if (source instanceof storage.SourceTable) {
+    return `${source.schema}.${source.table}`;
+  }
   return `${source.schema}.${source.name}`;
 }
 

modules/module-postgres-storage/src/storage/batch/PostgresBucketBatch.ts

@@ -448,6 +448,14 @@ export class PostgresBucketBatch
     });
   }
 
+  async updateTableProgress(
+    table: storage.SourceTable,
+    progress: Partial<storage.TableSnapshotStatus>
+  ): Promise<storage.SourceTable> {
+    // TODO: implement
+    return table;
+  }
+
   addCustomWriteCheckpoint(checkpoint: storage.BatchedCustomWriteCheckpointOptions): void {
     this.write_checkpoint_batch.push({
       ...checkpoint,

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

@@ -1,10 +1,16 @@
-import { ColumnDescriptor, SourceTable } from '@powersync/service-core';
+import { ColumnDescriptor, SourceTable, bson } from '@powersync/service-core';
 import { PgChunk, PgConnection, PgType, StatementParam, PgTypeOid } from '@powersync/service-jpgwire';
 import { escapeIdentifier } from '@powersync/lib-service-postgres';
 import { SqliteValue } from '@powersync/service-sync-rules';
+import { ServiceAssertionError } from '@powersync/lib-services-framework';
 
 export interface SnapshotQuery {
   initialize(): Promise<void>;
+  /**
+   * Returns an async iterable iterator that yields chunks of data.
+   *
+   * If the last chunk has 0 rows, it indicates that there are no more rows to fetch.
+   */
   nextChunk(): AsyncIterableIterator<PgChunk>;
 }
 
@@ -74,7 +80,7 @@ export class ChunkedSnapshotQuery implements SnapshotQuery {
   }
 
   private readonly key: ColumnDescriptor;
-  private lastKey: string | bigint | null = null;
+  lastKey: string | bigint | null = null;
 
   public constructor(
     private readonly connection: PgConnection,
@@ -88,6 +94,23 @@ export class ChunkedSnapshotQuery implements SnapshotQuery {
     // No-op
   }
 
+  public setLastKeySerialized(key: Uint8Array) {
+    const decoded = bson.deserialize(key, { useBigInt64: true });
+    const keys = Object.keys(decoded);
+    if (keys.length != 1) {
+      throw new ServiceAssertionError(`Multiple keys found: ${keys.join(', ')}`);
+    }
+    if (keys[0] != this.key.name) {
+      throw new ServiceAssertionError(`Key name mismatch: expected ${this.key.name}, got ${keys[0]}`);
+    }
+    const value = decoded[this.key.name];
+    this.lastKey = value;
+  }
+
+  public getLastKeySerialized(): Uint8Array {
+    return bson.serialize({ [this.key.name]: this.lastKey });
+  }
+
   public async *nextChunk(): AsyncIterableIterator<PgChunk> {
     let stream: AsyncIterableIterator<PgChunk>;
     if (this.lastKey == null) {