@powersync/service-module-mssql 0.10.3 → 0.10.4

package/test/src/CDCStream.test.ts

@@ -1,11 +1,19 @@
+import { MSSQLConnectionManager } from '@module/replication/MSSQLConnectionManager.js';
 import { getLatestLSN } from '@module/utils/mssql.js';
+import { ReplicationAbortedError } from '@powersync/lib-services-framework';
 import { storage } from '@powersync/service-core';
 import { METRICS_HELPER, putOp, removeOp } from '@powersync/service-core-tests';
 import { ReplicationMetric } from '@powersync/service-types';
 import sql from 'mssql';
-import { describe, expect, test } from 'vitest';
+import { describe, expect, test, vi } from 'vitest';
 import { CDCStreamTestContext } from './CDCStreamTestContext.js';
-import { createTestTable, describeWithStorage, insertTestData, waitForPendingCDCChanges } from './util.js';
+import {
+  createTestTable,
+  describeWithStorage,
+  enableCDCForTable,
+  insertTestData,
+  waitForPendingCDCChanges
+} from './util.js';
 
 const BASIC_SYNC_RULES = `
 bucket_definitions:
@@ -14,6 +22,23 @@ bucket_definitions:
       - SELECT id, description FROM "test_data"
 `;
 
+const COMPOSITE_KEY_SYNC_CONFIG = `
+config:
+  edition: 3
+streams:
+  global:
+    auto_subscribe: true
+    query: SELECT id_user AS id, id_user, id_visit FROM test_composite
+`;
+
+async function getRowsReplicated() {
+  return (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.ROWS_REPLICATED)) ?? 0;
+}
+
+async function getTransactionsReplicated() {
+  return (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.TRANSACTIONS_REPLICATED)) ?? 0;
+}
+
 describe('CDCStream tests', () => {
   describeWithStorage({ timeout: 20_000 }, defineCDCStreamTests);
 });
@@ -30,12 +55,12 @@ function defineCDCStreamTests(config: storage.TestStorageConfig) {
     const beforeLSN = await getLatestLSN(connectionManager);
     const testData = await insertTestData(connectionManager, 'test_data');
     await waitForPendingCDCChanges(beforeLSN, connectionManager);
-    const startRowCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.ROWS_REPLICATED)) ?? 0;
+    const startRowCount = await getRowsReplicated();
 
     await context.replicateSnapshot();
     await context.startStreaming();
 
-    const endRowCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.ROWS_REPLICATED)) ?? 0;
+    const endRowCount = await getRowsReplicated();
     const data = await context.getBucketData('global[]');
     expect(data).toMatchObject([putOp('test_data', testData)]);
     expect(endRowCount - startRowCount).toEqual(1);
@@ -49,8 +74,8 @@ function defineCDCStreamTests(config: storage.TestStorageConfig) {
     await createTestTable(connectionManager, 'test_data');
     await context.replicateSnapshot();
 
-    const startRowCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.ROWS_REPLICATED)) ?? 0;
-    const startTxCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.TRANSACTIONS_REPLICATED)) ?? 0;
+    const startRowCount = await getRowsReplicated();
+    const startTxCount = await getTransactionsReplicated();
 
     await context.startStreaming();
 
@@ -59,8 +84,8 @@ function defineCDCStreamTests(config: storage.TestStorageConfig) {
     const data = await context.getBucketData('global[]');
 
     expect(data).toMatchObject([putOp('test_data', testData)]);
-    const endRowCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.ROWS_REPLICATED)) ?? 0;
-    const endTxCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.TRANSACTIONS_REPLICATED)) ?? 0;
+    const endRowCount = await getRowsReplicated();
+    const endTxCount = await getTransactionsReplicated();
     expect(endRowCount - startRowCount).toEqual(1);
     expect(endTxCount - startTxCount).toEqual(1);
   });
@@ -156,8 +181,8 @@ function defineCDCStreamTests(config: storage.TestStorageConfig) {
 
     await context.replicateSnapshot();
 
-    const startRowCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.ROWS_REPLICATED)) ?? 0;
-    const startTxCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.TRANSACTIONS_REPLICATED)) ?? 0;
+    const startRowCount = await getRowsReplicated();
+    const startTxCount = await getTransactionsReplicated();
 
     await context.startStreaming();
 
@@ -165,8 +190,8 @@ function defineCDCStreamTests(config: storage.TestStorageConfig) {
     const data = await context.getBucketData('global[]');
 
     expect(data).toMatchObject([]);
-    const endRowCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.ROWS_REPLICATED)) ?? 0;
-    const endTxCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.TRANSACTIONS_REPLICATED)) ?? 0;
+    const endRowCount = await getRowsReplicated();
+    const endTxCount = await getTransactionsReplicated();
 
     // There was a transaction, but it is not counted since it is not for a table in the sync config
     expect(endRowCount - startRowCount).toEqual(0);
@@ -187,8 +212,8 @@ function defineCDCStreamTests(config: storage.TestStorageConfig) {
 
     await context.replicateSnapshot();
 
-    const startRowCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.ROWS_REPLICATED)) ?? 0;
-    const startTxCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.TRANSACTIONS_REPLICATED)) ?? 0;
+    const startRowCount = await getRowsReplicated();
+    const startTxCount = await getTransactionsReplicated();
 
     await context.startStreaming();
 
@@ -196,9 +221,74 @@ function defineCDCStreamTests(config: storage.TestStorageConfig) {
     const data = await context.getBucketData('global[]');
 
     expect(data).toMatchObject([putOp('test_DATA', testData)]);
-    const endRowCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.ROWS_REPLICATED)) ?? 0;
-    const endTxCount = (await METRICS_HELPER.getMetricValueForTests(ReplicationMetric.TRANSACTIONS_REPLICATED)) ?? 0;
+    const endRowCount = await getRowsReplicated();
+    const endTxCount = await getTransactionsReplicated();
     expect(endRowCount - startRowCount).toEqual(1);
     expect(endTxCount - startTxCount).toBeGreaterThanOrEqual(1);
   });
+
+  // 250 rows end with a partial batch, 300 rows end with an empty batch.
+  test.each([250, 300])(
+    'Initial snapshot for a SimpleSnapshotQuery table with %i rows, larger than the snapshot batch size.',
+    async (rowCount) => {
+      await using context = await CDCStreamTestContext.open(factory, { cdcStreamOptions: { snapshotBatchSize: 100 } });
+      const { connectionManager } = context;
+      await context.updateSyncRules(COMPOSITE_KEY_SYNC_CONFIG);
+
+      await createCompositeKeyTable(connectionManager);
+      await insertCompositeKeyTableRows(connectionManager, rowCount);
+      const startRowCount = await getRowsReplicated();
+
+      await context.replicateSnapshot();
+      await context.startStreaming();
+
+      const endRowCount = await getRowsReplicated();
+      expect(endRowCount - startRowCount).toEqual(rowCount);
+      expect(await context.getBucketData('global|0[]')).toHaveLength(rowCount);
+    }
+  );
+
+  test('Interrupting initial snapshot for a SimpleSnapshotQuery table between batches', async () => {
+    // The snapshot query is still in progress when the snapshot is interrupted between batches.
+    // It must be cancelled before the snapshot transaction can be rolled back.
+    await using context = await CDCStreamTestContext.open(factory, { cdcStreamOptions: { snapshotBatchSize: 100 } });
+    const { connectionManager } = context;
+    await context.updateSyncRules(COMPOSITE_KEY_SYNC_CONFIG);
+    await createCompositeKeyTable(connectionManager);
+    await insertCompositeKeyTableRows(connectionManager, 5000);
+    const startRowCount = await getRowsReplicated();
+
+    const snapshotError = context.replicateSnapshot().then(
+      () => null,
+      (e) => e
+    );
+    await vi.waitFor(async () => expect((await getRowsReplicated()) - startRowCount).toBeGreaterThanOrEqual(200), {
+      timeout: 10_000,
+      interval: 1
+    });
+    await context.dispose();
+
+    expect(await snapshotError).toBeInstanceOf(ReplicationAbortedError);
+  });
+}
+
+/**
+ * A composite primary key cannot use BatchedSnapshotQuery, so this table is snapshotted with SimpleSnapshotQuery.
+ */
+async function createCompositeKeyTable(connectionManager: MSSQLConnectionManager) {
+  await connectionManager.query(`
+    CREATE TABLE test_composite (
+      id_user INT NOT NULL,
+      id_visit INT NOT NULL,
+      CONSTRAINT PK_test_composite PRIMARY KEY (id_user, id_visit)
+    )`);
+  await enableCDCForTable({ connectionManager, table: 'test_composite' });
+}
+
+async function insertCompositeKeyTableRows(connectionManager: MSSQLConnectionManager, rowCount: number) {
+  const beforeLSN = await getLatestLSN(connectionManager);
+  await connectionManager.query(`
+    WITH n AS (SELECT 1 AS i UNION ALL SELECT i + 1 FROM n WHERE i < ${rowCount})
+    INSERT INTO test_composite (id_user, id_visit) SELECT i, i % 10 FROM n OPTION (MAXRECURSION ${rowCount})`);
+  await waitForPendingCDCChanges(beforeLSN, connectionManager);
 }

package/test/src/CDCStreamTestContext.ts

@@ -64,7 +64,9 @@ export class CDCStreamTestContext implements AsyncDisposable {
 
   async dispose() {
     this.abortController.abort();
-    await this.snapshotPromise;
+    // Let an interrupted snapshot finish before closing its connections.
+    // Its error is handled by the caller of replicateSnapshot().
+    await this.snapshotPromise?.catch(() => {});
     await this.streamPromise;
     await this.connectionManager.end();
     await this.factory?.[Symbol.asyncDispose]();
@@ -154,7 +156,8 @@ export class CDCStreamTestContext implements AsyncDisposable {
   }
 
   async replicateSnapshot() {
-    await this.cdcStream.initReplication();
+    this.snapshotPromise = this.cdcStream.initReplication();
+    await this.snapshotPromise;
     this.replicationDone = true;
   }