@powersync/service-module-mssql 0.10.2 → 0.10.4

package/src/replication/MSSQLSnapshotQuery.ts

@@ -5,61 +5,102 @@ import { MSSQLBaseType } from '../types/mssql-data-types.js';
 import { escapeIdentifier } from '../utils/mssql.js';
 
 export interface MSSQLSnapshotQuery {
-  initialize(): Promise<void>;
+  /**
+   *  Returns the column metadata for the rows that would be returned for this query.
+   */
+  getColumnMetadata(): Promise<sql.IColumnMetadata>;
 
   /**
-   *  Returns an async iterable iterator that yields the column metadata for the query followed by rows of data.
+   *  Returns an async iterator for the next batch of rows, yielding at most the batch size.
+   *  A batch smaller than the batch size indicates that there are no rows left.
    */
-  next(): AsyncIterableIterator<sql.IColumnMetadata | sql.IRecordSet<any>>;
+  nextBatch(): AsyncIterableIterator<Record<string, any>>;
+
+  /**
+   *  Cancels the query if it is still in progress.
+   */
+  close(): Promise<void>;
 }
 
 /**
- * Snapshot query using a plain SELECT * FROM table
- *
- * This supports all tables but does not efficiently resume the snapshot
- * if the process is restarted.
+ *  Helper class that encapsulates the streaming logic of the queries used for snapshots.
+ */
+class StreamingQuery {
+  readonly rows: AsyncIterableIterator<any>;
+  private readonly completed: Promise<void>;
+  private finished = false;
+
+  constructor(
+    private readonly request: sql.Request,
+    query: string
+  ) {
+    const stream = request.toReadableStream();
+    // Errors are surfaced by iterating the rows. This listener prevents an unhandled 'error' event
+    // when the query fails after iteration has stopped, such as the cancellation error from cancel().
+    stream.on('error', () => {});
+    // In stream mode, errors are emitted on the stream instead of rejecting this promise.
+    this.completed = request.query(query).then(
+      () => {
+        this.finished = true;
+      },
+      () => {
+        this.finished = true;
+      }
+    );
+    this.rows = stream[Symbol.asyncIterator]();
+  }
+
+  async cancel(): Promise<void> {
+    if (!this.finished) {
+      this.request.cancel();
+    }
+    // Resolves once the driver has released the transaction's connection.
+    await this.completed;
+  }
+}
+
+/**
+ * Snapshot query using a plain SELECT * FROM table.
+ * This supports all tables but cannot resume the snapshot if the process is restarted.
  */
 export class SimpleSnapshotQuery implements MSSQLSnapshotQuery {
+  private query: StreamingQuery | null = null;
+
   public constructor(
     private readonly transaction: sql.Transaction,
-    private readonly qualifiedTableName: string
+    private readonly qualifiedTableName: string,
+    private readonly batchSize: number = 10_000
   ) {}
 
-  public async initialize(): Promise<void> {}
-
-  public async *next(): AsyncIterableIterator<sql.IColumnMetadata | sql.IRecordSet<any>> {
-    const metadataRequest = this.transaction.request();
-    metadataRequest.stream = true;
-    const metadataPromise = new Promise<sql.IColumnMetadata>((resolve, reject) => {
-      metadataRequest.on('recordset', resolve);
-      metadataRequest.on('error', reject);
-    });
-
-    metadataRequest.query(`SELECT TOP(0) * FROM ${this.qualifiedTableName}`);
-
-    const columnMetadata: sql.IColumnMetadata = await metadataPromise;
-    yield columnMetadata;
-
-    const request = this.transaction.request();
-    const stream = request.toReadableStream();
-
-    request.query(`SELECT * FROM ${this.qualifiedTableName}`);
+  public getColumnMetadata(): Promise<sql.IColumnMetadata> {
+    return queryColumnMetadata(this.transaction, this.qualifiedTableName);
+  }
 
-    // MSSQL only streams one row at a time
-    for await (const row of stream) {
-      yield row;
+  public async *nextBatch(): AsyncIterableIterator<Record<string, any>> {
+    // Opened on the first call, and iterated straight away so that the row iterator receives any query error.
+    this.query ??= new StreamingQuery(this.transaction.request(), `SELECT * FROM ${this.qualifiedTableName}`);
+    for (let i = 0; i < this.batchSize; i++) {
+      // MSSQL only streams one row at a time
+      const result = await this.query.rows.next();
+      if (result.done) {
+        return;
+      }
+      yield result.value;
     }
   }
+
+  public async close(): Promise<void> {
+    await this.query?.cancel();
+  }
 }
 
 /**
  * Performs a table snapshot query, batching by ranges of primary key data.
  *
  * This may miss some rows if they are modified during the snapshot query.
- * In that case, replication will pick up those rows afterward,
- * possibly resulting in an IdSnapshotQuery.
+ * In that case, replication will pick up those rows afterward.
  *
- * Currently, this only supports a table with a single primary key column,
+ * Currently, this only supports tables with a single primary key column,
  * of a select few types.
  */
 export class BatchedSnapshotQuery implements MSSQLSnapshotQuery {
@@ -93,6 +134,7 @@ export class BatchedSnapshotQuery implements MSSQLSnapshotQuery {
   }
 
   private readonly key: ColumnDescriptor;
+  private query: StreamingQuery | null = null;
   lastKey: string | bigint | null = null;
 
   public constructor(
@@ -109,57 +151,52 @@ export class BatchedSnapshotQuery implements MSSQLSnapshotQuery {
     }
   }
 
-  public async initialize(): Promise<void> {
-    // No-op
-  }
-
-  public getLastKeySerialized(): Uint8Array {
-    return bson.serialize({ [this.key.name]: this.lastKey });
-  }
-
-  public async *next(): AsyncIterableIterator<sql.IColumnMetadata | sql.IRecordSet<any>> {
-    const escapedKeyName = escapeIdentifier(this.key.name);
-    const metadataRequest = this.transaction.request();
-    metadataRequest.stream = true;
-    const metadataPromise = new Promise<sql.IColumnMetadata>((resolve, reject) => {
-      metadataRequest.on('recordset', resolve);
-      metadataRequest.on('error', reject);
-    });
-    metadataRequest.query(`SELECT TOP(0) * FROM ${this.qualifiedTableName}`);
-
-    const columnMetadata: sql.IColumnMetadata = await metadataPromise;
+  public async getColumnMetadata(): Promise<sql.IColumnMetadata> {
+    const columnMetadata = await queryColumnMetadata(this.transaction, this.qualifiedTableName);
 
     const foundPrimaryKey = columnMetadata[this.key.name];
     if (!foundPrimaryKey) {
       throw new Error(
-        `Cannot find primary key column ${this.key.name} in results. Keys: ${Object.keys(columnMetadata.columns).join(', ')}`
+        `Cannot find primary key column ${this.key.name} in results. Keys: ${Object.keys(columnMetadata).join(', ')}`
       );
     }
+    return columnMetadata;
+  }
 
-    yield columnMetadata;
+  public getLastKeySerialized(): Uint8Array {
+    return bson.serialize({ [this.key.name]: this.lastKey });
+  }
 
+  public async *nextBatch(): AsyncIterableIterator<Record<string, any>> {
+    const escapedKeyName = escapeIdentifier(this.key.name);
     const request = this.transaction.request();
-    const stream = request.toReadableStream();
     if (this.lastKey == null) {
-      request.query(`SELECT TOP(${this.batchSize}) * FROM ${this.qualifiedTableName} ORDER BY ${escapedKeyName}`);
+      this.query = new StreamingQuery(
+        request,
+        `SELECT TOP(${this.batchSize}) * FROM ${this.qualifiedTableName} ORDER BY ${escapedKeyName}`
+      );
     } else {
       if (this.key.typeId == null) {
         throw new Error(`typeId required for primary key ${this.key.name}`);
       }
-      request
-        .input('lastKey', this.lastKey)
-        .query(
-          `SELECT TOP(${this.batchSize}) * FROM ${this.qualifiedTableName} WHERE ${escapedKeyName} > @lastKey ORDER BY ${escapedKeyName}`
-        );
+      request.input('lastKey', this.lastKey);
+      this.query = new StreamingQuery(
+        request,
+        `SELECT TOP(${this.batchSize}) * FROM ${this.qualifiedTableName} WHERE ${escapedKeyName} > @lastKey ORDER BY ${escapedKeyName}`
+      );
     }
 
     // MSSQL only streams one row at a time
-    for await (const row of stream) {
+    for await (const row of this.query.rows) {
       this.lastKey = row[this.key.name];
       yield row;
     }
   }
 
+  public async close(): Promise<void> {
+    await this.query?.cancel();
+  }
+
   private deserializeKey(key: Uint8Array) {
     const decoded = bson.deserialize(key, { useBigInt64: true });
     const keys = Object.keys(decoded);
@@ -173,3 +210,11 @@ export class BatchedSnapshotQuery implements MSSQLSnapshotQuery {
     return decoded[this.key.name];
   }
 }
+
+async function queryColumnMetadata(
+  transaction: sql.Transaction,
+  qualifiedTableName: string
+): Promise<sql.IColumnMetadata> {
+  const { recordset } = await transaction.request().query(`SELECT TOP(0) * FROM ${qualifiedTableName}`);
+  return recordset.columns;
+}

package/src/utils/mssql.ts

@@ -251,28 +251,6 @@ export function toQualifiedTableName(schema: string, tableName: string): string
   return `${escapeIdentifier(schema)}.${escapeIdentifier(tableName)}`;
 }
 
-export function isIColumnMetadata(obj: any): obj is sql.IColumnMetadata {
-  if (obj === null || typeof obj !== 'object' || Array.isArray(obj)) {
-    return false;
-  }
-
-  let propertiesMatched = true;
-  for (const value of Object.values(obj)) {
-    const property = value as any;
-    propertiesMatched =
-      typeof property.index === 'number' &&
-      typeof property.name === 'string' &&
-      (typeof property.length === 'number' || typeof property.length === 'undefined') &&
-      (typeof property.type === 'function' || typeof property.type === 'object') &&
-      typeof property.nullable === 'boolean' &&
-      typeof property.caseSensitive === 'boolean' &&
-      typeof property.identity === 'boolean' &&
-      typeof property.readOnly === 'boolean';
-  }
-
-  return propertiesMatched;
-}
-
 export function addParameters(request: sql.Request, parameters: MSSQLParameter[]): sql.Request {
   for (const param of parameters) {
     if (param.type) {

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

@@ -11,7 +11,7 @@ import {
   SyncRulesBucketStorage,
   updateSyncRulesFromYaml
 } from '@powersync/service-core';
-import { bucketRequest, METRICS_HELPER, test_utils } from '@powersync/service-core-tests';
+import { bucketRequest, getTestStorage, METRICS_HELPER, test_utils } from '@powersync/service-core-tests';
 import timers from 'timers/promises';
 import { clearTestDb, getClientCheckpoint, TEST_CONNECTION_OPTIONS } from './util.js';
 
@@ -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]();
@@ -86,7 +88,7 @@ export class CDCStreamTestContext implements AsyncDisposable {
       updateSyncRulesFromYaml(content, { validate: true, storageVersion: LEGACY_STORAGE_VERSION })
     );
     this.syncRulesContent = replicationStream.syncConfigContent[0];
-    this.storage = this.factory.getInstance(replicationStream);
+    this.storage = await getTestStorage(this.factory, replicationStream);
     return this.storage!;
   }
 
@@ -97,7 +99,7 @@ export class CDCStreamTestContext implements AsyncDisposable {
     }
 
     this.syncRulesContent = syncConfig.content;
-    this.storage = syncConfig.storage;
+    this.storage = await getTestStorage(this.factory, syncConfig.replicationStream);
     return this.storage!;
   }
 
@@ -108,7 +110,7 @@ export class CDCStreamTestContext implements AsyncDisposable {
     }
 
     this.syncRulesContent = syncConfig.content;
-    this.storage = syncConfig.storage;
+    this.storage = await getTestStorage(this.factory, syncConfig.replicationStream);
     return this.storage!;
   }
 
@@ -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;
   }