@powersync/service-module-postgres-storage 0.7.5 → 0.8.1

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));
@@ -279,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();
@@ -305,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();
@@ -333,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 }},
@@ -371,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);
     }
 
@@ -379,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() }}
@@ -399,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
@@ -409,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
@@ -448,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,
@@ -657,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;
@@ -677,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}`
           );
         }
@@ -771,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}`
           );
         }
@@ -811,7 +871,7 @@ export class PostgresBucketBatch
               }
             }
           );
-          logger.error(
+          this.logger.error(
             `Failed to evaluate parameter query on ${record.sourceTable.qualifiedName}.${after.id}: ${error.error}`
           );
         }

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.
    *

package/test/src/migrations.test.ts

@@ -1,15 +1,23 @@
-import { describe, expect, it } from 'vitest';
+import { beforeEach, describe, expect, it } from 'vitest';
 
+import { Direction } from '@powersync/lib-services-framework';
 import { register } from '@powersync/service-core-tests';
 import { PostgresMigrationAgent } from '../../src/migrations/PostgresMigrationAgent.js';
 import { env } from './env.js';
-import { POSTGRES_STORAGE_FACTORY } from './util.js';
+import { POSTGRES_STORAGE_FACTORY, POSTGRES_STORAGE_SETUP } from './util.js';
 
 const MIGRATION_AGENT_FACTORY = () => {
   return new PostgresMigrationAgent({ type: 'postgresql', uri: env.PG_STORAGE_TEST_URL, sslmode: 'disable' });
 };
 
 describe('Migrations', () => {
+  beforeEach(async () => {
+    // The migration tests clear the migration store, without running the down migrations.
+    // This ensures all the down migrations have been run before.
+    const setup = POSTGRES_STORAGE_SETUP;
+    await setup.migrate(Direction.Down);
+  });
+
   register.registerMigrationTests(MIGRATION_AGENT_FACTORY);
 
   it('Should have tables declared', async () => {

package/test/src/util.ts

@@ -2,7 +2,10 @@ import path from 'path';
 import { fileURLToPath } from 'url';
 import { normalizePostgresStorageConfig } from '../../src//types/types.js';
 import { PostgresMigrationAgent } from '../../src/migrations/PostgresMigrationAgent.js';
-import { PostgresTestStorageFactoryGenerator } from '../../src/storage/PostgresTestStorageFactoryGenerator.js';
+import {
+  postgresTestSetup,
+  PostgresTestStorageFactoryGenerator
+} from '../../src/storage/PostgresTestStorageFactoryGenerator.js';
 import { env } from './env.js';
 
 const __filename = fileURLToPath(import.meta.url);
@@ -28,7 +31,9 @@ class TestPostgresMigrationAgent extends PostgresMigrationAgent {
   }
 }
 
-export const POSTGRES_STORAGE_FACTORY = PostgresTestStorageFactoryGenerator({
+export const POSTGRES_STORAGE_SETUP = postgresTestSetup({
   url: env.PG_STORAGE_TEST_URL,
   migrationAgent: (config) => new TestPostgresMigrationAgent(config)
 });
+
+export const POSTGRES_STORAGE_FACTORY = POSTGRES_STORAGE_SETUP.factory;