@powersync/service-core 0.0.0-dev-20240620165206

package/src/storage/mongo/MongoIdSequence.ts

@@ -0,0 +1,24 @@
+/**
+ * Manages op_id or similar sequence in memory.
+ *
+ * This is typically used within a transaction, with the last value persisted
+ * at the end of the transaction.
+ */
+export class MongoIdSequence {
+  private _last: bigint;
+
+  constructor(last: bigint) {
+    if (typeof last != 'bigint') {
+      throw new Error(`BigInt required, got ${last} ${typeof last}`);
+    }
+    this._last = last;
+  }
+
+  next() {
+    return ++this._last;
+  }
+
+  last() {
+    return this._last;
+  }
+}

package/src/storage/mongo/MongoPersistedSyncRules.ts

@@ -0,0 +1,16 @@
+import { SqlSyncRules } from '@powersync/service-sync-rules';
+
+import { PersistedSyncRules } from '../BucketStorage.js';
+
+export class MongoPersistedSyncRules implements PersistedSyncRules {
+  public readonly slot_name: string;
+
+  constructor(
+    public readonly id: number,
+    public readonly sync_rules: SqlSyncRules,
+    public readonly checkpoint_lsn: string | null,
+    slot_name: string | null
+  ) {
+    this.slot_name = slot_name ?? `powersync_${id}`;
+  }
+}

package/src/storage/mongo/MongoPersistedSyncRulesContent.ts

@@ -0,0 +1,47 @@
+import { SqlSyncRules } from '@powersync/service-sync-rules';
+import * as mongo from 'mongodb';
+
+import { PersistedSyncRulesContent } from '../BucketStorage.js';
+import { MongoPersistedSyncRules } from './MongoPersistedSyncRules.js';
+import { MongoSyncRulesLock } from './MongoSyncRulesLock.js';
+import { PowerSyncMongo } from './db.js';
+import { SyncRuleDocument } from './models.js';
+
+export class MongoPersistedSyncRulesContent implements PersistedSyncRulesContent {
+  public readonly slot_name: string;
+
+  public readonly id: number;
+  public readonly sync_rules_content: string;
+  public readonly last_checkpoint_lsn: string | null;
+  public readonly last_fatal_error: string | null;
+  public readonly last_keepalive_ts: Date | null;
+  public readonly last_checkpoint_ts: Date | null;
+
+  public current_lock: MongoSyncRulesLock | null = null;
+
+  constructor(private db: PowerSyncMongo, doc: mongo.WithId<SyncRuleDocument>) {
+    this.id = doc._id;
+    this.sync_rules_content = doc.content;
+    this.last_checkpoint_lsn = doc.last_checkpoint_lsn;
+    // Handle legacy values
+    this.slot_name = doc.slot_name ?? `powersync_${this.id}`;
+    this.last_fatal_error = doc.last_fatal_error;
+    this.last_checkpoint_ts = doc.last_checkpoint_ts;
+    this.last_keepalive_ts = doc.last_keepalive_ts;
+  }
+
+  parsed() {
+    return new MongoPersistedSyncRules(
+      this.id,
+      SqlSyncRules.fromYaml(this.sync_rules_content),
+      this.last_checkpoint_lsn,
+      this.slot_name
+    );
+  }
+
+  async lock() {
+    const lock = await MongoSyncRulesLock.createLock(this.db, this);
+    this.current_lock = lock;
+    return lock;
+  }
+}

package/src/storage/mongo/MongoSyncBucketStorage.ts

@@ -0,0 +1,533 @@
+import { SqliteJsonRow, SqliteJsonValue, SqlSyncRules } from '@powersync/service-sync-rules';
+import * as bson from 'bson';
+import * as mongo from 'mongodb';
+
+import * as db from '../../db/db-index.js';
+import * as replication from '../../replication/WalStream.js';
+import * as util from '../../util/util-index.js';
+import {
+  BucketDataBatchOptions,
+  BucketStorageBatch,
+  DEFAULT_DOCUMENT_BATCH_LIMIT,
+  DEFAULT_DOCUMENT_CHUNK_LIMIT_BYTES,
+  FlushedResult,
+  ResolveTableOptions,
+  ResolveTableResult,
+  SyncRulesBucketStorage,
+  SyncRuleStatus
+} from '../BucketStorage.js';
+import { MongoBucketStorage } from '../MongoBucketStorage.js';
+import { SourceTable } from '../SourceTable.js';
+import { PowerSyncMongo } from './db.js';
+import { BucketDataDocument, BucketDataKey, SourceKey, SyncRuleState } from './models.js';
+import { MongoBucketBatch } from './MongoBucketBatch.js';
+import { BSON_DESERIALIZE_OPTIONS, idPrefixFilter, readSingleBatch, serializeLookup } from './util.js';
+import { ChecksumCache, FetchPartialBucketChecksum } from '../ChecksumCache.js';
+
+export class MongoSyncBucketStorage implements SyncRulesBucketStorage {
+  private readonly db: PowerSyncMongo;
+  private checksumCache = new ChecksumCache({
+    fetchChecksums: (batch) => {
+      return this.getChecksumsInternal(batch);
+    }
+  });
+
+  constructor(
+    public readonly factory: MongoBucketStorage,
+    public readonly group_id: number,
+    public readonly sync_rules: SqlSyncRules,
+    public readonly slot_name: string
+  ) {
+    this.db = factory.db;
+  }
+
+  async getCheckpoint() {
+    const doc = await this.db.sync_rules.findOne(
+      { _id: this.group_id },
+      {
+        projection: { last_checkpoint: 1, last_checkpoint_lsn: 1 }
+      }
+    );
+    return {
+      checkpoint: util.timestampToOpId(doc?.last_checkpoint ?? 0n),
+      lsn: doc?.last_checkpoint_lsn ?? replication.ZERO_LSN
+    };
+  }
+
+  async startBatch(options: {}, callback: (batch: BucketStorageBatch) => Promise<void>): Promise<FlushedResult | null> {
+    const doc = await this.db.sync_rules.findOne(
+      {
+        _id: this.group_id
+      },
+      { projection: { last_checkpoint_lsn: 1, no_checkpoint_before: 1 } }
+    );
+    const checkpoint_lsn = doc?.last_checkpoint_lsn ?? null;
+
+    const batch = new MongoBucketBatch(
+      this.db,
+      this.sync_rules,
+      this.group_id,
+      this.slot_name,
+      checkpoint_lsn,
+      doc?.no_checkpoint_before ?? null
+    );
+    try {
+      await callback(batch);
+      await batch.flush();
+      await batch.abort();
+      if (batch.last_flushed_op) {
+        return { flushed_op: String(batch.last_flushed_op) };
+      } else {
+        return null;
+      }
+    } catch (e) {
+      await batch.abort();
+      throw e;
+    }
+  }
+
+  async resolveTable(options: ResolveTableOptions): Promise<ResolveTableResult> {
+    const { group_id, connection_id, connection_tag, relation } = options;
+
+    const { schema, name: table, relationId, replicationColumns } = relation;
+
+    const columns = replicationColumns.map((column) => ({ name: column.name, type_oid: column.typeOid }));
+    let result: ResolveTableResult | null = null;
+    await this.db.client.withSession(async (session) => {
+      const col = this.db.source_tables;
+      let doc = await col.findOne(
+        {
+          group_id: group_id,
+          connection_id: connection_id,
+          relation_id: relationId,
+          schema_name: schema,
+          table_name: table,
+          replica_id_columns2: columns
+        },
+        { session }
+      );
+      if (doc == null) {
+        doc = {
+          _id: new bson.ObjectId(),
+          group_id: group_id,
+          connection_id: connection_id,
+          relation_id: relationId,
+          schema_name: schema,
+          table_name: table,
+          replica_id_columns: null,
+          replica_id_columns2: columns,
+          snapshot_done: false
+        };
+
+        await col.insertOne(doc, { session });
+      }
+      const sourceTable = new SourceTable(
+        doc._id,
+        connection_tag,
+        relationId,
+        schema,
+        table,
+        replicationColumns,
+        doc.snapshot_done ?? true
+      );
+      sourceTable.syncData = options.sync_rules.tableSyncsData(sourceTable);
+      sourceTable.syncParameters = options.sync_rules.tableSyncsParameters(sourceTable);
+
+      const truncate = await col
+        .find(
+          {
+            group_id: group_id,
+            connection_id: connection_id,
+            _id: { $ne: doc._id },
+            $or: [{ relation_id: relationId }, { schema_name: schema, table_name: table }]
+          },
+          { session }
+        )
+        .toArray();
+      result = {
+        table: sourceTable,
+        dropTables: truncate.map(
+          (doc) =>
+            new SourceTable(
+              doc._id,
+              connection_tag,
+              doc.relation_id ?? 0,
+              doc.schema_name,
+              doc.table_name,
+              doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid })) ?? [],
+              doc.snapshot_done ?? true
+            )
+        )
+      };
+    });
+    return result!;
+  }
+
+  async getParameterSets(checkpoint: util.OpId, lookups: SqliteJsonValue[][]): Promise<SqliteJsonRow[]> {
+    const lookupFilter = lookups.map((lookup) => {
+      return serializeLookup(lookup);
+    });
+    const rows = await this.db.bucket_parameters
+      .aggregate([
+        {
+          $match: {
+            'key.g': this.group_id,
+            lookup: { $in: lookupFilter },
+            _id: { $lte: BigInt(checkpoint) }
+          }
+        },
+        {
+          $sort: {
+            _id: -1
+          }
+        },
+        {
+          $group: {
+            _id: '$key',
+            bucket_parameters: {
+              $first: '$bucket_parameters'
+            }
+          }
+        }
+      ])
+      .toArray();
+    const groupedParameters = rows.map((row) => {
+      return row.bucket_parameters;
+    });
+    return groupedParameters.flat();
+  }
+
+  async *getBucketDataBatch(
+    checkpoint: util.OpId,
+    dataBuckets: Map<string, string>,
+    options?: BucketDataBatchOptions
+  ): AsyncIterable<util.SyncBucketData> {
+    if (dataBuckets.size == 0) {
+      return;
+    }
+    let filters: mongo.Filter<BucketDataDocument>[] = [];
+
+    const end = checkpoint ? BigInt(checkpoint) : new bson.MaxKey();
+    for (let [name, start] of dataBuckets.entries()) {
+      filters.push({
+        _id: {
+          $gt: {
+            g: this.group_id,
+            b: name,
+            o: BigInt(start)
+          },
+          $lte: {
+            g: this.group_id,
+            b: name,
+            o: end as any
+          }
+        }
+      });
+    }
+
+    const limit = options?.limit ?? DEFAULT_DOCUMENT_BATCH_LIMIT;
+    const sizeLimit = options?.chunkLimitBytes ?? DEFAULT_DOCUMENT_CHUNK_LIMIT_BYTES;
+
+    const cursor = this.db.bucket_data.find(
+      {
+        $or: filters
+      },
+      {
+        session: undefined,
+        sort: { _id: 1 },
+        limit: limit,
+        // Increase batch size above the default 101, so that we can fill an entire batch in
+        // one go.
+        batchSize: limit,
+        // Raw mode is returns an array of Buffer instead of parsed documents.
+        // We use it so that:
+        // 1. We can calculate the document size accurately without serializing again.
+        // 2. We can delay parsing the results until it's needed.
+        // We manually use bson.deserialize below
+        raw: true,
+
+        // Since we're using raw: true and parsing ourselves later, we don't need bigint
+        // support here.
+        // Disabling due to https://jira.mongodb.org/browse/NODE-6165, and the fact that this
+        // is one of our most common queries.
+        useBigInt64: false
+      }
+    ) as unknown as mongo.FindCursor<Buffer>;
+
+    // We want to limit results to a single batch to avoid high memory usage.
+    // This approach uses MongoDB's batch limits to limit the data here, which limits
+    // to the lower of the batch count and size limits.
+    // This is similar to using `singleBatch: true` in the find options, but allows
+    // detecting "hasMore".
+    let { data, hasMore } = await readSingleBatch(cursor);
+    if (data.length == limit) {
+      // Limit reached - could have more data, despite the cursor being drained.
+      hasMore = true;
+    }
+
+    let batchSize = 0;
+    let currentBatch: util.SyncBucketData | null = null;
+
+    // Ordered by _id, meaning buckets are grouped together
+    for (let rawData of data) {
+      const row = bson.deserialize(rawData, BSON_DESERIALIZE_OPTIONS) as BucketDataDocument;
+      const bucket = row._id.b;
+
+      if (currentBatch == null || currentBatch.bucket != bucket || batchSize >= sizeLimit) {
+        let start: string | undefined = undefined;
+        if (currentBatch != null) {
+          if (currentBatch.bucket == bucket) {
+            currentBatch.has_more = true;
+          }
+
+          const yieldBatch = currentBatch;
+          start = currentBatch.after;
+          currentBatch = null;
+          batchSize = 0;
+          yield yieldBatch;
+        }
+
+        start ??= dataBuckets.get(bucket);
+        if (start == null) {
+          throw new Error(`data for unexpected bucket: ${bucket}`);
+        }
+        currentBatch = {
+          bucket,
+          after: start,
+          has_more: hasMore,
+          data: [],
+          next_after: start
+        };
+      }
+
+      const entry: util.OplogEntry = {
+        op_id: util.timestampToOpId(row._id.o),
+        op: row.op,
+        object_type: row.table,
+        object_id: row.row_id,
+        checksum: Number(row.checksum),
+        subkey: `${row.source_table}/${row.source_key.toHexString()}`,
+        data: row.data
+      };
+      currentBatch.data.push(entry);
+      currentBatch.next_after = entry.op_id;
+
+      batchSize += rawData.byteLength;
+    }
+
+    if (currentBatch != null) {
+      const yieldBatch = currentBatch;
+      currentBatch = null;
+      yield yieldBatch;
+    }
+  }
+
+  async getChecksums(checkpoint: util.OpId, buckets: string[]): Promise<util.ChecksumMap> {
+    return this.checksumCache.getChecksumMap(checkpoint, buckets);
+  }
+
+  private async getChecksumsInternal(batch: FetchPartialBucketChecksum[]): Promise<util.ChecksumMap> {
+    if (batch.length == 0) {
+      return new Map();
+    }
+
+    const filters: any[] = [];
+    for (let request of batch) {
+      filters.push({
+        _id: {
+          $gt: {
+            g: this.group_id,
+            b: request.bucket,
+            o: request.start ? BigInt(request.start) : new bson.MinKey()
+          },
+          $lte: {
+            g: this.group_id,
+            b: request.bucket,
+            o: BigInt(request.end)
+          }
+        }
+      });
+    }
+
+    const aggregate = await this.db.bucket_data
+      .aggregate(
+        [
+          {
+            $match: {
+              $or: filters
+            }
+          },
+          {
+            $group: { _id: '$_id.b', checksum_total: { $sum: '$checksum' }, count: { $sum: 1 } }
+          }
+        ],
+        { session: undefined }
+      )
+      .toArray();
+
+    return new Map<string, util.BucketChecksum>(
+      aggregate.map((doc) => {
+        return [
+          doc._id,
+          {
+            bucket: doc._id,
+            count: doc.count,
+            checksum: Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff
+          } satisfies util.BucketChecksum
+        ];
+      })
+    );
+  }
+
+  async terminate() {
+    await this.clear();
+
+    await this.db.sync_rules.updateOne(
+      {
+        _id: this.group_id
+      },
+      {
+        $set: {
+          state: SyncRuleState.TERMINATED,
+          persisted_lsn: null,
+          snapshot_done: false
+        }
+      }
+    );
+  }
+
+  async getStatus(): Promise<SyncRuleStatus> {
+    const doc = await this.db.sync_rules.findOne(
+      {
+        _id: this.group_id
+      },
+      {
+        projection: {
+          snapshot_done: 1,
+          last_checkpoint_lsn: 1,
+          state: 1
+        }
+      }
+    );
+    if (doc == null) {
+      throw new Error('Cannot find sync rules status');
+    }
+
+    return {
+      snapshot_done: doc.snapshot_done,
+      active: doc.state == 'ACTIVE',
+      checkpoint_lsn: doc.last_checkpoint_lsn
+    };
+  }
+
+  async clear(): Promise<void> {
+    // Individual operations here may time out with the maxTimeMS option.
+    // It is expected to still make progress, and continue on the next try.
+
+    // TODO: Transactional?
+    await this.db.sync_rules.updateOne(
+      {
+        _id: this.group_id
+      },
+      {
+        $set: {
+          snapshot_done: false,
+          persisted_lsn: null,
+          last_checkpoint_lsn: null,
+          last_checkpoint: null,
+          no_checkpoint_before: null
+        }
+      },
+      { maxTimeMS: db.mongo.MONGO_OPERATION_TIMEOUT_MS }
+    );
+    await this.db.bucket_data.deleteMany(
+      {
+        _id: idPrefixFilter<BucketDataKey>({ g: this.group_id }, ['b', 'o'])
+      },
+      { maxTimeMS: db.mongo.MONGO_OPERATION_TIMEOUT_MS }
+    );
+    await this.db.bucket_parameters.deleteMany(
+      {
+        key: idPrefixFilter<SourceKey>({ g: this.group_id }, ['t', 'k'])
+      },
+      { maxTimeMS: db.mongo.MONGO_OPERATION_TIMEOUT_MS }
+    );
+
+    await this.db.current_data.deleteMany(
+      {
+        _id: idPrefixFilter<SourceKey>({ g: this.group_id }, ['t', 'k'])
+      },
+      { maxTimeMS: db.mongo.MONGO_OPERATION_TIMEOUT_MS }
+    );
+
+    await this.db.source_tables.deleteMany(
+      {
+        group_id: this.group_id
+      },
+      { maxTimeMS: db.mongo.MONGO_OPERATION_TIMEOUT_MS }
+    );
+  }
+
+  async setSnapshotDone(lsn: string): Promise<void> {
+    await this.db.sync_rules.updateOne(
+      {
+        _id: this.group_id
+      },
+      {
+        $set: {
+          snapshot_done: true,
+          persisted_lsn: lsn,
+          last_checkpoint_ts: new Date()
+        }
+      }
+    );
+  }
+
+  async autoActivate(): Promise<void> {
+    await this.db.client.withSession(async (session) => {
+      await session.withTransaction(async () => {
+        const doc = await this.db.sync_rules.findOne({ _id: this.group_id }, { session });
+        if (doc && doc.state == 'PROCESSING') {
+          await this.db.sync_rules.updateOne(
+            {
+              _id: this.group_id
+            },
+            {
+              $set: {
+                state: SyncRuleState.ACTIVE
+              }
+            },
+            { session }
+          );
+
+          await this.db.sync_rules.updateMany(
+            {
+              _id: { $ne: this.group_id },
+              state: SyncRuleState.ACTIVE
+            },
+            {
+              $set: {
+                state: SyncRuleState.STOP
+              }
+            },
+            { session }
+          );
+        }
+      });
+    });
+  }
+
+  async reportError(e: any): Promise<void> {
+    const message = String(e.message ?? 'Replication failure');
+    await this.db.sync_rules.updateOne(
+      {
+        _id: this.group_id
+      },
+      {
+        $set: {
+          last_fatal_error: message
+        }
+      }
+    );
+  }
+}

package/src/storage/mongo/MongoSyncRulesLock.ts

@@ -0,0 +1,81 @@
+import crypto from 'crypto';
+
+import { PersistedSyncRulesContent, ReplicationLock } from '../BucketStorage.js';
+import { PowerSyncMongo } from './db.js';
+import { logger } from '@powersync/lib-services-framework';
+
+/**
+ * Manages a lock on a sync rules document, so that only one process
+ * replicates those sync rules at a time.
+ */
+export class MongoSyncRulesLock implements ReplicationLock {
+  private readonly refreshInterval: NodeJS.Timer;
+
+  static async createLock(db: PowerSyncMongo, sync_rules: PersistedSyncRulesContent): Promise<MongoSyncRulesLock> {
+    const lockId = crypto.randomBytes(8).toString('hex');
+    const doc = await db.sync_rules.findOneAndUpdate(
+      { _id: sync_rules.id, $or: [{ lock: null }, { 'lock.expires_at': { $lt: new Date() } }] },
+      {
+        $set: {
+          lock: {
+            id: lockId,
+            expires_at: new Date(Date.now() + 60 * 1000)
+          }
+        }
+      },
+      {
+        projection: { lock: 1 },
+        returnDocument: 'before'
+      }
+    );
+
+    if (doc == null) {
+      throw new Error(`Replication slot ${sync_rules.slot_name} is locked by another process`);
+    }
+    return new MongoSyncRulesLock(db, sync_rules.id, lockId);
+  }
+
+  constructor(private db: PowerSyncMongo, public sync_rules_id: number, private lock_id: string) {
+    this.refreshInterval = setInterval(async () => {
+      try {
+        await this.refresh();
+      } catch (e) {
+        logger.error('Failed to refresh lock', e);
+        clearInterval(this.refreshInterval);
+      }
+    }, 30_130);
+  }
+
+  async release(): Promise<void> {
+    clearInterval(this.refreshInterval);
+    const result = await this.db.sync_rules.updateOne(
+      {
+        _id: this.sync_rules_id,
+        'lock.id': this.lock_id
+      },
+      {
+        $unset: { lock: 1 }
+      }
+    );
+    if (result.modifiedCount == 0) {
+      // Log and ignore
+      logger.warn(`Lock already released: ${this.sync_rules_id}/${this.lock_id}`);
+    }
+  }
+
+  private async refresh(): Promise<void> {
+    const result = await this.db.sync_rules.findOneAndUpdate(
+      {
+        _id: this.sync_rules_id,
+        'lock.id': this.lock_id
+      },
+      {
+        $set: { 'lock.expires_at': new Date(Date.now() + 60 * 1000) }
+      },
+      { returnDocument: 'after' }
+    );
+    if (result == null) {
+      throw new Error(`Lock not held anymore: ${this.sync_rules_id}/${this.lock_id}`);
+    }
+  }
+}