[POC] Bucket storage report (operations vs rows)

.changeset/bucket-storage-report.md

@@ -0,0 +1,8 @@
+---
+'@powersync/service-core': minor
+'@powersync/service-types': minor
+'@powersync/service-module-mongodb-storage': minor
+'@powersync/service-core-tests': minor
+---
+
+Add a `POST /api/admin/v1/bucket-report` admin endpoint reporting operations vs rows per bucket (MongoDB storage).

modules/module-mongodb-storage/src/storage/implementation/v1/MongoSyncBucketStorageV1.ts

@@ -32,7 +32,14 @@ import { MongoChecksums } from '../MongoChecksums.js';
 import { MongoCompactOptions, MongoCompactor } from '../MongoCompactor.js';
 import { MongoParameterCompactor } from '../MongoParameterCompactor.js';
 import { MongoPersistedReplicationStream } from '../MongoPersistedReplicationStream.js';
-import { MongoSyncBucketStorage, MongoSyncBucketStorageOptions } from '../MongoSyncBucketStorage.js';
+import {
+  BucketRowEstimate,
+  MongoSyncBucketStorage,
+  MongoSyncBucketStorageOptions,
+  TopBucketCandidate,
+  TopBucketSelection,
+  TopDefinitionCandidate
+} from '../MongoSyncBucketStorage.js';
 import {
   BucketDataDocumentV1,
   BucketDataKeyV1,
@@ -190,6 +197,75 @@ export class MongoSyncBucketStorageV1 extends MongoSyncBucketStorage {
     return new MongoCompactorV1(this, this.db, options);
   }
 
+  // For storage v1/v2, bucket state and bucket data are shared collections scoped by group (replication stream).
+  protected async collectTopBuckets(limit: number): Promise<TopBucketSelection> {
+    const { buckets, definitions, definitionsTruncated, totals } = await this.aggregateTopBuckets(
+      this.db.bucketStateV1,
+      { '_id.g': this.replicationStreamId },
+      limit
+    );
+    return {
+      buckets: buckets.map((b) => ({ bucket: b.id.b, operations: b.operations, operationBytes: b.operationBytes })),
+      definitions,
+      definitionsTruncated,
+      totals
+    };
+  }
+
+  protected estimateBucketRows(candidate: TopBucketCandidate): Promise<BucketRowEstimate> {
+    // v1/v2 store one document per operation, so a bucket's ops are an id-prefix range that can be sampled directly.
+    const sampled = this.shouldSampleBucketRows(candidate.operations);
+    const buildPrefix = (applySample: boolean): mongo.Document[] => {
+      // Range-match on the whole `_id` (g, b, o) so the {_id} index is used; a dotted `{'_id.g','_id.b'}` match
+      // cannot use the compound-object index and would scan the whole collection per bucket.
+      const prefix: mongo.Document[] = [
+        {
+          $match: {
+            _id: idPrefixFilter<{ g: number; b: string; o: unknown }>(
+              { g: this.replicationStreamId, b: candidate.bucket },
+              ['o']
+            )
+          }
+        }
+      ];
+      if (applySample) {
+        prefix.push({ $match: { $sampleRate: this.bucketRowSampleRate(candidate.operations) } });
+      }
+      return prefix;
+    };
+    return this.estimateRowsFromOperationSample(this.db.bucketDataV1, buildPrefix, candidate.operations, sampled);
+  }
+
+  protected estimateDefinitionRows(candidate: TopDefinitionCandidate): Promise<BucketRowEstimate> {
+    const sampled = this.shouldSampleBucketRows(candidate.operations);
+    const buildPrefix = (applySample: boolean): mongo.Document[] => {
+      // All of a definition's bucket names start with `<definition>[`, so an `_id` range on that string
+      // prefix selects exactly the definition's operations via the index. `\\` (0x5C) is the character
+      // after `[` (0x5B), so [`<definition>[`, `<definition>\\`) cannot include any other definition:
+      // a longer definition name would have to differ at or before the `[`.
+      const prefix: mongo.Document[] = [
+        {
+          $match: {
+            _id: {
+              $gte: { g: this.replicationStreamId, b: `${candidate.definition}[`, o: new bson.MinKey() },
+              $lt: { g: this.replicationStreamId, b: `${candidate.definition}\\`, o: new bson.MinKey() }
+            }
+          }
+        }
+      ];
+      if (applySample) {
+        prefix.push({ $match: { $sampleRate: this.bucketRowSampleRate(candidate.operations) } });
+      }
+      return prefix;
+    };
+    // Include the bucket name in the row key: at definition grain a row counts once per bucket holding it.
+    return this.estimateRowsFromOperationSample(this.db.bucketDataV1, buildPrefix, candidate.operations, sampled, {
+      b: '$_id.b',
+      table: '$table',
+      row_id: '$row_id'
+    });
+  }
+
   protected createMongoParameterCompactor(
     checkpoint: InternalOpId,
     options: storage.CompactOptions

modules/module-mongodb-storage/src/storage/implementation/v3/MongoSyncBucketStorageV3.ts

@@ -15,15 +15,22 @@ import {
 import { JSONBig } from '@powersync/service-jsonbig';
 import { ParameterLookupRows, ScopedParameterLookup, SqliteJsonRow } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
-import { mapOpEntry, readSingleBatch, setSessionSnapshotTime } from '../../../utils/util.js';
+import { idPrefixFilter, mapOpEntry, readSingleBatch, setSessionSnapshotTime } from '../../../utils/util.js';
 import { MongoBucketStorage } from '../../MongoBucketStorage.js';
 import { BucketDataDoc } from '../common/BucketDataDoc.js';
 import { MongoSyncBucketStorageCheckpoint } from '../common/MongoSyncBucketStorageCheckpoint.js';
 import { MongoChecksums } from '../MongoChecksums.js';
 import { MongoCompactOptions, MongoCompactor } from '../MongoCompactor.js';
 import { MongoParameterCompactor } from '../MongoParameterCompactor.js';
 import { MongoPersistedReplicationStream } from '../MongoPersistedReplicationStream.js';
-import { MongoSyncBucketStorage, MongoSyncBucketStorageOptions } from '../MongoSyncBucketStorage.js';
+import {
+  BucketRowEstimate,
+  MongoSyncBucketStorage,
+  MongoSyncBucketStorageOptions,
+  TopBucketCandidate,
+  TopBucketSelection,
+  TopDefinitionCandidate
+} from '../MongoSyncBucketStorage.js';
 import { loadBucketDataDocument } from './bucket-format.js';
 import {
   BucketDataDocumentV3,
@@ -209,6 +216,73 @@ export class MongoSyncBucketStorageV3 extends MongoSyncBucketStorage {
     return new MongoCompactorV3(this, this.db, options);
   }
 
+  // For storage v3, bucket state is a per-stream collection and bucket data is split into per-definition collections.
+  // A replication stream can host multiple sync configs (active + processing + stopped, until cleanup runs), all
+  // sharing these collections. Scope to the active config's definition ids so the report excludes stale buckets
+  // from old/stopped definitions. `this.storageIds` is derived from the active config only (see getActiveSyncConfig).
+  protected async collectTopBuckets(limit: number): Promise<TopBucketSelection> {
+    const { buckets, definitions, definitionsTruncated, totals } = await this.aggregateTopBuckets(
+      this.db.bucketState(this.replicationStreamId),
+      { '_id.d': { $in: this.storageIds.bucketDefinitionIds } },
+      limit
+    );
+    return {
+      buckets: buckets.map((b) => ({
+        bucket: b.id.b,
+        operations: b.operations,
+        operationBytes: b.operationBytes,
+        defId: b.id.d
+      })),
+      definitions,
+      definitionsTruncated,
+      totals
+    };
+  }
+
+  protected estimateBucketRows(candidate: TopBucketCandidate): Promise<BucketRowEstimate> {
+    // v3 batches operations into documents (one doc holds an `ops` array), in a per-definition collection.
+    // Sample whole batch documents, then unwind to operation level so the shared estimator sees one doc per op.
+    const sampled = this.shouldSampleBucketRows(candidate.operations);
+    const collection = this.db.bucketData(this.replicationStreamId, candidate.defId!);
+    const buildPrefix = (applySample: boolean): mongo.Document[] => {
+      // Range-match on the whole `_id` (b, o) so the {_id} index is used; a dotted `{'_id.b': ...}` match
+      // cannot use the compound-object index and would scan the whole collection per bucket.
+      const prefix: mongo.Document[] = [
+        { $match: { _id: idPrefixFilter<{ b: string; o: unknown }>({ b: candidate.bucket }, ['o']) } }
+      ];
+      if (applySample) {
+        prefix.push({ $match: { $sampleRate: this.bucketRowSampleRate(candidate.operations) } });
+      }
+      prefix.push({ $unwind: '$ops' }, { $replaceRoot: { newRoot: '$ops' } });
+      return prefix;
+    };
+    return this.estimateRowsFromOperationSample(collection, buildPrefix, candidate.operations, sampled);
+  }
+
+  protected estimateDefinitionRows(candidate: TopDefinitionCandidate): Promise<BucketRowEstimate> {
+    // A definition's operations are exactly its per-definition bucket_data collection, so no match stage is
+    // needed. Keep the bucket name alongside each unwound operation: at definition grain a row counts once
+    // per bucket holding it.
+    const sampled = this.shouldSampleBucketRows(candidate.operations);
+    const collection = this.db.bucketData(this.replicationStreamId, candidate.defId!);
+    const buildPrefix = (applySample: boolean): mongo.Document[] => {
+      const prefix: mongo.Document[] = [];
+      if (applySample) {
+        prefix.push({ $match: { $sampleRate: this.bucketRowSampleRate(candidate.operations) } });
+      }
+      prefix.push(
+        { $unwind: '$ops' },
+        { $project: { b: '$_id.b', op: '$ops.op', table: '$ops.table', row_id: '$ops.row_id' } }
+      );
+      return prefix;
+    };
+    return this.estimateRowsFromOperationSample(collection, buildPrefix, candidate.operations, sampled, {
+      b: '$b',
+      table: '$table',
+      row_id: '$row_id'
+    });
+  }
+
   protected createMongoParameterCompactor(
     checkpoint: InternalOpId,
     options: storage.CompactOptions

modules/module-mongodb-storage/test/src/bucket-report-scoping.test.ts

@@ -0,0 +1,128 @@
+import { MongoSyncBucketStorageV3 } from '@module/storage/implementation/v3/MongoSyncBucketStorageV3.js';
+import { storage, updateSyncRulesFromYaml } from '@powersync/service-core';
+import { test_utils } from '@powersync/service-core-tests';
+import * as bson from 'bson';
+import { describe, expect, test } from 'vitest';
+import { INITIALIZED_MONGO_STORAGE_FACTORY } from './util.js';
+
+function sourceDescriptor(name: string, objectId: string): storage.SourceEntityDescriptor {
+  return {
+    connectionTag: storage.SourceTable.DEFAULT_TAG,
+    objectId,
+    schema: 'public',
+    name,
+    replicaIdColumns: [{ name: 'id', type: 'VARCHAR', typeId: 25 }]
+  };
+}
+
+function objectIdGenerator(id: string) {
+  let used = false;
+  return () => {
+    if (used) {
+      throw new Error(`Can only generate a single id using ${id}`);
+    }
+    used = true;
+    return new bson.ObjectId(id);
+  };
+}
+
+/**
+ * In V3 a replication stream can host multiple sync configs (active + stopped, until cleanup runs), all sharing
+ * the per-stream bucket_state and source_records collections. The report must only include the active config's
+ * bucket definitions, not stale ones from a previous (now stopped) config.
+ */
+describe('bucket report scoping - mongodb v3', () => {
+  test('excludes buckets from stopped/old sync configs sharing the replication stream', async () => {
+    await using factory = await INITIALIZED_MONGO_STORAGE_FACTORY.factory();
+
+    // Config 1: replicate todos. Writing a row creates a data bucket for this config.
+    const first = await factory.updateSyncRules(
+      updateSyncRulesFromYaml(
+        `
+config:
+  edition: 3
+
+streams:
+  by_owner:
+    query: SELECT * FROM todos WHERE owner_id = subscription.parameter('owner_id')
+`,
+        { storageVersion: 3 }
+      )
+    );
+    const firstStorage = factory.getInstance(first) as MongoSyncBucketStorageV3;
+    await using firstWriter = await firstStorage.createWriter(test_utils.BATCH_OPTIONS);
+    const todosTable = (
+      await firstWriter.resolveTables({
+        connection_id: 1,
+        source: sourceDescriptor('todos', 'todos-relation'),
+        idGenerator: objectIdGenerator('6544e3899293153fa7b38360')
+      })
+    ).tables[0];
+    await firstWriter.save({
+      sourceTable: todosTable,
+      tag: storage.SaveOperationTag.INSERT,
+      after: { id: 'todo-1', owner_id: 'user-1' },
+      afterReplicaId: test_utils.rid('todo-1')
+    });
+    await firstWriter.markAllSnapshotDone('1/1');
+    await firstWriter.commit('1/1');
+    await firstWriter.flush();
+
+    // While config 1 is active, its bucket(s) show up in the report.
+    const firstReport = await firstStorage.getBucketReport();
+    expect(firstReport.totals.bucketCount).toBeGreaterThan(0);
+
+    // Config 2: a different stream over a different table. Config 1 transitions to STOP, but its bucket_state
+    // and source_records rows remain in the shared collections until cleanup runs (which we deliberately skip).
+    const second = await factory.updateSyncRules(
+      updateSyncRulesFromYaml(
+        `
+config:
+  edition: 3
+
+streams:
+  by_project:
+    query: SELECT * FROM scenes WHERE project_id = subscription.parameter('project_id')
+`,
+        { storageVersion: 3 }
+      )
+    );
+    expect(second.replicationStreamId).toBe(first.replicationStreamId);
+
+    // Drive config 2 to snapshot-done so it becomes ACTIVE and config 1 transitions to STOP (config 1 keeps
+    // serving until the new config finishes processing). Config 1's stale rows remain until cleanup, which we skip.
+    const replicatingStreams = await factory.getReplicatingReplicationStreams();
+    expect(replicatingStreams).toHaveLength(1);
+    const secondStorage = factory.getInstance(replicatingStreams[0]) as MongoSyncBucketStorageV3;
+    await using secondWriter = await secondStorage.createWriter(test_utils.BATCH_OPTIONS);
+    // Give config 2 its own replicated row, so the report has an active-config bucket to include.
+    const scenesTable = (
+      await secondWriter.resolveTables({
+        connection_id: 1,
+        source: sourceDescriptor('scenes', 'scenes-relation'),
+        idGenerator: objectIdGenerator('6544e3899293153fa7b38361')
+      })
+    ).tables[0];
+    await secondWriter.save({
+      sourceTable: scenesTable,
+      tag: storage.SaveOperationTag.INSERT,
+      after: { id: 'scene-1', project_id: 'project-1' },
+      afterReplicaId: test_utils.rid('scene-1')
+    });
+    await secondWriter.markAllSnapshotDone('2/1');
+    await secondWriter.commit('2/1');
+    await secondWriter.flush();
+
+    const activeConfig = await factory.getActiveSyncConfig();
+    expect(activeConfig).not.toBeNull();
+    const activeStorage = activeConfig!.storage as MongoSyncBucketStorageV3;
+    const secondReport = await activeStorage.getBucketReport();
+
+    // The active config's own bucket is included (include-active), while config 1's stale buckets, which still
+    // exist in the shared collections, are excluded (exclude-stale). Without scoping to the active config's
+    // definition ids, config 1's bucket would leak in here.
+    expect(secondReport.totals.bucketCount).toBeGreaterThan(0);
+    const firstBucketNames = new Set(firstReport.buckets.map((b) => b.bucket));
+    expect(secondReport.buckets.some((b) => firstBucketNames.has(b.bucket))).toBe(false);
+  });
+});

modules/module-mongodb-storage/test/src/storage.test.ts

@@ -20,6 +20,8 @@ for (let storageVersion of TEST_STORAGE_VERSIONS) {
   describe(`Mongo Sync Bucket Storage - Checkpoints - v${storageVersion}`, () =>
     register.registerDataStorageCheckpointTests({ ...INITIALIZED_MONGO_STORAGE_FACTORY, storageVersion }));
 
+  describe(`Mongo Sync Bucket Storage - Bucket report - v${storageVersion}`, () =>
+    register.registerBucketReportTests({ ...INITIALIZED_MONGO_STORAGE_FACTORY, storageVersion }));
   describe(`Mongo Sync Bucket Storage - write checkpoint metadata - v${storageVersion}`, () => {
     test('uses checkpoint_requested_at as the client-requested checkpoint marker', async () => {
       await using factory = await INITIALIZED_MONGO_STORAGE_FACTORY.factory();