@powersync/service-module-mongodb-storage 0.15.3 → 0.16.0

package/src/storage/implementation/v3/PersistedBatchV3.ts

@@ -0,0 +1,318 @@
+import { mongo } from '@powersync/lib-service-mongodb';
+import { ReplicationAssertionError } from '@powersync/lib-services-framework';
+import { InternalOpId, storage } from '@powersync/service-core';
+import { BucketDataSource } from '@powersync/service-sync-rules';
+import * as bson from 'bson';
+import { mongoTableId } from '../../../utils/util.js';
+import { BucketDefinitionId } from '../BucketDefinitionMapping.js';
+import {
+  BucketStateUpdate,
+  PersistedBatch,
+  SaveParameterDataOptions,
+  UpsertCurrentDataOptions
+} from '../common/PersistedBatch.js';
+import { SourceTableKey } from '../models.js';
+import {
+  BucketParameterDocumentV3,
+  BucketStateDocumentV3,
+  CurrentDataDocumentV3,
+  serializeBucketDataV3,
+  SourceTableDocumentV3,
+  taggedBucketParameterDocumentToV3
+} from './models.js';
+import { serializeParameterLookupV3 } from './MongoParameterLookupV3.js';
+import { VersionedPowerSyncMongoV3 } from './VersionedPowerSyncMongoV3.js';
+
+export class PersistedBatchV3 extends PersistedBatch {
+  declare protected readonly db: VersionedPowerSyncMongoV3;
+
+  currentData: { sourceTableId: bson.ObjectId; operation: mongo.AnyBulkWriteOperation<CurrentDataDocumentV3> }[] = [];
+  sourceTablePendingDeletes = new Map<string, InternalOpId>();
+
+  protected checkDefinitionId(definitionId: BucketDefinitionId | null): BucketDefinitionId {
+    if (definitionId == null) {
+      // This is required for V3 storage.
+      throw new ReplicationAssertionError('Expected v3 bucket when incrementalReprocessing is enabled');
+    }
+    return definitionId;
+  }
+
+  protected getBucketDefinitionId(bucketSource: BucketDataSource): BucketDefinitionId {
+    return this.mapping.bucketSourceId(bucketSource);
+  }
+
+  saveParameterData(data: SaveParameterDataOptions) {
+    const { sourceTable, sourceKey, evaluated } = data;
+    const remaining_lookups = new Map<string, SaveParameterDataOptions['existing_lookups'][number]>();
+
+    for (let lookup of data.existing_lookups) {
+      if (lookup.indexId == null) {
+        throw new ReplicationAssertionError('Expected v3 lookup when incrementalReprocessing is enabled');
+      }
+      remaining_lookups.set(`${lookup.indexId}.${lookup.lookup.toString('base64')}`, lookup);
+    }
+
+    for (let result of evaluated) {
+      const sourceDefinitionId = this.mapping.parameterLookupId(result.lookup.source);
+      const binLookup = serializeParameterLookupV3(result.lookup);
+      remaining_lookups.delete(`${sourceDefinitionId}.${binLookup.toString('base64')}`);
+
+      const op_id = data.op_seq.next();
+      this.debugLastOpId = op_id;
+      const values: BucketParameterDocumentV3 = {
+        _id: op_id,
+        key: {
+          t: mongoTableId(sourceTable.id),
+          k: sourceKey
+        } satisfies SourceTableKey,
+        lookup: binLookup,
+        bucket_parameters: result.bucketParameters
+      };
+      this.bucketParameters.push({
+        ...values,
+        index: sourceDefinitionId
+      });
+
+      this.currentSize += 200;
+    }
+
+    for (let lookup of remaining_lookups.values()) {
+      const op_id = data.op_seq.next();
+      this.debugLastOpId = op_id;
+      const indexId = lookup.indexId;
+      if (indexId == null) {
+        throw new ReplicationAssertionError('Expected v3 lookup when incrementalReprocessing is enabled');
+      }
+      const values: BucketParameterDocumentV3 = {
+        _id: op_id,
+        key: {
+          t: mongoTableId(sourceTable.id),
+          k: sourceKey
+        } satisfies SourceTableKey,
+        lookup: lookup.lookup,
+        bucket_parameters: []
+      };
+      this.bucketParameters.push({
+        ...values,
+        index: indexId
+      });
+
+      this.currentSize += 200;
+    }
+  }
+
+  hardDeleteCurrentData(sourceTableId: bson.ObjectId, replicaId: storage.ReplicaId) {
+    this.currentData.push({
+      sourceTableId,
+      operation: {
+        deleteOne: {
+          filter: { _id: replicaId }
+        }
+      }
+    });
+    this.currentSize += 50;
+  }
+
+  softDeleteCurrentData(
+    sourceTableId: bson.ObjectId,
+    replicaId: storage.ReplicaId,
+    checkpointGreaterThan: InternalOpId
+  ) {
+    this.currentData.push({
+      sourceTableId,
+      operation: {
+        updateOne: {
+          filter: { _id: replicaId },
+          update: {
+            $set: {
+              data: null,
+              buckets: [] as CurrentDataDocumentV3['buckets'],
+              lookups: [] as CurrentDataDocumentV3['lookups'],
+              pending_delete: checkpointGreaterThan
+            }
+          },
+          upsert: true
+        }
+      }
+    });
+    const sourceTableKey = sourceTableId.toHexString();
+    const existingPendingDelete = this.sourceTablePendingDeletes.get(sourceTableKey);
+    if (existingPendingDelete == null || checkpointGreaterThan > existingPendingDelete) {
+      this.sourceTablePendingDeletes.set(sourceTableKey, checkpointGreaterThan);
+    }
+
+    this.currentSize += 50;
+  }
+
+  upsertCurrentData(values: UpsertCurrentDataOptions) {
+    const buckets = values.buckets.map((bucket) => {
+      if (bucket.definitionId == null) {
+        throw new ReplicationAssertionError('Expected v3 bucket when incrementalReprocessing is enabled');
+      }
+      return {
+        def: bucket.definitionId,
+        bucket: bucket.bucket,
+        table: bucket.table,
+        id: bucket.id
+      };
+    });
+    const lookups = values.lookups.map((lookup) => {
+      if (lookup.indexId == null) {
+        throw new ReplicationAssertionError('Expected v3 lookup when incrementalReprocessing is enabled');
+      }
+      return {
+        i: lookup.indexId,
+        l: lookup.lookup
+      };
+    });
+
+    this.currentData.push({
+      sourceTableId: values.sourceTableId,
+      operation: {
+        updateOne: {
+          filter: { _id: values.replicaId },
+          update: {
+            $set: {
+              data: values.data,
+              buckets,
+              lookups
+            },
+            $unset: { pending_delete: 1 }
+          },
+          upsert: true
+        }
+      }
+    });
+    this.currentSize += (values.data?.length() ?? 0) + 100;
+  }
+
+  protected get currentDataCount() {
+    return this.currentData.length;
+  }
+
+  protected async flushBucketData(session: mongo.ClientSession) {
+    const operationsByDefinition = new Map<BucketDefinitionId, typeof this.bucketData>();
+    for (const document of this.bucketData) {
+      const existing = operationsByDefinition.get(document.bucketKey.definitionId) ?? [];
+      existing.push(document);
+      operationsByDefinition.set(document.bucketKey.definitionId, existing);
+    }
+
+    for (const [definitionId, documents] of operationsByDefinition.entries()) {
+      await this.db.bucketDataV3(this.group_id, definitionId).bulkWrite(
+        documents.map((document) => ({
+          insertOne: {
+            document: serializeBucketDataV3(document)
+          }
+        })),
+        {
+          session,
+          ordered: false
+        }
+      );
+    }
+  }
+
+  protected async flushBucketParameters(session: mongo.ClientSession) {
+    const operationsByIndex = new Map<string, typeof this.bucketParameters>();
+    for (const document of this.bucketParameters) {
+      const existing = operationsByIndex.get(document.index) ?? [];
+      existing.push(document);
+      operationsByIndex.set(document.index, existing);
+    }
+
+    for (const [indexId, documents] of operationsByIndex.entries()) {
+      await this.db.parameterIndexV3(this.group_id, indexId).bulkWrite(
+        documents.map((document) => ({
+          insertOne: {
+            document: taggedBucketParameterDocumentToV3(document)
+          }
+        })),
+        {
+          session,
+          ordered: false
+        }
+      );
+    }
+  }
+
+  protected async flushCurrentData(session: mongo.ClientSession) {
+    const operationsBySourceTable = new Map<string, typeof this.currentData>();
+    for (const operation of this.currentData) {
+      const sourceTableId = operation.sourceTableId.toHexString();
+      const existing = operationsBySourceTable.get(sourceTableId) ?? [];
+      existing.push(operation);
+      operationsBySourceTable.set(sourceTableId, existing);
+    }
+
+    const sourceTableUpdates: mongo.AnyBulkWriteOperation<SourceTableDocumentV3>[] = [
+      ...this.sourceTablePendingDeletes.entries()
+    ].map(([key, value]) => {
+      return {
+        updateOne: {
+          filter: { _id: new bson.ObjectId(key) },
+          update: {
+            $max: {
+              latest_pending_delete: value
+            }
+          }
+        }
+      };
+    });
+
+    if (sourceTableUpdates.length > 0) {
+      await this.db.sourceTablesV3(this.group_id).bulkWrite(sourceTableUpdates, { session, ordered: false });
+    }
+
+    for (const operations of operationsBySourceTable.values()) {
+      const sourceTableId = operations[0]!.sourceTableId;
+      await this.db.sourceRecordsV3(this.group_id, sourceTableId).bulkWrite(
+        operations.map((entry) => entry.operation),
+        {
+          session,
+          ordered: true
+        }
+      );
+    }
+  }
+
+  protected async flushBucketStates(session: mongo.ClientSession) {
+    await this.db.bucketStateV3(this.group_id).bulkWrite(this.getBucketStateUpdates(), {
+      session,
+      ordered: false
+    });
+  }
+
+  protected resetCurrentData() {
+    this.currentData = [];
+    this.sourceTablePendingDeletes.clear();
+  }
+
+  private getBucketStateUpdates(): mongo.AnyBulkWriteOperation<BucketStateDocumentV3>[] {
+    return Array.from(this.bucketStates.values()).map((state: BucketStateUpdate) => {
+      if (state.definitionId == null) {
+        throw new ReplicationAssertionError('Expected bucket definition id when incrementalReprocessing is enabled');
+      }
+      return {
+        updateOne: {
+          filter: {
+            _id: {
+              d: state.definitionId,
+              b: state.bucket
+            }
+          },
+          update: {
+            $set: {
+              last_op: state.lastOp
+            },
+            $inc: {
+              'estimate_since_compact.count': state.incrementCount,
+              'estimate_since_compact.bytes': state.incrementBytes
+            }
+          },
+          upsert: true
+        }
+      } satisfies mongo.AnyBulkWriteOperation<BucketStateDocumentV3>;
+    });
+  }
+}

package/src/storage/implementation/v3/SingleBucketStoreV3.ts

@@ -0,0 +1,68 @@
+import { mongo } from '@powersync/lib-service-mongodb';
+import { InternalOpId } from '@powersync/service-core';
+import { BucketDataDoc, BucketKey } from '../common/BucketDataDoc.js';
+import {
+  BucketDataDocumentGeneric,
+  BucketDataDocumentGenericId,
+  SingleBucketStore
+} from '../common/SingleBucketStore.js';
+import { BucketDataProperties } from '../models.js';
+import { VersionedPowerSyncMongoV3 } from './VersionedPowerSyncMongoV3.js';
+import { BucketDataDocumentV3, BucketDataKeyV3, loadBucketDataDocumentV3, serializeBucketDataV3 } from './models.js';
+
+export class SingleBucketStoreV3 implements SingleBucketStore {
+  public readonly collection: mongo.Collection<BucketDataDocumentGeneric>;
+
+  constructor(
+    private db: VersionedPowerSyncMongoV3,
+    public readonly key: BucketKey
+  ) {
+    this.collection = db.bucketDataV3(
+      key.replicationStreamId,
+      key.definitionId
+    ) as unknown as mongo.Collection<BucketDataDocumentGeneric>;
+  }
+
+  docId(o: InternalOpId): BucketDataDocumentGenericId {
+    // `satisfies BucketDataKeyV3` checks that we use the correct type for V3 storage
+    // `as BucketDataDocumentGenericId` does a cast to get the interface virtual type
+    return {
+      b: this.key.bucket,
+      o
+    } satisfies BucketDataKeyV3 as BucketDataDocumentGenericId;
+  }
+
+  get minId(): BucketDataDocumentGenericId {
+    return {
+      b: this.key.bucket,
+      o: new mongo.MinKey()
+    } as any; // No way to properly type this
+  }
+
+  get maxId(): BucketDataDocumentGenericId {
+    return {
+      b: this.key.bucket,
+      o: new mongo.MaxKey()
+    } as any; // No way to properly type this
+  }
+
+  toPersistedDocument(source: Omit<BucketDataDoc, 'bucketKey'>): BucketDataDocumentGeneric {
+    return serializeBucketDataV3({ bucketKey: this.key, ...source }) as BucketDataDocumentGeneric;
+  }
+
+  fromPersistedDocument(doc: BucketDataDocumentGeneric): BucketDataDoc {
+    return loadBucketDataDocumentV3(this.key, doc as BucketDataDocumentV3);
+  }
+
+  fromPartialPersistedDocument<T extends keyof BucketDataProperties>(
+    doc: Pick<BucketDataDocumentGeneric, '_id' | T>
+  ): Pick<BucketDataDoc, 'bucketKey' | 'o' | T> {
+    const document = doc as Pick<BucketDataDocumentV3, '_id' | T>;
+    const { _id, ...rest } = document;
+    return {
+      bucketKey: this.key,
+      o: _id.o,
+      ...rest
+    } as Pick<BucketDataDoc, 'bucketKey' | 'o' | T>;
+  }
+}

package/src/storage/implementation/v3/SourceRecordStoreV3.ts

@@ -0,0 +1,226 @@
+import * as lib_mongo from '@powersync/lib-service-mongodb';
+import { mongo } from '@powersync/lib-service-mongodb';
+import { Logger } from '@powersync/lib-services-framework';
+import { storage } from '@powersync/service-core';
+import { EvaluatedParameters, EvaluatedRow } from '@powersync/service-sync-rules';
+import * as bson from 'bson';
+import { retryOnMongoMaxTimeMSExpired } from '../../../utils/util.js';
+import { BucketDefinitionMapping } from '../BucketDefinitionMapping.js';
+import { cacheKey } from '../OperationBatch.js';
+import { LoadedSourceRecord, SourceRecordLookupEntry, SourceRecordStore } from '../common/SourceRecordStore.js';
+import { serializeParameterLookupV3 } from './MongoParameterLookupV3.js';
+import { VersionedPowerSyncMongoV3 } from './VersionedPowerSyncMongoV3.js';
+import { CurrentDataDocumentV3, SourceTableDocumentV3 } from './models.js';
+
+export class SourceRecordStoreV3 implements SourceRecordStore {
+  constructor(
+    private readonly db: VersionedPowerSyncMongoV3,
+    private readonly groupId: number,
+    private readonly mapping: BucketDefinitionMapping
+  ) {}
+
+  mapEvaluatedBuckets(evaluated: EvaluatedRow[]): LoadedSourceRecord['buckets'] {
+    return evaluated.map((entry) => ({
+      definitionId: this.mapping.bucketSourceId(entry.source),
+      bucket: entry.bucket,
+      table: entry.table,
+      id: entry.id
+    }));
+  }
+
+  mapParameterLookups(paramEvaluated: EvaluatedParameters[]): LoadedSourceRecord['lookups'] {
+    return paramEvaluated.map((entry) => ({
+      indexId: this.mapping.parameterLookupId(entry.lookup.source),
+      lookup: serializeParameterLookupV3(entry.lookup)
+    }));
+  }
+
+  private createLoadedDocument(
+    sourceTableId: bson.ObjectId,
+    id: storage.ReplicaId,
+    data: bson.Binary | null,
+    buckets: CurrentDataDocumentV3['buckets'],
+    lookups: CurrentDataDocumentV3['lookups']
+  ): LoadedSourceRecord {
+    return {
+      sourceTableId,
+      replicaId: id,
+      data,
+      buckets: buckets.map((bucket) => ({
+        definitionId: bucket.def,
+        bucket: bucket.bucket,
+        table: bucket.table,
+        id: bucket.id
+      })),
+      lookups: lookups.map((lookup) => ({
+        indexId: lookup.i,
+        lookup: lookup.l
+      })),
+      cacheKey: cacheKey(sourceTableId, id)
+    };
+  }
+
+  async loadSizes(session: mongo.ClientSession, entries: SourceRecordLookupEntry[]): Promise<Map<string, number>> {
+    const sizes = new Map<string, number>();
+    for (const [sourceTableId, replicaIds] of this.groupEntries(entries)) {
+      const filter = {
+        _id: { $in: replicaIds as any[] }
+      } as unknown as mongo.Filter<CurrentDataDocumentV3>;
+      const sizeCursor: mongo.AggregationCursor<CurrentDataDocumentV3 & { size: number }> = this.db
+        .sourceRecordsV3(this.groupId, sourceTableId)
+        .aggregate(
+          [
+            {
+              $match: filter
+            },
+            {
+              $project: {
+                _id: 1,
+                size: { $bsonSize: '$$ROOT' }
+              }
+            }
+          ],
+          { session }
+        );
+      for await (const doc of sizeCursor.stream()) {
+        sizes.set(cacheKey(sourceTableId, doc._id), doc.size);
+      }
+    }
+    return sizes;
+  }
+
+  async loadDocuments(
+    session: mongo.ClientSession,
+    entries: SourceRecordLookupEntry[],
+    idsOnly: boolean
+  ): Promise<Map<string, LoadedSourceRecord>> {
+    const documents = new Map<string, LoadedSourceRecord>();
+    const projection = idsOnly ? { _id: 1 } : undefined;
+    for (const [sourceTableId, replicaIds] of this.groupEntries(entries)) {
+      const filter = {
+        _id: { $in: replicaIds as any[] }
+      } as unknown as mongo.Filter<CurrentDataDocumentV3>;
+      const cursor = this.db.sourceRecordsV3(this.groupId, sourceTableId).find(filter, { session, projection });
+      for await (const doc of cursor.stream()) {
+        const loaded = this.createLoadedDocument(
+          sourceTableId,
+          doc._id,
+          idsOnly ? null : doc.data,
+          idsOnly ? [] : doc.buckets,
+          idsOnly ? [] : doc.lookups
+        );
+        documents.set(loaded.cacheKey, loaded);
+      }
+    }
+    return documents;
+  }
+
+  async loadTruncateBatch(
+    session: mongo.ClientSession,
+    sourceTableId: bson.ObjectId,
+    limit: number
+  ): Promise<LoadedSourceRecord[]> {
+    const cursor = this.db.sourceRecordsV3(this.groupId, sourceTableId).find(
+      {
+        pending_delete: { $exists: false }
+      },
+      {
+        projection: {
+          _id: 1,
+          buckets: 1,
+          lookups: 1
+        },
+        limit,
+        session
+      }
+    );
+    return (await cursor.toArray()).map((doc) =>
+      this.createLoadedDocument(sourceTableId, doc._id, null, doc.buckets, doc.lookups)
+    );
+  }
+
+  async postCommitCleanup(lastCheckpoint: bigint, logger: Logger): Promise<void> {
+    // This cleans up soft deletes in source_records collections.
+    // Since there may be a lot (100+) of these collections in some cases, we track which
+    // ones have dirty deletes in source_tables.
+
+    const dirtySourceTables = await this.db
+      .sourceTablesV3(this.groupId)
+      .find(
+        {
+          latest_pending_delete: { $exists: true }
+        },
+        {
+          projection: { _id: 1, latest_pending_delete: 1 }
+        }
+      )
+      .toArray();
+
+    let deletedCount = 0;
+    const sourceTableUpdates: mongo.AnyBulkWriteOperation<SourceTableDocumentV3>[] = [];
+    for (const sourceTable of dirtySourceTables) {
+      const collection = this.db.sourceRecordsV3(this.groupId, sourceTable._id);
+      const result = await this.deletePendingDeletes(collection, sourceTable._id, lastCheckpoint, logger);
+      deletedCount += result.deletedCount;
+
+      if (sourceTable.latest_pending_delete != null && sourceTable.latest_pending_delete <= lastCheckpoint) {
+        sourceTableUpdates.push({
+          updateOne: {
+            filter: {
+              _id: sourceTable._id,
+              // If the source table received more writes in the meantime, this will filter it out
+              latest_pending_delete: sourceTable.latest_pending_delete
+            },
+            update: {
+              $unset: {
+                latest_pending_delete: 1
+              }
+            }
+          }
+        });
+      }
+    }
+
+    if (sourceTableUpdates.length > 0) {
+      await this.db.sourceTablesV3(this.groupId).bulkWrite(sourceTableUpdates, { ordered: false });
+    }
+    if (deletedCount > 0) {
+      logger.info(`Cleaned up ${deletedCount} pending delete current_data records for checkpoint ${lastCheckpoint}`);
+    }
+  }
+
+  private async deletePendingDeletes(
+    collection: mongo.Collection<CurrentDataDocumentV3>,
+    sourceTableId: bson.ObjectId,
+    lastCheckpoint: bigint,
+    logger: Logger
+  ) {
+    return retryOnMongoMaxTimeMSExpired(
+      () =>
+        collection.deleteMany(
+          {
+            pending_delete: { $exists: true, $lte: lastCheckpoint }
+          },
+          {
+            maxTimeMS: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS
+          }
+        ),
+      {
+        retryDelayMs: lib_mongo.db.MONGO_OPERATION_TIMEOUT_MS / 5,
+        onRetry: (n: number) => {
+          logger.warn(`Cleared batch ${n} of pending deletes for source table ${sourceTableId}, continuing...`);
+        }
+      }
+    );
+  }
+
+  private groupEntries(entries: SourceRecordLookupEntry[]): Map<bson.ObjectId, storage.ReplicaId[]> {
+    const grouped = new Map<bson.ObjectId, storage.ReplicaId[]>();
+    for (const entry of entries) {
+      const existing = grouped.get(entry.sourceTableId) ?? [];
+      existing.push(entry.replicaId);
+      grouped.set(entry.sourceTableId, existing);
+    }
+    return grouped;
+  }
+}