SERVER-42298 Use SortPattern in the implementation of SortKeyGenerator

3 files changed, 19 insertions, 8 deletions

diff --git a/src/mongo/dbtests/query_stage_ensure_sorted.cpp b/src/mongo/dbtests/query_stage_ensure_sorted.cpp
index 49e9383442f..20d5750c958 100644
--- a/src/mongo/dbtests/query_stage_ensure_sorted.cpp
+++ b/src/mongo/dbtests/query_stage_ensure_sorted.cpp
@@ -78,10 +78,15 @@ public:
             queuedDataStage->pushBack(id);
         }
 
+        // Create a mock ExpressionContext.
+        boost::intrusive_ptr<ExpressionContext> pExpCtx(
+            new ExpressionContext(opCtx.get(), collator));
+        pExpCtx->setCollator(collator);
+
         // Initialization.
         BSONObj pattern = fromjson(patternStr);
         auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>(
-            opCtx.get(), queuedDataStage.release(), &ws, pattern, collator);
+            pExpCtx, queuedDataStage.release(), &ws, pattern);
         EnsureSortedStage ess(opCtx.get(), pattern, &ws, sortKeyGen.release());
         WorkingSetID id = WorkingSet::INVALID_ID;
         PlanStage::StageState state = PlanStage::NEED_TIME;
@@ -115,10 +120,13 @@ protected:
 TEST_F(QueryStageEnsureSortedTest, EnsureSortedEmptyWorkingSet) {
     auto opCtx = _serviceContext.makeOperationContext();
 
+    // Create a mock ExpressionContext.
+    boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(opCtx.get(), nullptr));
+
     WorkingSet ws;
     auto queuedDataStage = std::make_unique<QueuedDataStage>(opCtx.get(), &ws);
-    auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>(
-        opCtx.get(), queuedDataStage.release(), &ws, BSONObj(), nullptr);
+    auto sortKeyGen =
+        std::make_unique<SortKeyGeneratorStage>(pExpCtx, queuedDataStage.release(), &ws, BSONObj());
     EnsureSortedStage ess(opCtx.get(), BSONObj(), &ws, sortKeyGen.release());
 
     WorkingSetID id = WorkingSet::INVALID_ID;
diff --git a/src/mongo/dbtests/query_stage_sort.cpp b/src/mongo/dbtests/query_stage_sort.cpp
index 5b855933793..0919fdea1ad 100644
--- a/src/mongo/dbtests/query_stage_sort.cpp
+++ b/src/mongo/dbtests/query_stage_sort.cpp
@@ -120,7 +120,7 @@ public:
         params.limit = limit();
 
         auto keyGenStage = std::make_unique<SortKeyGeneratorStage>(
-            &_opCtx, queuedDataStage.release(), ws.get(), params.pattern, nullptr);
+            _pExpCtx, queuedDataStage.release(), ws.get(), params.pattern);
 
         auto ss = std::make_unique<SortStage>(&_opCtx, params, ws.get(), keyGenStage.release());
 
@@ -158,7 +158,7 @@ public:
         params.limit = limit();
 
         auto keyGenStage = std::make_unique<SortKeyGeneratorStage>(
-            &_opCtx, queuedDataStage.release(), ws.get(), params.pattern, nullptr);
+            _pExpCtx, queuedDataStage.release(), ws.get(), params.pattern);
 
         auto sortStage =
             std::make_unique<SortStage>(&_opCtx, params, ws.get(), keyGenStage.release());
@@ -227,6 +227,7 @@ public:
 protected:
     const ServiceContext::UniqueOperationContext _txnPtr = cc().makeOperationContext();
     OperationContext& _opCtx = *_txnPtr;
+    boost::intrusive_ptr<ExpressionContext> _pExpCtx = new ExpressionContext(&_opCtx, nullptr);
     DBDirectClient _client;
 };
 
@@ -558,7 +559,7 @@ public:
         params.pattern = BSON("b" << -1 << "c" << 1 << "a" << 1);
 
         auto keyGenStage = std::make_unique<SortKeyGeneratorStage>(
-            &_opCtx, queuedDataStage.release(), ws.get(), params.pattern, nullptr);
+            _pExpCtx, queuedDataStage.release(), ws.get(), params.pattern);
 
         auto sortStage =
             std::make_unique<SortStage>(&_opCtx, params, ws.get(), keyGenStage.release());
diff --git a/src/mongo/dbtests/query_stage_sort_key_generator.cpp b/src/mongo/dbtests/query_stage_sort_key_generator.cpp
index daa667313e4..8675a3e7a51 100644
--- a/src/mongo/dbtests/query_stage_sort_key_generator.cpp
+++ b/src/mongo/dbtests/query_stage_sort_key_generator.cpp
@@ -65,6 +65,7 @@ BSONObj extractKeyFromKeyGenStage(SortKeyGeneratorStage* sortKeyGen, WorkingSet*
 BSONObj extractSortKey(const char* sortSpec, const char* doc, const CollatorInterface* collator) {
     QueryTestServiceContext serviceContext;
     auto opCtx = serviceContext.makeOperationContext();
+    boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(opCtx.get(), collator));
 
     WorkingSet workingSet;
 
@@ -77,7 +78,7 @@ BSONObj extractSortKey(const char* sortSpec, const char* doc, const CollatorInte
 
     BSONObj sortPattern = fromjson(sortSpec);
     SortKeyGeneratorStage sortKeyGen{
-        opCtx.get(), mockStage.release(), &workingSet, std::move(sortPattern), collator};
+        pExpCtx, mockStage.release(), &workingSet, std::move(sortPattern)};
     return extractKeyFromKeyGenStage(&sortKeyGen, &workingSet);
 }
 
@@ -93,6 +94,7 @@ BSONObj extractSortKeyCovered(const char* sortSpec,
                               const CollatorInterface* collator) {
     QueryTestServiceContext serviceContext;
     auto opCtx = serviceContext.makeOperationContext();
+    boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(opCtx.get(), collator));
 
     WorkingSet workingSet;
 
@@ -105,7 +107,7 @@ BSONObj extractSortKeyCovered(const char* sortSpec,
 
     BSONObj sortPattern = fromjson(sortSpec);
     SortKeyGeneratorStage sortKeyGen{
-        opCtx.get(), mockStage.release(), &workingSet, std::move(sortPattern), collator};
+        pExpCtx, mockStage.release(), &workingSet, std::move(sortPattern)};
     return extractKeyFromKeyGenStage(&sortKeyGen, &workingSet);
 }