Corrected overwritten fields in remove_conditions_from_selection_set (

#6116)

apollo-federation/src/query_plan/conditions.rs

@@ -10,8 +10,10 @@ use serde::Serialize;
 
 use crate::error::FederationError;
 use crate::operation::DirectiveList;
+use crate::operation::NamedFragments;
 use crate::operation::Selection;
 use crate::operation::SelectionMap;
+use crate::operation::SelectionMapperReturn;
 use crate::operation::SelectionSet;
 use crate::query_graph::graph_path::OpPathElement;
 
@@ -176,37 +178,37 @@ pub(crate) fn remove_conditions_from_selection_set(
             Ok(selection_set.clone())
         }
         Conditions::Variables(variable_conditions) => {
-            let mut selection_map = SelectionMap::new();
-
-            for selection in selection_set.selections.values() {
+            selection_set.lazy_map(&NamedFragments::default(), |selection| {
                 let element = selection.element()?;
                 // We remove any of the conditions on the element and recurse.
                 let updated_element =
                     remove_conditions_of_element(element.clone(), variable_conditions);
-                let new_selection = if let Some(selection_set) = selection.selection_set() {
+                if let Some(selection_set) = selection.selection_set() {
                     let updated_selection_set =
                         remove_conditions_from_selection_set(selection_set, conditions)?;
                     if updated_element == element {
                         if *selection_set == updated_selection_set {
-                            selection.clone()
+                            Ok(SelectionMapperReturn::Selection(selection.clone()))
                         } else {
-                            selection.with_updated_selection_set(Some(updated_selection_set))?
+                            Ok(SelectionMapperReturn::Selection(
+                                selection
+                                    .with_updated_selection_set(Some(updated_selection_set))?,
+                            ))
                         }
                     } else {
-                        Selection::from_element(updated_element, Some(updated_selection_set))?
+                        Ok(SelectionMapperReturn::Selection(Selection::from_element(
+                            updated_element,
+                            Some(updated_selection_set),
+                        )?))
                     }
                 } else if updated_element == element {
-                    selection.clone()
+                    Ok(SelectionMapperReturn::Selection(selection.clone()))
                 } else {
-                    Selection::from_element(updated_element, None)?
-                };
-                selection_map.insert(new_selection);
-            }
-
-            Ok(SelectionSet {
-                schema: selection_set.schema.clone(),
-                type_position: selection_set.type_position.clone(),
-                selections: Arc::new(selection_map),
+                    Ok(SelectionMapperReturn::Selection(Selection::from_element(
+                        updated_element,
+                        None,
+                    )?))
+                }
             })
         }
     }

apollo-federation/tests/query_plan/build_query_plan_tests/field_merging_with_skip_and_include.rs

@@ -146,7 +146,7 @@ fn merging_skip_and_include_directives_multiple_applications_differing_order() {
               hello: Hello!
               extraFieldToPreventSkipIncludeNodes: String!
           }
-    
+
           type Hello {
               world: String!
               goodbye: String!
@@ -227,3 +227,85 @@ fn merging_skip_and_include_directives_multiple_applications_differing_quantity(
         "###
     );
 }
+
+#[test]
+fn fields_are_not_overwritten_when_directives_are_removed() {
+    let planner = planner!(
+        SubgraphSkip: r#"
+          type Query {
+            foo: Foo
+          }
+
+          type Foo {
+            bar: Bar
+          }
+
+          type Bar {
+            things: String
+            name: String
+          }
+        "#,
+    );
+    assert_plan!(
+        &planner,
+        r#"
+          query Test($b: Boolean!) {
+            foo @include(if: $b) {
+              bar {
+                name
+              }
+              bar @include(if: $b) {
+                things
+              }
+            }
+          }
+        "#,
+        @r###"
+          QueryPlan {
+            Include(if: $b) {
+              Fetch(service: "SubgraphSkip") {
+                {
+                  foo {
+                    bar {
+                      name
+                      things
+                    }
+                  }
+                }
+              },
+            },
+          }
+        "###
+    );
+    assert_plan!(
+        &planner,
+        r#"
+          query Test($b: Boolean!) {
+            foo @skip(if: $b) {
+              bar {
+                name
+              }
+              bar @skip(if: $b) {
+                things
+              }
+            }
+          }
+        "#,
+        @r###"
+          QueryPlan {
+            Skip(if: $b) {
+              Fetch(service: "SubgraphSkip") {
+                {
+                  foo {
+                    bar {
+                      name
+                      things
+                    }
+                  }
+                }
+              },
+            },
+          }
+        "###
+    );
+}

apollo-federation/tests/query_plan/supergraphs/fields_are_not_overwritten_when_directives_are_removed.graphql

@@ -0,0 +1,64 @@
+# Composed from subgraphs with hash: 654c7bcba86c6f5845a7cb520710cfa37b678e3d
+schema
+  @link(url: "https://specs.apollo.dev/link/v1.0")
+  @link(url: "https://specs.apollo.dev/join/v0.4", for: EXECUTION)
+{
+  query: Query
+}
+
+directive @join__directive(graphs: [join__Graph!], name: String!, args: join__DirectiveArguments) repeatable on SCHEMA | OBJECT | INTERFACE | FIELD_DEFINITION
+
+directive @join__enumValue(graph: join__Graph!) repeatable on ENUM_VALUE
+
+directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean, overrideLabel: String) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
+
+directive @join__graph(name: String!, url: String!) on ENUM_VALUE
+
+directive @join__implements(graph: join__Graph!, interface: String!) repeatable on OBJECT | INTERFACE
+
+directive @join__type(graph: join__Graph!, key: join__FieldSet, extension: Boolean! = false, resolvable: Boolean! = true, isInterfaceObject: Boolean! = false) repeatable on OBJECT | INTERFACE | UNION | ENUM | INPUT_OBJECT | SCALAR
+
+directive @join__unionMember(graph: join__Graph!, member: String!) repeatable on UNION
+
+directive @link(url: String, as: String, for: link__Purpose, import: [link__Import]) repeatable on SCHEMA
+
+type Bar
+  @join__type(graph: SUBGRAPHSKIP)
+{
+  things: String
+  name: String
+}
+
+type Foo
+  @join__type(graph: SUBGRAPHSKIP)
+{
+  bar: Bar
+}
+
+scalar join__DirectiveArguments
+
+scalar join__FieldSet
+
+enum join__Graph {
+  SUBGRAPHSKIP @join__graph(name: "SubgraphSkip", url: "none")
+}
+
+scalar link__Import
+
+enum link__Purpose {
+  """
+  `SECURITY` features provide metadata necessary to securely resolve fields.
+  """
+  SECURITY
+
+  """
+  `EXECUTION` features provide metadata necessary for operation execution.
+  """
+  EXECUTION
+}
+
+type Query
+  @join__type(graph: SUBGRAPHSKIP)
+{
+  foo: Foo
+}

apollo-federation/tests/query_plan/supergraphs/merging_skip_and_include_directives_multiple_applications_differing_order.graphql

@@ -1,4 +1,4 @@
-# Composed from subgraphs with hash: 1cce0a8dc674e651027f8d368504edb388f9b239
+# Composed from subgraphs with hash: fbbe470b5e40af130de42043f74772ec30ee60ff
 schema
   @link(url: "https://specs.apollo.dev/link/v1.0")
   @link(url: "https://specs.apollo.dev/join/v0.4", for: EXECUTION)