SERVER-36123: Reject $out with mode: replaceCollection if the output collection is sharded

6 files changed, 40 insertions, 14 deletions

diff --git a/src/mongo/db/pipeline/document_source_out.cpp b/src/mongo/db/pipeline/document_source_out.cpp
index 02d2caf3374..8d8d3db49e8 100644
--- a/src/mongo/db/pipeline/document_source_out.cpp
+++ b/src/mongo/db/pipeline/document_source_out.cpp
@@ -49,8 +49,10 @@ std::unique_ptr<LiteParsedDocumentSourceForeignCollections> DocumentSourceOut::l
             spec.type() == BSONType::String || spec.type() == BSONType::Object);
 
     NamespaceString targetNss;
+    bool allowSharded;
     if (spec.type() == BSONType::String) {
         targetNss = NamespaceString(request.getNamespaceString().db(), spec.valueStringData());
+        allowSharded = false;
     } else if (spec.type() == BSONType::Object) {
         auto outSpec =
             DocumentSourceOutSpec::parse(IDLParserErrorContext("$out"), spec.embeddedObject());
@@ -61,6 +63,9 @@ std::unique_ptr<LiteParsedDocumentSourceForeignCollections> DocumentSourceOut::l
             targetNss =
                 NamespaceString(request.getNamespaceString().db(), outSpec.getTargetCollection());
         }
+
+        // Sharded output collections are not allowed with mode "replaceCollection".
+        allowSharded = outSpec.getMode() != WriteModeEnum::kModeReplaceCollection;
     }
 
     uassert(ErrorCodes::InvalidNamespace,
@@ -74,7 +79,6 @@ std::unique_ptr<LiteParsedDocumentSourceForeignCollections> DocumentSourceOut::l
 
     PrivilegeVector privileges{Privilege(ResourcePattern::forExactNamespace(targetNss), actions)};
 
-    constexpr bool allowSharded = true;
     return stdx::make_unique<LiteParsedDocumentSourceForeignCollections>(
         std::move(targetNss), std::move(privileges), allowSharded);
 }
@@ -192,13 +196,21 @@ intrusive_ptr<DocumentSource> DocumentSourceOut::createFromBson(
         } else {
             outputNs = NamespaceString(expCtx->ns.db(), spec.getTargetCollection());
         }
-
     } else {
         uasserted(16990,
                   str::stream() << "$out only supports a string or object argument, not "
                                 << typeName(elem.type()));
     }
 
+    // Although we perform a check for "replaceCollection" mode with a sharded output collection
+    // during lite parsing, we need to do it here as well in case mongos is stale or the command is
+    // sent directly to the shard.
+    uassert(17017,
+            str::stream() << "$out with mode " << WriteMode_serializer(mode)
+                          << " is not supported to an existing *sharded* output collection.",
+            !(mode == WriteModeEnum::kModeReplaceCollection &&
+              expCtx->mongoProcessInterface->isSharded(expCtx->opCtx, outputNs)));
+
     uassert(17385, "Can't $out to special collection: " + outputNs.coll(), !outputNs.isSpecial());
 
     switch (mode) {
diff --git a/src/mongo/db/pipeline/document_source_out_replace_coll.cpp b/src/mongo/db/pipeline/document_source_out_replace_coll.cpp
index 7b0be7fd252..7b0ab47f256 100644
--- a/src/mongo/db/pipeline/document_source_out_replace_coll.cpp
+++ b/src/mongo/db/pipeline/document_source_out_replace_coll.cpp
@@ -46,14 +46,9 @@ void DocumentSourceOutReplaceColl::initializeWriteNs() {
     _originalOutOptions = pExpCtx->mongoProcessInterface->getCollectionOptions(outputNs);
     _originalIndexes = conn->getIndexSpecs(outputNs.ns());
 
-    // Check if it's sharded or capped to make sure we have a chance of succeeding before we do
-    // all the work. If the collection becomes capped during processing, the collection options will
-    // have changed, and the $out will fail. If it becomes sharded during processing, the final
-    // rename will fail.
-    uassert(17017,
-            str::stream() << "namespace '" << outputNs.ns()
-                          << "' is sharded so it can't be used for $out'",
-            !pExpCtx->mongoProcessInterface->isSharded(pExpCtx->opCtx, outputNs));
+    // Check if it's capped to make sure we have a chance of succeeding before we do all the work.
+    // If the collection becomes capped during processing, the collection options will have changed,
+    // and the $out will fail.
     uassert(17152,
             str::stream() << "namespace '" << outputNs.ns()
                           << "' is capped so it can't be used for $out",
diff --git a/src/mongo/db/pipeline/document_source_out_test.cpp b/src/mongo/db/pipeline/document_source_out_test.cpp
index 1c01b88ab35..a8e4a1bec78 100644
--- a/src/mongo/db/pipeline/document_source_out_test.cpp
+++ b/src/mongo/db/pipeline/document_source_out_test.cpp
@@ -44,8 +44,24 @@ StringData kModeFieldName = DocumentSourceOutSpec::kModeFieldName;
 StringData kUniqueKeyFieldName = DocumentSourceOutSpec::kUniqueKeyFieldName;
 StringData kDefaultMode = WriteMode_serializer(WriteModeEnum::kModeReplaceCollection);
 
+/**
+ * For the purpsoses of this test, assume every collection is unsharded. Stages may ask this during
+ * setup. For example, to compute its constraints, the $out stage needs to know if the output
+ * collection is sharded.
+ */
+class MongoProcessInterfaceForTest : public StubMongoProcessInterface {
+public:
+    bool isSharded(OperationContext* opCtx, const NamespaceString& ns) override {
+        return false;
+    }
+};
+
 class DocumentSourceOutTest : public AggregationContextFixture {
 public:
+    DocumentSourceOutTest() : AggregationContextFixture() {
+        getExpCtx()->mongoProcessInterface = std::make_shared<MongoProcessInterfaceForTest>();
+    }
+
     intrusive_ptr<DocumentSourceOut> createOutStage(BSONObj spec) {
         auto specElem = spec.firstElement();
         intrusive_ptr<DocumentSourceOut> outStage = dynamic_cast<DocumentSourceOut*>(
diff --git a/src/mongo/db/pipeline/mongos_process_interface.cpp b/src/mongo/db/pipeline/mongos_process_interface.cpp
index 7e3e451a63c..1e25b490f9c 100644
--- a/src/mongo/db/pipeline/mongos_process_interface.cpp
+++ b/src/mongo/db/pipeline/mongos_process_interface.cpp
@@ -204,4 +204,9 @@ std::vector<GenericCursor> MongoSInterface::getCursors(
     return cursorManager->getAllCursors();
 }
 
+bool MongoSInterface::isSharded(OperationContext* opCtx, const NamespaceString& nss) {
+    auto routingInfo = Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss);
+    return routingInfo.isOK() && routingInfo.getValue().cm();
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/mongos_process_interface.h b/src/mongo/db/pipeline/mongos_process_interface.h
index 1c654d83e2f..41af61d9b1f 100644
--- a/src/mongo/db/pipeline/mongos_process_interface.h
+++ b/src/mongo/db/pipeline/mongos_process_interface.h
@@ -59,9 +59,7 @@ public:
         MONGO_UNREACHABLE;
     }
 
-    bool isSharded(OperationContext* opCtx, const NamespaceString& nss) final {
-        MONGO_UNREACHABLE;
-    }
+    bool isSharded(OperationContext* opCtx, const NamespaceString& nss) final;
 
     void insert(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                 const NamespaceString& ns,
diff --git a/src/mongo/db/pipeline/stub_mongo_process_interface.h b/src/mongo/db/pipeline/stub_mongo_process_interface.h
index bda78c03d93..7e3d05a9fce 100644
--- a/src/mongo/db/pipeline/stub_mongo_process_interface.h
+++ b/src/mongo/db/pipeline/stub_mongo_process_interface.h
@@ -54,7 +54,7 @@ public:
     }
 
     bool isSharded(OperationContext* opCtx, const NamespaceString& ns) override {
-        MONGO_UNREACHABLE;
+        return false;
     }
 
     void insert(const boost::intrusive_ptr<ExpressionContext>& expCtx,