SERVER-73417 Change method signature of produceDocumentForUpsert() to include immutablePaths

6 files changed, 83 insertions, 97 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 4bc93760ab4..86494b576dd 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -1456,7 +1456,7 @@ env.Library(
         'stats/serveronly_stats',
         'storage/remove_saver',
         'storage/storage_options',
-        'update/produce_document_for_upsert',
+        'update/update_driver',
     ],
     LIBDEPS_PRIVATE=[
         '$BUILD_DIR/mongo/db/catalog/collection_crud',
diff --git a/src/mongo/db/exec/upsert_stage.cpp b/src/mongo/db/exec/upsert_stage.cpp
index cf7121e8e92..c4917b191ab 100644
--- a/src/mongo/db/exec/upsert_stage.cpp
+++ b/src/mongo/db/exec/upsert_stage.cpp
@@ -101,14 +101,7 @@ PlanStage::StageState UpsertStage::doWork(WorkingSetID* out) {
     _specificStats.nUpserted = 1;
 
     // Generate the new document to be inserted.
-    const auto& collDesc = _cachedShardingCollectionDescription.getCollectionDescription(opCtx());
-    _specificStats.objInserted = update::produceDocumentForUpsert(opCtx(),
-                                                                  _params.request,
-                                                                  _params.driver,
-                                                                  _params.canonicalQuery,
-                                                                  _isUserInitiatedWrite,
-                                                                  _doc,
-                                                                  collDesc);
+    _specificStats.objInserted = _produceNewDocumentForInsert();
 
     // If this is an explain, skip performing the actual insert.
     if (!_params.request->explain()) {
@@ -206,4 +199,69 @@ void UpsertStage::_performInsert(BSONObj newDocument) {
     });
 }
 
+BSONObj UpsertStage::_produceNewDocumentForInsert() {
+    FieldRefSet shardKeyPaths, immutablePaths;
+    if (_isUserInitiatedWrite) {
+        // Obtain the collection description. This will be needed to compute the shardKey paths.
+        const auto& collDesc =
+            _cachedShardingCollectionDescription.getCollectionDescription(opCtx());
+
+        // If the collection is sharded, add all fields from the shard key to the 'shardKeyPaths'
+        // set.
+        if (collDesc.isSharded()) {
+            shardKeyPaths.fillFrom(collDesc.getKeyPatternFields());
+        }
+
+        // An unversioned request cannot update the shard key, so all shardKey paths are immutable.
+        if (!OperationShardingState::isComingFromRouter(opCtx())) {
+            for (auto&& shardKeyPath : shardKeyPaths) {
+                immutablePaths.insert(shardKeyPath);
+            }
+        }
+
+        // The _id field is always immutable to user requests, even if the shard key is mutable.
+        immutablePaths.keepShortest(&idFieldRef);
+    }
+
+    // Generate the new document to be inserted.
+    update::produceDocumentForUpsert(
+        opCtx(), _params.request, _params.driver, _params.canonicalQuery, immutablePaths, _doc);
+
+    // Assert that the finished document has all required fields and is valid for storage.
+    _assertDocumentToBeInsertedIsValid(_doc, shardKeyPaths);
+
+    auto newDocument = _doc.getObject();
+    if (!DocumentValidationSettings::get(opCtx()).isInternalValidationDisabled()) {
+        uassert(17420,
+                str::stream() << "Document to upsert is larger than " << BSONObjMaxUserSize,
+                newDocument.objsize() <= BSONObjMaxUserSize);
+    }
+
+    return newDocument;
+}
+
+void UpsertStage::_assertDocumentToBeInsertedIsValid(const mutablebson::Document& document,
+                                                     const FieldRefSet& shardKeyPaths) {
+    // For a non-internal operation, we assert that the document contains all required paths, that
+    // no shard key fields have arrays at any point along their paths, and that the document is
+    // valid for storage. Skip all such checks for an internal operation.
+    if (_isUserInitiatedWrite) {
+        // Shard key values are permitted to be missing, and so the only required field is _id. We
+        // should always have an _id here, since we generated one earlier if not already present.
+        invariant(document.root().ok() && document.root()[idFieldName].ok());
+        bool containsDotsAndDollarsField = false;
+
+        storage_validation::scanDocument(document,
+                                         true, /* allowTopLevelDollarPrefixes */
+                                         true, /* Should validate for storage */
+                                         &containsDotsAndDollarsField);
+        if (containsDotsAndDollarsField)
+            _params.driver->setContainsDotsAndDollarsField(true);
+
+        //  Neither _id nor the shard key fields may have arrays at any point along their paths.
+        update::assertPathsNotArray(document, {{&idFieldRef}});
+        update::assertPathsNotArray(document, shardKeyPaths);
+    }
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/exec/upsert_stage.h b/src/mongo/db/exec/upsert_stage.h
index 34660c37984..d3a1b1671c4 100644
--- a/src/mongo/db/exec/upsert_stage.h
+++ b/src/mongo/db/exec/upsert_stage.h
@@ -61,7 +61,10 @@ public:
     StageState doWork(WorkingSetID* out) final;
 
 private:
+    BSONObj _produceNewDocumentForInsert();
     void _performInsert(BSONObj newDocument);
+    void _assertDocumentToBeInsertedIsValid(const mutablebson::Document& document,
+                                            const FieldRefSet& shardKeyPaths);
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/update/SConscript b/src/mongo/db/update/SConscript
index 3c1b1db5b29..e4806030657 100644
--- a/src/mongo/db/update/SConscript
+++ b/src/mongo/db/update/SConscript
@@ -90,6 +90,7 @@ env.Library(
     target='update_driver',
     source=[
         'update_driver.cpp',
+        'produce_document_for_upsert.cpp',
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/db/common',
@@ -102,18 +103,6 @@ env.Library(
 )
 
 env.Library(
-    target='produce_document_for_upsert',
-    source=[
-        'produce_document_for_upsert.cpp',
-    ],
-    LIBDEPS=[
-        '$BUILD_DIR/mongo/db/catalog/document_validation',
-        '$BUILD_DIR/mongo/db/shard_role_api',
-        '$BUILD_DIR/mongo/db/update/update_driver',
-    ],
-)
-
-env.Library(
     target='update_test_helpers',
     source=[
         'document_diff_test_helpers.cpp',
diff --git a/src/mongo/db/update/produce_document_for_upsert.cpp b/src/mongo/db/update/produce_document_for_upsert.cpp
index 144db75bc43..b541d2b3aa1 100644
--- a/src/mongo/db/update/produce_document_for_upsert.cpp
+++ b/src/mongo/db/update/produce_document_for_upsert.cpp
@@ -30,7 +30,6 @@
 #include "mongo/db/update/produce_document_for_upsert.h"
 
 #include "mongo/bson/mutable/algorithm.h"
-#include "mongo/db/catalog/document_validation.h"
 #include "mongo/db/s/operation_sharding_state.h"
 #include "mongo/db/update/storage_validation.h"
 
@@ -87,60 +86,12 @@ void generateNewDocumentFromSuppliedDoc(OperationContext* opCtx,
         opCtx, {}, &document, validateForStorage, immutablePaths, isInsert));
 }
 
-void assertDocumentToBeInsertedIsValid(const mutablebson::Document& document,
-                                       const FieldRefSet& shardKeyPaths,
-                                       bool isUserInitiatedWrite,
-                                       UpdateDriver* driver) {
-    // For a non-internal operation, we assert that the document contains all required paths, that
-    // no shard key fields have arrays at any point along their paths, and that the document is
-    // valid for storage. Skip all such checks for an internal operation.
-    if (isUserInitiatedWrite) {
-        // Shard key values are permitted to be missing, and so the only required field is _id. We
-        // should always have an _id here, since we generated one earlier if not already present.
-        invariant(document.root().ok() && document.root()[idFieldName].ok());
-        bool containsDotsAndDollarsField = false;
-
-        storage_validation::scanDocument(document,
-                                         true, /* allowTopLevelDollarPrefixes */
-                                         true, /* Should validate for storage */
-                                         &containsDotsAndDollarsField);
-        if (containsDotsAndDollarsField)
-            driver->setContainsDotsAndDollarsField(true);
-
-        //  Neither _id nor the shard key fields may have arrays at any point along their paths.
-        assertPathsNotArray(document, {{&idFieldRef}});
-        assertPathsNotArray(document, shardKeyPaths);
-    }
-}
-
-BSONObj produceDocumentForUpsert(OperationContext* opCtx,
-                                 const UpdateRequest* request,
-                                 UpdateDriver* driver,
-                                 const CanonicalQuery* canonicalQuery,
-                                 bool isUserInitiatedWrite,
-                                 mutablebson::Document& doc,
-                                 const ScopedCollectionDescription& collDesc) {
-    // Obtain the collection description. This will be needed to compute the shardKey paths.
-    FieldRefSet shardKeyPaths, immutablePaths;
-
-    if (isUserInitiatedWrite) {
-        // If the collection is sharded, add all fields from the shard key to the 'shardKeyPaths'
-        // set.
-        if (collDesc.isSharded()) {
-            shardKeyPaths.fillFrom(collDesc.getKeyPatternFields());
-        }
-
-        // An unversioned request cannot update the shard key, so all shardKey paths are immutable.
-        if (!OperationShardingState::isComingFromRouter(opCtx)) {
-            for (auto&& shardKeyPath : shardKeyPaths) {
-                immutablePaths.insert(shardKeyPath);
-            }
-        }
-
-        // The _id field is always immutable to user requests, even if the shard key is mutable.
-        immutablePaths.keepShortest(&idFieldRef);
-    }
-
+void produceDocumentForUpsert(OperationContext* opCtx,
+                              const UpdateRequest* request,
+                              UpdateDriver* driver,
+                              const CanonicalQuery* canonicalQuery,
+                              const FieldRefSet& immutablePaths,
+                              mutablebson::Document& doc) {
     // Reset the document into which we will be writing.
     doc.reset();
 
@@ -163,19 +114,6 @@ BSONObj produceDocumentForUpsert(OperationContext* opCtx,
 
     // Third: ensure _id is first if it exists, and generate a new OID otherwise.
     ensureIdFieldIsFirst(&doc, true);
-
-    // Fourth: assert that the finished document has all required fields and is valid for storage.
-    assertDocumentToBeInsertedIsValid(doc, shardKeyPaths, isUserInitiatedWrite, driver);
-
-    // Fifth: validate that the newly-produced document does not exceed the maximum BSON user size.
-    auto newDocument = doc.getObject();
-    if (!DocumentValidationSettings::get(opCtx).isInternalValidationDisabled()) {
-        uassert(17420,
-                str::stream() << "Document to upsert is larger than " << BSONObjMaxUserSize,
-                newDocument.objsize() <= BSONObjMaxUserSize);
-    }
-
-    return newDocument;
 }
 
 void assertPathsNotArray(const mutablebson::Document& document, const FieldRefSet& paths) {
diff --git a/src/mongo/db/update/produce_document_for_upsert.h b/src/mongo/db/update/produce_document_for_upsert.h
index 938205676bc..3056d8222a0 100644
--- a/src/mongo/db/update/produce_document_for_upsert.h
+++ b/src/mongo/db/update/produce_document_for_upsert.h
@@ -29,7 +29,6 @@
 
 #include "mongo/db/operation_context.h"
 #include "mongo/db/ops/update_request.h"
-#include "mongo/db/s/scoped_collection_metadata.h"
 #include "mongo/db/update/update_driver.h"
 #include "mongo/util/safe_num.h"
 
@@ -39,13 +38,12 @@ namespace update {
 /**
  * Use an UpdateDriver and UpdateRequest to produce the document to insert.
  **/
-BSONObj produceDocumentForUpsert(OperationContext* opCtx,
-                                 const UpdateRequest* request,
-                                 UpdateDriver* driver,
-                                 const CanonicalQuery* cq,
-                                 bool isUserInitiatedWrite,
-                                 mutablebson::Document& doc,
-                                 const ScopedCollectionDescription& collDesc);
+void produceDocumentForUpsert(OperationContext* opCtx,
+                              const UpdateRequest* request,
+                              UpdateDriver* driver,
+                              const CanonicalQuery* cq,
+                              const FieldRefSet& immutablePaths,
+                              mutablebson::Document& doc);
 
 void ensureIdFieldIsFirst(mutablebson::Document* doc, bool generateOIDIfMissing);
 void assertPathsNotArray(const mutablebson::Document& document, const FieldRefSet& paths);