refactor: use $createModifiedPathsSnapshot for transaction wrapper

Re: #14699
Automattic · Jul 14, 2024 · 93d33c6 · 93d33c6
1 parent 1972f7f
commit 93d33c6
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 43 deletions.
diff --git a/lib/connection.js b/lib/connection.js
@@ -574,32 +574,7 @@ async function _wrapUserTransaction(fn, session, mongoose) {
 function _resetSessionDocuments(session) {
   for (const doc of session[sessionNewDocuments].keys()) {
     const state = session[sessionNewDocuments].get(doc);
-    if (state.hasOwnProperty('isNew')) {
-      doc.$isNew = state.isNew;
-    }
-    if (state.hasOwnProperty('versionKey')) {
-      doc.set(doc.schema.options.versionKey, state.versionKey);
-    }
-
-    if (state.modifiedPaths.length > 0 && doc.$__.activePaths.states.modify == null) {
-      doc.$__.activePaths.states.modify = {};
-    }
-    for (const path of state.modifiedPaths) {
-      const currentState = doc.$__.activePaths.paths[path];
-      if (currentState != null) {
-        delete doc.$__.activePaths[currentState][path];
-      }
-      doc.$__.activePaths.paths[path] = 'modify';
-      doc.$__.activePaths.states.modify[path] = true;
-    }
-
-    for (const path of state.atomics.keys()) {
-      const val = doc.$__getValue(path);
-      if (val == null) {
-        continue;
-      }
-      val[arrayAtomicsSymbol] = state.atomics.get(path);
-    }
+    doc.$restoreModifiedPathsSnapshot(state);
   }
 }
 

diff --git a/lib/document.js b/lib/document.js
@@ -5249,7 +5249,8 @@ Document.prototype.$createModifiedPathsSnapshot = function $createModifiedPathsS
   return new ModifiedPathsSnapshot(
     subdocSnapshot,
     this.$__.activePaths.clone(),
-    this.$__.version
+    this.$__.version,
+    this.$isNew
   );
 };
 
@@ -5281,6 +5282,7 @@ Document.prototype.$createModifiedPathsSnapshot = function $createModifiedPathsS
 Document.prototype.$restoreModifiedPathsSnapshot = function $restoreModifiedPathsSnapshot(snapshot) {
   this.$__.activePaths = snapshot.activePaths.clone();
   this.$__.version = snapshot.version;
+  this.$isNew = snapshot.isNew;
   if (!this.$isSubdocument) {
     const subdocs = this.$getAllSubdocs();
     for (const child of subdocs) {

diff --git a/lib/modifiedPathsSnapshot.js b/lib/modifiedPathsSnapshot.js
@@ -1,9 +1,10 @@
 'use strict';
 
 module.exports = class ModifiedPathsSnapshot {
-  constructor(subdocSnapshot, activePaths, version) {
+  constructor(subdocSnapshot, activePaths, version, isNew) {
     this.subdocSnapshot = subdocSnapshot;
     this.activePaths = activePaths;
     this.version = version;
+    this.isNew = isNew;
   }
 };
diff --git a/lib/plugins/trackTransaction.js b/lib/plugins/trackTransaction.js
@@ -15,24 +15,11 @@ module.exports = function trackTransaction(schema) {
     }
 
     if (!session[sessionNewDocuments].has(this)) {
-      const initialState = {};
-      if (this.isNew) {
-        initialState.isNew = true;
-      }
-      if (this.$__schema.options.versionKey) {
-        initialState.versionKey = this.get(this.$__schema.options.versionKey);
-      }
-
-      initialState.modifiedPaths = new Set(Object.keys(this.$__.activePaths.getStatePaths('modify')));
-      initialState.atomics = _getAtomics(this);
-
-      session[sessionNewDocuments].set(this, initialState);
+      session[sessionNewDocuments].set(this, this.$createModifiedPathsSnapshot());
     } else {
       const state = session[sessionNewDocuments].get(this);
 
-      for (const path of Object.keys(this.$__.activePaths.getStatePaths('modify'))) {
-        state.modifiedPaths.add(path);
-      }
+      state.activePaths.states.modify = { ...this.$__.activePaths.getStatePaths('modify') };
       state.atomics = _getAtomics(this, state.atomics);
     }
   });

diff --git a/test/document.test.js b/test/document.test.js
@@ -13631,6 +13631,7 @@ describe('document', function() {
       docArr: [new Schema({ subprop3: String }, { _id: false })]
     });
     const Test = db.model('Test', schema);
+    await Test.deleteMany({});
 
     const doc = new Test({});
     await doc.save();
@@ -13651,14 +13652,17 @@ describe('document', function() {
     assert.deepStrictEqual(doc.subdoc.getChanges(), { $set: { subprop2: 'test3' } });
     assert.deepStrictEqual(doc.docArr[0].getChanges(), { $set: { subprop3: 'test4' } });
 
+    assert.ok(!doc.$isNew);
     const snapshot = doc.$createModifiedPathsSnapshot();
     doc.$clearModifiedPaths();
+    doc.$isNew = true;
 
     assert.deepStrictEqual(doc.getChanges(), {});
     assert.deepStrictEqual(doc.subdoc.getChanges(), {});
     assert.deepStrictEqual(doc.docArr[0].getChanges(), {});
 
     doc.$restoreModifiedPathsSnapshot(snapshot);
+    assert.ok(!doc.$isNew);
     assert.deepStrictEqual(doc.getChanges().$set, {
       name: 'test1',
       nested: { subprop1: 'test2' },