@powersync/service-module-mongodb 0.0.0-dev-20250102111825 → 0.0.0-dev-20250108073049

package/dist/storage/implementation/MongoBucketBatch.js

@@ -1,681 +0,0 @@
-import * as bson from 'bson';
-import * as mongo from 'mongodb';
-import { container, DisposableObserver, errors, logger } from '@powersync/lib-services-framework';
-import { SaveOperationTag, storage, utils } from '@powersync/service-core';
-import * as timers from 'node:timers/promises';
-import { MongoIdSequence } from './MongoIdSequence.js';
-import { batchCreateCustomWriteCheckpoints } from './MongoWriteCheckpointAPI.js';
-import { cacheKey, OperationBatch, RecordOperation } from './OperationBatch.js';
-import { PersistedBatch } from './PersistedBatch.js';
-import { idPrefixFilter } from './util.js';
-/**
- * 15MB
- */
-const MAX_ROW_SIZE = 15 * 1024 * 1024;
-// Currently, we can only have a single flush() at a time, since it locks the op_id sequence.
-// While the MongoDB transaction retry mechanism handles this okay, using an in-process Mutex
-// makes it more fair and has less overhead.
-//
-// In the future, we can investigate allowing multiple replication streams operating independently.
-const replicationMutex = new utils.Mutex();
-export class MongoBucketBatch extends DisposableObserver {
-    constructor(options) {
-        super();
-        this.batch = null;
-        this.write_checkpoint_batch = [];
-        /**
-         * Last LSN received associated with a checkpoint.
-         *
-         * This could be either:
-         * 1. A commit LSN.
-         * 2. A keepalive message LSN.
-         */
-        this.last_checkpoint_lsn = null;
-        this.persisted_op = null;
-        /**
-         * For tests only - not for persistence logic.
-         */
-        this.last_flushed_op = null;
-        this.lastWaitingLogThottled = 0;
-        this.client = options.db.client;
-        this.db = options.db;
-        this.group_id = options.groupId;
-        this.last_checkpoint_lsn = options.lastCheckpointLsn;
-        this.no_checkpoint_before_lsn = options.noCheckpointBeforeLsn;
-        this.session = this.client.startSession();
-        this.slot_name = options.slotName;
-        this.sync_rules = options.syncRules;
-        this.storeCurrentData = options.storeCurrentData;
-        this.skipExistingRows = options.skipExistingRows;
-        this.batch = new OperationBatch();
-        if (options.keepaliveOp) {
-            this.persisted_op = BigInt(options.keepaliveOp);
-        }
-    }
-    addCustomWriteCheckpoint(checkpoint) {
-        this.write_checkpoint_batch.push({
-            ...checkpoint,
-            sync_rules_id: this.group_id
-        });
-    }
-    get lastCheckpointLsn() {
-        return this.last_checkpoint_lsn;
-    }
-    async flush() {
-        let result = null;
-        // One flush may be split over multiple transactions.
-        // Each flushInner() is one transaction.
-        while (this.batch != null) {
-            let r = await this.flushInner();
-            if (r) {
-                result = r;
-            }
-        }
-        await batchCreateCustomWriteCheckpoints(this.db, this.write_checkpoint_batch);
-        this.write_checkpoint_batch = [];
-        return result;
-    }
-    async flushInner() {
-        const batch = this.batch;
-        if (batch == null) {
-            return null;
-        }
-        let last_op = null;
-        let resumeBatch = null;
-        await this.withReplicationTransaction(`Flushing ${batch.length} ops`, async (session, opSeq) => {
-            resumeBatch = await this.replicateBatch(session, batch, opSeq);
-            last_op = opSeq.last();
-        });
-        // null if done, set if we need another flush
-        this.batch = resumeBatch;
-        if (last_op == null) {
-            throw new Error('Unexpected last_op == null');
-        }
-        this.persisted_op = last_op;
-        this.last_flushed_op = last_op;
-        return { flushed_op: String(last_op) };
-    }
-    async replicateBatch(session, batch, op_seq) {
-        let sizes = undefined;
-        if (this.storeCurrentData && !this.skipExistingRows) {
-            // We skip this step if we don't store current_data, since the sizes will
-            // always be small in that case.
-            // With skipExistingRows, we don't load the full documents into memory,
-            // so we can also skip the size lookup step.
-            // Find sizes of current_data documents, to assist in intelligent batching without
-            // exceeding memory limits.
-            //
-            // A previous attempt tried to do batching by the results of the current_data query
-            // (automatically limited to 48MB(?) per batch by MongoDB). The issue is that it changes
-            // the order of processing, which then becomes really tricky to manage.
-            // This now takes 2+ queries, but doesn't have any issues with order of operations.
-            const sizeLookups = batch.batch.map((r) => {
-                return { g: this.group_id, t: r.record.sourceTable.id, k: r.beforeId };
-            });
-            sizes = new Map();
-            const sizeCursor = this.db.current_data.aggregate([
-                {
-                    $match: {
-                        _id: { $in: sizeLookups }
-                    }
-                },
-                {
-                    $project: {
-                        _id: 1,
-                        size: { $bsonSize: '$$ROOT' }
-                    }
-                }
-            ], { session });
-            for await (let doc of sizeCursor.stream()) {
-                const key = cacheKey(doc._id.t, doc._id.k);
-                sizes.set(key, doc.size);
-            }
-        }
-        // If set, we need to start a new transaction with this batch.
-        let resumeBatch = null;
-        let transactionSize = 0;
-        // Now batch according to the sizes
-        // This is a single batch if storeCurrentData == false
-        for await (let b of batch.batched(sizes)) {
-            if (resumeBatch) {
-                for (let op of b) {
-                    resumeBatch.push(op);
-                }
-                continue;
-            }
-            const lookups = b.map((r) => {
-                return { g: this.group_id, t: r.record.sourceTable.id, k: r.beforeId };
-            });
-            let current_data_lookup = new Map();
-            // With skipExistingRows, we only need to know whether or not the row exists.
-            const projection = this.skipExistingRows ? { _id: 1 } : undefined;
-            const cursor = this.db.current_data.find({
-                _id: { $in: lookups }
-            }, { session, projection });
-            for await (let doc of cursor.stream()) {
-                current_data_lookup.set(cacheKey(doc._id.t, doc._id.k), doc);
-            }
-            let persistedBatch = new PersistedBatch(this.group_id, transactionSize);
-            for (let op of b) {
-                if (resumeBatch) {
-                    resumeBatch.push(op);
-                    continue;
-                }
-                const currentData = current_data_lookup.get(op.internalBeforeKey) ?? null;
-                if (currentData != null) {
-                    // If it will be used again later, it will be set again using nextData below
-                    current_data_lookup.delete(op.internalBeforeKey);
-                }
-                const nextData = this.saveOperation(persistedBatch, op, currentData, op_seq);
-                if (nextData != null) {
-                    // Update our current_data and size cache
-                    current_data_lookup.set(op.internalAfterKey, nextData);
-                    sizes?.set(op.internalAfterKey, nextData.data.length());
-                }
-                if (persistedBatch.shouldFlushTransaction()) {
-                    // Transaction is getting big.
-                    // Flush, and resume in a new transaction.
-                    await persistedBatch.flush(this.db, this.session);
-                    persistedBatch = null;
-                    // Computing our current progress is a little tricky here, since
-                    // we're stopping in the middle of a batch.
-                    // We create a new batch, and push any remaining operations to it.
-                    resumeBatch = new OperationBatch();
-                }
-            }
-            if (persistedBatch) {
-                transactionSize = persistedBatch.currentSize;
-                await persistedBatch.flush(this.db, this.session);
-            }
-        }
-        return resumeBatch;
-    }
-    saveOperation(batch, operation, current_data, opSeq) {
-        const record = operation.record;
-        const beforeId = operation.beforeId;
-        const afterId = operation.afterId;
-        let after = record.after;
-        const sourceTable = record.sourceTable;
-        let existing_buckets = [];
-        let new_buckets = [];
-        let existing_lookups = [];
-        let new_lookups = [];
-        const before_key = { g: this.group_id, t: record.sourceTable.id, k: beforeId };
-        if (this.skipExistingRows) {
-            if (record.tag == SaveOperationTag.INSERT) {
-                if (current_data != null) {
-                    // Initial replication, and we already have the record.
-                    // This may be a different version of the record, but streaming replication
-                    // will take care of that.
-                    // Skip the insert here.
-                    return null;
-                }
-            }
-            else {
-                throw new Error(`${record.tag} not supported with skipExistingRows: true`);
-            }
-        }
-        if (record.tag == SaveOperationTag.UPDATE) {
-            const result = current_data;
-            if (result == null) {
-                // Not an error if we re-apply a transaction
-                existing_buckets = [];
-                existing_lookups = [];
-                // Log to help with debugging if there was a consistency issue
-                if (this.storeCurrentData) {
-                    logger.warn(`Cannot find previous record for update on ${record.sourceTable.qualifiedName}: ${beforeId} / ${record.before?.id}`);
-                }
-            }
-            else {
-                existing_buckets = result.buckets;
-                existing_lookups = result.lookups;
-                if (this.storeCurrentData) {
-                    const data = bson.deserialize(result.data.buffer, storage.BSON_DESERIALIZE_OPTIONS);
-                    after = storage.mergeToast(after, data);
-                }
-            }
-        }
-        else if (record.tag == SaveOperationTag.DELETE) {
-            const result = current_data;
-            if (result == null) {
-                // Not an error if we re-apply a transaction
-                existing_buckets = [];
-                existing_lookups = [];
-                // Log to help with debugging if there was a consistency issue
-                if (this.storeCurrentData) {
-                    logger.warn(`Cannot find previous record for delete on ${record.sourceTable.qualifiedName}: ${beforeId} / ${record.before?.id}`);
-                }
-            }
-            else {
-                existing_buckets = result.buckets;
-                existing_lookups = result.lookups;
-            }
-        }
-        let afterData;
-        if (afterId != null && !this.storeCurrentData) {
-            afterData = new bson.Binary(bson.serialize({}));
-        }
-        else if (afterId != null) {
-            try {
-                // This will fail immediately if the record is > 16MB.
-                afterData = new bson.Binary(bson.serialize(after));
-                // We additionally make sure it's <= 15MB - we need some margin for metadata.
-                if (afterData.length() > MAX_ROW_SIZE) {
-                    throw new Error(`Row too large: ${afterData.length()}`);
-                }
-            }
-            catch (e) {
-                // Replace with empty values, equivalent to TOAST values
-                after = Object.fromEntries(Object.entries(after).map(([key, value]) => {
-                    return [key, undefined];
-                }));
-                afterData = new bson.Binary(bson.serialize(after));
-                container.reporter.captureMessage(`Data too big on ${record.sourceTable.qualifiedName}.${record.after?.id}: ${e.message}`, {
-                    level: errors.ErrorSeverity.WARNING,
-                    metadata: {
-                        replication_slot: this.slot_name,
-                        table: record.sourceTable.qualifiedName
-                    }
-                });
-            }
-        }
-        // 2. Save bucket data
-        if (beforeId != null && (afterId == null || !storage.replicaIdEquals(beforeId, afterId))) {
-            // Source ID updated
-            if (sourceTable.syncData) {
-                // Delete old record
-                batch.saveBucketData({
-                    op_seq: opSeq,
-                    sourceKey: beforeId,
-                    table: sourceTable,
-                    before_buckets: existing_buckets,
-                    evaluated: []
-                });
-                // Clear this, so we don't also try to REMOVE for the new id
-                existing_buckets = [];
-            }
-            if (sourceTable.syncParameters) {
-                // Delete old parameters
-                batch.saveParameterData({
-                    op_seq: opSeq,
-                    sourceKey: beforeId,
-                    sourceTable,
-                    evaluated: [],
-                    existing_lookups
-                });
-                existing_lookups = [];
-            }
-        }
-        // If we re-apply a transaction, we can end up with a partial row.
-        //
-        // We may end up with toasted values, which means the record is not quite valid.
-        // However, it will be valid by the end of the transaction.
-        //
-        // In this case, we don't save the op, but we do save the current data.
-        if (afterId && after && utils.isCompleteRow(after)) {
-            // Insert or update
-            if (sourceTable.syncData) {
-                const { results: evaluated, errors: syncErrors } = this.sync_rules.evaluateRowWithErrors({
-                    record: after,
-                    sourceTable
-                });
-                for (let error of syncErrors) {
-                    container.reporter.captureMessage(`Failed to evaluate data query on ${record.sourceTable.qualifiedName}.${record.after?.id}: ${error.error}`, {
-                        level: errors.ErrorSeverity.WARNING,
-                        metadata: {
-                            replication_slot: this.slot_name,
-                            table: record.sourceTable.qualifiedName
-                        }
-                    });
-                    logger.error(`Failed to evaluate data query on ${record.sourceTable.qualifiedName}.${record.after?.id}: ${error.error}`);
-                }
-                // Save new one
-                batch.saveBucketData({
-                    op_seq: opSeq,
-                    sourceKey: afterId,
-                    evaluated,
-                    table: sourceTable,
-                    before_buckets: existing_buckets
-                });
-                new_buckets = evaluated.map((e) => {
-                    return {
-                        bucket: e.bucket,
-                        table: e.table,
-                        id: e.id
-                    };
-                });
-            }
-            if (sourceTable.syncParameters) {
-                // Parameters
-                const { results: paramEvaluated, errors: paramErrors } = this.sync_rules.evaluateParameterRowWithErrors(sourceTable, after);
-                for (let error of paramErrors) {
-                    container.reporter.captureMessage(`Failed to evaluate parameter query on ${record.sourceTable.qualifiedName}.${record.after?.id}: ${error.error}`, {
-                        level: errors.ErrorSeverity.WARNING,
-                        metadata: {
-                            replication_slot: this.slot_name,
-                            table: record.sourceTable.qualifiedName
-                        }
-                    });
-                    logger.error(`Failed to evaluate parameter query on ${record.sourceTable.qualifiedName}.${after.id}: ${error.error}`);
-                }
-                batch.saveParameterData({
-                    op_seq: opSeq,
-                    sourceKey: afterId,
-                    sourceTable,
-                    evaluated: paramEvaluated,
-                    existing_lookups
-                });
-                new_lookups = paramEvaluated.map((p) => {
-                    return storage.serializeLookup(p.lookup);
-                });
-            }
-        }
-        let result = null;
-        // 5. TOAST: Update current data and bucket list.
-        if (afterId) {
-            // Insert or update
-            const after_key = { g: this.group_id, t: sourceTable.id, k: afterId };
-            batch.upsertCurrentData(after_key, {
-                data: afterData,
-                buckets: new_buckets,
-                lookups: new_lookups
-            });
-            result = {
-                _id: after_key,
-                data: afterData,
-                buckets: new_buckets,
-                lookups: new_lookups
-            };
-        }
-        if (afterId == null || !storage.replicaIdEquals(beforeId, afterId)) {
-            // Either a delete (afterId == null), or replaced the old replication id
-            batch.deleteCurrentData(before_key);
-        }
-        return result;
-    }
-    async withTransaction(cb) {
-        await replicationMutex.exclusiveLock(async () => {
-            await this.session.withTransaction(async () => {
-                try {
-                    await cb();
-                }
-                catch (e) {
-                    if (e instanceof mongo.MongoError && e.hasErrorLabel('TransientTransactionError')) {
-                        // Likely write conflict caused by concurrent write stream replicating
-                    }
-                    else {
-                        logger.warn('Transaction error', e);
-                    }
-                    await timers.setTimeout(Math.random() * 50);
-                    throw e;
-                }
-            }, { maxCommitTimeMS: 10000 });
-        });
-    }
-    async withReplicationTransaction(description, callback) {
-        let flushTry = 0;
-        const start = Date.now();
-        const lastTry = start + 90000;
-        const session = this.session;
-        await this.withTransaction(async () => {
-            flushTry += 1;
-            if (flushTry % 10 == 0) {
-                logger.info(`${this.slot_name} ${description} - try ${flushTry}`);
-            }
-            if (flushTry > 20 && Date.now() > lastTry) {
-                throw new Error('Max transaction tries exceeded');
-            }
-            const next_op_id_doc = await this.db.op_id_sequence.findOneAndUpdate({
-                _id: 'main'
-            }, {
-                $setOnInsert: { op_id: 0n },
-                $set: {
-                    // Force update to ensure we get a mongo lock
-                    ts: Date.now()
-                }
-            }, {
-                upsert: true,
-                returnDocument: 'after',
-                session
-            });
-            const opSeq = new MongoIdSequence(next_op_id_doc?.op_id ?? 0n);
-            await callback(session, opSeq);
-            await this.db.op_id_sequence.updateOne({
-                _id: 'main'
-            }, {
-                $set: {
-                    op_id: opSeq.last()
-                }
-            }, {
-                session
-            });
-            await this.db.sync_rules.updateOne({
-                _id: this.group_id
-            }, {
-                $set: {
-                    last_keepalive_ts: new Date()
-                }
-            }, { session });
-        });
-    }
-    async [Symbol.asyncDispose]() {
-        await this.session.endSession();
-        super[Symbol.dispose]();
-    }
-    async commit(lsn) {
-        await this.flush();
-        if (this.last_checkpoint_lsn != null && lsn < this.last_checkpoint_lsn) {
-            // When re-applying transactions, don't create a new checkpoint until
-            // we are past the last transaction.
-            logger.info(`Re-applied transaction ${lsn} - skipping checkpoint`);
-            return false;
-        }
-        if (lsn < this.no_checkpoint_before_lsn) {
-            if (Date.now() - this.lastWaitingLogThottled > 5000) {
-                logger.info(`Waiting until ${this.no_checkpoint_before_lsn} before creating checkpoint, currently at ${lsn}. Persisted op: ${this.persisted_op}`);
-                this.lastWaitingLogThottled = Date.now();
-            }
-            // Edge case: During initial replication, we have a no_checkpoint_before_lsn set,
-            // and don't actually commit the snapshot.
-            // The first commit can happen from an implicit keepalive message.
-            // That needs the persisted_op to get an accurate checkpoint, so
-            // we persist that in keepalive_op.
-            await this.db.sync_rules.updateOne({
-                _id: this.group_id
-            }, {
-                $set: {
-                    keepalive_op: this.persisted_op == null ? null : String(this.persisted_op)
-                }
-            }, { session: this.session });
-            return false;
-        }
-        const now = new Date();
-        const update = {
-            last_checkpoint_lsn: lsn,
-            last_checkpoint_ts: now,
-            last_keepalive_ts: now,
-            snapshot_done: true,
-            last_fatal_error: null,
-            keepalive_op: null
-        };
-        if (this.persisted_op != null) {
-            update.last_checkpoint = this.persisted_op;
-        }
-        await this.db.sync_rules.updateOne({
-            _id: this.group_id
-        }, {
-            $set: update
-        }, { session: this.session });
-        this.persisted_op = null;
-        this.last_checkpoint_lsn = lsn;
-        return true;
-    }
-    async keepalive(lsn) {
-        if (this.last_checkpoint_lsn != null && lsn <= this.last_checkpoint_lsn) {
-            // No-op
-            return false;
-        }
-        if (lsn < this.no_checkpoint_before_lsn) {
-            return false;
-        }
-        if (this.persisted_op != null) {
-            // The commit may have been skipped due to "no_checkpoint_before_lsn".
-            // Apply it now if relevant
-            logger.info(`Commit due to keepalive at ${lsn} / ${this.persisted_op}`);
-            return await this.commit(lsn);
-        }
-        await this.db.sync_rules.updateOne({
-            _id: this.group_id
-        }, {
-            $set: {
-                last_checkpoint_lsn: lsn,
-                snapshot_done: true,
-                last_fatal_error: null,
-                last_keepalive_ts: new Date()
-            }
-        }, { session: this.session });
-        this.last_checkpoint_lsn = lsn;
-        return true;
-    }
-    async save(record) {
-        const { after, before, sourceTable, tag } = record;
-        for (const event of this.getTableEvents(sourceTable)) {
-            this.iterateListeners((cb) => cb.replicationEvent?.({
-                batch: this,
-                table: sourceTable,
-                data: {
-                    op: tag,
-                    after: after && utils.isCompleteRow(after) ? after : undefined,
-                    before: before && utils.isCompleteRow(before) ? before : undefined
-                },
-                event
-            }));
-        }
-        /**
-         * Return if the table is just an event table
-         */
-        if (!sourceTable.syncData && !sourceTable.syncParameters) {
-            return null;
-        }
-        logger.debug(`Saving ${record.tag}:${record.before?.id}/${record.after?.id}`);
-        this.batch ?? (this.batch = new OperationBatch());
-        this.batch.push(new RecordOperation(record));
-        if (this.batch.shouldFlush()) {
-            const r = await this.flush();
-            // HACK: Give other streams a  chance to also flush
-            await timers.setTimeout(5);
-            return r;
-        }
-        return null;
-    }
-    /**
-     * Drop is equivalent to TRUNCATE, plus removing our record of the table.
-     */
-    async drop(sourceTables) {
-        await this.truncate(sourceTables);
-        const result = await this.flush();
-        await this.withTransaction(async () => {
-            for (let table of sourceTables) {
-                await this.db.source_tables.deleteOne({ _id: table.id });
-            }
-        });
-        return result;
-    }
-    async truncate(sourceTables) {
-        await this.flush();
-        let last_op = null;
-        for (let table of sourceTables) {
-            last_op = await this.truncateSingle(table);
-        }
-        if (last_op) {
-            this.persisted_op = last_op;
-        }
-        return {
-            flushed_op: String(last_op)
-        };
-    }
-    async truncateSingle(sourceTable) {
-        let last_op = null;
-        // To avoid too large transactions, we limit the amount of data we delete per transaction.
-        // Since we don't use the record data here, we don't have explicit size limits per batch.
-        const BATCH_LIMIT = 2000;
-        let lastBatchCount = BATCH_LIMIT;
-        while (lastBatchCount == BATCH_LIMIT) {
-            await this.withReplicationTransaction(`Truncate ${sourceTable.qualifiedName}`, async (session, opSeq) => {
-                const current_data_filter = {
-                    _id: idPrefixFilter({ g: this.group_id, t: sourceTable.id }, ['k'])
-                };
-                const cursor = this.db.current_data.find(current_data_filter, {
-                    projection: {
-                        _id: 1,
-                        buckets: 1,
-                        lookups: 1
-                    },
-                    limit: BATCH_LIMIT,
-                    session: session
-                });
-                const batch = await cursor.toArray();
-                const persistedBatch = new PersistedBatch(this.group_id, 0);
-                for (let value of batch) {
-                    persistedBatch.saveBucketData({
-                        op_seq: opSeq,
-                        before_buckets: value.buckets,
-                        evaluated: [],
-                        table: sourceTable,
-                        sourceKey: value._id.k
-                    });
-                    persistedBatch.saveParameterData({
-                        op_seq: opSeq,
-                        existing_lookups: value.lookups,
-                        evaluated: [],
-                        sourceTable: sourceTable,
-                        sourceKey: value._id.k
-                    });
-                    persistedBatch.deleteCurrentData(value._id);
-                }
-                await persistedBatch.flush(this.db, session);
-                lastBatchCount = batch.length;
-                last_op = opSeq.last();
-            });
-        }
-        return last_op;
-    }
-    async markSnapshotDone(tables, no_checkpoint_before_lsn) {
-        const session = this.session;
-        const ids = tables.map((table) => table.id);
-        await this.withTransaction(async () => {
-            await this.db.source_tables.updateMany({ _id: { $in: ids } }, {
-                $set: {
-                    snapshot_done: true
-                }
-            }, { session });
-            if (no_checkpoint_before_lsn > this.no_checkpoint_before_lsn) {
-                this.no_checkpoint_before_lsn = no_checkpoint_before_lsn;
-                await this.db.sync_rules.updateOne({
-                    _id: this.group_id
-                }, {
-                    $set: {
-                        no_checkpoint_before: no_checkpoint_before_lsn,
-                        last_keepalive_ts: new Date()
-                    }
-                }, { session: this.session });
-            }
-        });
-        return tables.map((table) => {
-            const copy = new storage.SourceTable(table.id, table.connectionTag, table.objectId, table.schema, table.table, table.replicaIdColumns, table.snapshotComplete);
-            copy.syncData = table.syncData;
-            copy.syncParameters = table.syncParameters;
-            return copy;
-        });
-    }
-    /**
-     * Gets relevant {@link SqlEventDescriptor}s for the given {@link SourceTable}
-     */
-    getTableEvents(table) {
-        return this.sync_rules.event_descriptors.filter((evt) => [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table)));
-    }
-}
-export function currentBucketKey(b) {
-    return `${b.bucket}/${b.table}/${b.id}`;
-}
-//# sourceMappingURL=MongoBucketBatch.js.map