fix selecting nested field with name that begins with dollar sign (#108)

This PR updates the logic for building selection documents for the `$replaceWith` pipeline stage to be more rigorous. It uses an expanded `ColumnRef` enum to keep track of whether the selected reference has a name that needs to be escaped, is a variable, etc.

CHANGELOG.md

@@ -4,6 +4,14 @@ This changelog documents the changes between release versions.
 
 ## [Unreleased]
 
+### Added
+
+### Fixed
+
+- Selecting nested fields with names that begin with a dollar sign ([#108](https://github.com/hasura/ndc-mongodb/pull/108))
+
+### Changed
+
 ## [1.2.0] - 2024-09-12
 
 ### Added
@@ -117,10 +125,6 @@ definition:
     typeName: InstitutionStaffComparisonExp
 ```
 
-### Fixed
-
-### Changed
-
 ## [1.1.0] - 2024-08-16
 
 - Accept predicate arguments in native mutations and native queries ([#92](https://github.com/hasura/ndc-mongodb/pull/92))

crates/integration-tests/src/tests/basic.rs

@@ -70,3 +70,23 @@ async fn selects_array_within_array() -> anyhow::Result<()> {
     );
     Ok(())
 }
+
+#[tokio::test]
+async fn selects_nested_field_with_dollar_sign_in_name() -> anyhow::Result<()> {
+    assert_yaml_snapshot!(
+        graphql_query(
+            r#"
+            query {
+              testCases_nestedFieldWithDollar(order_by: { configuration: Asc }) {
+                configuration {
+                  schema
+                }
+              }
+            }
+            "#
+        )
+        .run()
+        .await?
+    );
+    Ok(())
+}

crates/integration-tests/src/tests/snapshots/integration_tests__tests__basic__selects_nested_field_with_dollar_sign_in_name.snap

@@ -0,0 +1,13 @@
+---
+source: crates/integration-tests/src/tests/basic.rs
+expression: "graphql_query(r#\"\n            query {\n              testCases_nestedFieldWithDollar(order_by: { configuration: Asc }) {\n                configuration {\n                  schema\n                }\n              }\n            }\n            \"#).run().await?"
+---
+data:
+  testCases_nestedFieldWithDollar:
+    - configuration:
+        schema: ~
+    - configuration:
+        schema: schema1
+    - configuration:
+        schema: schema3
+errors: ~

crates/mongodb-agent-common/src/mongodb/selection.rs

@@ -1,11 +1,13 @@
 use indexmap::IndexMap;
 use mongodb::bson::{self, doc, Bson, Document};
+use ndc_models::FieldName;
 use serde::{Deserialize, Serialize};
 
 use crate::{
     interface_types::MongoAgentError,
     mongo_query_plan::{Field, NestedArray, NestedField, NestedObject, QueryPlan},
     mongodb::sanitize::get_field,
+    query::column_ref::ColumnRef,
 };
 
 /// Wraps a BSON document that represents a MongoDB "expression" that constructs a document based
@@ -32,51 +34,50 @@ impl Selection {
         } else {
             &empty_map
         };
-        let doc = from_query_request_helper(&[], fields)?;
+        let doc = from_query_request_helper(None, fields)?;
         Ok(Selection(doc))
     }
 }
 
 fn from_query_request_helper(
-    parent_columns: &[&str],
+    parent: Option<ColumnRef<'_>>,
     field_selection: &IndexMap<ndc_models::FieldName, Field>,
 ) -> Result<Document, MongoAgentError> {
     field_selection
         .iter()
-        .map(|(key, value)| Ok((key.to_string(), selection_for_field(parent_columns, value)?)))
+        .map(|(key, value)| Ok((key.to_string(), selection_for_field(parent.clone(), value)?)))
         .collect()
 }
 
 /// Wraps column reference with an `$isNull` check. That catches cases where a field is missing
 /// from a document, and substitutes a concrete null value. Otherwise the field would be omitted
 /// from query results which leads to an error in the engine.
-fn value_or_null(col_path: String) -> Bson {
-    doc! { "$ifNull": [col_path, Bson::Null] }.into()
+fn value_or_null(value: Bson) -> Bson {
+    doc! { "$ifNull": [value, Bson::Null] }.into()
 }
 
-fn selection_for_field(parent_columns: &[&str], field: &Field) -> Result<Bson, MongoAgentError> {
+fn selection_for_field(
+    parent: Option<ColumnRef<'_>>,
+    field: &Field,
+) -> Result<Bson, MongoAgentError> {
     match field {
         Field::Column {
             column,
             fields: None,
             ..
         } => {
-            let col_path = match parent_columns {
-                [] => format!("${column}"),
-                _ => format!("${}.{}", parent_columns.join("."), column),
-            };
-            let bson_col_path = value_or_null(col_path);
-            Ok(bson_col_path)
+            let col_ref = nested_column_reference(parent, column);
+            let col_ref_or_null = value_or_null(col_ref.into_aggregate_expression());
+            Ok(col_ref_or_null)
         }
         Field::Column {
             column,
             fields: Some(NestedField::Object(NestedObject { fields })),
             ..
         } => {
-            let nested_parent_columns = append_to_path(parent_columns, column.as_str());
-            let nested_parent_col_path = format!("${}", nested_parent_columns.join("."));
-            let nested_selection = from_query_request_helper(&nested_parent_columns, fields)?;
-            Ok(doc! {"$cond": {"if": nested_parent_col_path, "then": nested_selection, "else": Bson::Null}}.into())
+            let col_ref = nested_column_reference(parent, column);
+            let nested_selection = from_query_request_helper(Some(col_ref.clone()), fields)?;
+            Ok(doc! {"$cond": {"if": col_ref.into_aggregate_expression(), "then": nested_selection, "else": Bson::Null}}.into())
         }
         Field::Column {
             column,
@@ -85,11 +86,7 @@ fn selection_for_field(parent_columns: &[&str], field: &Field) -> Result<Bson, M
                     fields: nested_field,
                 })),
             ..
-        } => selection_for_array(
-            &append_to_path(parent_columns, column.as_str()),
-            nested_field,
-            0,
-        ),
+        } => selection_for_array(nested_column_reference(parent, column), nested_field, 0),
         Field::Relationship {
             relationship,
             aggregates,
@@ -161,31 +158,34 @@ fn selection_for_field(parent_columns: &[&str], field: &Field) -> Result<Bson, M
 }
 
 fn selection_for_array(
-    parent_columns: &[&str],
+    parent: ColumnRef<'_>,
     field: &NestedField,
     array_nesting_level: usize,
 ) -> Result<Bson, MongoAgentError> {
     match field {
         NestedField::Object(NestedObject { fields }) => {
-            let nested_parent_col_path = format!("${}", parent_columns.join("."));
-            let mut nested_selection = from_query_request_helper(&["$this"], fields)?;
+            let mut nested_selection =
+                from_query_request_helper(Some(ColumnRef::variable("this")), fields)?;
             for _ in 0..array_nesting_level {
                 nested_selection = doc! {"$map": {"input": "$$this", "in": nested_selection}}
             }
-            let map_expression =
-                doc! {"$map": {"input": &nested_parent_col_path, "in": nested_selection}};
-            Ok(doc! {"$cond": {"if": &nested_parent_col_path, "then": map_expression, "else": Bson::Null}}.into())
+            let map_expression = doc! {"$map": {"input": parent.clone().into_aggregate_expression(), "in": nested_selection}};
+            Ok(doc! {"$cond": {"if": parent.into_aggregate_expression(), "then": map_expression, "else": Bson::Null}}.into())
         }
         NestedField::Array(NestedArray {
             fields: nested_field,
-        }) => selection_for_array(parent_columns, nested_field, array_nesting_level + 1),
+        }) => selection_for_array(parent, nested_field, array_nesting_level + 1),
     }
 }
-fn append_to_path<'a, 'b, 'c>(parent_columns: &'a [&'b str], column: &'c str) -> Vec<&'c str>
-where
-    'b: 'c,
-{
-    parent_columns.iter().copied().chain(Some(column)).collect()
+
+fn nested_column_reference<'a>(
+    parent: Option<ColumnRef<'a>>,
+    column: &'a FieldName,
+) -> ColumnRef<'a> {
+    match parent {
+        Some(parent) => parent.into_nested_field(column),
+        None => ColumnRef::from_field_path([column]),
+    }
 }
 
 /// The extend implementation provides a shallow merge.

crates/mongodb-agent-common/src/query/column_ref.rs

@@ -31,7 +31,13 @@ use crate::{
 /// caller to switch contexts in the second case.
 #[derive(Clone, Debug, PartialEq)]
 pub enum ColumnRef<'a> {
+    /// Reference that can be used as a key in a match document. For example, "$imdb.rating".
     MatchKey(Cow<'a, str>),
+
+    /// Just like MatchKey, except that this form can reference variables. For example,
+    /// "$$this.title". Can only be used in aggregation expressions, is not used as a key.
+    ExpressionStringShorthand(Cow<'a, str>),
+
     Expression(Bson),
 }
 
@@ -68,13 +74,19 @@ impl<'a> ColumnRef<'a> {
         fold_path_element(None, field_name.as_ref())
     }
 
+    /// Get a reference to a pipeline variable
+    pub fn variable(variable_name: impl std::fmt::Display) -> Self {
+        Self::ExpressionStringShorthand(format!("$${variable_name}").into())
+    }
+
     pub fn into_nested_field<'b: 'a>(self, field_name: &'b ndc_models::FieldName) -> ColumnRef<'b> {
         fold_path_element(Some(self), field_name.as_ref())
     }
 
     pub fn into_aggregate_expression(self) -> Bson {
         match self {
             ColumnRef::MatchKey(key) => format!("${key}").into(),
+            ColumnRef::ExpressionStringShorthand(key) => key.to_string().into(),
             ColumnRef::Expression(expr) => expr,
         }
     }
@@ -107,10 +119,8 @@ fn from_comparison_target(column: &ComparisonTarget) -> ColumnRef<'_> {
             // "$$ROOT" is not actually a valid match key, but cheating here makes the
             // implementation much simpler. This match branch produces a ColumnRef::Expression
             // in all cases.
-            let init = ColumnRef::MatchKey(format!("${}", name_from_scope(scope)).into());
-            // The None case won't come up if the input to [from_target_helper] has at least
-            // one element, and we know it does because we start the iterable with `name`
-            let col_ref = from_path(
+            let init = ColumnRef::variable(name_from_scope(scope));
+            from_path(
                 Some(init),
                 once(name.as_ref() as &str).chain(
                     field_path
@@ -119,12 +129,9 @@ fn from_comparison_target(column: &ComparisonTarget) -> ColumnRef<'_> {
                         .map(|field_name| field_name.as_ref() as &str),
                 ),
             )
-            .unwrap();
-            match col_ref {
-                // move from MatchKey to Expression because "$$ROOT" is not valid in a match key
-                ColumnRef::MatchKey(key) => ColumnRef::Expression(format!("${key}").into()),
-                e @ ColumnRef::Expression(_) => e,
-            }
+            // The None case won't come up if the input to [from_target_helper] has at least
+            // one element, and we know it does because we start the iterable with `name`
+            .unwrap()
         }
     }
 }
@@ -186,28 +193,13 @@ fn fold_path_element<'a>(
         (Some(ColumnRef::MatchKey(parent)), true) => {
             ColumnRef::MatchKey(format!("{parent}.{path_element}").into())
         }
-        (Some(ColumnRef::MatchKey(parent)), false) => ColumnRef::Expression(
-            doc! {
-                "$getField": {
-                    "input": format!("${parent}"),
-                    "field": { "$literal": path_element },
-                }
-            }
-            .into(),
-        ),
-        (Some(ColumnRef::Expression(parent)), true) => ColumnRef::Expression(
-            doc! {
-                "$getField": {
-                    "input": parent,
-                    "field": path_element,
-                }
-            }
-            .into(),
-        ),
-        (Some(ColumnRef::Expression(parent)), false) => ColumnRef::Expression(
+        (Some(ColumnRef::ExpressionStringShorthand(parent)), true) => {
+            ColumnRef::ExpressionStringShorthand(format!("{parent}.{path_element}").into())
+        }
+        (Some(parent), _) => ColumnRef::Expression(
             doc! {
                 "$getField": {
-                    "input": parent,
+                    "input": parent.into_aggregate_expression(),
                     "field": { "$literal": path_element },
                 }
             }
@@ -293,7 +285,7 @@ mod tests {
             doc! {
                 "$getField": {
                     "input": { "$getField": { "$literal": "meta.subtitles" } },
-                    "field": "english_us",
+                    "field": { "$literal": "english_us" },
                 }
             }
             .into(),
@@ -360,7 +352,7 @@ mod tests {
             scope: Scope::Root,
         };
         let actual = ColumnRef::from_comparison_target(&target);
-        let expected = ColumnRef::Expression("$$scope_root.field.prop1.prop2".into());
+        let expected = ColumnRef::ExpressionStringShorthand("$$scope_root.field.prop1.prop2".into());
         assert_eq!(actual, expected);
         Ok(())
     }