Fixed a subtle bug in add_at_path method (#5200)

Co-authored-by: Iryna Shestak <[email protected]>
apollographql · May 21, 2024 · 8b22c40 · 8b22c40
1 parent 18a5b6a
commit 8b22c40
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 9 deletions.
diff --git a/apollo-federation/src/query_graph/graph_path.rs b/apollo-federation/src/query_graph/graph_path.rs
@@ -333,6 +333,15 @@ impl OpPathElement {
         }
     }
 
+    pub(crate) fn sub_selection_type_position(
+        &self,
+    ) -> Result<Option<CompositeTypeDefinitionPosition>, FederationError> {
+        match self {
+            OpPathElement::Field(field) => Ok(field.data().output_base_type()?.try_into().ok()),
+            OpPathElement::InlineFragment(inline) => Ok(Some(inline.data().casted_type())),
+        }
+    }
+
     pub(crate) fn extract_operation_conditionals(
         &self,
     ) -> Result<Vec<OperationConditional>, FederationError> {

diff --git a/apollo-federation/src/query_plan/operation.rs b/apollo-federation/src/query_plan/operation.rs
@@ -1781,7 +1781,7 @@ mod normalized_inline_fragment_selection {
             }
         }
 
-        pub(super) fn casted_type(&self) -> CompositeTypeDefinitionPosition {
+        pub(crate) fn casted_type(&self) -> CompositeTypeDefinitionPosition {
             self.type_condition_position
                 .clone()
                 .unwrap_or_else(|| self.parent_type_position.clone())
@@ -2531,10 +2531,8 @@ impl SelectionSet {
                 "Unable to rebase selection updates",
             ));
         };
-        let sub_selection_parent_type: Option<CompositeTypeDefinitionPosition> = match element {
-            OpPathElement::Field(ref field) => field.data().output_base_type()?.try_into().ok(),
-            OpPathElement::InlineFragment(ref inline) => Some(inline.data().casted_type()),
-        };
+        let sub_selection_parent_type: Option<CompositeTypeDefinitionPosition> =
+            element.sub_selection_type_position()?;
 
         let Some(ref sub_selection_parent_type) = sub_selection_parent_type else {
             // This is a leaf, so all updates should correspond ot the same field and we just use the first.
@@ -2786,17 +2784,17 @@ impl SelectionSet {
         match path.split_first() {
             // If we have a sub-path, recurse.
             Some((ele, path @ &[_, ..])) => {
+                let Some(sub_selection_type) = ele.sub_selection_type_position()? else {
+                    return Err(FederationError::internal("unexpected error: add_at_path encountered a field that is not of a composite type".to_string()));
+                };
                 let mut selection = Arc::make_mut(&mut self.selections)
                     .entry(ele.key())
                     .or_insert(|| {
                         Selection::from_element(
                             OpPathElement::clone(ele),
                             // We immediately add a selection afterward to make this selection set
                             // valid.
-                            Some(SelectionSet::empty(
-                                self.schema.clone(),
-                                self.type_position.clone(),
-                            )),
+                            Some(SelectionSet::empty(self.schema.clone(), sub_selection_type)),
                         )
                     })?;
                 match &mut selection {