@powersync/service-module-postgres-storage 0.0.0-dev-20250507154604 → 0.0.0-dev-20250617143654

package/src/storage/PostgresSyncRulesStorage.ts

@@ -107,12 +107,6 @@ export class PostgresSyncRulesStorage
     return new PostgresCompactor(this.db, this.group_id, options).compact();
   }
 
-  batchCreateCustomWriteCheckpoints(checkpoints: storage.BatchedCustomWriteCheckpointOptions[]): Promise<void> {
-    return this.writeCheckpointAPI.batchCreateCustomWriteCheckpoints(
-      checkpoints.map((c) => ({ ...c, sync_rules_id: this.group_id }))
-    );
-  }
-
   lastWriteCheckpoint(filters: storage.SyncStorageLastWriteCheckpointFilters): Promise<bigint | null> {
     return this.writeCheckpointAPI.lastWriteCheckpoint({
       ...filters,
@@ -233,6 +227,13 @@ export class PostgresSyncRulesStorage
         replicationColumns,
         sourceTableRow!.snapshot_done ?? true
       );
+      if (!sourceTable.snapshotComplete) {
+        sourceTable.snapshotStatus = {
+          totalEstimatedCount: Number(sourceTableRow!.snapshot_total_estimated_count ?? -1n),
+          replicatedCount: Number(sourceTableRow!.snapshot_replicated_count ?? 0n),
+          lastKey: sourceTableRow!.snapshot_last_key
+        };
+      }
       sourceTable.syncEvent = options.sync_rules.tableTriggersEvent(sourceTable);
       sourceTable.syncData = options.sync_rules.tableSyncsData(sourceTable);
       sourceTable.syncParameters = options.sync_rules.tableSyncsParameters(sourceTable);
@@ -321,6 +322,7 @@ export class PostgresSyncRulesStorage
     const checkpoint_lsn = syncRules?.last_checkpoint_lsn ?? null;
 
     const batch = new PostgresBucketBatch({
+      logger: options.logger ?? framework.logger,
       db: this.db,
       sync_rules: this.sync_rules.parsed(options).sync_rules,
       group_id: this.group_id,
@@ -330,13 +332,14 @@ export class PostgresSyncRulesStorage
       no_checkpoint_before_lsn: syncRules?.no_checkpoint_before ?? options.zeroLSN,
       store_current_data: options.storeCurrentData,
       skip_existing_rows: options.skipExistingRows ?? false,
-      batch_limits: this.options.batchLimits
+      batch_limits: this.options.batchLimits,
+      markRecordUnavailable: options.markRecordUnavailable
     });
     this.iterateListeners((cb) => cb.batchStarted?.(batch));
 
     await callback(batch);
     await batch.flush();
-    if (batch.last_flushed_op) {
+    if (batch.last_flushed_op != null) {
       return { flushed_op: batch.last_flushed_op };
     } else {
       return null;
@@ -580,6 +583,7 @@ export class PostgresSyncRulesStorage
     const syncRulesRow = await this.db.sql`
       SELECT
         snapshot_done,
+        snapshot_lsn,
         last_checkpoint_lsn,
         state
       FROM
@@ -587,7 +591,7 @@ export class PostgresSyncRulesStorage
       WHERE
         id = ${{ type: 'int4', value: this.group_id }}
     `
-      .decoded(pick(models.SyncRules, ['snapshot_done', 'last_checkpoint_lsn', 'state']))
+      .decoded(pick(models.SyncRules, ['snapshot_done', 'last_checkpoint_lsn', 'state', 'snapshot_lsn']))
       .first();
 
     if (syncRulesRow == null) {
@@ -597,7 +601,8 @@ export class PostgresSyncRulesStorage
     return {
       snapshot_done: syncRulesRow.snapshot_done,
       active: syncRulesRow.state == storage.SyncRuleState.ACTIVE,
-      checkpoint_lsn: syncRulesRow.last_checkpoint_lsn ?? null
+      checkpoint_lsn: syncRulesRow.last_checkpoint_lsn ?? null,
+      snapshot_lsn: syncRulesRow.snapshot_lsn ?? null
     };
   }
 

package/src/storage/PostgresTestStorageFactoryGenerator.ts

@@ -12,50 +12,62 @@ export type PostgresTestStorageOptions = {
   migrationAgent?: (config: PostgresStorageConfigDecoded) => PostgresMigrationAgent;
 };
 
-export const PostgresTestStorageFactoryGenerator = (factoryOptions: PostgresTestStorageOptions) => {
-  return async (options?: TestStorageOptions) => {
-    try {
-      const migrationManager: PowerSyncMigrationManager = new framework.MigrationManager();
-
-      const BASE_CONFIG = {
-        type: 'postgresql' as const,
-        uri: factoryOptions.url,
-        sslmode: 'disable' as const
-      };
-
-      const TEST_CONNECTION_OPTIONS = normalizePostgresStorageConfig(BASE_CONFIG);
-
-      await using migrationAgent = factoryOptions.migrationAgent
-        ? factoryOptions.migrationAgent(BASE_CONFIG)
-        : new PostgresMigrationAgent(BASE_CONFIG);
-      migrationManager.registerMigrationAgent(migrationAgent);
-
-      const mockServiceContext = { configuration: { storage: BASE_CONFIG } } as unknown as ServiceContext;
-
-      if (!options?.doNotClear) {
-        await migrationManager.migrate({
-          direction: framework.migrations.Direction.Down,
-          migrationContext: {
-            service_context: mockServiceContext
-          }
-        });
+export const postgresTestSetup = (factoryOptions: PostgresTestStorageOptions) => {
+  const BASE_CONFIG = {
+    type: 'postgresql' as const,
+    uri: factoryOptions.url,
+    sslmode: 'disable' as const
+  };
+
+  const TEST_CONNECTION_OPTIONS = normalizePostgresStorageConfig(BASE_CONFIG);
+
+  const migrate = async (direction: framework.migrations.Direction) => {
+    await using migrationManager: PowerSyncMigrationManager = new framework.MigrationManager();
+    await using migrationAgent = factoryOptions.migrationAgent
+      ? factoryOptions.migrationAgent(BASE_CONFIG)
+      : new PostgresMigrationAgent(BASE_CONFIG);
+    migrationManager.registerMigrationAgent(migrationAgent);
+
+    const mockServiceContext = { configuration: { storage: BASE_CONFIG } } as unknown as ServiceContext;
+
+    await migrationManager.migrate({
+      direction: framework.migrations.Direction.Down,
+      migrationContext: {
+        service_context: mockServiceContext
       }
+    });
 
+    if (direction == framework.migrations.Direction.Up) {
       await migrationManager.migrate({
         direction: framework.migrations.Direction.Up,
         migrationContext: {
           service_context: mockServiceContext
         }
       });
-
-      return new PostgresBucketStorageFactory({
-        config: TEST_CONNECTION_OPTIONS,
-        slot_name_prefix: 'test_'
-      });
-    } catch (ex) {
-      // Vitest does not display these errors nicely when using the `await using` syntx
-      console.error(ex, ex.cause);
-      throw ex;
     }
   };
+
+  return {
+    factory: async (options?: TestStorageOptions) => {
+      try {
+        if (!options?.doNotClear) {
+          await migrate(framework.migrations.Direction.Up);
+        }
+
+        return new PostgresBucketStorageFactory({
+          config: TEST_CONNECTION_OPTIONS,
+          slot_name_prefix: 'test_'
+        });
+      } catch (ex) {
+        // Vitest does not display these errors nicely when using the `await using` syntx
+        console.error(ex, ex.cause);
+        throw ex;
+      }
+    },
+    migrate
+  };
+};
+
+export const PostgresTestStorageFactoryGenerator = (factoryOptions: PostgresTestStorageOptions) => {
+  return postgresTestSetup(factoryOptions).factory;
 };

package/src/storage/batch/PostgresBucketBatch.ts

@@ -4,12 +4,12 @@ import {
   container,
   ErrorCode,
   errors,
-  logger,
+  Logger,
   ReplicationAssertionError,
   ServiceAssertionError,
   ServiceError
 } from '@powersync/lib-services-framework';
-import { InternalOpId, storage, utils } from '@powersync/service-core';
+import { BucketStorageMarkRecordUnavailable, InternalOpId, storage, utils } from '@powersync/service-core';
 import * as sync_rules from '@powersync/service-sync-rules';
 import * as timers from 'timers/promises';
 import * as t from 'ts-codec';
@@ -22,6 +22,7 @@ import { cacheKey, encodedCacheKey, OperationBatch, RecordOperation } from './Op
 import { PostgresPersistedBatch } from './PostgresPersistedBatch.js';
 
 export interface PostgresBucketBatchOptions {
+  logger: Logger;
   db: lib_postgres.DatabaseClient;
   sync_rules: sync_rules.SqlSyncRules;
   group_id: number;
@@ -35,6 +36,8 @@ export interface PostgresBucketBatchOptions {
    */
   skip_existing_rows: boolean;
   batch_limits: RequiredOperationBatchLimits;
+
+  markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
 }
 
 /**
@@ -54,6 +57,8 @@ export class PostgresBucketBatch
   extends BaseObserver<storage.BucketBatchStorageListener>
   implements storage.BucketStorageBatch
 {
+  private logger: Logger;
+
   public last_flushed_op: InternalOpId | null = null;
 
   protected db: lib_postgres.DatabaseClient;
@@ -67,15 +72,18 @@ export class PostgresBucketBatch
   protected readonly sync_rules: sync_rules.SqlSyncRules;
   protected batch: OperationBatch | null;
   private lastWaitingLogThrottled = 0;
+  private markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
 
   constructor(protected options: PostgresBucketBatchOptions) {
     super();
+    this.logger = options.logger;
     this.db = options.db;
     this.group_id = options.group_id;
     this.last_checkpoint_lsn = options.last_checkpoint_lsn;
     this.no_checkpoint_before_lsn = options.no_checkpoint_before_lsn;
     this.write_checkpoint_batch = [];
     this.sync_rules = options.sync_rules;
+    this.markRecordUnavailable = options.markRecordUnavailable;
     this.batch = null;
     this.persisted_op = null;
     if (options.keep_alive_op) {
@@ -115,7 +123,7 @@ export class PostgresBucketBatch
       return null;
     }
 
-    logger.debug(`Saving ${record.tag}:${record.before?.id}/${record.after?.id}`);
+    this.logger.debug(`Saving ${record.tag}:${record.before?.id}/${record.after?.id}`);
 
     this.batch ??= new OperationBatch(this.options.batch_limits);
     this.batch.push(new RecordOperation(record));
@@ -208,13 +216,7 @@ export class PostgresBucketBatch
       return null;
     }
 
-    const currentSequence = await this.db.sql`
-      SELECT
-        LAST_VALUE AS value
-      FROM
-        op_id_sequence;
-    `.first<{ value: bigint }>();
-    return currentSequence!.value;
+    return this.getLastOpIdSequence(this.db);
   }
 
   async drop(sourceTables: storage.SourceTable[]): Promise<storage.FlushedResult | null> {
@@ -262,13 +264,7 @@ export class PostgresBucketBatch
     const lastOp = await this.withReplicationTransaction(async (db) => {
       resumeBatch = await this.replicateBatch(db, batch);
 
-      const sequence = await db.sql`
-        SELECT
-          LAST_VALUE AS value
-        FROM
-          op_id_sequence;
-      `.first<{ value: bigint }>();
-      return sequence!.value;
+      return this.getLastOpIdSequence(db);
     });
 
     // null if done, set if we need another flush
@@ -291,13 +287,14 @@ export class PostgresBucketBatch
     if (this.last_checkpoint_lsn != null && lsn < this.last_checkpoint_lsn) {
       // When re-applying transactions, don't create a new checkpoint until
       // we are past the last transaction.
-      logger.info(`Re-applied transaction ${lsn} - skipping checkpoint`);
+      this.logger.info(`Re-applied transaction ${lsn} - skipping checkpoint`);
+      // Cannot create a checkpoint yet - return false
       return false;
     }
 
     if (lsn < this.no_checkpoint_before_lsn) {
       if (Date.now() - this.lastWaitingLogThrottled > 5_000) {
-        logger.info(
+        this.logger.info(
           `Waiting until ${this.no_checkpoint_before_lsn} before creating checkpoint, currently at ${lsn}. Persisted op: ${this.persisted_op}`
         );
         this.lastWaitingLogThrottled = Date.now();
@@ -317,12 +314,14 @@ export class PostgresBucketBatch
           id = ${{ type: 'int4', value: this.group_id }}
       `.execute();
 
+      // Cannot create a checkpoint yet - return false
       return false;
     }
 
     // Don't create a checkpoint if there were no changes
     if (!createEmptyCheckpoints && this.persisted_op == null) {
-      return false;
+      // Nothing to commit - return true
+      return true;
     }
 
     const now = new Date().toISOString();
@@ -345,6 +344,7 @@ export class PostgresBucketBatch
         keepalive_op = ${{ type: 'int8', value: update.keepalive_op }},
         last_fatal_error = ${{ type: 'varchar', value: update.last_fatal_error }},
         snapshot_done = ${{ type: 'bool', value: update.snapshot_done }},
+        snapshot_lsn = NULL,
         last_keepalive_ts = ${{ type: 1184, value: update.last_keepalive_ts }},
         last_checkpoint = COALESCE(
           ${{ type: 'int8', value: update.last_checkpoint }},
@@ -383,7 +383,7 @@ export class PostgresBucketBatch
     if (this.persisted_op != null) {
       // The commit may have been skipped due to "no_checkpoint_before_lsn".
       // Apply it now if relevant
-      logger.info(`Commit due to keepalive at ${lsn} / ${this.persisted_op}`);
+      this.logger.info(`Commit due to keepalive at ${lsn} / ${this.persisted_op}`);
       return await this.commit(lsn);
     }
 
@@ -391,6 +391,7 @@ export class PostgresBucketBatch
       UPDATE sync_rules
       SET
         snapshot_done = ${{ type: 'bool', value: true }},
+        snapshot_lsn = NULL,
         last_checkpoint_lsn = ${{ type: 'varchar', value: lsn }},
         last_fatal_error = ${{ type: 'varchar', value: null }},
         last_keepalive_ts = ${{ type: 1184, value: new Date().toISOString() }}
@@ -411,6 +412,16 @@ export class PostgresBucketBatch
     return true;
   }
 
+  async setSnapshotLsn(lsn: string): Promise<void> {
+    await this.db.sql`
+      UPDATE sync_rules
+      SET
+        snapshot_lsn = ${{ type: 'varchar', value: lsn }}
+      WHERE
+        id = ${{ type: 'int4', value: this.group_id }}
+    `.execute();
+  }
+
   async markSnapshotDone(
     tables: storage.SourceTable[],
     no_checkpoint_before_lsn: string
@@ -421,7 +432,10 @@ export class PostgresBucketBatch
       await db.sql`
         UPDATE source_tables
         SET
-          snapshot_done = ${{ type: 'bool', value: true }}
+          snapshot_done = ${{ type: 'bool', value: true }},
+          snapshot_total_estimated_count = NULL,
+          snapshot_replicated_count = NULL,
+          snapshot_last_key = NULL
         WHERE
           id IN (
             SELECT
@@ -460,6 +474,31 @@ export class PostgresBucketBatch
     });
   }
 
+  async updateTableProgress(
+    table: storage.SourceTable,
+    progress: Partial<storage.TableSnapshotStatus>
+  ): Promise<storage.SourceTable> {
+    const copy = table.clone();
+    const snapshotStatus = {
+      totalEstimatedCount: progress.totalEstimatedCount ?? copy.snapshotStatus?.totalEstimatedCount ?? 0,
+      replicatedCount: progress.replicatedCount ?? copy.snapshotStatus?.replicatedCount ?? 0,
+      lastKey: progress.lastKey ?? copy.snapshotStatus?.lastKey ?? null
+    };
+    copy.snapshotStatus = snapshotStatus;
+
+    await this.db.sql`
+      UPDATE source_tables
+      SET
+        snapshot_total_estimated_count = ${{ type: 'int4', value: snapshotStatus.totalEstimatedCount }},
+        snapshot_replicated_count = ${{ type: 'int4', value: snapshotStatus.replicatedCount }},
+        snapshot_last_key = ${{ type: 'bytea', value: snapshotStatus.lastKey }}
+      WHERE
+        id = ${{ type: 'varchar', value: table.id }}
+    `.execute();
+
+    return copy;
+  }
+
   addCustomWriteCheckpoint(checkpoint: storage.BatchedCustomWriteCheckpointOptions): void {
     this.write_checkpoint_batch.push({
       ...checkpoint,
@@ -669,10 +708,19 @@ export class PostgresBucketBatch
         existingBuckets = [];
         existingLookups = [];
         // Log to help with debugging if there was a consistency issue
+
         if (this.options.store_current_data) {
-          logger.warn(
-            `Cannot find previous record for update on ${record.sourceTable.qualifiedName}: ${beforeId} / ${record.before?.id}`
-          );
+          if (this.markRecordUnavailable != null) {
+            // This will trigger a "resnapshot" of the record.
+            // This is not relevant if storeCurrentData is false, since we'll get the full row
+            // directly in the replication stream.
+            this.markRecordUnavailable(record);
+          } else {
+            // Log to help with debugging if there was a consistency issue
+            this.logger.warn(
+              `Cannot find previous record for update on ${record.sourceTable.qualifiedName}: ${beforeId} / ${record.before?.id}`
+            );
+          }
         }
       } else {
         existingBuckets = result.buckets;
@@ -689,8 +737,8 @@ export class PostgresBucketBatch
         existingBuckets = [];
         existingLookups = [];
         // Log to help with debugging if there was a consistency issue
-        if (this.options.store_current_data) {
-          logger.warn(
+        if (this.options.store_current_data && this.markRecordUnavailable == null) {
+          this.logger.warn(
             `Cannot find previous record for delete on ${record.sourceTable.qualifiedName}: ${beforeId} / ${record.before?.id}`
           );
         }
@@ -783,7 +831,7 @@ export class PostgresBucketBatch
               }
             }
           );
-          logger.error(
+          this.logger.error(
             `Failed to evaluate data query on ${record.sourceTable.qualifiedName}.${record.after?.id}: ${error.error}`
           );
         }
@@ -823,7 +871,7 @@ export class PostgresBucketBatch
               }
             }
           );
-          logger.error(
+          this.logger.error(
             `Failed to evaluate parameter query on ${record.sourceTable.qualifiedName}.${after.id}: ${error.error}`
           );
         }
@@ -873,7 +921,7 @@ export class PostgresBucketBatch
    * TODO maybe share this with an abstract class
    */
   protected getTableEvents(table: storage.SourceTable): sync_rules.SqlEventDescriptor[] {
-    return this.sync_rules.event_descriptors.filter((evt) =>
+    return this.sync_rules.eventDescriptors.filter((evt) =>
       [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table))
     );
   }
@@ -895,6 +943,23 @@ export class PostgresBucketBatch
       `.execute();
     }
   }
+
+  private async getLastOpIdSequence(db: lib_postgres.AbstractPostgresConnection) {
+    // When no op_id has been generated, last_value = 1 and nextval() will be 1.
+    // To cater for this case, we check is_called, and default to 0 if no value has been generated.
+    const sequence = await db.sql`
+      SELECT
+        (
+          CASE
+            WHEN is_called THEN last_value
+            ELSE 0
+          END
+        ) AS value
+      FROM
+        op_id_sequence;
+    `.first<{ value: bigint }>();
+    return sequence!.value;
+  }
 }
 
 /**

package/src/storage/batch/PostgresPersistedBatch.ts

@@ -152,7 +152,7 @@ export class PostgresPersistedBatch {
       const base64 = binLookup.toString('base64');
       remaining_lookups.delete(base64);
       const hexLookup = binLookup.toString('hex');
-      const serializedBucketParameters = JSONBig.stringify(result.bucket_parameters);
+      const serializedBucketParameters = JSONBig.stringify(result.bucketParameters);
       this.parameterDataInserts.push({
         group_id: this.group_id,
         source_table: table.id,

package/src/storage/checkpoints/PostgresWriteCheckpointAPI.ts

@@ -1,7 +1,6 @@
 import * as lib_postgres from '@powersync/lib-service-postgres';
 import * as framework from '@powersync/lib-services-framework';
-import { storage, sync } from '@powersync/service-core';
-import { JSONBig, JsonContainer } from '@powersync/service-jsonbig';
+import { InternalOpId, storage } from '@powersync/service-core';
 import { models } from '../../types/types.js';
 
 export type PostgresCheckpointAPIOptions = {
@@ -26,7 +25,10 @@ export class PostgresWriteCheckpointAPI implements storage.WriteCheckpointAPI {
     this._mode = mode;
   }
 
-  async batchCreateCustomWriteCheckpoints(checkpoints: storage.CustomWriteCheckpointOptions[]): Promise<void> {
+  async batchCreateCustomWriteCheckpoints(
+    checkpoints: storage.CustomWriteCheckpointOptions[],
+    op_id: InternalOpId
+  ): Promise<void> {
     return batchCreateCustomWriteCheckpoints(this.db, checkpoints);
   }
 
@@ -58,13 +60,6 @@ export class PostgresWriteCheckpointAPI implements storage.WriteCheckpointAPI {
     return row!.write_checkpoint;
   }
 
-  watchUserWriteCheckpoint(
-    options: storage.WatchUserWriteCheckpointOptions
-  ): AsyncIterable<storage.WriteCheckpointResult> {
-    // Not used for Postgres currently
-    throw new Error('Method not implemented.');
-  }
-
   async lastWriteCheckpoint(filters: storage.LastWriteCheckpointFilters): Promise<bigint | null> {
     switch (this.writeCheckpointMode) {
       case storage.WriteCheckpointMode.CUSTOM:

package/src/storage/sync-rules/PostgresPersistedSyncRulesContent.ts

@@ -14,6 +14,7 @@ export class PostgresPersistedSyncRulesContent implements storage.PersistedSyncR
   public readonly last_fatal_error: string | null;
   public readonly last_keepalive_ts: Date | null;
   public readonly last_checkpoint_ts: Date | null;
+  public readonly active: boolean;
   current_lock: storage.ReplicationLock | null = null;
 
   constructor(
@@ -27,6 +28,7 @@ export class PostgresPersistedSyncRulesContent implements storage.PersistedSyncR
     this.last_fatal_error = row.last_fatal_error;
     this.last_checkpoint_ts = row.last_checkpoint_ts ? new Date(row.last_checkpoint_ts) : null;
     this.last_keepalive_ts = row.last_keepalive_ts ? new Date(row.last_keepalive_ts) : null;
+    this.active = row.state == 'ACTIVE';
   }
 
   parsed(options: storage.ParseSyncRulesOptions): storage.PersistedSyncRules {

package/src/types/models/SourceTable.ts

@@ -1,5 +1,5 @@
 import * as t from 'ts-codec';
-import { bigint, jsonb, jsonb_raw, pgwire_number } from '../codecs.js';
+import { bigint, hexBuffer, jsonb, jsonb_raw, pgwire_number } from '../codecs.js';
 
 export type StoredRelationId = {
   object_id: string | number | undefined;
@@ -25,7 +25,10 @@ export const SourceTable = t.object({
   schema_name: t.string,
   table_name: t.string,
   replica_id_columns: t.Null.or(jsonb(t.array(ColumnDescriptor))),
-  snapshot_done: t.boolean
+  snapshot_done: t.boolean,
+  snapshot_total_estimated_count: t.Null.or(bigint),
+  snapshot_replicated_count: t.Null.or(bigint),
+  snapshot_last_key: t.Null.or(hexBuffer)
 });
 
 export type SourceTable = t.Encoded<typeof SourceTable>;

package/src/types/models/SyncRules.ts

@@ -11,6 +11,10 @@ export const SyncRules = t.object({
    * Can only be false if state == PROCESSING.
    */
   snapshot_done: t.boolean,
+  /**
+   * May be set if snapshot_done = false, if the replication stream requires it.
+   */
+  snapshot_lsn: t.Null.or(t.string),
   /**
    * The last consistent checkpoint.
    *