fix: support @include and @skip in initial fetch node (#591)

Fixed an issue where `@skip` and `@include` directives were incorrectly
removed from the initial Fetch of the Query Plan.

References #579

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: kamilkisiela

.changeset/fix-cond.md

@@ -0,0 +1,7 @@
+---
+node-addon: patch
+router: patch
+query-planner: patch
+---
+
+Fixed an issue where `@skip` and `@include` directives were incorrectly removed from the initial Fetch of the Query Plan.

lib/query-planner/src/ast/merge_path.rs

@@ -13,6 +13,21 @@ pub enum Condition {
     Include(String),
 }
 
+impl Condition {
+    pub fn to_skip_if(&self) -> Option<String> {
+        match self {
+            Condition::Skip(var) => Some(var.clone()),
+            _ => None,
+        }
+    }
+    pub fn to_include_if(&self) -> Option<String> {
+        match self {
+            Condition::Include(var) => Some(var.clone()),
+            _ => None,
+        }
+    }
+}
+
 impl Display for Condition {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {

lib/query-planner/src/ast/selection_item.rs

@@ -1,5 +1,6 @@
 use crate::{
-    ast::normalization::utils::extract_type_condition, utils::pretty_display::PrettyDisplay,
+    ast::normalization::utils::extract_type_condition,
+    utils::pretty_display::{get_indent, PrettyDisplay},
 };
 use graphql_parser::query as query_ast;
 
@@ -49,7 +50,10 @@ impl PrettyDisplay for SelectionItem {
             SelectionItem::InlineFragment(fragment_selection) => {
                 fragment_selection.pretty_fmt(f, depth)?
             }
-            SelectionItem::FragmentSpread(name) => write!(f, "...{}", name)?,
+            SelectionItem::FragmentSpread(name) => {
+                let indent = get_indent(depth);
+                writeln!(f, "{indent}...{}", name)?
+            }
         }
 
         Ok(())

lib/query-planner/src/ast/selection_set.rs

@@ -303,6 +303,8 @@ impl PrettyDisplay for FieldSelection {
             None => String::new(),
         };
 
+        write!(f, "{indent}{}{}{}", alias_str, self.name, args_str)?;
+
         if let Some(skip_if) = &self.skip_if {
             write!(f, " @skip(if: ${})", skip_if)?;
         }
@@ -312,10 +314,10 @@ impl PrettyDisplay for FieldSelection {
         }
 
         if self.is_leaf() {
-            return writeln!(f, "{indent}{}{}{}", alias_str, self.name, args_str);
+            return writeln!(f);
         }
 
-        writeln!(f, "{indent}{}{}{} {{", alias_str, self.name, args_str)?;
+        writeln!(f, " {{")?;
         self.selections.pretty_fmt(f, depth + 1)?;
         writeln!(f, "{indent}}}")
     }

lib/query-planner/src/graph/edge.rs

@@ -159,6 +159,10 @@ pub enum Edge {
     EntityMove(EntityMove),
     /// join__implements
     AbstractMove(String),
+    /// A special edge representing a case where an inline fragment is used with a condition,
+    /// and we don't really move anywhere, but we need to ensure that
+    /// the condition is preserved when planning the query.
+    Selfie(String),
     /// Represents a special case where going from @interfaceObject
     /// to an object type due to the `__typename` field usage,
     /// or usage of a type condition (fragment),
@@ -227,6 +231,7 @@ impl Edge {
             Self::FieldMove(fm) => &fm.name,
             Self::EntityMove(EntityMove { key, .. }) => key,
             Self::AbstractMove(id) => id,
+            Self::Selfie(_) => "selfie",
             Self::SubgraphEntrypoint { name, .. } => &name.0,
             Self::InterfaceObjectTypeMove(InterfaceObjectTypeMove {
                 object_type_name, ..
@@ -264,6 +269,7 @@ impl Display for Edge {
             Edge::SubgraphEntrypoint { name, .. } => write!(f, "{}", name.0),
             Edge::EntityMove(EntityMove { .. }) => write!(f, "🔑"),
             Edge::AbstractMove(_) => write!(f, "🔮"),
+            Edge::Selfie(_) => write!(f, "🤳"),
             Edge::FieldMove(field_move) => write!(f, "{}", field_move.name),
             Edge::InterfaceObjectTypeMove(m) => write!(f, "🔎 {}", m.object_type_name),
         }?;
@@ -318,6 +324,7 @@ impl Debug for Edge {
                 write!(f, "🔑 {}", key)
             }
             Edge::AbstractMove(name) => write!(f, "🔮 {}", name),
+            Edge::Selfie(_) => write!(f, "🤳"),
             Edge::InterfaceObjectTypeMove(m) => write!(f, "🔎 {}", m.object_type_name),
         }
     }

lib/query-planner/src/graph/mod.rs

@@ -643,6 +643,18 @@ impl Graph {
                     );
                 }
 
+                trace!(
+                    "[x] Creating self-referencing edge for '{}/{}'",
+                    def_name,
+                    graph_id
+                );
+                let head = self.upsert_node(Node::new_node(
+                    def_name,
+                    state.resolve_graph_id(graph_id)?,
+                    state.is_interface_object_in_subgraph(def_name, graph_id),
+                ));
+                self.upsert_edge(head, head, Edge::Selfie(def_name.clone()));
+
                 for (field_name, field_definition) in definition.fields().iter() {
                     let (is_available, maybe_join_field) =
                         FederationRules::check_field_subgraph_availability(

lib/query-planner/src/graph/tests.rs

@@ -287,8 +287,8 @@ mod graph_tests {
             .assert_field_edge("createdBy", "User/products")
             .assert_field_edge("__typename", "String/products")
             .no_field_edge("reviews");
-        assert_eq!(incoming.edges.len(), 2);
-        assert_eq!(outgoing.edges.len(), 11);
+        assert_eq!(incoming.edges.len(), 3);
+        assert_eq!(outgoing.edges.len(), 12);
 
         // requires preserves selection set in the graph
         let outgoing = find_node(&graph, "Product/inventory").1;
@@ -315,8 +315,8 @@ mod graph_tests {
         assert_eq!(graph.root_subscription_node(), None);
 
         let (incoming, outgoing) = find_node(&graph, "Product/products");
-        assert_eq!(incoming.edges.len(), 11);
-        assert_eq!(outgoing.edges.len(), 15);
+        assert_eq!(incoming.edges.len(), 14);
+        assert_eq!(outgoing.edges.len(), 18);
 
         incoming
             .assert_key_edge("id", "Product/inventory")

lib/query-planner/src/planner/fetch/fetch_graph.rs

@@ -186,6 +186,26 @@ impl FetchGraph {
 
         Ok(())
     }
+
+    /// Checks whether the given step is an ancestor of a step that has the given condition
+    pub fn is_ancestor_of_condition(&self, step_index: NodeIndex, condition: &Condition) -> bool {
+        let mut bfs = Bfs::new(&self.graph, step_index);
+
+        while let Some(current_index) = bfs.next(&self.graph) {
+            let current_step = match self.get_step_data(current_index) {
+                Ok(s) => s,
+                Err(_) => continue,
+            };
+
+            if let Some(step_condition) = &current_step.condition {
+                if step_condition == condition {
+                    return true;
+                }
+            }
+        }
+
+        false
+    }
 }
 
 impl Display for FetchGraph {
@@ -1004,6 +1024,92 @@ fn process_subgraph_entrypoint_edge(
     )
 }
 
+// TODO: simplfy args
+#[allow(clippy::too_many_arguments)]
+#[instrument(level = "trace",skip_all, fields(
+  parent_fetch_step_index = parent_fetch_step_index.index(),
+  requiring_fetch_step_index = requiring_fetch_step_index.map(|f| f.index()),
+  type_name = target_type_name,
+  response_path = response_path.to_string(),
+  fetch_path = fetch_path.to_string()
+))]
+fn process_selfie_edge(
+    graph: &Graph,
+    fetch_graph: &mut FetchGraph,
+    override_context: &PlannerOverrideContext,
+    query_node: &QueryTreeNode,
+    parent_fetch_step_index: NodeIndex,
+    requiring_fetch_step_index: Option<NodeIndex>,
+    response_path: &MergePath,
+    fetch_path: &MergePath,
+    target_type_name: &String,
+    condition: Option<&Condition>,
+) -> Result<Vec<NodeIndex>, FetchGraphError> {
+    let is_ancestor_of_condition = match condition {
+        Some(c) => fetch_graph.is_ancestor_of_condition(parent_fetch_step_index, c),
+        None => false,
+    };
+
+    let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?;
+    trace!(
+        "starting an inline fragment for type '{}' to fetch step [{}]",
+        parent_fetch_step_index.index(),
+        target_type_name,
+    );
+    parent_fetch_step.output.add_at_path(
+        &TypeAwareSelection {
+            selection_set: SelectionSet {
+                items: vec![SelectionItem::InlineFragment(InlineFragmentSelection {
+                    type_condition: target_type_name.clone(),
+                    selections: SelectionSet::default(),
+                    skip_if: condition.and_then(|c| {
+                        if is_ancestor_of_condition {
+                            None
+                        } else {
+                            c.to_skip_if()
+                        }
+                    }),
+                    include_if: condition.and_then(|c| {
+                        if is_ancestor_of_condition {
+                            None
+                        } else {
+                            c.to_include_if()
+                        }
+                    }),
+                })],
+            },
+            type_name: target_type_name.clone(),
+        },
+        fetch_path.clone(),
+        false,
+    )?;
+
+    let segment_condition = if is_ancestor_of_condition {
+        None
+    } else {
+        condition.cloned()
+    };
+    let child_response_path = response_path.push(Segment::Cast(
+        target_type_name.clone(),
+        segment_condition.clone(),
+    ));
+    let child_fetch_path =
+        fetch_path.push(Segment::Cast(target_type_name.clone(), segment_condition));
+
+    process_children_for_fetch_steps(
+        graph,
+        fetch_graph,
+        override_context,
+        query_node,
+        parent_fetch_step_index,
+        &child_response_path,
+        &child_fetch_path,
+        requiring_fetch_step_index,
+        condition,
+        false,
+    )
+}
+
 // TODO: simplfy args
 #[allow(clippy::too_many_arguments)]
 #[instrument(level = "trace",skip_all, fields(
@@ -1111,6 +1217,25 @@ fn process_plain_field_edge(
         fetch_graph.connect(parent_fetch_step_index, requiring_fetch_step_index);
     }
 
+    // If we are inserting into "... | [Product] @include(if: $bool)", we don't need the condition on the field.
+    let condition_in_path = if let Some(condition) = condition {
+        fetch_path.inner.iter().any(|segment| {
+            matches!(
+                segment,
+                Segment::Cast(_, Some(c)) | Segment::Field(_, _, Some(c)) if c == condition
+            )
+        })
+    } else {
+        false
+    };
+
+    let ancestor_of_condition = match condition {
+        Some(c) => fetch_graph.is_ancestor_of_condition(parent_fetch_step_index, c),
+        None => false,
+    };
+
+    let should_strip_condition = condition_in_path || ancestor_of_condition;
+
     let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?;
     trace!(
         "adding output field '{}' to fetch step [{}]",
@@ -1126,8 +1251,20 @@ fn process_plain_field_edge(
                     alias: query_node.selection_alias().map(|a| a.to_string()),
                     selections: SelectionSet::default(),
                     arguments: query_node.selection_arguments().cloned(),
-                    skip_if: None,
-                    include_if: None,
+                    skip_if: condition.and_then(|c| {
+                        if should_strip_condition {
+                            None
+                        } else {
+                            c.to_skip_if()
+                        }
+                    }),
+                    include_if: condition.and_then(|c| {
+                        if should_strip_condition {
+                            None
+                        } else {
+                            c.to_include_if()
+                        }
+                    }),
                 })],
             },
             type_name: field_move.type_name.to_string(),
@@ -1137,19 +1274,24 @@ fn process_plain_field_edge(
     )?;
 
     let child_segment = query_node.selection_alias().unwrap_or(&field_move.name);
+    let segment_condition = if ancestor_of_condition {
+        None
+    } else {
+        condition.cloned()
+    };
     let segment_args_hash = query_node
         .selection_arguments()
         .map(|a| a.hash_u64())
         .unwrap_or(0);
     let mut child_response_path = response_path.push(Segment::Field(
         child_segment.to_string(),
         segment_args_hash,
-        None,
+        segment_condition.clone(),
     ));
     let mut child_fetch_path = fetch_path.push(Segment::Field(
         child_segment.to_string(),
         segment_args_hash,
-        None,
+        segment_condition,
     ));
 
     if field_move.is_list {
@@ -1595,6 +1737,18 @@ fn process_query_node(
                     created_from_requires,
                 ),
             },
+            Edge::Selfie(type_name) => process_selfie_edge(
+                graph,
+                fetch_graph,
+                override_context,
+                query_node,
+                parent_fetch_step_index,
+                requiring_fetch_step_index,
+                response_path,
+                fetch_path,
+                type_name,
+                condition,
+            ),
             Edge::AbstractMove(type_name) => process_abstract_edge(
                 graph,
                 fetch_graph,