@360-l/mongo-bulk-data-migration
Version:
MongoDB bulk data migration for node scripts
292 lines (291 loc) • 13 kB
JavaScript
"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.FETCH_ALL = exports.DELETE_COLLECTION = void 0;
const lodash_1 = __importDefault(require("lodash"));
const p_limit_1 = __importDefault(require("p-limit"));
const MigrationBulk_1 = require("./lib/MigrationBulk");
const BackupBulk_1 = require("./lib/BackupBulk");
const AbstractBulkOperationResults_1 = require("./lib/AbstractBulkOperationResults");
const RollbackBulk_1 = require("./lib/RollbackBulk");
const computeRollbackQuery_1 = require("./lib/computeRollbackQuery");
const DEFAULT_BULK_SIZE = 5000;
const COUNT_TOO_LONG_WARNING_THRESHOLD_MS = 30000;
const COLLECTION_VALIDATION_LEVEL = 'moderate';
exports.DELETE_COLLECTION = Symbol();
exports.FETCH_ALL = Symbol();
const defaultLogger = {
info: (...args) => {
if (process.env.NODE_ENV === 'test') {
return;
}
console.log({ ENV: process.env.NODE_ENV });
console.log.call(console, args);
},
warn: (...args) => {
if (process.env.NODE_ENV === 'test') {
return;
}
console.warn.call(console, args);
},
};
class MongoBulkDataMigration {
options = {
arrayFilters: [],
bypassRollbackValidation: false,
bypassUpdateValidation: false,
continueOnBulkWriteError: false,
dontCount: false,
maxBulkSize: DEFAULT_BULK_SIZE,
maxConcurrentUpdateCalls: 10,
rollbackable: true,
throttle: 0,
};
collectionName;
id;
migrationInfos;
logger;
constructor(...args) {
const config = args[0];
this.id = config.id;
this.collectionName = config.collectionName;
Object.assign(this.options, { ...config.options });
this.migrationInfos = {
db: config.db,
operation: config.operation,
projection: config.projection,
rollback: config.rollback,
query: config.query,
update: config.update,
};
this.logger = config.logger ?? defaultLogger;
}
setLogger(logger) {
this.logger = logger;
}
getInfos() {
return {
id: this.id,
migrationInfos: this.migrationInfos,
collection: this.getCollection(),
collectionName: this.collectionName,
};
}
async getRollbackCollection() {
return this.migrationInfos.db.collection(this.getRollbackCollectionName());
}
getRollbackCollectionName() {
return `_rollback_${this.collectionName}_${this.id}`;
}
getCollection() {
return this.migrationInfos.db.collection(this.collectionName);
}
async update() {
const migrationCollection = this.getCollection();
const rollbackCollection = await this.getRollbackCollection();
if (this.migrationInfos.operation === exports.DELETE_COLLECTION) {
const status = await this.renameCollection(this.collectionName, this.getRollbackCollectionName());
return { ok: status ? 1 : 0 };
}
await this.lowerValidationLevel('update');
const { cursor, totalEntries } = await this.getCursorAndCount(migrationCollection, rollbackCollection);
const formattedTotalEntries = totalEntries === AbstractBulkOperationResults_1.NO_COUNT_AVAILABLE
? 'N/A (dontCount option ON)'
: totalEntries;
this.logger.info({
collectionName: this.collectionName,
id: this.id,
totalEntries: formattedTotalEntries,
}, 'Starting migration UPDATE process');
const bulkMigration = new MigrationBulk_1.MigrationBulk(migrationCollection, this.logger, totalEntries);
const bulkBackup = new BackupBulk_1.BackupBulk(rollbackCollection, this.logger, totalEntries);
const updatePromiseLimiter = (0, p_limit_1.default)(this.options.maxConcurrentUpdateCalls);
let updatePromises = [];
let document = (await cursor.next());
while (document !== null) {
const bulkUpdateWrappedPromise = updatePromiseLimiter(this.buildBulkUpdater(document, bulkBackup, bulkMigration));
updatePromises.push(bulkUpdateWrappedPromise);
document = (await cursor.next());
if (!document || updatePromises.length >= this.options.maxBulkSize) {
await Promise.all(updatePromises);
const backupRes = (await bulkBackup.execute()).getResults();
const updateRes = (await bulkMigration.execute(this.options.continueOnBulkWriteError)).getResults();
if (this.options.rollbackable) {
const totalNewBackupDocs = backupRes.nUpserted + backupRes.nInserted;
const totalUpdatedDocument = updateRes.nModified + updateRes.nRemoved;
if (totalNewBackupDocs < totalUpdatedDocument) {
this.logger.warn({ totalNewBackupDocs, totalUpdatedDocument }, "The number of backup documents should be equal to the total updated documents. Check your query is idempotent or ensure you don't use a same migration id for different migrations.");
}
}
await this.throttle();
updatePromises = [];
}
}
this.logger.info({ collectionName: this.collectionName, id: this.id }, 'Ending migration UPDATE process');
await this.restoreValidationLevel('update');
return bulkMigration.getResults();
}
async getCursorAndCount(migrationCollection, rollbackCollection) {
const resolvedQuery = await this.resolveQuery(rollbackCollection);
const cursor = getCursor(resolvedQuery, this.migrationInfos);
const countTakingTooLongTimeout = setTimeout(() => this.logger.warn({ COUNT_TOO_LONG_WARNING_THRESHOLD_MS }, 'Count is taking a significant amount of time, consider using dontCount:true option'), COUNT_TOO_LONG_WARNING_THRESHOLD_MS);
const totalEntries = this.options.dontCount
? AbstractBulkOperationResults_1.NO_COUNT_AVAILABLE
: await getTotalEntries(resolvedQuery);
clearTimeout(countTakingTooLongTimeout);
return { cursor, totalEntries };
function getCursor(query, { projection }) {
if (isPipeline(query)) {
const pipelineWithProjection = query.concat(lodash_1.default.isEmpty(projection) ? [] : [{ $project: projection }]);
return migrationCollection.aggregate(pipelineWithProjection);
}
return migrationCollection.find(query, { projection });
}
async function getTotalEntries(query) {
if (isPipeline(query)) {
const pipelineComputeTotal = query.concat({ $count: 'totalEntries' });
const cursorComputeTotal = migrationCollection.aggregate(pipelineComputeTotal);
const total = (await cursorComputeTotal.next());
return total === null ? 0 : total.totalEntries;
}
return migrationCollection.countDocuments(query);
}
function isPipeline(query) {
return Array.isArray(query);
}
}
buildBulkUpdater(document, bulkBackup, bulkMigration) {
return async () => {
const updateQuery = lodash_1.default.isFunction(this.migrationInfos.update)
? await this.migrationInfos.update(lodash_1.default.cloneDeep(document))
: this.migrationInfos.update;
if (this.options.rollbackable) {
const backupDocument = this.buildBackupDocument(document, updateQuery);
bulkBackup.addInsertOperation(document, backupDocument);
}
bulkMigration.addUpdateOrRemoveOperation(updateQuery, document._id, this.options.arrayFilters);
};
}
buildBackupDocument(document, updateQuery) {
return {
_id: document._id,
backup: this.options.projectionBackupFilter
? lodash_1.default.pick(document, this.options.projectionBackupFilter)
: document,
date: new Date(),
updateQuery: JSON.stringify(updateQuery),
};
}
async throttle() {
if (this.options.throttle > 0) {
await new Promise((resolve) => setTimeout(resolve, this.options.throttle));
}
}
async resolveQuery(rollbackCollection) {
if (this.migrationInfos.query !== exports.FETCH_ALL) {
return this.migrationInfos.query;
}
const lastBackup = await rollbackCollection
.find({}, { projection: { _id: 1 } })
.sort({ _id: -1 })
.limit(1)
.next();
return lastBackup
? { _id: { $gt: lastBackup._id } }
: {};
}
async rollback() {
if (!this.options.rollbackable) {
this.logger.warn('Calling rollback() on a non rollbackable script');
return { ok: 1 };
}
const collection = this.getCollection();
const rollbackCollection = await this.getRollbackCollection();
if (this.migrationInfos.operation === exports.DELETE_COLLECTION) {
const status = await this.renameCollection(this.getRollbackCollectionName(), this.collectionName);
return { ok: status ? 1 : 0 };
}
const cursor = rollbackCollection.find({});
const totalEntries = await rollbackCollection.countDocuments({});
this.logger.info({ collectionName: this.collectionName, id: this.id, totalEntries }, 'Starting migration ROLLBACK process');
const bulkRollback = new RollbackBulk_1.RollbackBulk(collection, this.logger, totalEntries);
await this.lowerValidationLevel('rollback');
let rollbackDocument = (await cursor.next());
while (rollbackDocument !== null) {
const updateQuery = await this.getRollbackUpdateQuery(rollbackDocument);
if (this.migrationInfos.update === MigrationBulk_1.DELETE_OPERATION) {
bulkRollback.addRollbackFullDocumentOperation(rollbackDocument.backup);
}
else {
const { arrayFilters, ...cleanedUpdateQuery } = updateQuery;
bulkRollback.addRollbackOperation(cleanedUpdateQuery, rollbackDocument._id, arrayFilters ?? []);
}
rollbackDocument =
(await cursor.next());
if (!rollbackDocument || bulkRollback.size >= this.options.maxBulkSize) {
await bulkRollback.execute();
}
}
this.logger.info({ collectionName: this.collectionName, id: this.id }, 'Ending migration ROLLBACK process');
await this.clean();
await this.restoreValidationLevel('rollback');
return bulkRollback.getResults();
}
async getRollbackUpdateQuery(obj) {
if (this.migrationInfos.rollback) {
return this.migrationInfos.rollback(obj.backup);
}
const updateOperation = JSON.parse(obj.updateQuery);
if (updateOperation === null) {
return { $set: obj.backup };
}
return (0, computeRollbackQuery_1.computeRollbackQuery)(updateOperation, obj.backup);
}
async lowerValidationLevel(action) {
if (this.canUpdateValidation(action)) {
await this.setValidationLevel('off');
}
}
async restoreValidationLevel(action) {
if (this.canUpdateValidation(action)) {
await this.setValidationLevel(COLLECTION_VALIDATION_LEVEL);
}
}
canUpdateValidation(action) {
return ((action === 'update' && this.options.bypassUpdateValidation) ||
(action === 'rollback' && this.options.bypassRollbackValidation));
}
async setValidationLevel(validationLevel) {
await this.migrationInfos.db.command({
collMod: this.collectionName,
validationLevel,
});
}
async renameCollection(nameBefore, nameAfter) {
try {
await this.migrationInfos.db.renameCollection(nameBefore, nameAfter);
return true;
}
catch (err) {
this.logger.warn({ err }, "Couldn't rename collection");
return false;
}
}
async clean() {
const migrationCollectionName = this.getRollbackCollectionName();
const migrationCollection = await this.getRollbackCollection();
this.logger.info({ migrationCollectionName }, 'Deleting migration rollback collection');
try {
await migrationCollection.drop();
}
catch (err) {
if (err?.message === 'ns not found') {
return;
}
throw err;
}
}
}
exports.default = MongoBulkDataMigration;