SERVER-64875 Enable `collectionUUID` parameter on `aggregate` command more generally on mongos

3 files changed, 79 insertions, 31 deletions

diff --git a/jstests/aggregation/aggregation_with_uuids.js b/jstests/aggregation/aggregation_with_uuids.js
index b5f70c54c89..eebe1ac189f 100644
--- a/jstests/aggregation/aggregation_with_uuids.js
+++ b/jstests/aggregation/aggregation_with_uuids.js
@@ -24,18 +24,9 @@ const validateErrorResponse = function(
     assert.eq(res.actualCollection, actualCollection);
 };
 
-// On mongos, collectionUUID is only supported for $collStats and $indexStats aggregations.
-if (FixtureHelpers.isMongos(db)) {
-    assert.commandFailedWithCode(
-        testDB.runCommand(
-            {aggregate: 1, collectionUUID: UUID(), pipeline: [{$match: {}}], cursor: {}}),
-        6256300);
-    return;
-}
-
 const docs = [{_id: 1}, {_id: 2}];
 
-testColl.drop({writeConcern: {w: "majority"}});
+assert.commandWorked(testDB.dropDatabase());
 assert.commandWorked(testColl.insert(docs));
 
 // Get the namespace's initial UUID.
@@ -99,10 +90,9 @@ assert.sameMembers(collNameRes.cursor.firstBatch, docs);
 
 // An aggregation with a collectionUUID should fail with CollectionUUIDMismatch 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"}}));
+const viewName = 'view';
+assert.commandWorked(testDB.runCommand(
+    {create: viewName, viewOn: testColl.getName(), pipeline: [], writeConcern: {w: "majority"}}));
 
 res = assert.commandFailedWithCode(
     testDB.runCommand(
@@ -128,7 +118,7 @@ assert.commandFailedWithCode(
 
 // collectionUUID is not allowed with collectionless aggregations.
 assert.commandFailedWithCode(
-    testDB.runCommand(
-        {aggregate: 1, collectionUUID: uuid, pipeline: [{$listLocalSessions: {}}], cursor: {}}),
+    testDB.adminCommand(
+        {aggregate: 1, collectionUUID: uuid, pipeline: [{$currentOp: {}}], cursor: {}}),
     4928901);
 })();
diff --git a/jstests/sharding/collection_uuid_aggregate.js b/jstests/sharding/collection_uuid_aggregate.js
index 11a864cc80c..3607ff6bb61 100644
--- a/jstests/sharding/collection_uuid_aggregate.js
+++ b/jstests/sharding/collection_uuid_aggregate.js
@@ -9,16 +9,18 @@
 (function() {
 'use strict';
 
-const st = new ShardingTest({shards: 2});
+const st = new ShardingTest({shards: 3});
 
 const db = st.s.getDB(jsTestName());
 assert.commandWorked(st.s.adminCommand({enableSharding: db.getName()}));
 st.ensurePrimaryShard(db.getName(), st.shard0.shardName);
 
-const shardedColl = db.sharded;
+const shardedColl1 = db.sharded_1;
+const shardedColl2 = db.sharded_2;
 const unshardedColl = db.unsharded;
 
-assert.commandWorked(shardedColl.insert({_id: 0}));
+assert.commandWorked(shardedColl1.insert({_id: 0}));
+assert.commandWorked(shardedColl2.insert({_id: 1}));
 assert.commandWorked(unshardedColl.insert({_id: 2}));
 
 const uuid = function(coll) {
@@ -28,15 +30,28 @@ const uuid = function(coll) {
 };
 
 assert.commandWorked(
-    st.s.adminCommand({shardCollection: shardedColl.getFullName(), key: {_id: 1}}));
+    st.s.adminCommand({shardCollection: shardedColl1.getFullName(), key: {_id: 1}}));
+assert.commandWorked(
+    st.s.adminCommand({shardCollection: shardedColl2.getFullName(), key: {_id: 1}}));
 
-// Move the chunk to shard1.
+// Move shardedColl1's chunk to shard1 and shardedColl2's chunk to shard2.
+assert.commandWorked(st.s.adminCommand(
+    {moveChunk: shardedColl1.getFullName(), find: {_id: 0}, to: st.shard1.shardName}));
 assert.commandWorked(st.s.adminCommand(
-    {moveChunk: shardedColl.getFullName(), find: {_id: 0}, to: st.shard1.shardName}));
+    {moveChunk: shardedColl2.getFullName(), find: {_id: 0}, to: st.shard2.shardName}));
+
+// Cannot use the collectionUUID parameter with a pipeline that executes entirely on mongos.
+assert.commandFailedWithCode(db.adminCommand({
+    aggregate: 1,
+    pipeline: [{$currentOp: {localOps: true}}],
+    cursor: {},
+    collectionUUID: uuid(unshardedColl),
+}),
+                             6487500);
 
-// Run an aggregate which will only target shard1, since shard0 does not own any chunks.
+// Run an aggregate which will only target shard1, since shard0 and shard2 do not own any chunks.
 let res = assert.commandFailedWithCode(db.runCommand({
-    aggregate: shardedColl.getName(),
+    aggregate: shardedColl1.getName(),
     pipeline: [{$indexStats: {}}],
     cursor: {},
     collectionUUID: uuid(unshardedColl),
@@ -45,8 +60,49 @@ let res = assert.commandFailedWithCode(db.runCommand({
 jsTestLog('$indexStats result: ' + tojson(res));
 assert.eq(res.db, db.getName());
 assert.eq(res.collectionUUID, uuid(unshardedColl));
-assert.eq(res.expectedCollection, shardedColl.getName());
+assert.eq(res.expectedCollection, shardedColl1.getName());
 assert.eq(res.actualCollection, unshardedColl.getName());
 
+res = assert.commandFailedWithCode(db.runCommand({
+    aggregate: shardedColl1.getName(),
+    pipeline: [{$indexStats: {}}],
+    cursor: {},
+    collectionUUID: uuid(shardedColl2),
+}),
+                                   ErrorCodes.CollectionUUIDMismatch);
+jsTestLog('$indexStats result: ' + tojson(res));
+assert.eq(res.db, db.getName());
+assert.eq(res.collectionUUID, uuid(shardedColl2));
+assert.eq(res.expectedCollection, shardedColl1.getName());
+assert.eq(res.actualCollection, shardedColl2.getName());
+
+// Run an aggregate which will only target shard1, since that is the shard on which the chunk
+// containing the matching document lives.
+res = assert.commandFailedWithCode(db.runCommand({
+    aggregate: shardedColl1.getName(),
+    pipeline: [{$match: {_id: 0}}],
+    cursor: {},
+    collectionUUID: uuid(unshardedColl),
+}),
+                                   ErrorCodes.CollectionUUIDMismatch);
+jsTestLog('$match result: ' + tojson(res));
+assert.eq(res.db, db.getName());
+assert.eq(res.collectionUUID, uuid(unshardedColl));
+assert.eq(res.expectedCollection, shardedColl1.getName());
+assert.eq(res.actualCollection, unshardedColl.getName());
+
+res = assert.commandFailedWithCode(db.runCommand({
+    aggregate: shardedColl1.getName(),
+    pipeline: [{$match: {_id: 0}}],
+    cursor: {},
+    collectionUUID: uuid(shardedColl2),
+}),
+                                   ErrorCodes.CollectionUUIDMismatch);
+jsTestLog('$match result: ' + tojson(res));
+assert.eq(res.db, db.getName());
+assert.eq(res.collectionUUID, uuid(shardedColl2));
+assert.eq(res.expectedCollection, shardedColl1.getName());
+assert.eq(res.actualCollection, shardedColl2.getName());
+
 st.stop();
 })();
diff --git a/src/mongo/s/query/cluster_aggregate.cpp b/src/mongo/s/query/cluster_aggregate.cpp
index bd291ed6f70..2034489bb49 100644
--- a/src/mongo/s/query/cluster_aggregate.cpp
+++ b/src/mongo/s/query/cluster_aggregate.cpp
@@ -302,12 +302,6 @@ Status ClusterAggregate::runAggregate(OperationContext* opCtx,
     auto involvedNamespaces = liteParsedPipeline.getInvolvedNamespaces();
     auto shouldDoFLERewrite = ::mongo::shouldDoFLERewrite(request);
 
-    uassert(6256300,
-            str::stream() << "On mongos, " << AggregateCommandRequest::kCollectionUUIDFieldName
-                          << " is only supported for $collStats and $indexStats aggregation.",
-            !request.getCollectionUUID() || liteParsedPipeline.startsWithCollStats() ||
-                liteParsedPipeline.startsWithIndexStats());
-
     // If the routing table is not already taken by the higher level, fill it now.
     if (!cm) {
         // If the routing table is valid, we obtain a reference to it. If the table is not valid,
@@ -396,6 +390,14 @@ Status ClusterAggregate::runAggregate(OperationContext* opCtx,
         allowedToPassthrough,
         request.getPassthroughToShard().has_value());
 
+    uassert(
+        6487500,
+        fmt::format("Cannot use {} with an aggregation that executes entirely on mongos",
+                    AggregateCommandRequest::kCollectionUUIDFieldName),
+        !request.getCollectionUUID() ||
+            targeter.policy !=
+                cluster_aggregation_planner::AggregationTargeter::TargetingPolicy::kMongosRequired);
+
     if (!expCtx) {
         // When the AggregationTargeter chooses a "passthrough" policy, it does not call the
         // 'pipelineBuilder' function, so we never get an expression context. Because this is a