@powersync/service-module-postgres 0.0.0-dev-20241209071534 → 0.0.0-dev-20241219091224

package/src/replication/WalStream.ts

@@ -192,6 +192,15 @@ export class WalStream {
     if (slotExists) {
       // This checks that the slot is still valid
       const r = await this.checkReplicationSlot();
+      if (snapshotDone && r.needsNewSlot) {
+        // We keep the current snapshot, and create a new replication slot
+        throw new MissingReplicationSlotError(`Replication slot ${slotName} is not valid anymore`);
+      }
+      // We can have:
+      //   needsInitialSync: true, needsNewSlot: true -> initial sync from scratch
+      //   needsInitialSync: true, needsNewSlot: false -> resume initial sync
+      //   needsInitialSync: false, needsNewSlot: true -> handled above
+      //   needsInitialSync: false, needsNewSlot: false -> resume streaming replication
       return {
         needsInitialSync: !snapshotDone,
         needsNewSlot: r.needsNewSlot
@@ -204,7 +213,7 @@ export class WalStream {
   /**
    * If a replication slot exists, check that it is healthy.
    */
-  private async checkReplicationSlot(): Promise<InitResult> {
+  private async checkReplicationSlot(): Promise<{ needsNewSlot: boolean }> {
     let last_error = null;
     const slotName = this.slot_name;
 
@@ -244,7 +253,7 @@ export class WalStream {
 
         // Success
         logger.info(`Slot ${slotName} appears healthy`);
-        return { needsInitialSync: false, needsNewSlot: false };
+        return { needsNewSlot: false };
       } catch (e) {
         last_error = e;
         logger.warn(`${slotName} Replication slot error`, e);
@@ -274,9 +283,9 @@ export class WalStream {
           // Sample: publication "powersync" does not exist
           //   Happens when publication deleted or never created.
           //   Slot must be re-created in this case.
-          logger.info(`${slotName} does not exist anymore, will create new slot`);
+          logger.info(`${slotName} is not valid anymore`);
 
-          return { needsInitialSync: true, needsNewSlot: true };
+          return { needsNewSlot: true };
         }
         // Try again after a pause
         await new Promise((resolve) => setTimeout(resolve, 1000));
@@ -349,18 +358,10 @@ WHERE  oid = $1::regclass`,
               logger.info(`${this.slot_name} Skipping ${table.qualifiedName} - snapshot already done`);
               continue;
             }
-            let tableLsnNotBefore: string;
-            await db.query('BEGIN');
-            try {
-              await this.snapshotTable(batch, db, table);
-
-              const rs = await db.query(`select pg_current_wal_lsn() as lsn`);
-              tableLsnNotBefore = rs.rows[0][0];
-            } finally {
-              // Read-only transaction, commit does not actually do anything.
-              await db.query('COMMIT');
-            }
+            await this.snapshotTable(batch, db, table);
 
+            const rs = await db.query(`select pg_current_wal_lsn() as lsn`);
+            const tableLsnNotBefore = rs.rows[0][0];
             await batch.markSnapshotDone([table], tableLsnNotBefore);
             await touch();
           }
@@ -386,70 +387,51 @@ WHERE  oid = $1::regclass`,
     const estimatedCount = await this.estimatedCount(db, table);
     let at = 0;
     let lastLogIndex = 0;
-
-    // We do streaming on two levels:
-    // 1. Coarse level: DELCARE CURSOR, FETCH 10000 at a time.
-    // 2. Fine level: Stream chunks from each fetch call.
-    await db.query(`DECLARE powersync_cursor CURSOR FOR SELECT * FROM ${table.escapedIdentifier}`);
-
+    const cursor = db.stream({ statement: `SELECT * FROM ${table.escapedIdentifier}` });
     let columns: { i: number; name: string }[] = [];
-    let hasRemainingData = true;
-    while (hasRemainingData) {
-      // Fetch 10k at a time.
-      // The balance here is between latency overhead per FETCH call,
-      // and not spending too much time on each FETCH call.
-      // We aim for a couple of seconds on each FETCH call.
-      const cursor = db.stream({
-        statement: `FETCH 10000 FROM powersync_cursor`
-      });
-      hasRemainingData = false;
-      // pgwire streams rows in chunks.
-      // These chunks can be quite small (as little as 16KB), so we don't flush chunks automatically.
-      // There are typically 100-200 rows per chunk.
-      for await (let chunk of cursor) {
-        if (chunk.tag == 'RowDescription') {
-          // We get a RowDescription for each FETCH call, but they should
-          // all be the same.
-          let i = 0;
-          columns = chunk.payload.map((c) => {
-            return { i: i++, name: c.name };
-          });
-          continue;
-        }
-
-        const rows = chunk.rows.map((row) => {
-          let q: DatabaseInputRow = {};
-          for (let c of columns) {
-            q[c.name] = row[c.i];
-          }
-          return q;
+    // pgwire streams rows in chunks.
+    // These chunks can be quite small (as little as 16KB), so we don't flush chunks automatically.
+
+    for await (let chunk of cursor) {
+      if (chunk.tag == 'RowDescription') {
+        let i = 0;
+        columns = chunk.payload.map((c) => {
+          return { i: i++, name: c.name };
         });
-        if (rows.length > 0 && at - lastLogIndex >= 5000) {
-          logger.info(`${this.slot_name} Replicating ${table.qualifiedName} ${at}/${estimatedCount}`);
-          lastLogIndex = at;
-          hasRemainingData = true;
-        }
-        if (this.abort_signal.aborted) {
-          throw new Error(`Aborted initial replication of ${this.slot_name}`);
-        }
+        continue;
+      }
 
-        for (const record of WalStream.getQueryData(rows)) {
-          // This auto-flushes when the batch reaches its size limit
-          await batch.save({
-            tag: storage.SaveOperationTag.INSERT,
-            sourceTable: table,
-            before: undefined,
-            beforeReplicaId: undefined,
-            after: record,
-            afterReplicaId: getUuidReplicaIdentityBson(record, table.replicaIdColumns)
-          });
+      const rows = chunk.rows.map((row) => {
+        let q: DatabaseInputRow = {};
+        for (let c of columns) {
+          q[c.name] = row[c.i];
         }
+        return q;
+      });
+      if (rows.length > 0 && at - lastLogIndex >= 5000) {
+        logger.info(`${this.slot_name} Replicating ${table.qualifiedName} ${at}/${estimatedCount}`);
+        lastLogIndex = at;
+      }
+      if (this.abort_signal.aborted) {
+        throw new Error(`Aborted initial replication of ${this.slot_name}`);
+      }
 
-        at += rows.length;
-        Metrics.getInstance().rows_replicated_total.add(rows.length);
-
-        await touch();
+      for (const record of WalStream.getQueryData(rows)) {
+        // This auto-flushes when the batch reaches its size limit
+        await batch.save({
+          tag: storage.SaveOperationTag.INSERT,
+          sourceTable: table,
+          before: undefined,
+          beforeReplicaId: undefined,
+          after: record,
+          afterReplicaId: getUuidReplicaIdentityBson(record, table.replicaIdColumns)
+        });
       }
+
+      at += rows.length;
+      Metrics.getInstance().rows_replicated_total.add(rows.length);
+
+      await touch();
     }
 
     await batch.flush();

package/src/types/types.ts

@@ -133,6 +133,12 @@ export function normalizeConnectionConfig(options: PostgresConnectionConfig): No
   };
 }
 
+export function isPostgresConfig(
+  config: service_types.configFile.DataSourceConfig
+): config is PostgresConnectionConfig {
+  return config.type == POSTGRES_CONNECTION_TYPE;
+}
+
 /**
  * Check whether the port is in a "safe" range.
  *

package/test/src/env.ts

@@ -2,6 +2,7 @@ import { utils } from '@powersync/lib-services-framework';
 
 export const env = utils.collectEnvironmentVariables({
   PG_TEST_URL: utils.type.string.default('postgres://postgres:postgres@localhost:5432/powersync_test'),
+  MONGO_TEST_URL: utils.type.string.default('mongodb://localhost:27017/powersync_test'),
   CI: utils.type.boolean.default('false'),
   SLOW_TESTS: utils.type.boolean.default('false')
 });

package/test/src/large_batch.test.ts

@@ -1,17 +1,17 @@
-import { MONGO_STORAGE_FACTORY, StorageFactory, StorageOptions } from '@core-tests/util.js';
+import { Metrics } from '@powersync/service-core';
+import { test_utils } from '@powersync/service-core-tests';
+import * as timers from 'timers/promises';
 import { describe, expect, test } from 'vitest';
 import { populateData } from '../../dist/utils/populate_test_data.js';
 import { env } from './env.js';
-import { TEST_CONNECTION_OPTIONS } from './util.js';
+import { INITIALIZED_MONGO_STORAGE_FACTORY, TEST_CONNECTION_OPTIONS } from './util.js';
 import { WalStreamTestContext } from './wal_stream_utils.js';
-import * as timers from 'timers/promises';
-import { Metrics } from '@powersync/service-core';
 
 describe('batch replication tests - mongodb', { timeout: 120_000 }, function () {
   // These are slow but consistent tests.
   // Not run on every test run, but we do run on CI, or when manually debugging issues.
   if (env.CI || env.SLOW_TESTS) {
-    defineBatchTests(MONGO_STORAGE_FACTORY);
+    defineBatchTests(INITIALIZED_MONGO_STORAGE_FACTORY);
   } else {
     // Need something in this file.
     test('no-op', () => {});
@@ -23,7 +23,7 @@ const BASIC_SYNC_RULES = `bucket_definitions:
     data:
       - SELECT id, description, other FROM "test_data"`;
 
-function defineBatchTests(factory: StorageFactory) {
+function defineBatchTests(factory: test_utils.StorageFactory) {
   test('update large record', async () => {
     await using context = await WalStreamTestContext.open(factory);
     // This test generates a large transaction in MongoDB, despite the replicated data
@@ -176,6 +176,71 @@ function defineBatchTests(factory: StorageFactory) {
     console.log(`Truncated ${truncateCount} ops in ${truncateDuration}ms ${truncatePerSecond} ops/s. ${used}MB heap`);
   });
 
+  test('large number of bucket_data docs', async () => {
+    // This tests that we don't run into this error:
+    //   MongoBulkWriteError: BSONObj size: 16814023 (0x1008FC7) is invalid. Size must be between 0 and 16793600(16MB) First element: insert: "bucket_data"
+    // The test is quite sensitive to internals, since we need to
+    // generate an internal batch that is just below 16MB.
+    //
+    // For the test to work, we need a:
+    // 1. Large number of documents in the batch.
+    // 2. More bucket_data documents than current_data documents,
+    //    otherwise other batch limiting thresholds are hit.
+    // 3. A large document to make sure we get to just below the 16MB
+    //    limit.
+    // 4. Another document to make sure the internal batching overflows
+    //    to a second batch.
+
+    await using context = await WalStreamTestContext.open(factory);
+    await context.updateSyncRules(`bucket_definitions:
+  global:
+    data:
+      # Sync 4x so we get more bucket_data documents
+      - SELECT * FROM test_data
+      - SELECT * FROM test_data
+      - SELECT * FROM test_data
+      - SELECT * FROM test_data
+      `);
+    const { pool } = context;
+
+    await pool.query(`CREATE TABLE test_data(id serial primary key, description text)`);
+
+    const numDocs = 499;
+    let description = '';
+    while (description.length < 2650) {
+      description += '.';
+    }
+
+    await pool.query({
+      statement: `INSERT INTO test_data(description) SELECT $2 FROM generate_series(1, $1) i`,
+      params: [
+        { type: 'int4', value: numDocs },
+        { type: 'varchar', value: description }
+      ]
+    });
+
+    let largeDescription = '';
+
+    while (largeDescription.length < 2_768_000) {
+      largeDescription += '.';
+    }
+    await pool.query({
+      statement: 'INSERT INTO test_data(description) VALUES($1)',
+      params: [{ type: 'varchar', value: largeDescription }]
+    });
+    await pool.query({
+      statement: 'INSERT INTO test_data(description) VALUES($1)',
+      params: [{ type: 'varchar', value: 'testingthis' }]
+    });
+    await context.replicateSnapshot();
+
+    context.startStreaming();
+
+    const checkpoint = await context.getCheckpoint({ timeout: 50_000 });
+    const checksum = await context.storage!.getChecksums(checkpoint, ['global[]']);
+    expect(checksum.get('global[]')!.count).toEqual((numDocs + 2) * 4);
+  });
+
   test('resuming initial replication (1)', async () => {
     // Stop early - likely to not include deleted row in first replication attempt.
     await testResumingReplication(2000);

package/test/src/schema_changes.test.ts

@@ -1,10 +1,8 @@
-import { compareIds, putOp, removeOp } from '@core-tests/stream_utils.js';
-import { reduceBucket } from '@powersync/service-core';
-import { setTimeout } from 'node:timers/promises';
+import { compareIds, putOp, reduceBucket, removeOp, test_utils } from '@powersync/service-core-tests';
+import * as timers from 'timers/promises';
 import { describe, expect, test } from 'vitest';
 import { INITIALIZED_MONGO_STORAGE_FACTORY, StorageFactory } from './util.js';
 import { WalStreamTestContext } from './wal_stream_utils.js';
-
 describe('schema changes', { timeout: 20_000 }, function () {
   defineTests(INITIALIZED_MONGO_STORAGE_FACTORY);
 });
@@ -16,12 +14,12 @@ bucket_definitions:
       - SELECT id, * FROM "test_data"
 `;
 
-const PUT_T1 = putOp('test_data', { id: 't1', description: 'test1' });
-const PUT_T2 = putOp('test_data', { id: 't2', description: 'test2' });
-const PUT_T3 = putOp('test_data', { id: 't3', description: 'test3' });
+const PUT_T1 = test_utils.putOp('test_data', { id: 't1', description: 'test1' });
+const PUT_T2 = test_utils.putOp('test_data', { id: 't2', description: 'test2' });
+const PUT_T3 = test_utils.putOp('test_data', { id: 't3', description: 'test3' });
 
-const REMOVE_T1 = removeOp('test_data', 't1');
-const REMOVE_T2 = removeOp('test_data', 't2');
+const REMOVE_T1 = test_utils.removeOp('test_data', 't1');
+const REMOVE_T2 = test_utils.removeOp('test_data', 't2');
 
 function defineTests(factory: StorageFactory) {
   test('re-create table', async () => {
@@ -544,7 +542,7 @@ function defineTests(factory: StorageFactory) {
     );
 
     // Need some delay for the snapshot to be triggered
-    await setTimeout(5);
+    await timers.setTimeout(5);
 
     let stop = false;
 

package/test/src/slow_tests.test.ts

@@ -2,15 +2,14 @@ import * as bson from 'bson';
 import { afterEach, describe, expect, test } from 'vitest';
 import { WalStream, WalStreamOptions } from '../../src/replication/WalStream.js';
 import { env } from './env.js';
-import { clearTestDb, connectPgPool, getClientCheckpoint, TEST_CONNECTION_OPTIONS } from './util.js';
+import { clearTestDb, connectPgPool, getClientCheckpoint, INITIALIZED_MONGO_STORAGE_FACTORY, TEST_CONNECTION_OPTIONS } from './util.js';
 
 import * as pgwire from '@powersync/service-jpgwire';
 import { SqliteRow } from '@powersync/service-sync-rules';
 
-import { mapOpEntry, MongoBucketStorage } from '@/storage/storage-index.js';
-import { validateCompactedBucket } from '@core-tests/bucket_validation.js';
-import { MONGO_STORAGE_FACTORY, StorageFactory } from '@core-tests/util.js';
 import { PgManager } from '@module/replication/PgManager.js';
+import { test_utils } from '@powersync/service-core-tests';
+import * as mongo_module from '@powersync/service-module-mongodb';
 import * as timers from 'node:timers/promises';
 import { reduceBucket } from '@powersync/service-core';
 
@@ -18,14 +17,14 @@ describe('slow tests - mongodb', function () {
   // These are slow, inconsistent tests.
   // Not run on every test run, but we do run on CI, or when manually debugging issues.
   if (env.CI || env.SLOW_TESTS) {
-    defineSlowTests(MONGO_STORAGE_FACTORY);
+    defineSlowTests(INITIALIZED_MONGO_STORAGE_FACTORY);
   } else {
     // Need something in this file.
     test('no-op', () => {});
   }
 });
 
-function defineSlowTests(factory: StorageFactory) {
+function defineSlowTests(factory: test_utils.StorageFactory) {
   let walStream: WalStream | undefined;
   let connections: PgManager | undefined;
   let abortController: AbortController | undefined;
@@ -74,7 +73,7 @@ function defineSlowTests(factory: StorageFactory) {
     const replicationConnection = await connections.replicationConnection();
     const pool = connections.pool;
     await clearTestDb(pool);
-    const f = (await factory()) as MongoBucketStorage;
+    const f = (await factory()) as mongo_module.storage.MongoBucketStorage;
 
     const syncRuleContent = `
 bucket_definitions:
@@ -171,13 +170,13 @@ bucket_definitions:
           const checkpoint = BigInt((await storage.getCheckpoint()).checkpoint);
           const opsBefore = (await f.db.bucket_data.find().sort({ _id: 1 }).toArray())
             .filter((row) => row._id.o <= checkpoint)
-            .map(mapOpEntry);
+            .map(mongo_module.storage.mapOpEntry);
           await storage.compact({ maxOpId: checkpoint });
           const opsAfter = (await f.db.bucket_data.find().sort({ _id: 1 }).toArray())
             .filter((row) => row._id.o <= checkpoint)
-            .map(mapOpEntry);
+            .map(mongo_module.storage.mapOpEntry);
 
-          validateCompactedBucket(opsBefore, opsAfter);
+          test_utils.validateCompactedBucket(opsBefore, opsAfter);
         }
       };
 
@@ -202,8 +201,8 @@ bucket_definitions:
       const ops = await f.db.bucket_data.find().sort({ _id: 1 }).toArray();
 
       // All a single bucket in this test
-      const bucket = ops.map((op) => mapOpEntry(op));
-      const reduced = reduceBucket(bucket);
+      const bucket = ops.map((op) => mongo_module.storage.mapOpEntry(op));
+      const reduced = test_utils.reduceBucket(bucket);
       expect(reduced).toMatchObject([
         {
           op_id: '0',

package/test/src/util.ts

@@ -1,12 +1,12 @@
-import { connectMongo } from '@core-tests/util.js';
+import { PostgresRouteAPIAdapter } from '@module/api/PostgresRouteAPIAdapter.js';
 import * as types from '@module/types/types.js';
 import * as pg_utils from '@module/utils/pgwire_utils.js';
 import { logger } from '@powersync/lib-services-framework';
-import { BucketStorageFactory, Metrics, MongoBucketStorage, OpId } from '@powersync/service-core';
+import { BucketStorageFactory, Metrics, OpId } from '@powersync/service-core';
+import { test_utils } from '@powersync/service-core-tests';
 import * as pgwire from '@powersync/service-jpgwire';
-import { pgwireRows } from '@powersync/service-jpgwire';
+import * as mongo_module from '@powersync/service-module-mongodb';
 import { env } from './env.js';
-import { PostgresRouteAPIAdapter } from '@module/api/PostgresRouteAPIAdapter.js';
 
 // The metrics need to be initialized before they can be used
 await Metrics.initialise({
@@ -26,7 +26,7 @@ export const TEST_CONNECTION_OPTIONS = types.normalizeConnectionConfig({
 
 export type StorageFactory = () => Promise<BucketStorageFactory>;
 
-export const INITIALIZED_MONGO_STORAGE_FACTORY: StorageFactory = async () => {
+export const INITIALIZED_MONGO_STORAGE_FACTORY: StorageFactory = async (options?: test_utils.StorageOptions) => {
   const db = await connectMongo();
 
   // None of the PG tests insert data into this collection, so it was never created
@@ -34,13 +34,26 @@ export const INITIALIZED_MONGO_STORAGE_FACTORY: StorageFactory = async () => {
     await db.db.createCollection('bucket_parameters');
   }
 
-  await db.clear();
+  if (!options?.doNotClear) {
+    await db.clear();
+  }
 
-  return new MongoBucketStorage(db, {
+  return new mongo_module.storage.MongoBucketStorage(db, {
     slot_name_prefix: 'test_'
   });
 };
 
+export async function connectMongo() {
+  // Short timeout for tests, to fail fast when the server is not available.
+  // Slightly longer timeouts for CI, to avoid arbitrary test failures
+  const client = mongo_module.storage.createMongoClient(env.MONGO_TEST_URL, {
+    connectTimeoutMS: env.CI ? 15_000 : 5_000,
+    socketTimeoutMS: env.CI ? 15_000 : 5_000,
+    serverSelectionTimeoutMS: env.CI ? 15_000 : 2_500
+  });
+  return new mongo_module.storage.PowerSyncMongo(client);
+}
+
 export async function clearTestDb(db: pgwire.PgClient) {
   await db.query(
     "select pg_drop_replication_slot(slot_name) from pg_replication_slots where active = false and slot_name like 'test_%'"

package/test/src/validation.test.ts

@@ -1,10 +1,12 @@
-import { MONGO_STORAGE_FACTORY } from '@core-tests/util.js';
 import { getDebugTablesInfo } from '@module/replication/replication-utils.js';
 import { expect, test } from 'vitest';
+
+// Not quite a walStreamTest, but it helps to manage the connection
+import { INITIALIZED_MONGO_STORAGE_FACTORY } from './util.js';
 import { WalStreamTestContext } from './wal_stream_utils.js';
 
 test('validate tables', async () => {
-  await using context = await WalStreamTestContext.open(MONGO_STORAGE_FACTORY);
+  await using context = await WalStreamTestContext.open(INITIALIZED_MONGO_STORAGE_FACTORY);
   const { pool } = context;
 
   await pool.query(`CREATE TABLE test_data(id uuid primary key default uuid_generate_v4(), description text)`);

package/test/src/wal_stream.test.ts

@@ -1,10 +1,11 @@
-import { putOp, removeOp } from '@core-tests/stream_utils.js';
-import { MONGO_STORAGE_FACTORY } from '@core-tests/util.js';
 import { BucketStorageFactory, Metrics } from '@powersync/service-core';
+import { putOp, removeOp } from '@powersync/service-core-tests';
 import { pgwireRows } from '@powersync/service-jpgwire';
 import * as crypto from 'crypto';
 import { describe, expect, test } from 'vitest';
+import { INITIALIZED_MONGO_STORAGE_FACTORY } from './util.js';
 import { WalStreamTestContext } from './wal_stream_utils.js';
+import { MissingReplicationSlotError } from '@module/replication/WalStream.js';
 
 type StorageFactory = () => Promise<BucketStorageFactory>;
 
@@ -16,7 +17,7 @@ bucket_definitions:
 `;
 
 describe('wal stream - mongodb', { timeout: 20_000 }, function () {
-  defineWalStreamTests(MONGO_STORAGE_FACTORY);
+  defineWalStreamTests(INITIALIZED_MONGO_STORAGE_FACTORY);
 });
 
 function defineWalStreamTests(factory: StorageFactory) {
@@ -291,4 +292,52 @@ bucket_definitions:
     expect(endRowCount - startRowCount).toEqual(0);
     expect(endTxCount - startTxCount).toEqual(1);
   });
+
+  test('reporting slot issues', async () => {
+    {
+      await using context = await WalStreamTestContext.open(factory);
+      const { pool } = context;
+      await context.updateSyncRules(`
+bucket_definitions:
+  global:
+    data:
+      - SELECT id, description FROM "test_data"`);
+
+      await pool.query(
+        `CREATE TABLE test_data(id uuid primary key default uuid_generate_v4(), description text, num int8)`
+      );
+      await pool.query(
+        `INSERT INTO test_data(id, description) VALUES('8133cd37-903b-4937-a022-7c8294015a3a', 'test1') returning id as test_id`
+      );
+      await context.replicateSnapshot();
+      await context.startStreaming();
+
+      const data = await context.getBucketData('global[]');
+
+      expect(data).toMatchObject([
+        putOp('test_data', {
+          id: '8133cd37-903b-4937-a022-7c8294015a3a',
+          description: 'test1'
+        })
+      ]);
+
+      expect(await context.storage!.getStatus()).toMatchObject({ active: true, snapshot_done: true });
+    }
+
+    {
+      await using context = await WalStreamTestContext.open(factory, { doNotClear: true });
+      const { pool } = context;
+      await pool.query('DROP PUBLICATION powersync');
+      await pool.query(`UPDATE test_data SET description = 'updated'`);
+      await pool.query('CREATE PUBLICATION powersync FOR ALL TABLES');
+
+      await context.loadActiveSyncRules();
+      await expect(async () => {
+        await context.replicateSnapshot();
+      }).rejects.toThrowError(MissingReplicationSlotError);
+
+      // The error is handled on a higher level, which triggers
+      // creating a new replication slot.
+    }
+  });
 }

package/test/src/wal_stream_utils.ts

@@ -1,10 +1,9 @@
-import { fromAsync } from '@core-tests/stream_utils.js';
 import { PgManager } from '@module/replication/PgManager.js';
 import { PUBLICATION_NAME, WalStream, WalStreamOptions } from '@module/replication/WalStream.js';
 import { BucketStorageFactory, OplogEntry, SyncRulesBucketStorage } from '@powersync/service-core';
+import { StorageOptions, test_utils } from '@powersync/service-core-tests';
 import * as pgwire from '@powersync/service-jpgwire';
 import { clearTestDb, getClientCheckpoint, TEST_CONNECTION_OPTIONS } from './util.js';
-import { StorageOptions } from '@core-tests/util.js';
 
 export class WalStreamTestContext implements AsyncDisposable {
   private _walStream?: WalStream;
@@ -77,6 +76,16 @@ export class WalStreamTestContext implements AsyncDisposable {
     return this.storage!;
   }
 
+  async loadActiveSyncRules() {
+    const syncRules = await this.factory.getActiveSyncRulesContent();
+    if (syncRules == null) {
+      throw new Error(`Active sync rules not available`);
+    }
+
+    this.storage = this.factory.getInstance(syncRules);
+    return this.storage!;
+  }
+
   get walStream() {
     if (this.storage == null) {
       throw new Error('updateSyncRules() first');
@@ -122,7 +131,7 @@ export class WalStreamTestContext implements AsyncDisposable {
   async getBucketsDataBatch(buckets: Record<string, string>, options?: { timeout?: number }) {
     let checkpoint = await this.getCheckpoint(options);
     const map = new Map<string, string>(Object.entries(buckets));
-    return fromAsync(this.storage!.getBucketDataBatch(checkpoint, map));
+    return test_utils.fromAsync(this.storage!.getBucketDataBatch(checkpoint, map));
   }
 
   /**
@@ -136,7 +145,7 @@ export class WalStreamTestContext implements AsyncDisposable {
     while (true) {
       const batch = this.storage!.getBucketDataBatch(checkpoint, map);
 
-      const batches = await fromAsync(batch);
+      const batches = await test_utils.fromAsync(batch);
       data = data.concat(batches[0]?.batch.data ?? []);
       if (batches.length == 0 || !batches[0]!.batch.has_more) {
         break;
@@ -154,7 +163,7 @@ export class WalStreamTestContext implements AsyncDisposable {
     const { checkpoint } = await this.storage!.getCheckpoint();
     const map = new Map<string, string>([[bucket, start]]);
     const batch = this.storage!.getBucketDataBatch(checkpoint, map);
-    const batches = await fromAsync(batch);
+    const batches = await test_utils.fromAsync(batch);
     return batches[0]?.batch.data ?? [];
   }
 }