@powersync/service-module-postgres-storage 0.8.4 → 0.9.0

package/src/storage/PostgresSyncRulesStorage.ts

@@ -7,6 +7,8 @@ import {
   InternalOpId,
   internalToExternalOpId,
   LastValueSink,
+  maxLsn,
+  ReplicationCheckpoint,
   storage,
   utils,
   WatchWriteCheckpointOptions
@@ -135,18 +137,19 @@ export class PostgresSyncRulesStorage
       .decoded(pick(models.SyncRules, ['last_checkpoint', 'last_checkpoint_lsn']))
       .first();
 
-    return {
-      checkpoint: checkpointRow?.last_checkpoint ?? 0n,
-      lsn: checkpointRow?.last_checkpoint_lsn ?? null
-    };
+    return new PostgresReplicationCheckpoint(
+      this,
+      checkpointRow?.last_checkpoint ?? 0n,
+      checkpointRow?.last_checkpoint_lsn ?? null
+    );
   }
 
   async resolveTable(options: storage.ResolveTableOptions): Promise<storage.ResolveTableResult> {
     const { group_id, connection_id, connection_tag, entity_descriptor } = options;
 
-    const { schema, name: table, objectId, replicationColumns } = entity_descriptor;
+    const { schema, name: table, objectId, replicaIdColumns } = entity_descriptor;
 
-    const columns = replicationColumns.map((column) => ({
+    const normalizedReplicaIdColumns = replicaIdColumns.map((column) => ({
       name: column.name,
       type: column.type,
       // The PGWire returns this as a BigInt. We want to store this as JSONB
@@ -166,7 +169,7 @@ export class PostgresSyncRulesStorage
             AND relation_id = ${{ type: 'jsonb', value: { object_id: objectId } satisfies StoredRelationId }}
             AND schema_name = ${{ type: 'varchar', value: schema }}
             AND table_name = ${{ type: 'varchar', value: table }}
-            AND replica_id_columns = ${{ type: 'jsonb', value: columns }}
+            AND replica_id_columns = ${{ type: 'jsonb', value: normalizedReplicaIdColumns }}
         `
           .decoded(models.SourceTable)
           .first();
@@ -181,7 +184,7 @@ export class PostgresSyncRulesStorage
             AND connection_id = ${{ type: 'int4', value: connection_id }}
             AND schema_name = ${{ type: 'varchar', value: schema }}
             AND table_name = ${{ type: 'varchar', value: table }}
-            AND replica_id_columns = ${{ type: 'jsonb', value: columns }}
+            AND replica_id_columns = ${{ type: 'jsonb', value: normalizedReplicaIdColumns }}
         `
           .decoded(models.SourceTable)
           .first();
@@ -208,7 +211,7 @@ export class PostgresSyncRulesStorage
               ${{ type: 'jsonb', value: { object_id: objectId } satisfies StoredRelationId }},
               ${{ type: 'varchar', value: schema }},
               ${{ type: 'varchar', value: table }},
-              ${{ type: 'jsonb', value: columns }}
+              ${{ type: 'jsonb', value: normalizedReplicaIdColumns }}
             )
           RETURNING
             *
@@ -218,15 +221,15 @@ export class PostgresSyncRulesStorage
         sourceTableRow = row;
       }
 
-      const sourceTable = new storage.SourceTable(
-        sourceTableRow!.id,
-        connection_tag,
-        objectId,
-        schema,
-        table,
-        replicationColumns,
-        sourceTableRow!.snapshot_done ?? true
-      );
+      const sourceTable = new storage.SourceTable({
+        id: sourceTableRow!.id,
+        connectionTag: connection_tag,
+        objectId: objectId,
+        schema: schema,
+        name: table,
+        replicaIdColumns: replicaIdColumns,
+        snapshotComplete: sourceTableRow!.snapshot_done ?? true
+      });
       if (!sourceTable.snapshotComplete) {
         sourceTable.snapshotStatus = {
           totalEstimatedCount: Number(sourceTableRow!.snapshot_total_estimated_count ?? -1n),
@@ -284,19 +287,20 @@ export class PostgresSyncRulesStorage
         table: sourceTable,
         dropTables: truncatedTables.map(
           (doc) =>
-            new storage.SourceTable(
-              doc.id,
-              connection_tag,
-              doc.relation_id?.object_id ?? 0,
-              doc.schema_name,
-              doc.table_name,
-              doc.replica_id_columns?.map((c) => ({
-                name: c.name,
-                typeOid: c.typeId,
-                type: c.type
-              })) ?? [],
-              doc.snapshot_done ?? true
-            )
+            new storage.SourceTable({
+              id: doc.id,
+              connectionTag: connection_tag,
+              objectId: doc.relation_id?.object_id ?? 0,
+              schema: doc.schema_name,
+              name: doc.table_name,
+              replicaIdColumns:
+                doc.replica_id_columns?.map((c) => ({
+                  name: c.name,
+                  typeOid: c.typeId,
+                  type: c.type
+                })) ?? [],
+              snapshotComplete: doc.snapshot_done ?? true
+            })
         )
       };
     });
@@ -310,13 +314,14 @@ export class PostgresSyncRulesStorage
       SELECT
         last_checkpoint_lsn,
         no_checkpoint_before,
-        keepalive_op
+        keepalive_op,
+        snapshot_lsn
       FROM
         sync_rules
       WHERE
         id = ${{ type: 'int4', value: this.group_id }}
     `
-      .decoded(pick(models.SyncRules, ['last_checkpoint_lsn', 'no_checkpoint_before', 'keepalive_op']))
+      .decoded(pick(models.SyncRules, ['last_checkpoint_lsn', 'no_checkpoint_before', 'keepalive_op', 'snapshot_lsn']))
       .first();
 
     const checkpoint_lsn = syncRules?.last_checkpoint_lsn ?? null;
@@ -330,6 +335,7 @@ export class PostgresSyncRulesStorage
       last_checkpoint_lsn: checkpoint_lsn,
       keep_alive_op: syncRules?.keepalive_op,
       no_checkpoint_before_lsn: syncRules?.no_checkpoint_before ?? options.zeroLSN,
+      resumeFromLsn: maxLsn(syncRules?.snapshot_lsn, checkpoint_lsn),
       store_current_data: options.storeCurrentData,
       skip_existing_rows: options.skipExistingRows ?? false,
       batch_limits: this.options.batchLimits,
@@ -347,7 +353,7 @@ export class PostgresSyncRulesStorage
   }
 
   async getParameterSets(
-    checkpoint: utils.InternalOpId,
+    checkpoint: ReplicationCheckpoint,
     lookups: sync_rules.ParameterLookup[]
   ): Promise<sync_rules.SqliteJsonRow[]> {
     const rows = await this.db.sql`
@@ -370,7 +376,7 @@ export class PostgresSyncRulesStorage
         value: lookups.map((l) => storage.serializeLookupBuffer(l).toString('hex'))
       }}) AS FILTER
         )
-        AND id <= ${{ type: 'int8', value: checkpoint }}
+        AND id <= ${{ type: 'int8', value: checkpoint.checkpoint }}
       ORDER BY
         lookup,
         source_table,
@@ -644,43 +650,6 @@ export class PostgresSyncRulesStorage
     `.execute();
   }
 
-  async autoActivate(): Promise<void> {
-    await this.db.transaction(async (db) => {
-      const syncRulesRow = await db.sql`
-        SELECT
-          state
-        FROM
-          sync_rules
-        WHERE
-          id = ${{ type: 'int4', value: this.group_id }}
-      `
-        .decoded(pick(models.SyncRules, ['state']))
-        .first();
-
-      if (syncRulesRow && syncRulesRow.state == storage.SyncRuleState.PROCESSING) {
-        await db.sql`
-          UPDATE sync_rules
-          SET
-            state = ${{ type: 'varchar', value: storage.SyncRuleState.ACTIVE }}
-          WHERE
-            id = ${{ type: 'int4', value: this.group_id }}
-        `.execute();
-      }
-
-      await db.sql`
-        UPDATE sync_rules
-        SET
-          state = ${{ type: 'varchar', value: storage.SyncRuleState.STOP }}
-        WHERE
-          (
-            state = ${{ value: storage.SyncRuleState.ACTIVE, type: 'varchar' }}
-            OR state = ${{ value: storage.SyncRuleState.ERRORED, type: 'varchar' }}
-          )
-          AND id != ${{ type: 'int4', value: this.group_id }}
-      `.execute();
-    });
-  }
-
   private async getChecksumsInternal(batch: storage.FetchPartialBucketChecksum[]): Promise<storage.PartialChecksumMap> {
     if (batch.length == 0) {
       return new Map();
@@ -867,9 +836,18 @@ export class PostgresSyncRulesStorage
   }
 
   private makeActiveCheckpoint(row: models.ActiveCheckpointDecoded | null) {
-    return {
-      checkpoint: row?.last_checkpoint ?? 0n,
-      lsn: row?.last_checkpoint_lsn ?? null
-    } satisfies storage.ReplicationCheckpoint;
+    return new PostgresReplicationCheckpoint(this, row?.last_checkpoint ?? 0n, row?.last_checkpoint_lsn ?? null);
+  }
+}
+
+class PostgresReplicationCheckpoint implements storage.ReplicationCheckpoint {
+  constructor(
+    private storage: PostgresSyncRulesStorage,
+    public readonly checkpoint: utils.InternalOpId,
+    public readonly lsn: string | null
+  ) {}
+
+  getParameterSets(lookups: sync_rules.ParameterLookup[]): Promise<sync_rules.SqliteJsonRow[]> {
+    return this.storage.getParameterSets(this, lookups);
   }
 }

package/src/storage/batch/PostgresBucketBatch.ts

@@ -31,6 +31,7 @@ export interface PostgresBucketBatchOptions {
   no_checkpoint_before_lsn: string;
   store_current_data: boolean;
   keep_alive_op?: InternalOpId | null;
+  resumeFromLsn: string | null;
   /**
    * Set to true for initial replication.
    */
@@ -61,6 +62,8 @@ export class PostgresBucketBatch
 
   public last_flushed_op: InternalOpId | null = null;
 
+  public resumeFromLsn: string | null;
+
   protected db: lib_postgres.DatabaseClient;
   protected group_id: number;
   protected last_checkpoint_lsn: string | null;
@@ -73,6 +76,7 @@ export class PostgresBucketBatch
   protected batch: OperationBatch | null;
   private lastWaitingLogThrottled = 0;
   private markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
+  private needsActivation = true;
 
   constructor(protected options: PostgresBucketBatchOptions) {
     super();
@@ -81,6 +85,7 @@ export class PostgresBucketBatch
     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.resumeFromLsn = options.resumeFromLsn;
     this.write_checkpoint_batch = [];
     this.sync_rules = options.sync_rules;
     this.markRecordUnavailable = options.markRecordUnavailable;
@@ -321,6 +326,7 @@ export class PostgresBucketBatch
     // Don't create a checkpoint if there were no changes
     if (!createEmptyCheckpoints && this.persisted_op == null) {
       // Nothing to commit - return true
+      await this.autoActivate(lsn);
       return true;
     }
 
@@ -363,6 +369,7 @@ export class PostgresBucketBatch
       .decoded(StatefulCheckpoint)
       .first();
 
+    await this.autoActivate(lsn);
     await notifySyncRulesUpdate(this.db, doc!);
 
     this.persisted_op = null;
@@ -406,13 +413,14 @@ export class PostgresBucketBatch
       .decoded(StatefulCheckpoint)
       .first();
 
+    await this.autoActivate(lsn);
     await notifySyncRulesUpdate(this.db, updated!);
 
     this.last_checkpoint_lsn = lsn;
     return true;
   }
 
-  async setSnapshotLsn(lsn: string): Promise<void> {
+  async setResumeLsn(lsn: string): Promise<void> {
     await this.db.sql`
       UPDATE sync_rules
       SET
@@ -459,15 +467,15 @@ export class PostgresBucketBatch
       }
     });
     return tables.map((table) => {
-      const copy = new storage.SourceTable(
-        table.id,
-        table.connectionTag,
-        table.objectId,
-        table.schema,
-        table.table,
-        table.replicaIdColumns,
-        table.snapshotComplete
-      );
+      const copy = new storage.SourceTable({
+        id: table.id,
+        connectionTag: table.connectionTag,
+        objectId: table.objectId,
+        schema: table.schema,
+        name: table.name,
+        replicaIdColumns: table.replicaIdColumns,
+        snapshotComplete: table.snapshotComplete
+      });
       copy.syncData = table.syncData;
       copy.syncParameters = table.syncParameters;
       return copy;
@@ -916,6 +924,59 @@ export class PostgresBucketBatch
     return result;
   }
 
+  /**
+   * Switch from processing -> active if relevant.
+   *
+   * Called on new commits.
+   */
+  private async autoActivate(lsn: string): Promise<void> {
+    if (!this.needsActivation) {
+      // Already activated
+      return;
+    }
+
+    let didActivate = false;
+    await this.db.transaction(async (db) => {
+      const syncRulesRow = await db.sql`
+        SELECT
+          state
+        FROM
+          sync_rules
+        WHERE
+          id = ${{ type: 'int4', value: this.group_id }}
+      `
+        .decoded(pick(models.SyncRules, ['state']))
+        .first();
+
+      if (syncRulesRow && syncRulesRow.state == storage.SyncRuleState.PROCESSING) {
+        await db.sql`
+          UPDATE sync_rules
+          SET
+            state = ${{ type: 'varchar', value: storage.SyncRuleState.ACTIVE }}
+          WHERE
+            id = ${{ type: 'int4', value: this.group_id }}
+        `.execute();
+        didActivate = true;
+      }
+
+      await db.sql`
+        UPDATE sync_rules
+        SET
+          state = ${{ type: 'varchar', value: storage.SyncRuleState.STOP }}
+        WHERE
+          (
+            state = ${{ value: storage.SyncRuleState.ACTIVE, type: 'varchar' }}
+            OR state = ${{ value: storage.SyncRuleState.ERRORED, type: 'varchar' }}
+          )
+          AND id != ${{ type: 'int4', value: this.group_id }}
+      `.execute();
+    });
+    if (didActivate) {
+      this.logger.info(`Activated new sync rules at ${lsn}`);
+    }
+    this.needsActivation = false;
+  }
+
   /**
    * Gets relevant {@link SqlEventDescriptor}s for the given {@link SourceTable}
    * TODO maybe share this with an abstract class