SERVER-58581 Add SEARCH_META variable

10 files changed, 128 insertions, 12 deletions

diff --git a/src/mongo/db/pipeline/variables.cpp b/src/mongo/db/pipeline/variables.cpp
index 6adc0e280c1..676c32262be 100644
--- a/src/mongo/db/pipeline/variables.cpp
+++ b/src/mongo/db/pipeline/variables.cpp
@@ -51,6 +51,7 @@ constexpr StringData kNowName = "NOW"_sd;
 constexpr StringData kClusterTimeName = "CLUSTER_TIME"_sd;
 constexpr StringData kJsScopeName = "JS_SCOPE"_sd;
 constexpr StringData kIsMapReduceName = "IS_MR"_sd;
+constexpr StringData kSearchMetaName = "SEARCH_META"_sd;
 
 const StringMap<Variables::Id> Variables::kBuiltinVarNameToId = {
     {kRootName.rawData(), kRootId},
@@ -58,7 +59,8 @@ const StringMap<Variables::Id> Variables::kBuiltinVarNameToId = {
     {kNowName.rawData(), kNowId},
     {kClusterTimeName.rawData(), kClusterTimeId},
     {kJsScopeName.rawData(), kJsScopeId},
-    {kIsMapReduceName.rawData(), kIsMapReduceId}};
+    {kIsMapReduceName.rawData(), kIsMapReduceId},
+    {kSearchMetaName.rawData(), kSearchMetaId}};
 
 const std::map<Variables::Id, std::string> Variables::kIdToBuiltinVarName = {
     {kRootId, kRootName.rawData()},
@@ -66,7 +68,8 @@ const std::map<Variables::Id, std::string> Variables::kIdToBuiltinVarName = {
     {kNowId, kNowName.rawData()},
     {kClusterTimeId, kClusterTimeName.rawData()},
     {kJsScopeId, kJsScopeName.rawData()},
-    {kIsMapReduceId, kIsMapReduceName.rawData()}};
+    {kIsMapReduceId, kIsMapReduceName.rawData()},
+    {kSearchMetaId, kSearchMetaName.rawData()}};
 
 const std::map<StringData, std::function<void(const Value&)>> Variables::kSystemVarValidators = {
     {kNowName,
@@ -106,6 +109,24 @@ void Variables::setValue(Id id, const Value& value, bool isConstant) {
     _definitions[id] = {value, isConstant};
 }
 
+void Variables::setReservedValue(Id id, const Value& value, bool isConstant) {
+    // If a value has already been set for 'id', and that value was marked as constant, then it
+    // is illegal to modify.
+    switch (id) {
+        case Variables::kSearchMetaId:
+            tassert(5858101,
+                    "Can't set a variable that has been set to be constant ",
+                    !hasConstantValue(id));
+            _definitions[id] = {value, isConstant};
+            break;
+        default:
+            // Currently it is only allowed to manually set the SEARCH_META builtin variable.
+            tasserted(5858102,
+                      str::stream() << "Attempted to set '$$" << getBuiltinVariableName(id)
+                                    << "' which is not permitted");
+    }
+}
+
 void Variables::setValue(Variables::Id id, const Value& value) {
     const bool isConstant = false;
     setValue(id, value, isConstant);
@@ -142,6 +163,15 @@ Value Variables::getValue(Id id, const Document& root) const {
                 uasserted(51144,
                           str::stream() << "Builtin variable '$$" << getBuiltinVariableName(id)
                                         << "' is not available");
+            case Variables::kSearchMetaId: {
+                uassert(5858105,
+                        str::stream() << "Must enable 'featureFlagSearchMeta' to access '$$"
+                                      << getBuiltinVariableName(id),
+                        ::mongo::feature_flags::gFeatureFlagSearchMeta.isEnabled(
+                            serverGlobalParams.featureCompatibility));
+                auto metaIt = _definitions.find(id);
+                return metaIt == _definitions.end() ? Value() : metaIt->second.value;
+            }
             default:
                 MONGO_UNREACHABLE;
         }
diff --git a/src/mongo/db/pipeline/variables.h b/src/mongo/db/pipeline/variables.h
index d1fcadea9f7..03407882adf 100644
--- a/src/mongo/db/pipeline/variables.h
+++ b/src/mongo/db/pipeline/variables.h
@@ -84,6 +84,7 @@ public:
     static constexpr auto kClusterTimeId = Id(-4);
     static constexpr auto kJsScopeId = Id(-5);
     static constexpr auto kIsMapReduceId = Id(-6);
+    static constexpr auto kSearchMetaId = Id(-7);
 
     // Map from builtin var name to reserved id number.
     static const StringMap<Id> kBuiltinVarNameToId;
@@ -103,6 +104,12 @@ public:
     void setConstantValue(Variables::Id id, const Value& value);
 
     /**
+     * Same as 'setValue' but is only allowed on reserved, builtin, variables. Should not be used
+     * when setting from user input.
+     */
+    void setReservedValue(Variables::Id id, const Value& value, bool isConstant);
+
+    /**
      * Gets the value of a user-defined or system variable. If the 'id' provided represents the
      * special ROOT variable, then we return 'root' in Value form.
      */
@@ -175,6 +182,15 @@ public:
 
     LegacyRuntimeConstants transitionalExtractRuntimeConstants() const;
 
+    static auto getBuiltinVariableName(Variables::Id variable) {
+        for (auto& [name, id] : kBuiltinVarNameToId) {
+            if (variable == id) {
+                return name;
+            }
+        }
+        MONGO_UNREACHABLE_TASSERT(5858104);
+    }
+
 private:
     struct ValueAndState {
         ValueAndState() = default;
@@ -187,15 +203,6 @@ private:
 
     void setValue(Id id, const Value& value, bool isConstant);
 
-    static auto getBuiltinVariableName(Variables::Id variable) {
-        for (auto& [name, id] : kBuiltinVarNameToId) {
-            if (variable == id) {
-                return name;
-            }
-        }
-        return std::string();
-    }
-
     IdGenerator _idGenerator;
     stdx::unordered_map<Id, ValueAndState> _definitions;
 };
diff --git a/src/mongo/db/query/cursor_response.cpp b/src/mongo/db/query/cursor_response.cpp
index 06c62c3091f..00c609acd91 100644
--- a/src/mongo/db/query/cursor_response.cpp
+++ b/src/mongo/db/query/cursor_response.cpp
@@ -44,6 +44,7 @@ const char kCursorsField[] = "cursors";
 const char kCursorField[] = "cursor";
 const char kIdField[] = "id";
 const char kNsField[] = "ns";
+const char kVarsField[] = "vars";
 const char kAtClusterTimeField[] = "atClusterTime";
 const char kBatchField[] = "nextBatch";
 const char kBatchFieldInitial[] = "firstBatch";
@@ -129,6 +130,7 @@ CursorResponse::CursorResponse(NamespaceString nss,
                                boost::optional<long long> numReturnedSoFar,
                                boost::optional<BSONObj> postBatchResumeToken,
                                boost::optional<BSONObj> writeConcernError,
+                               boost::optional<BSONObj> varsField,
                                bool partialResultsReturned,
                                bool invalidated)
     : _nss(std::move(nss)),
@@ -138,6 +140,7 @@ CursorResponse::CursorResponse(NamespaceString nss,
       _numReturnedSoFar(numReturnedSoFar),
       _postBatchResumeToken(std::move(postBatchResumeToken)),
       _writeConcernError(std::move(writeConcernError)),
+      _varsField(std::move(varsField)),
       _partialResultsReturned(partialResultsReturned),
       _invalidated(invalidated) {}
 
@@ -200,6 +203,13 @@ StatusWith<CursorResponse> CursorResponse::parseFromBSON(const BSONObj& cmdRespo
     }
     fullns = nsElt.String();
 
+    BSONElement varsElt = cmdResponse[kVarsField];
+    if (!varsElt.eoo() && varsElt.type() != BSONType::Object) {
+        return {ErrorCodes::TypeMismatch,
+                str::stream() << "Field '" << kVarsField
+                              << "' must be of type object in: " << cmdResponse};
+    }
+
     BSONElement batchElt = cursorObj[kBatchField];
     if (batchElt.eoo()) {
         batchElt = cursorObj[kBatchFieldInitial];
@@ -280,6 +290,7 @@ StatusWith<CursorResponse> CursorResponse::parseFromBSON(const BSONObj& cmdRespo
              postBatchResumeTokenElem ? postBatchResumeTokenElem.Obj().getOwned()
                                       : boost::optional<BSONObj>{},
              writeConcernError ? writeConcernError.Obj().getOwned() : boost::optional<BSONObj>{},
+             varsElt ? varsElt.Obj().getOwned() : boost::optional<BSONObj>{},
              partialResultsReturned.trueValue(),
              invalidatedElem.trueValue()}};
 }
diff --git a/src/mongo/db/query/cursor_response.h b/src/mongo/db/query/cursor_response.h
index 0bd028825ed..b9473f5eb5b 100644
--- a/src/mongo/db/query/cursor_response.h
+++ b/src/mongo/db/query/cursor_response.h
@@ -199,6 +199,7 @@ public:
                    boost::optional<long long> numReturnedSoFar = boost::none,
                    boost::optional<BSONObj> postBatchResumeToken = boost::none,
                    boost::optional<BSONObj> writeConcernError = boost::none,
+                   boost::optional<BSONObj> varsField = boost::none,
                    bool partialResultsReturned = false,
                    bool invalidated = false);
 
@@ -241,6 +242,10 @@ public:
         return _atClusterTime;
     }
 
+    boost::optional<BSONObj> getVarsField() const {
+        return _varsField;
+    }
+
     bool getPartialResultsReturned() const {
         return _partialResultsReturned;
     }
@@ -266,6 +271,7 @@ private:
     boost::optional<long long> _numReturnedSoFar;
     boost::optional<BSONObj> _postBatchResumeToken;
     boost::optional<BSONObj> _writeConcernError;
+    boost::optional<BSONObj> _varsField;
     bool _partialResultsReturned = false;
     bool _invalidated = false;
 };
diff --git a/src/mongo/db/query/cursor_response.idl b/src/mongo/db/query/cursor_response.idl
index 66c118ff5f6..639a3cf3fe4 100644
--- a/src/mongo/db/query/cursor_response.idl
+++ b/src/mongo/db/query/cursor_response.idl
@@ -79,6 +79,10 @@ structs:
             cursor:
                 description: "A response cursor object."
                 type: InitialResponseCursor
+            vars:
+                description: "An optional field containing additional response information for the query."
+                type: object
+                optional: true
 
     GetMoreResponseCursor:
         description: "A struct representing a subsequent response cursor."
diff --git a/src/mongo/db/query/cursor_response_test.cpp b/src/mongo/db/query/cursor_response_test.cpp
index 2a3424f2029..57912f4ca9b 100644
--- a/src/mongo/db/query/cursor_response_test.cpp
+++ b/src/mongo/db/query/cursor_response_test.cpp
@@ -211,6 +211,53 @@ TEST(CursorResponseTest, parseFromBSONPartialResultsReturnedFieldWrongType) {
     ASSERT_NOT_OK(result.getStatus());
 }
 
+TEST(CursorResponseTest, parseFromBSONVarsFieldCorrect) {
+    BSONObj varsContents = BSON("randomVar" << 7);
+    StatusWith<CursorResponse> result = CursorResponse::parseFromBSON(BSON(
+        "cursor" << BSON("id" << CursorId(123) << "ns"
+                              << "db.coll"
+                              << "firstBatch" << BSON_ARRAY(BSON("_id" << 1) << BSON("_id" << 2)))
+                 << "vars" << varsContents << "ok" << 1));
+    ASSERT_OK(result.getStatus());
+
+    CursorResponse response = std::move(result.getValue());
+    ASSERT_EQ(response.getCursorId(), CursorId(123));
+    ASSERT_EQ(response.getNSS().ns(), "db.coll");
+    ASSERT_EQ(response.getBatch().size(), 2U);
+    ASSERT_BSONOBJ_EQ(response.getBatch()[0], BSON("_id" << 1));
+    ASSERT_BSONOBJ_EQ(response.getBatch()[1], BSON("_id" << 2));
+    ASSERT_TRUE(response.getVarsField());
+    ASSERT_BSONOBJ_EQ(response.getVarsField().get(), varsContents);
+}
+
+TEST(CursorResponseTest, parseFromBSONVarsFieldWrongType) {
+    StatusWith<CursorResponse> result = CursorResponse::parseFromBSON(BSON(
+        "cursor" << BSON("id" << CursorId(123) << "ns"
+                              << "db.coll"
+                              << "firstBatch" << BSON_ARRAY(BSON("_id" << 1) << BSON("_id" << 2)))
+                 << "vars" << 2 << "ok" << 1));
+    ASSERT_NOT_OK(result.getStatus());
+}
+
+TEST(CursorResponseTest, parseFromBSONMultipleVars) {
+    BSONObj varsContents = BSON("randomVar" << 7 << "otherVar" << BSON("nested" << 2));
+    StatusWith<CursorResponse> result = CursorResponse::parseFromBSON(BSON(
+        "cursor" << BSON("id" << CursorId(123) << "ns"
+                              << "db.coll"
+                              << "firstBatch" << BSON_ARRAY(BSON("_id" << 1) << BSON("_id" << 2)))
+                 << "vars" << varsContents << "ok" << 1));
+    ASSERT_OK(result.getStatus());
+
+    CursorResponse response = std::move(result.getValue());
+    ASSERT_EQ(response.getCursorId(), CursorId(123));
+    ASSERT_EQ(response.getNSS().ns(), "db.coll");
+    ASSERT_EQ(response.getBatch().size(), 2U);
+    ASSERT_BSONOBJ_EQ(response.getBatch()[0], BSON("_id" << 1));
+    ASSERT_BSONOBJ_EQ(response.getBatch()[1], BSON("_id" << 2));
+    ASSERT_TRUE(response.getVarsField());
+    ASSERT_BSONOBJ_EQ(response.getVarsField().get(), varsContents);
+}
+
 TEST(CursorResponseTest, roundTripThroughCursorResponseBuilderWithPartialResultsReturned) {
     CursorResponseBuilder::Options options;
     options.isInitialResponse = true;
@@ -301,6 +348,7 @@ TEST(CursorResponseTest, toBSONPartialResultsReturned) {
                             boost::none,
                             boost::none,
                             boost::none,
+                            boost::none,
                             true);
     BSONObj responseObj = response.toBSON(CursorResponse::ResponseType::InitialResponse);
     BSONObj expectedResponse = BSON(
diff --git a/src/mongo/executor/task_executor_cursor.cpp b/src/mongo/executor/task_executor_cursor.cpp
index 7e2e4b5c5e0..d9576fe521a 100644
--- a/src/mongo/executor/task_executor_cursor.cpp
+++ b/src/mongo/executor/task_executor_cursor.cpp
@@ -159,6 +159,8 @@ void TaskExecutorCursor::_getNextBatch(OperationContext* opCtx) {
     if (_cursorId == kUnitializedCursorId) {
         _ns = cr.getNSS();
         _rcr.dbname = _ns.db().toString();
+        // 'vars' are only included in the first batch.
+        _cursorVars = cr.getVarsField();
     }
 
     _cursorId = cr.getCursorId();
diff --git a/src/mongo/executor/task_executor_cursor.h b/src/mongo/executor/task_executor_cursor.h
index 020d4986039..8d6bd5dc0a0 100644
--- a/src/mongo/executor/task_executor_cursor.h
+++ b/src/mongo/executor/task_executor_cursor.h
@@ -107,6 +107,10 @@ public:
         return toRet;
     }
 
+    boost::optional<BSONObj> getCursorVars() {
+        return _cursorVars;
+    }
+
 private:
     /**
      * Runs a remote command and pipes the output back to this object
@@ -138,6 +142,9 @@ private:
 
     CursorId _cursorId = kUnitializedCursorId;
 
+    // Variables sent alongside the results in the cursor.
+    boost::optional<BSONObj> _cursorVars = boost::none;
+
     // This is a sum of the time spent waiting on remote calls.
     Milliseconds _millisecondsWaiting = Milliseconds(0);
 
diff --git a/src/mongo/s/query/cluster_find.cpp b/src/mongo/s/query/cluster_find.cpp
index c6173e6e58b..7ed9694090c 100644
--- a/src/mongo/s/query/cluster_find.cpp
+++ b/src/mongo/s/query/cluster_find.cpp
@@ -885,6 +885,7 @@ StatusWith<CursorResponse> ClusterFind::runGetMore(OperationContext* opCtx,
                           startingFrom,
                           postBatchResumeToken,
                           boost::none,
+                          boost::none,
                           partialResultsReturned);
 }
 
diff --git a/src/mongo/s/query/establish_cursors.cpp b/src/mongo/s/query/establish_cursors.cpp
index 4c44a93b2c9..3ac04e4da84 100644
--- a/src/mongo/s/query/establish_cursors.cpp
+++ b/src/mongo/s/query/establish_cursors.cpp
@@ -248,7 +248,7 @@ void CursorEstablisher::_handleFailure(const AsyncRequestsSender::Response& resp
         // This exception is eligible to be swallowed. Add an entry with a cursorID of 0, an
         // empty HostAndPort, and which has the 'partialResultsReturned' flag set to true.
         _remoteCursors.push_back(
-            {response.shardId.toString(), {}, {_nss, CursorId{0}, {}, {}, {}, {}, {}, true}});
+            {response.shardId.toString(), {}, {_nss, CursorId{0}, {}, {}, {}, {}, {}, {}, true}});
         return;
     }