SERVER-58583 Expressive projection in FLE for find

2 files changed, 25 insertions, 6 deletions

diff --git a/src/mongo/db/pipeline/document_source_project.cpp b/src/mongo/db/pipeline/document_source_project.cpp
index 657350dc656..e5f048a8e3d 100644
--- a/src/mongo/db/pipeline/document_source_project.cpp
+++ b/src/mongo/db/pipeline/document_source_project.cpp
@@ -37,7 +37,6 @@
 #include "mongo/db/exec/projection_executor.h"
 #include "mongo/db/exec/projection_executor_builder.h"
 #include "mongo/db/pipeline/lite_parsed_document_source.h"
-#include "mongo/db/query/projection_parser.h"
 
 namespace mongo {
 
@@ -62,7 +61,9 @@ BSONObj buildExclusionProjectionSpecification(const std::vector<BSONElement>& un
 }  // namespace
 
 intrusive_ptr<DocumentSource> DocumentSourceProject::create(
-    BSONObj projectSpec, const intrusive_ptr<ExpressionContext>& expCtx, StringData specifiedName) {
+    projection_ast::Projection projection,
+    const intrusive_ptr<ExpressionContext>& expCtx,
+    StringData specifiedName) {
     const bool isIndependentOfAnyCollection = false;
     intrusive_ptr<DocumentSource> project(new DocumentSourceSingleDocumentTransformation(
         expCtx,
@@ -72,8 +73,6 @@ intrusive_ptr<DocumentSource> DocumentSourceProject::create(
             // here so we can add the name that was actually specified by the user, be it $project
             // or an alias.
             try {
-                auto policies = ProjectionPolicies::aggregateProjectionPolicies();
-                auto projection = projection_ast::parse(expCtx, projectSpec, policies);
                 // We won't optimize the executor on creation, and will do it as part of the
                 // pipeline optimization process when requested via the 'optimize()' method on
                 // 'DocumentSourceSingleDocumentTransformation', so we won't pass the
@@ -87,7 +86,10 @@ intrusive_ptr<DocumentSource> DocumentSourceProject::create(
                     projection_executor::kDefaultBuilderParams};
                 builderParams.reset(projection_executor::kOptimizeExecutor);
                 return projection_executor::buildProjectionExecutor(
-                    expCtx, &projection, policies, builderParams);
+                    expCtx,
+                    &projection,
+                    ProjectionPolicies::aggregateProjectionPolicies(),
+                    builderParams);
             } catch (DBException& ex) {
                 ex.addContext("Invalid " + specifiedName.toString());
                 throw;
diff --git a/src/mongo/db/pipeline/document_source_project.h b/src/mongo/db/pipeline/document_source_project.h
index 43ee90fce0f..dc26b2e1c41 100644
--- a/src/mongo/db/pipeline/document_source_project.h
+++ b/src/mongo/db/pipeline/document_source_project.h
@@ -30,6 +30,7 @@
 #pragma once
 
 #include "mongo/db/pipeline/document_source_single_document_transformation.h"
+#include "mongo/db/query/projection_parser.h"
 
 namespace mongo {
 
@@ -45,12 +46,28 @@ public:
     static constexpr StringData kAliasNameUnset = "$unset"_sd;
 
     /**
+     * Method to create a $project stage from a Projection AST.
+     */
+    static boost::intrusive_ptr<DocumentSource> create(
+        projection_ast::Projection projection,
+        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+        StringData specifiedName);
+
+    /**
      * Convenience method to create a $project stage from 'projectSpec'.
      */
     static boost::intrusive_ptr<DocumentSource> create(
         BSONObj projectSpec,
         const boost::intrusive_ptr<ExpressionContext>& expCtx,
-        StringData specifiedName);
+        StringData specifiedName) try {
+
+        auto projection = projection_ast::parse(
+            expCtx, projectSpec, ProjectionPolicies::aggregateProjectionPolicies());
+        return create(projection, expCtx, specifiedName);
+    } catch (DBException& ex) {
+        ex.addContext("Invalid " + specifiedName.toString());
+        throw;
+    }
 
     /**
      * Parses a $project stage from the user-supplied BSON.