SERVER-40406 Add support for $$NOW and $$CLUSTER_TIME in the update command

5 files changed, 35 insertions, 4 deletions

diff --git a/src/mongo/db/ops/SConscript b/src/mongo/db/ops/SConscript
index 8f1df866c99..dff8b81dd47 100644
--- a/src/mongo/db/ops/SConscript
+++ b/src/mongo/db/ops/SConscript
@@ -36,6 +36,7 @@ env.Library(
         '$BUILD_DIR/mongo/base',
         '$BUILD_DIR/mongo/db/commands/test_commands_enabled',
         '$BUILD_DIR/mongo/db/dbmessage',
+        '$BUILD_DIR/mongo/db/pipeline/runtime_constants_idl',
         '$BUILD_DIR/mongo/idl/idl_parser',
     ],
 )
diff --git a/src/mongo/db/ops/parsed_update.cpp b/src/mongo/db/ops/parsed_update.cpp
index acedfc48d8a..d47d973c94b 100644
--- a/src/mongo/db/ops/parsed_update.cpp
+++ b/src/mongo/db/ops/parsed_update.cpp
@@ -44,7 +44,7 @@ ParsedUpdate::ParsedUpdate(OperationContext* opCtx,
                            const ExtensionsCallback& extensionsCallback)
     : _opCtx(opCtx),
       _request(request),
-      _driver(new ExpressionContext(opCtx, nullptr)),
+      _driver(new ExpressionContext(opCtx, nullptr, _request->getRuntimeConstants())),
       _canonicalQuery(),
       _extensionsCallback(extensionsCallback) {}
 
@@ -123,7 +123,8 @@ Status ParsedUpdate::parseQueryToCQ() {
         allowedMatcherFeatures &= ~MatchExpressionParser::AllowedFeatures::kExpr;
     }
 
-    boost::intrusive_ptr<ExpressionContext> expCtx;
+    auto expCtx =
+        make_intrusive<ExpressionContext>(_opCtx, _collator.get(), _request->getRuntimeConstants());
     auto statusWithCQ = CanonicalQuery::canonicalize(
         _opCtx, std::move(qr), std::move(expCtx), _extensionsCallback, allowedMatcherFeatures);
     if (statusWithCQ.isOK()) {
diff --git a/src/mongo/db/ops/update_request.h b/src/mongo/db/ops/update_request.h
index 587fdb1b718..2e84c575f11 100644
--- a/src/mongo/db/ops/update_request.h
+++ b/src/mongo/db/ops/update_request.h
@@ -34,6 +34,7 @@
 #include "mongo/db/logical_session_id.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/ops/write_ops_parsers.h"
+#include "mongo/db/pipeline/runtime_constants_gen.h"
 #include "mongo/db/query/explain.h"
 #include "mongo/util/str.h"
 
@@ -110,6 +111,14 @@ public:
         return _updateMod;
     }
 
+    inline void setRuntimeConstants(const RuntimeConstants& runtimeConstants) {
+        _runtimeConstants = runtimeConstants;
+    }
+
+    inline const boost::optional<RuntimeConstants>& getRuntimeConstants() const {
+        return _runtimeConstants;
+    }
+
     inline void setArrayFilters(const std::vector<BSONObj>& arrayFilters) {
         _arrayFilters = arrayFilters;
     }
@@ -210,6 +219,10 @@ public:
         builder << " updateModification: " << _updateMod.toString();
         builder << " stmtId: " << _stmtId;
 
+        if (_runtimeConstants) {
+            builder << "runtimeConstants: " << _runtimeConstants->toBSON().toString();
+        }
+
         builder << " arrayFilters: [";
         bool first = true;
         for (auto arrayFilter : _arrayFilters) {
@@ -248,6 +261,9 @@ private:
     // Contains the modifiers to apply to matched objects, or a replacement document.
     write_ops::UpdateModification _updateMod;
 
+    // Contains any constant values which may be required by the query or update operation.
+    boost::optional<RuntimeConstants> _runtimeConstants;
+
     // Filters to specify which array elements should be updated.
     std::vector<BSONObj> _arrayFilters;
 
diff --git a/src/mongo/db/ops/write_ops.idl b/src/mongo/db/ops/write_ops.idl
index c570b799b36..95ef7f5e71d 100644
--- a/src/mongo/db/ops/write_ops.idl
+++ b/src/mongo/db/ops/write_ops.idl
@@ -33,6 +33,7 @@ global:
 
 imports:
     - "mongo/db/logical_session_id.idl"
+    - "mongo/db/pipeline/runtime_constants.idl"
 
 types:
     multi_delete_bool:
@@ -163,6 +164,11 @@ commands:
                 description: "An array of one or more update statements to perform."
                 type: array<UpdateOpEntry>
                 supports_doc_sequence: true
+            runtimeConstants:
+                description: "A collection of values that do not change once computed. These are
+                              used by pipeline-style update operations."
+                type: RuntimeConstants
+                optional: true
 
     delete:
         description: "Parser for the 'delete' command."
diff --git a/src/mongo/db/ops/write_ops_exec.cpp b/src/mongo/db/ops/write_ops_exec.cpp
index 3bab66af9ae..acd90695b4f 100644
--- a/src/mongo/db/ops/write_ops_exec.cpp
+++ b/src/mongo/db/ops/write_ops_exec.cpp
@@ -685,7 +685,8 @@ static SingleWriteResult performSingleUpdateOp(OperationContext* opCtx,
 static SingleWriteResult performSingleUpdateOpWithDupKeyRetry(OperationContext* opCtx,
                                                               const NamespaceString& ns,
                                                               StmtId stmtId,
-                                                              const write_ops::UpdateOpEntry& op) {
+                                                              const write_ops::UpdateOpEntry& op,
+                                                              RuntimeConstants runtimeConstants) {
     globalOpCounters.gotUpdate();
     ServerWriteConcernMetrics::get(opCtx)->recordWriteConcernForUpdate(opCtx->getWriteConcern());
     auto& curOp = *CurOp::get(opCtx);
@@ -707,6 +708,7 @@ static SingleWriteResult performSingleUpdateOpWithDupKeyRetry(OperationContext*
     UpdateRequest request(ns);
     request.setQuery(op.getQ());
     request.setUpdateModification(op.getU());
+    request.setRuntimeConstants(std::move(runtimeConstants));
     request.setCollation(write_ops::collationOf(op));
     request.setStmtId(stmtId);
     request.setArrayFilters(write_ops::arrayFiltersOf(op));
@@ -770,6 +772,11 @@ WriteResult performUpdates(OperationContext* opCtx, const write_ops::Update& who
     WriteResult out;
     out.results.reserve(wholeOp.getUpdates().size());
 
+    // If the update command specified runtime constants, we adopt them. Otherwise, we set them to
+    // the current local and cluster time. These constants are applied to each update in the batch.
+    const auto& runtimeConstants =
+        wholeOp.getRuntimeConstants().value_or(Variables::generateRuntimeConstants(opCtx));
+
     for (auto&& singleOp : wholeOp.getUpdates()) {
         const auto stmtId = getStmtIdForWriteOp(opCtx, wholeOp, stmtIdIndex++);
         if (opCtx->getTxnNumber()) {
@@ -796,7 +803,7 @@ WriteResult performUpdates(OperationContext* opCtx, const write_ops::Update& who
         try {
             lastOpFixer.startingOp();
             out.results.emplace_back(performSingleUpdateOpWithDupKeyRetry(
-                opCtx, wholeOp.getNamespace(), stmtId, singleOp));
+                opCtx, wholeOp.getNamespace(), stmtId, singleOp, runtimeConstants));
             lastOpFixer.finishedOpSuccessfully();
         } catch (const DBException& ex) {
             const bool canContinue =