SERVER-49289 Add collectionUUID option to aggregate

6 files changed, 188 insertions, 2 deletions

diff --git a/jstests/aggregation/aggregation_with_uuids.js b/jstests/aggregation/aggregation_with_uuids.js
new file mode 100644
index 00000000000..a394489ed3b
--- /dev/null
+++ b/jstests/aggregation/aggregation_with_uuids.js
@@ -0,0 +1,122 @@
+/**
+ * Tests for aggregation requests with the collectionUUID parameter.
+ * @tags: [
+ *   requires_fcv_47,
+ *   # Change stream aggregations don't support read concerns other than 'majority'
+ *   assumes_read_concern_unchanged,
+ * ]
+ */
+(function() {
+"use strict";
+
+load("jstests/libs/fixture_helpers.js");  // For 'isMongos'
+
+const dbName = jsTestName();
+const collName = "foo";
+
+const testDB = db.getSiblingDB(dbName);
+const testColl = testDB.getCollection(collName);
+
+if (FixtureHelpers.isMongos(db)) {
+    // collectionUUID is not supported on mongos.
+    assert.commandFailedWithCode(
+        testDB.runCommand(
+            {aggregate: 1, collectionUUID: UUID(), pipeline: [{$match: {}}], cursor: {}}),
+        4928902);
+    return;
+}
+
+const docs = [{_id: 1}, {_id: 2}];
+
+testColl.drop({writeConcern: {w: "majority"}});
+assert.commandWorked(testColl.insert(docs));
+
+// Get the namespace's initial UUID.
+let collectionInfos = testDB.getCollectionInfos({name: collName});
+let uuid = collectionInfos[0].info.uuid;
+assert(uuid, "Expected collection " + collName + " to have a UUID.");
+
+// An aggregation with the UUID should succeed and find the same documents as an aggregation with
+// the collection name.
+let uuidRes = assert.commandWorked(testDB.runCommand(
+    {aggregate: collName, collectionUUID: uuid, pipeline: [{$match: {}}], cursor: {}}));
+assert.sameMembers(uuidRes.cursor.firstBatch, docs);
+
+let collNameRes = assert.commandWorked(
+    testDB.runCommand({aggregate: collName, pipeline: [{$match: {}}], cursor: {}}));
+assert.sameMembers(collNameRes.cursor.firstBatch, uuidRes.cursor.firstBatch);
+
+// getMore should work with cursors created by an aggregation with a uuid.
+uuidRes = assert.commandWorked(testDB.runCommand(
+    {aggregate: collName, pipeline: [{$match: {}}, {$sort: {_id: 1}}], cursor: {batchSize: 1}}));
+assert.eq(1, uuidRes.cursor.firstBatch.length, tojson(uuidRes));
+assert.eq(docs[0], uuidRes.cursor.firstBatch[0], tojson(uuidRes));
+
+const getMoreRes =
+    assert.commandWorked(testDB.runCommand({getMore: uuidRes.cursor.id, collection: collName}));
+assert.eq(1, getMoreRes.cursor.nextBatch.length, tojson(getMoreRes));
+assert.eq(docs[1], getMoreRes.cursor.nextBatch[0], tojson(getMoreRes));
+assert.eq(0, getMoreRes.cursor.id, tojson(getMoreRes));
+
+// An aggregation with collectionUUID throws NamespaceNotFound if the namespace does not exist, even
+// if a collection does exist with the given uuid.
+assert.commandFailedWithCode(
+    testDB.runCommand(
+        {aggregate: "doesNotExist", collectionUUID: uuid, pipeline: [{$match: {}}], cursor: {}}),
+    ErrorCodes.NamespaceNotFound);
+
+// Drop the collection.
+testColl.drop({writeConcern: {w: "majority"}});
+
+// An aggregation with the initial UUID should fail since the namespace doesn't exist.
+assert.commandFailedWithCode(
+    testDB.runCommand(
+        {aggregate: collName, collectionUUID: uuid, pipeline: [{$match: {}}], cursor: {}}),
+    ErrorCodes.NamespaceNotFound);
+
+// Now recreate the collection.
+assert.commandWorked(testColl.insert(docs));
+
+// An aggregation with the initial UUID should still fail despite the namespace existing.
+assert.commandFailedWithCode(
+    testDB.runCommand(
+        {aggregate: collName, collectionUUID: uuid, pipeline: [{$match: {}}], cursor: {}}),
+    ErrorCodes.NamespaceNotFound);
+
+collNameRes = assert.commandWorked(
+    testDB.runCommand({aggregate: collName, pipeline: [{$match: {}}], cursor: {}}));
+assert.sameMembers(collNameRes.cursor.firstBatch, docs);
+
+//
+// Tests for rejecting invalid collectionUUIDs and cases where collectionUUID is not allowed.
+//
+
+// collectionUUID must be a UUID.
+assert.commandFailedWithCode(
+    testDB.runCommand(
+        {aggregate: collName, collectionUUID: "NotAUUID", pipeline: [{$match: {}}], cursor: {}}),
+    ErrorCodes.InvalidUUID);
+
+// collectionUUID is not allowed with change streams.
+assert.commandFailedWithCode(
+    testDB.runCommand(
+        {aggregate: collName, collectionUUID: uuid, pipeline: [{$changeStream: {}}], cursor: {}}),
+    4928900);
+
+// collectionUUID is not allowed with collectionless aggregations.
+assert.commandFailedWithCode(
+    testDB.runCommand(
+        {aggregate: 1, collectionUUID: uuid, pipeline: [{$listLocalSessions: {}}], cursor: {}}),
+    4928901);
+
+// Aggregation with collectionUUID throws OptionNotSupportedOnView if the namespace is a view.
+const testView = testDB.getCollection("viewCollection");
+testView.drop({writeConcern: {w: "majority"}});
+assert.commandWorked(testView.runCommand(
+    "create", {viewOn: testColl.getName(), pipeline: [], writeConcern: {w: "majority"}}));
+
+assert.commandFailedWithCode(
+    testDB.runCommand(
+        {aggregate: "viewCollection", collectionUUID: uuid, pipeline: [{$match: {}}], cursor: {}}),
+    ErrorCodes.OptionNotSupportedOnView);
+})();
diff --git a/src/mongo/db/commands/run_aggregate.cpp b/src/mongo/db/commands/run_aggregate.cpp
index 74c118f5b8f..13121ed3b15 100644
--- a/src/mongo/db/commands/run_aggregate.cpp
+++ b/src/mongo/db/commands/run_aggregate.cpp
@@ -514,6 +514,11 @@ Status runAggregate(OperationContext* opCtx,
 
         // If this is a change stream, perform special checks and change the execution namespace.
         if (liteParsedPipeline.hasChangeStream()) {
+            uassert(4928900,
+                    str::stream() << AggregationRequest::kCollectionUUIDName
+                                  << " is not supported for a change stream",
+                    !request.getCollectionUUID());
+
             // Replace the execution namespace with that of the oplog.
             nss = NamespaceString::kRsOplogNamespace;
 
@@ -535,6 +540,11 @@ Status runAggregate(OperationContext* opCtx,
             // Obtain collection locks on the execution namespace; that is, the oplog.
             ctx.emplace(opCtx, nss, AutoGetCollectionViewMode::kViewsForbidden);
         } else if (nss.isCollectionlessAggregateNS() && pipelineInvolvedNamespaces.empty()) {
+            uassert(4928901,
+                    str::stream() << AggregationRequest::kCollectionUUIDName
+                                  << " is not supported for a collectionless aggregation",
+                    !request.getCollectionUUID());
+
             // If this is a collectionless agg with no foreign namespaces, don't acquire any locks.
             statsTracker.emplace(opCtx,
                                  nss,
@@ -563,6 +573,10 @@ Status runAggregate(OperationContext* opCtx,
         if (ctx && ctx->getView() && !liteParsedPipeline.startsWithCollStats()) {
             invariant(nss != NamespaceString::kRsOplogNamespace);
             invariant(!nss.isCollectionlessAggregateNS());
+            uassert(ErrorCodes::OptionNotSupportedOnView,
+                    str::stream() << AggregationRequest::kCollectionUUIDName
+                                  << " is not supported against a view",
+                    !request.getCollectionUUID());
 
             // Check that the default collation of 'view' is compatible with the operation's
             // collation. The check is skipped if the request did not specify a collation.
@@ -600,6 +614,14 @@ Status runAggregate(OperationContext* opCtx,
             return status;
         }
 
+        if (request.getCollectionUUID()) {
+            // If the namespace is not a view and collectionUUID was provided, verify the collection
+            // exists and has the expected UUID.
+            uassert(ErrorCodes::NamespaceNotFound,
+                    "No collection found with the given namespace and UUID",
+                    uuid && uuid == *request.getCollectionUUID());
+        }
+
         invariant(collatorToUse);
         expCtx = makeExpressionContext(opCtx, request, std::move(*collatorToUse), uuid);
 
diff --git a/src/mongo/db/pipeline/aggregation_request.cpp b/src/mongo/db/pipeline/aggregation_request.cpp
index 60fe025772a..a948696b526 100644
--- a/src/mongo/db/pipeline/aggregation_request.cpp
+++ b/src/mongo/db/pipeline/aggregation_request.cpp
@@ -205,6 +205,13 @@ StatusWith<AggregationRequest> AggregationRequest::parseFromBSON(
             auto bob = BSONObjBuilder{request.getLetParameters()};
             bob.appendElementsUnique(elem.embeddedObject());
             request._letParameters = bob.obj();
+        } else if (kCollectionUUIDName == fieldName) {
+            auto collectionUUIDSW = UUID::parse(elem);
+            if (!collectionUUIDSW.isOK()) {
+                return collectionUUIDSW.getStatus();
+            }
+
+            request.setCollectionUUID(collectionUUIDSW.getValue());
         } else if (fieldName == kUse44SortKeysName) {
             if (elem.type() != BSONType::Bool) {
                 return {ErrorCodes::TypeMismatch,
@@ -331,6 +338,8 @@ Document AggregationRequest::serializeToCommandObj() const {
         {kRuntimeConstantsName, _runtimeConstants ? Value(_runtimeConstants->toBSON()) : Value()},
         {kIsMapReduceCommandName, _isMapReduceCommand ? Value(true) : Value()},
         {kLetName, !_letParameters.isEmpty() ? Value(_letParameters) : Value()},
+        // Only serialize collection UUID if one was specified.
+        {kCollectionUUIDName, _collectionUUID ? Value(*_collectionUUID) : Value()},
     };
 }
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/aggregation_request.h b/src/mongo/db/pipeline/aggregation_request.h
index 454c246feb2..66e4d95df57 100644
--- a/src/mongo/db/pipeline/aggregation_request.h
+++ b/src/mongo/db/pipeline/aggregation_request.h
@@ -66,6 +66,7 @@ public:
     static constexpr StringData kUse44SortKeysName = "use44SortKeys"_sd;
     static constexpr StringData kIsMapReduceCommandName = "isMapReduceCommand"_sd;
     static constexpr StringData kLetName = "let"_sd;
+    static constexpr StringData kCollectionUUIDName = "collectionUUID"_sd;
 
     static constexpr long long kDefaultBatchSize = 101;
 
@@ -227,6 +228,10 @@ public:
         return _isMapReduceCommand;
     }
 
+    const auto& getCollectionUUID() const {
+        return _collectionUUID;
+    }
+
     //
     // Setters for optional fields.
     //
@@ -299,6 +304,10 @@ public:
         _isMapReduceCommand = isMapReduce;
     }
 
+    void setCollectionUUID(UUID collectionUUID) {
+        _collectionUUID = std::move(collectionUUID);
+    }
+
 private:
     // Required fields.
     const NamespaceString _nss;
@@ -349,6 +358,9 @@ private:
     // $$NOW).
     boost::optional<RuntimeConstants> _runtimeConstants;
 
+    // The expected UUID of the namespace the aggregation executes on.
+    boost::optional<UUID> _collectionUUID;
+
     // A document containing user-specified let parameter constants; i.e. values that do not change
     // once computed.
     BSONObj _letParameters;
diff --git a/src/mongo/db/pipeline/aggregation_request_test.cpp b/src/mongo/db/pipeline/aggregation_request_test.cpp
index f16b7a5891b..04a93a21ea1 100644
--- a/src/mongo/db/pipeline/aggregation_request_test.cpp
+++ b/src/mongo/db/pipeline/aggregation_request_test.cpp
@@ -56,12 +56,17 @@ const Document kDefaultCursorOptionDocument{
 
 TEST(AggregationRequestTest, ShouldParseAllKnownOptions) {
     NamespaceString nss("a.collection");
-    const BSONObj inputBson = fromjson(
+    BSONObj inputBson = fromjson(
         "{pipeline: [{$match: {a: 'abc'}}], explain: false, allowDiskUse: true, fromMongos: true, "
         "needsMerge: true, bypassDocumentValidation: true, collation: {locale: 'en_US'}, cursor: "
         "{batchSize: 10}, hint: {a: 1}, maxTimeMS: 100, readConcern: {level: 'linearizable'}, "
         "$queryOptions: {$readPreference: 'nearest'}, exchange: {policy: "
         "'roundrobin', consumers:NumberInt(2)}, isMapReduceCommand: true}");
+    auto uuid = UUID::gen();
+    BSONObjBuilder uuidBob;
+    uuid.appendToBuilder(&uuidBob, AggregationRequest::kCollectionUUIDName);
+    inputBson = inputBson.addField(uuidBob.obj().firstElement());
+
     auto request = unittest::assertGet(AggregationRequest::parseFromBSON(nss, inputBson));
     ASSERT_FALSE(request.getExplain());
     ASSERT_TRUE(request.shouldAllowDiskUse());
@@ -82,6 +87,7 @@ TEST(AggregationRequestTest, ShouldParseAllKnownOptions) {
                            << "nearest"));
     ASSERT_TRUE(request.getExchangeSpec().is_initialized());
     ASSERT_TRUE(request.getIsMapReduceCommand());
+    ASSERT_EQ(*request.getCollectionUUID(), uuid);
 }
 
 TEST(AggregationRequestTest, ShouldParseExplicitExplainTrue) {
@@ -193,6 +199,8 @@ TEST(AggregationRequestTest, ShouldSerializeOptionalValuesIfSet) {
     const auto letParamsObj = BSON("foo"
                                    << "bar");
     request.setLetParameters(letParamsObj);
+    auto uuid = UUID::gen();
+    request.setCollectionUUID(uuid);
 
     auto expectedSerialization =
         Document{{AggregationRequest::kCommandName, nss.coll()},
@@ -209,7 +217,8 @@ TEST(AggregationRequestTest, ShouldSerializeOptionalValuesIfSet) {
                  {QueryRequest::kUnwrappedReadPrefField, readPrefObj},
                  {QueryRequest::cmdOptionMaxTimeMS, 10},
                  {AggregationRequest::kIsMapReduceCommandName, true},
-                 {AggregationRequest::kLetName, letParamsObj}};
+                 {AggregationRequest::kLetName, letParamsObj},
+                 {AggregationRequest::kCollectionUUIDName, uuid}};
     ASSERT_DOCUMENT_EQ(request.serializeToCommandObj(), expectedSerialization);
 }
 
@@ -503,6 +512,14 @@ TEST(AggregationRequestTest, ShouldRejectInvalidWriteConcern) {
         fromjson("{pipeline: [{$match: {a: 'abc'}}], cursor: {}, writeConcern: 'invalid'}");
     ASSERT_NOT_OK(AggregationRequest::parseFromBSON(nss, inputBson).getStatus());
 }
+
+TEST(AggregationRequestTest, ShouldRejectInvalidCollectionUUID) {
+    NamespaceString nss("a.collection");
+    const BSONObj inputBSON = fromjson("{pipeline: [{$match: {}}], collectionUUID: 2}");
+    ASSERT_EQUALS(AggregationRequest::parseFromBSON(nss, inputBSON).getStatus().code(),
+                  ErrorCodes::InvalidUUID);
+}
+
 //
 // Ignore fields parsed elsewhere.
 //
diff --git a/src/mongo/s/query/cluster_aggregate.cpp b/src/mongo/s/query/cluster_aggregate.cpp
index ec2419f1aa3..56103395066 100644
--- a/src/mongo/s/query/cluster_aggregate.cpp
+++ b/src/mongo/s/query/cluster_aggregate.cpp
@@ -213,6 +213,10 @@ Status ClusterAggregate::runAggregate(OperationContext* opCtx,
                           << ", " << AggregationRequest::kFromMongosName
                           << "] cannot be set to 'true' when sent to mongos",
             !request.needsMerge() && !request.isFromMongos());
+    uassert(4928902,
+            str::stream() << AggregationRequest::kCollectionUUIDName
+                          << " is not supported on a mongos",
+            !request.getCollectionUUID());
 
     const auto isSharded = [](OperationContext* opCtx, const NamespaceString& nss) {
         const auto resolvedNsRoutingInfo =