@powersync/service-module-mongodb 0.0.0-dev-20261015071128 → 0.0.0-dev-20261015132441

package/test/src/snapshot_count.test.ts

@@ -0,0 +1,146 @@
+import { estimateFilteredSnapshotCount } from '@module/replication/estimateFilteredSnapshotCount.js';
+import { mongo } from '@powersync/lib-service-mongodb';
+import { createLogger } from '@powersync/lib-services-framework';
+import { afterEach, describe, expect, test, vi } from 'vitest';
+
+const FILTER = { $expr: { $eq: ['$label', 'live'] } };
+const LOGGER = createLogger({ silent: true });
+
+function fixture() {
+  const client = new mongo.MongoClient('mongodb://localhost:27017');
+  const collection = client.db('snapshot_count_unit').collection('documents');
+  const controller = new AbortController();
+  const count = vi.spyOn(collection, 'countDocuments').mockResolvedValue(42);
+  const metadata = vi.spyOn(collection, 'estimatedDocumentCount').mockResolvedValue(100_000);
+  const cursor = collection.aggregate();
+  const next = vi.spyOn(cursor, 'next').mockResolvedValue({ sampled: [{ count: 1_000 }], matches: [{ count: 250 }] });
+  const close = vi.spyOn(cursor, 'close').mockResolvedValue();
+  const aggregate = vi.spyOn(collection, 'aggregate').mockReturnValue(cursor);
+  const estimate = () =>
+    estimateFilteredSnapshotCount({ collection, filter: FILTER, signal: controller.signal, logger: LOGGER });
+  return { controller, count, metadata, next, close, aggregate, estimate };
+}
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe('filtered snapshot count', () => {
+  test('uses a bounded exact count with the snapshot collation', async () => {
+    const context = fixture();
+    expect(await context.estimate()).toBe(42);
+    expect(context.count).toHaveBeenCalledWith(FILTER, {
+      collation: { locale: 'simple' },
+      readConcern: 'majority',
+      timeoutMS: 10_000,
+      signal: context.controller.signal
+    });
+    expect(context.metadata).not.toHaveBeenCalled();
+    expect(context.aggregate).not.toHaveBeenCalled();
+  });
+
+  test('preserves an exact zero without sampling', async () => {
+    const context = fixture();
+    context.count.mockResolvedValue(0);
+    expect(await context.estimate()).toBe(0);
+    expect(context.metadata).not.toHaveBeenCalled();
+  });
+
+  test('samples before filtering and scales by collection metadata after a count timeout', async () => {
+    const context = fixture();
+    context.count.mockRejectedValue(new mongo.MongoOperationTimeoutError('count timed out'));
+    expect(await context.estimate()).toBe(25_000);
+    expect(context.aggregate).toHaveBeenCalledWith(
+      [
+        { $sample: { size: 1_000 } },
+        { $facet: { sampled: [{ $count: 'count' }], matches: [{ $match: FILTER }, { $count: 'count' }] } }
+      ],
+      expect.objectContaining({
+        collation: { locale: 'simple' },
+        readConcern: 'majority',
+        signal: context.controller.signal,
+        allowDiskUse: false
+      })
+    );
+    const timeoutMS = context.aggregate.mock.calls[0][1]!.timeoutMS!;
+    expect(timeoutMS).toBeGreaterThan(0);
+    expect(timeoutMS).toBeLessThanOrEqual(10_000);
+    expect(context.close).toHaveBeenCalled();
+  });
+
+  test.each([
+    [10, 1],
+    [100_000_000, 10_000]
+  ])('bounds the sample size for a collection of %s documents', async (total, sampleSize) => {
+    const context = fixture();
+    context.count.mockRejectedValue(new Error('count failed'));
+    context.metadata.mockResolvedValue(total);
+    await context.estimate();
+    expect(context.aggregate.mock.calls[0][0]![0]).toEqual({ $sample: { size: sampleSize } });
+  });
+
+  test('uses the actual number sampled when metadata or collection size has changed', async () => {
+    const context = fixture();
+    context.count.mockRejectedValue(new Error('count failed'));
+    context.next.mockResolvedValue({ sampled: [{ count: 80 }], matches: [{ count: 20 }] });
+    expect(await context.estimate()).toBe(25_000);
+  });
+
+  test.each([null, { sampled: [], matches: [] }, { sampled: [{ count: 1_000 }], matches: [] }])(
+    'leaves the total unknown for an empty sample or no matches: %j',
+    async (sample) => {
+      const context = fixture();
+      context.count.mockRejectedValue(new Error('count failed'));
+      context.next.mockResolvedValue(sample);
+      expect(await context.estimate()).toBe(-1);
+      expect(context.close).toHaveBeenCalled();
+    }
+  );
+
+  test.each(['metadata', 'sample', 'empty metadata', 'budget exhausted'])(
+    'leaves the total unknown after %s',
+    async (failure) => {
+      const context = fixture();
+      context.count.mockRejectedValue(new Error('count failed'));
+      if (failure == 'metadata') {
+        context.metadata.mockRejectedValue(new Error('metadata failed'));
+      } else if (failure == 'sample') {
+        context.next.mockRejectedValue(new mongo.MongoOperationTimeoutError('sample timed out'));
+      } else if (failure == 'empty metadata') {
+        context.metadata.mockResolvedValue(0);
+      } else {
+        const { performance } = await import('node:perf_hooks');
+        vi.spyOn(performance, 'now').mockReturnValueOnce(0).mockReturnValue(10_001);
+      }
+      expect(await context.estimate()).toBe(-1);
+      if (failure == 'sample') {
+        expect(context.close).toHaveBeenCalled();
+      } else {
+        expect(context.aggregate).not.toHaveBeenCalled();
+      }
+    }
+  );
+
+  test.each(['before count', 'during count', 'during sample'])('propagates cancellation %s', async (when) => {
+    const context = fixture();
+    const reason = new Error('replication aborted');
+    if (when == 'before count') {
+      context.controller.abort(reason);
+    } else if (when == 'during count') {
+      context.count.mockImplementation(async () => {
+        context.controller.abort(reason);
+        throw reason;
+      });
+    } else {
+      context.count.mockRejectedValue(new Error('count failed'));
+      context.next.mockImplementation(async () => {
+        context.controller.abort(reason);
+        throw reason;
+      });
+    }
+    await expect(context.estimate()).rejects.toBe(reason);
+    if (when != 'during sample') {
+      expect(context.metadata).not.toHaveBeenCalled();
+    }
+  });
+});

package/test/src/snapshot_count_integration.test.ts

@@ -0,0 +1,70 @@
+import { estimateFilteredSnapshotCount } from '@module/replication/estimateFilteredSnapshotCount.js';
+import { mongo } from '@powersync/lib-service-mongodb';
+import { createLogger } from '@powersync/lib-services-framework';
+import { randomUUID } from 'node:crypto';
+import { describe, expect, test, vi } from 'vitest';
+import { DATABASE_TYPE, DatabaseType } from './DatabaseType.js';
+import { connectMongoData } from './util.js';
+
+const LOGGER = createLogger({ silent: true });
+
+describe.skipIf(DATABASE_TYPE == DatabaseType.DOCUMENTDB)('filtered snapshot count integration', () => {
+  test('counts with simple collation even when the collection uses case-insensitive collation', async () => {
+    const { client, db } = await connectMongoData();
+    const collection = await db.createCollection(`snapshot_count_${randomUUID()}`, {
+      collation: { locale: 'en', strength: 2 }
+    });
+    try {
+      await collection.insertMany([{ label: 'live' }, { label: 'LIVE' }, { label: 'other' }]);
+      expect(
+        await estimateFilteredSnapshotCount({
+          collection,
+          filter: { $expr: { $eq: ['$label', 'live'] } },
+          signal: new AbortController().signal,
+          logger: LOGGER
+        })
+      ).toBe(1);
+    } finally {
+      await collection.drop();
+      await client.close();
+    }
+  });
+
+  test('executes the sample pipeline after a failed count and preserves binary filter semantics', async () => {
+    const { client, db } = await connectMongoData();
+    const collection = await db.createCollection(`snapshot_count_${randomUUID()}`, {
+      collation: { locale: 'en', strength: 2 }
+    });
+    try {
+      // For this small collection the sample contains one document. All rows have the same label,
+      // making the sampled result deterministic while exercising the real $sample/$facet pipeline.
+      await collection.insertMany(Array.from({ length: 50 }, () => ({ label: 'live' })));
+      const count = vi
+        .spyOn(collection, 'countDocuments')
+        .mockRejectedValue(new mongo.MongoOperationTimeoutError('timeout'));
+      try {
+        expect(
+          await estimateFilteredSnapshotCount({
+            collection,
+            filter: { $and: [{ label: 'live' }, { $expr: { $eq: ['$label', 'live'] } }] },
+            signal: new AbortController().signal,
+            logger: LOGGER
+          })
+        ).toBe(50);
+        expect(
+          await estimateFilteredSnapshotCount({
+            collection,
+            filter: { $expr: { $eq: ['$label', 'LIVE'] } },
+            signal: new AbortController().signal,
+            logger: LOGGER
+          })
+        ).toBe(-1);
+      } finally {
+        count.mockRestore();
+      }
+    } finally {
+      await collection.drop();
+      await client.close();
+    }
+  });
+});

package/test/src/snapshot_query_provider.test.ts

@@ -1,4 +1,5 @@
 import { DEFAULT_MONGO_REPLICATION_QUERY_PROVIDER } from '@module/replication/MongoReplicationQueryProvider.js';
+import { mongo } from '@powersync/lib-service-mongodb';
 import { describe, expect, test, vi } from 'vitest';
 import { openChangeStreamTestContext } from './change_stream_test_setup.js';
 import { DATABASE_TYPE, DatabaseType } from './DatabaseType.js';
@@ -6,7 +7,11 @@ import { describeWithStorage } from './util.js';
 
 describe.skipIf(DATABASE_TYPE == DatabaseType.DOCUMENTDB)('MongoDB snapshot query provider', () => {
   describeWithStorage({ timeout: 30_000 }, ({ factory, storageVersion }) => {
-    test.each([false, true])('pages snapshots with an expression filter: %s', async (filtered) => {
+    test.each([
+      { filtered: false, retryUnknown: false },
+      { filtered: true, retryUnknown: false },
+      { filtered: true, retryUnknown: true }
+    ])('pages snapshots: %j', async ({ filtered, retryUnknown }) => {
       // Run the same snapshot with and without a provider predicate. The filtered case selects even
       // ids between 2 and 12 whose label is exactly 'live'; null must preserve the full snapshot.
       const getSnapshotFilter = vi.fn(() =>
@@ -53,7 +58,21 @@ describe.skipIf(DATABASE_TYPE == DatabaseType.DOCUMENTDB)('MongoDB snapshot quer
       await collection.insertMany(
         Array.from({ length: 14 }, (_, i) => ({ _id: i + 1, label: (i + 1) % 3 == 0 ? 'LIVE' : 'live' }))
       );
-      await context.replicateSnapshot();
+      const count = vi.spyOn(mongo.Collection.prototype, 'countDocuments');
+      const metadata = vi.spyOn(mongo.Collection.prototype, 'estimatedDocumentCount');
+      try {
+        if (retryUnknown) {
+          // Leave the queued estimate unknown, then let the snapshot retry against the real source.
+          count.mockRejectedValueOnce(new Error('count unavailable'));
+          metadata.mockRejectedValueOnce(new Error('metadata unavailable'));
+        }
+        await context.replicateSnapshot();
+        expect(count).toHaveBeenCalledTimes(retryUnknown ? 2 : filtered ? 1 : 0);
+        expect(metadata).toHaveBeenCalledTimes(retryUnknown ? 1 : filtered ? 0 : 2);
+      } finally {
+        count.mockRestore();
+        metadata.mockRestore();
+      }
       const rows = (await context.getBucketData('global|0[]')).filter((op) => op.op == 'PUT');
       // Matching ids 4, 8 and 10 require two pages at a page size of two. After the first page ends
       // at id 8, pagination adds its own $expr requiring _id > 8. It must combine that condition