SERVER-23473 Add collation parameter to aggregate

6 files changed, 67 insertions, 4 deletions

diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index 72609a6d0cc..0eeb6aadaba 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -149,6 +149,8 @@ env.Library(
         'document_value',
         '$BUILD_DIR/mongo/db/auth/authorization_manager_global',
         '$BUILD_DIR/mongo/db/auth/authcore',
+        '$BUILD_DIR/mongo/db/query/collation/collation_serializer',
+        '$BUILD_DIR/mongo/db/query/collation/collator_factory_interface',
         '$BUILD_DIR/mongo/db/repl/read_concern_args',
         '$BUILD_DIR/mongo/db/storage/storage_options',
     ]
@@ -177,6 +179,8 @@ env.CppUnitTest(
     LIBDEPS=[
         'pipeline',
         '$BUILD_DIR/mongo/db/auth/authorization_manager_mock_init',
+        '$BUILD_DIR/mongo/db/query/collation/collator_interface_mock',
+        '$BUILD_DIR/mongo/db/query/query_test_service_context',
         '$BUILD_DIR/mongo/db/service_context',
         '$BUILD_DIR/mongo/db/service_context_noop_init',
         ],
diff --git a/src/mongo/db/pipeline/expression_context.h b/src/mongo/db/pipeline/expression_context.h
index d4e18dcc582..35b01eef9b8 100644
--- a/src/mongo/db/pipeline/expression_context.h
+++ b/src/mongo/db/pipeline/expression_context.h
@@ -32,6 +32,7 @@
 
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/operation_context.h"
+#include "mongo/db/query/collation/collator_interface.h"
 #include "mongo/util/intrusive_counter.h"
 
 namespace mongo {
@@ -60,6 +61,7 @@ public:
     std::string tempDir;  // Defaults to empty to prevent external sorting in mongos.
 
     OperationContext* opCtx;
+    std::unique_ptr<CollatorInterface> collator;
     static const int kInterruptCheckPeriod = 128;
     int interruptCounter = kInterruptCheckPeriod;  // when 0, check interruptStatus
 };
diff --git a/src/mongo/db/pipeline/pipeline.cpp b/src/mongo/db/pipeline/pipeline.cpp
index 535c05a4030..32285652e24 100644
--- a/src/mongo/db/pipeline/pipeline.cpp
+++ b/src/mongo/db/pipeline/pipeline.cpp
@@ -38,11 +38,14 @@
 #include "mongo/db/catalog/document_validation.h"
 #include "mongo/db/commands.h"
 #include "mongo/db/jsobj.h"
+#include "mongo/db/operation_context.h"
 #include "mongo/db/pipeline/accumulator.h"
 #include "mongo/db/pipeline/document.h"
 #include "mongo/db/pipeline/document_source.h"
 #include "mongo/db/pipeline/expression.h"
 #include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/query/collation/collation_serializer.h"
+#include "mongo/db/query/collation/collator_factory_interface.h"
 #include "mongo/db/repl/read_concern_args.h"
 #include "mongo/util/mongoutils/str.h"
 
@@ -56,6 +59,7 @@ using std::vector;
 
 const char Pipeline::commandName[] = "aggregate";
 const char Pipeline::pipelineName[] = "pipeline";
+const char Pipeline::collationName[] = "collation";
 const char Pipeline::explainName[] = "explain";
 const char Pipeline::fromRouterName[] = "fromRouter";
 const char Pipeline::serverPipelineName[] = "serverPipeline";
@@ -101,6 +105,19 @@ intrusive_ptr<Pipeline> Pipeline::parseCommand(string& errmsg,
             continue;
         }
 
+        if (str::equals(pFieldName, collationName)) {
+            uassert(ErrorCodes::BadValue,
+                    str::stream() << collationName << " must be an object, not a "
+                                  << typeName(cmdElement.type()),
+                    cmdElement.type() == BSONType::Object);
+            auto statusWithCollator =
+                CollatorFactoryInterface::get(pCtx->opCtx->getServiceContext())
+                    ->makeFromBSON(cmdElement.Obj());
+            uassertStatusOK(statusWithCollator.getStatus());
+            pCtx->collator = std::move(statusWithCollator.getValue());
+            continue;
+        }
+
         /* look for the aggregation command */
         if (!strcmp(pFieldName, commandName)) {
             continue;
@@ -437,6 +454,11 @@ Document Pipeline::serialize() const {
         serialized.setField(bypassDocumentValidationCommandOption(), Value(true));
     }
 
+    if (pCtx->collator.get()) {
+        serialized.setField(collationName,
+                            Value(CollationSerializer::specToBSON(pCtx->collator->getSpec())));
+    }
+
     return serialized.freeze();
 }
 
diff --git a/src/mongo/db/pipeline/pipeline.h b/src/mongo/db/pipeline/pipeline.h
index 61fab79c347..be9ff35dc70 100644
--- a/src/mongo/db/pipeline/pipeline.h
+++ b/src/mongo/db/pipeline/pipeline.h
@@ -212,6 +212,7 @@ private:
     friend class Optimizations::Sharded;
 
     static const char pipelineName[];
+    static const char collationName[];
     static const char explainName[];
     static const char fromRouterName[];
     static const char serverPipelineName[];
diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp
index 33f6cf104df..18309f709c6 100644
--- a/src/mongo/db/pipeline/pipeline_d.cpp
+++ b/src/mongo/db/pipeline/pipeline_d.cpp
@@ -48,6 +48,7 @@
 #include "mongo/db/matcher/extensions_callback_real.h"
 #include "mongo/db/pipeline/document_source.h"
 #include "mongo/db/pipeline/pipeline.h"
+#include "mongo/db/query/collation/collation_serializer.h"
 #include "mongo/db/query/get_executor.h"
 #include "mongo/db/query/query_planner.h"
 #include "mongo/db/service_context.h"
@@ -206,12 +207,15 @@ StatusWith<std::unique_ptr<PlanExecutor>> attemptToGetExecutor(
     BSONObj projectionObj,
     BSONObj sortObj,
     const size_t plannerOpts) {
-    const ExtensionsCallbackReal extensionsCallback(pExpCtx->opCtx, &pExpCtx->ns);
-
     auto lpq = stdx::make_unique<LiteParsedQuery>(pExpCtx->ns);
     lpq->setFilter(queryObj);
-    lpq->setSort(sortObj);
     lpq->setProj(projectionObj);
+    lpq->setSort(sortObj);
+    if (pExpCtx->collator) {
+        lpq->setCollation(CollationSerializer::specToBSON(pExpCtx->collator->getSpec()));
+    }
+
+    const ExtensionsCallbackReal extensionsCallback(pExpCtx->opCtx, &pExpCtx->ns);
 
     auto cq = CanonicalQuery::canonicalize(txn, std::move(lpq), extensionsCallback);
 
diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp
index 01e1dd2306a..ca31b0ab4e6 100644
--- a/src/mongo/db/pipeline/pipeline_test.cpp
+++ b/src/mongo/db/pipeline/pipeline_test.cpp
@@ -28,11 +28,13 @@
 
 #include "mongo/platform/basic.h"
 
+#include "mongo/db/operation_context_noop.h"
 #include "mongo/db/pipeline/document.h"
 #include "mongo/db/pipeline/expression_context.h"
 #include "mongo/db/pipeline/field_path.h"
 #include "mongo/db/pipeline/pipeline.h"
-#include "mongo/db/operation_context_noop.h"
+#include "mongo/db/query/collation/collator_interface_mock.h"
+#include "mongo/db/query/query_test_service_context.h"
 #include "mongo/dbtests/dbtests.h"
 
 namespace mongo {
@@ -1082,6 +1084,34 @@ TEST(PipelineInitialSource, MatchInitialQuery) {
     intrusive_ptr<Pipeline> pipe = Pipeline::parseCommand(errmsg, inputBson, ctx);
     ASSERT_EQ(pipe->getInitialQuery(), BSON("a" << 4));
 }
+
+TEST(PipelineInitialSource, CollationNotAnObjectFailsToParse) {
+    QueryTestServiceContext serviceContext;
+    auto txn = serviceContext.makeOperationContext();
+
+    const BSONObj inputBson = fromjson("{pipeline: [{$match: {a: 'abc'}}], collation: 1}");
+
+    intrusive_ptr<ExpressionContext> ctx =
+        new ExpressionContext(txn.get(), NamespaceString("a.collection"));
+    string errmsg;
+    ASSERT_THROWS(Pipeline::parseCommand(errmsg, inputBson, ctx), UserException);
+}
+
+TEST(PipelineInitialSource, ParseCollation) {
+    QueryTestServiceContext serviceContext;
+    auto txn = serviceContext.makeOperationContext();
+
+    const BSONObj inputBson =
+        fromjson("{pipeline: [{$match: {a: 'abc'}}], collation: {locale: 'reverse'}}");
+
+    intrusive_ptr<ExpressionContext> ctx =
+        new ExpressionContext(txn.get(), NamespaceString("a.collection"));
+    string errmsg;
+    intrusive_ptr<Pipeline> pipe = Pipeline::parseCommand(errmsg, inputBson, ctx);
+    ASSERT(ctx->collator.get());
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    ASSERT_TRUE(CollatorInterface::collatorsMatch(ctx->collator.get(), &collator));
+}
 }