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

2 files changed, 69 insertions, 8 deletions

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