hasura · hallettj · Nov 22, 2024 · Nov 11, 2024 · Nov 12, 2024 · Nov 13, 2024
diff --git a/crates/cli/src/native_query/pipeline/mod.rs b/crates/cli/src/native_query/pipeline/mod.rs
@@ -1,4 +1,5 @@
 mod match_stage;
+mod project_stage;
 
 use std::{collections::BTreeMap, iter::once};
 
@@ -54,7 +55,7 @@ pub fn infer_pipeline_types(
     if let TypeConstraint::Object(stage_type_name) = last_stage_type {
         if let Some(object_type) = context.get_object_type(&stage_type_name) {
             context.insert_object_type(object_type_name.clone(), object_type.into_owned());
-            context.set_stage_doc_type(TypeConstraint::Object(object_type_name))
+            context.set_stage_doc_type(TypeConstraint::Object(object_type_name));
         }
     }
 
@@ -93,9 +94,25 @@ fn infer_stage_output_type(
             None
         }
         Stage::Sort(_) => None,
-        Stage::Limit(_) => None,
+        Stage::Skip(expression) => {
+            infer_type_from_aggregation_expression(
+                context,
+                desired_object_type_name,
+                Some(&TypeConstraint::Scalar(BsonScalarType::Int)),
+                expression.clone(),
+            )?;
+            None
+        }
+        Stage::Limit(expression) => {
+            infer_type_from_aggregation_expression(
+                context,
+                desired_object_type_name,
+                Some(&TypeConstraint::Scalar(BsonScalarType::Int)),
+                expression.clone(),
+            )?;
+            None
+        }
         Stage::Lookup { .. } => todo!("lookup stage"),
-        Stage::Skip(_) => None,
         Stage::Group {
             key_expression,
             accumulators,
@@ -110,7 +127,18 @@ fn infer_stage_output_type(
         }
         Stage::Facet(_) => todo!("facet stage"),
         Stage::Count(_) => todo!("count stage"),
-        Stage::ReplaceWith(selection) => {
+        Stage::Project(doc) => {
+            let augmented_type = project_stage::infer_type_from_project_stage(
+                context,
+                &format!("{desired_object_type_name}_project"),
+                doc,
+            )?;
+            Some(augmented_type)
+        }
+        Stage::ReplaceRoot {
+            new_root: selection,
+        }
+        | Stage::ReplaceWith(selection) => {
             let selection: &Document = selection.into();
             Some(
                 aggregation_expression::infer_type_from_aggregation_expression(
@@ -291,7 +319,11 @@ fn infer_type_from_unwind_stage(
     Ok(TypeConstraint::WithFieldOverrides {
         augmented_object_type_name: format!("{desired_object_type_name}_unwind").into(),
         target_type: Box::new(context.get_input_document_type()?.clone()),
-        fields: unwind_stage_object_type.fields,
+        fields: unwind_stage_object_type
+            .fields
+            .into_iter()
+            .map(|(k, t)| (k, Some(t)))
+            .collect(),
     })
 }
 
@@ -360,7 +392,7 @@ mod tests {
         }))]);
         let config = mflix_config();
         let pipeline_types =
-            infer_pipeline_types(&config, "movies", Some(&("movies".into())), &pipeline).unwrap();
+            infer_pipeline_types(&config, "movies", Some(&("movies".into())), &pipeline)?;
         let expected = [(
             "movies_replaceWith".into(),
             ObjectType {
@@ -415,13 +447,18 @@ mod tests {
                 augmented_object_type_name: "unwind_stage_unwind".into(),
                 target_type: Box::new(TypeConstraint::Variable(input_doc_variable)),
                 fields: [
-                    ("idx".into(), TypeConstraint::Scalar(BsonScalarType::Long)),
+                    (
+                        "idx".into(),
+                        Some(TypeConstraint::Scalar(BsonScalarType::Long))
+                    ),
                     (
                         "words".into(),
-                        TypeConstraint::ElementOf(Box::new(TypeConstraint::FieldOf {
-                            target_type: Box::new(TypeConstraint::Variable(input_doc_variable)),
-                            path: nonempty!["words".into()],
-                        }))
+                        Some(TypeConstraint::ElementOf(Box::new(
+                            TypeConstraint::FieldOf {
+                                target_type: Box::new(TypeConstraint::Variable(input_doc_variable)),
+                                path: nonempty!["words".into()],
+                            }
+                        )))
                     )
                 ]
                 .into(),