From 65e2448c7e321d6f11f32e945329c819a369ddef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9e=20Kooi?= Date: Mon, 29 Jul 2024 17:33:50 +0200 Subject: [PATCH 1/4] fix(federation): pass on operation directives to subgraph queries --- .../src/query_plan/fetch_dependency_graph.rs | 10 ++++++++-- .../fetch_dependency_graph_processor.rs | 10 ++++++++-- .../src/query_plan/query_planner.rs | 3 ++- .../handles_operations_with_directives.rs | 16 ++++------------ 4 files changed, 22 insertions(+), 17 deletions(-) diff --git a/apollo-federation/src/query_plan/fetch_dependency_graph.rs b/apollo-federation/src/query_plan/fetch_dependency_graph.rs index 5fdfcb99d7..754a3eec78 100644 --- a/apollo-federation/src/query_plan/fetch_dependency_graph.rs +++ b/apollo-federation/src/query_plan/fetch_dependency_graph.rs @@ -14,6 +14,7 @@ use apollo_compiler::ast::Type; use apollo_compiler::collections::IndexMap; use apollo_compiler::collections::IndexSet; use apollo_compiler::executable; +use apollo_compiler::executable::DirectiveList; use apollo_compiler::executable::VariableDefinition; use apollo_compiler::name; use apollo_compiler::schema; @@ -2325,6 +2326,7 @@ impl FetchDependencyGraphNode { query_graph: &QueryGraph, handled_conditions: &Conditions, variable_definitions: &[Node], + operation_directives: &Arc, fragments: Option<&mut RebasedFragments>, operation_name: Option, ) -> Result, FederationError> { @@ -2358,6 +2360,7 @@ impl FetchDependencyGraphNode { subgraph_schema, selection, variable_definitions, + operation_directives, &operation_name, )? } else { @@ -2366,6 +2369,7 @@ impl FetchDependencyGraphNode { self.root_kind, selection, variable_definitions, + operation_directives, &operation_name, )? }; @@ -2536,6 +2540,7 @@ fn operation_for_entities_fetch( subgraph_schema: &ValidFederationSchema, selection_set: SelectionSet, all_variable_definitions: &[Node], + operation_directives: &Arc, operation_name: &Option, ) -> Result { let mut variable_definitions: Vec> = @@ -2611,7 +2616,7 @@ fn operation_for_entities_fetch( root_kind: SchemaRootDefinitionKind::Query, name: operation_name.clone(), variables: Arc::new(variable_definitions), - directives: Default::default(), + directives: Arc::clone(operation_directives), selection_set, named_fragments: Default::default(), }) @@ -2622,6 +2627,7 @@ fn operation_for_query_fetch( root_kind: SchemaRootDefinitionKind, selection_set: SelectionSet, variable_definitions: &[Node], + operation_directives: &Arc, operation_name: &Option, ) -> Result { let used_variables = selection_set.used_variables(); @@ -2636,7 +2642,7 @@ fn operation_for_query_fetch( root_kind, name: operation_name.clone(), variables: Arc::new(variable_definitions), - directives: Default::default(), + directives: Arc::clone(operation_directives), selection_set, named_fragments: Default::default(), }) diff --git a/apollo-federation/src/query_plan/fetch_dependency_graph_processor.rs b/apollo-federation/src/query_plan/fetch_dependency_graph_processor.rs index 4ee9b57da0..66132a4ce5 100644 --- a/apollo-federation/src/query_plan/fetch_dependency_graph_processor.rs +++ b/apollo-federation/src/query_plan/fetch_dependency_graph_processor.rs @@ -1,5 +1,7 @@ use std::collections::HashSet; +use std::sync::Arc; +use apollo_compiler::executable::DirectiveList; use apollo_compiler::executable::VariableDefinition; use apollo_compiler::Name; use apollo_compiler::Node; @@ -44,7 +46,8 @@ const FETCH_COST: QueryPlanCost = 1000.0; const PIPELINING_COST: QueryPlanCost = 100.0; pub(crate) struct FetchDependencyGraphToQueryPlanProcessor { - variable_definitions: Vec>, + variable_definitions: Arc>>, + operation_directives: Arc, fragments: Option, operation_name: Option, assigned_defer_labels: Option>, @@ -241,13 +244,15 @@ fn sequence_cost(values: impl IntoIterator) -> QueryPlanCo impl FetchDependencyGraphToQueryPlanProcessor { pub(crate) fn new( - variable_definitions: Vec>, + variable_definitions: Arc>>, + operation_directives: Arc, fragments: Option, operation_name: Option, assigned_defer_labels: Option>, ) -> Self { Self { variable_definitions, + operation_directives, fragments, operation_name, assigned_defer_labels, @@ -276,6 +281,7 @@ impl FetchDependencyGraphProcessor, DeferredDeferBlock> query_graph, handled_conditions, &self.variable_definitions, + &self.operation_directives, self.fragments.as_mut(), op_name, ) diff --git a/apollo-federation/src/query_plan/query_planner.rs b/apollo-federation/src/query_plan/query_planner.rs index 012dd65392..2a3cb21ef4 100644 --- a/apollo-federation/src/query_plan/query_planner.rs +++ b/apollo-federation/src/query_plan/query_planner.rs @@ -447,7 +447,8 @@ impl QueryPlanner { None }; let mut processor = FetchDependencyGraphToQueryPlanProcessor::new( - operation.variables.clone(), + normalized_operation.variables.clone(), + normalized_operation.directives.clone(), rebased_fragments, operation.name.clone(), assigned_defer_labels, diff --git a/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs b/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs index 986c43a0d2..a479a61863 100644 --- a/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs +++ b/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs @@ -39,8 +39,6 @@ const SUBGRAPH_B: &str = r#" "#; #[test] -#[should_panic(expected = "snapshot assertion")] -// TODO: investigate this failure (missing directives on fetch operation) fn test_if_directives_at_the_operation_level_are_passed_down_to_subgraph_queries() { let planner = planner!( subgraphA: SUBGRAPH_A, @@ -134,8 +132,8 @@ fn test_if_directives_at_the_operation_level_are_passed_down_to_subgraph_queries insta::assert_snapshot!(b_fetch_nodes[0].operation_document, @r#" query Operation__subgraphB__1($representations: [_Any!]!) @operation { _entities(representations: $representations) { - ... on Foo { - baz @field + ... on T { + f1 @field } } } @@ -144,8 +142,8 @@ fn test_if_directives_at_the_operation_level_are_passed_down_to_subgraph_queries insta::assert_snapshot!(b_fetch_nodes[1].operation_document, @r#" query Operation__subgraphB__2($representations: [_Any!]!) @operation { _entities(representations: $representations) { - ... on T { - f1 @field + ... on Foo { + baz @field } } } @@ -153,8 +151,6 @@ fn test_if_directives_at_the_operation_level_are_passed_down_to_subgraph_queries } #[test] -#[should_panic(expected = "snapshot assertion")] -// TODO: investigate this failure (missing `mutation` keyword and operation name) fn test_if_directives_on_mutations_are_passed_down_to_subgraph_queries() { let planner = planner!( subgraphA: SUBGRAPH_A, @@ -198,8 +194,6 @@ fn test_if_directives_on_mutations_are_passed_down_to_subgraph_queries() { } #[test] -#[should_panic(expected = "snapshot assertion")] -// TODO: investigate this failure (missing directives on fetch query) fn test_if_directives_with_arguments_applied_on_queries_are_ok() { let planner = planner!( Subgraph1: r#" @@ -244,8 +238,6 @@ fn test_if_directives_with_arguments_applied_on_queries_are_ok() { } #[test] -#[should_panic(expected = r#"unused variable: `$some_var`"#)] -// TODO: investigate this failure fn subgraph_query_retains_the_query_variables_used_in_the_directives_applied_to_the_query() { let planner = planner!( Subgraph1: r#" From 9ee9d011486ec75c75a689d3ec03390fa0150fd8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9e=20Kooi?= Date: Mon, 29 Jul 2024 18:08:20 +0200 Subject: [PATCH 2/4] fix(federation): detect variable usage in directive on operation --- apollo-federation/src/operation/mod.rs | 153 +++++++++--------- apollo-federation/src/operation/optimize.rs | 2 +- .../src/query_plan/fetch_dependency_graph.rs | 48 +++--- .../handles_operations_with_directives.rs | 2 +- 4 files changed, 104 insertions(+), 101 deletions(-) diff --git a/apollo-federation/src/operation/mod.rs b/apollo-federation/src/operation/mod.rs index c2441530a4..0f5378575d 100644 --- a/apollo-federation/src/operation/mod.rs +++ b/apollo-federation/src/operation/mod.rs @@ -3671,97 +3671,106 @@ impl RebasedFragments { // Collect used variables from operation types. -fn collect_variables_from_value<'selection>( - value: &'selection executable::Value, - variables: &mut HashSet<&'selection Name>, -) { - match value { - executable::Value::Variable(v) => { - variables.insert(v); - } - executable::Value::List(list) => { - for value in list { - collect_variables_from_value(value, variables); +pub(crate) struct VariableCollector<'s> { + variables: HashSet<&'s Name>, +} + +impl<'s> VariableCollector<'s> { + pub(crate) fn new() -> Self { + Self { variables: Default::default() } + } + + fn visit_value(&mut self, value: &'s executable::Value) { + match value { + executable::Value::Variable(v) => { + self.variables.insert(v); } - } - executable::Value::Object(object) => { - for (_key, value) in object { - collect_variables_from_value(value, variables); + executable::Value::List(list) => { + for value in list { + self.visit_value(value); + } } + executable::Value::Object(object) => { + for (_key, value) in object { + self.visit_value(value); + } + } + _ => {} } - _ => {} } -} -fn collect_variables_from_directive<'selection>( - directive: &'selection executable::Directive, - variables: &mut HashSet<&'selection Name>, -) { - for arg in directive.arguments.iter() { - collect_variables_from_value(&arg.value, variables); + fn visit_directive(&mut self, directive: &'s executable::Directive) { + for arg in directive.arguments.iter() { + self.visit_value(&arg.value); + } } -} -impl Field { - fn collect_variables<'selection>(&'selection self, variables: &mut HashSet<&'selection Name>) { - for arg in self.arguments.iter() { - collect_variables_from_value(&arg.value, variables); - } - for dir in self.directives.iter() { - collect_variables_from_directive(dir, variables); + pub(crate) fn visit_directive_list(&mut self, directives: &'s executable::DirectiveList) { + for dir in directives.iter() { + self.visit_directive(dir); } } -} -impl FieldSelection { - fn collect_variables<'selection>(&'selection self, variables: &mut HashSet<&'selection Name>) { - self.field.collect_variables(variables); - if let Some(set) = &self.selection_set { - set.collect_variables(variables); + fn visit_field(&mut self, field: &'s Field) { + for arg in field.arguments.iter() { + self.visit_value(&arg.value); } + self.visit_directive_list(&field.directives); } -} -impl InlineFragment { - fn collect_variables<'selection>(&'selection self, variables: &mut HashSet<&'selection Name>) { - for dir in self.directives.iter() { - collect_variables_from_directive(dir, variables); + fn visit_field_selection(&mut self, selection: &'s FieldSelection) { + self.visit_field(&selection.field); + if let Some(set) = &selection.selection_set { + self.visit_selection_set(set); } } -} -impl InlineFragmentSelection { - fn collect_variables<'selection>(&'selection self, variables: &mut HashSet<&'selection Name>) { - self.inline_fragment.collect_variables(variables); - self.selection_set.collect_variables(variables); + fn visit_inline_fragment(&mut self, fragment: &'s InlineFragment) { + self.visit_directive_list(&fragment.directives); } -} -impl FragmentSpread { - fn collect_variables<'selection>(&'selection self, variables: &mut HashSet<&'selection Name>) { - for dir in self.directives.iter() { - collect_variables_from_directive(dir, variables); + fn visit_inline_fragment_selection(&mut self, selection: &'s InlineFragmentSelection) { + self.visit_inline_fragment(&selection.inline_fragment); + self.visit_selection_set(&selection.selection_set); + } + + fn visit_fragment_spread(&mut self, fragment: &'s FragmentSpread) { + self.visit_directive_list(&fragment.directives); + self.visit_directive_list(&fragment.fragment_directives); + } + + fn visit_fragment_spread_selection(&mut self, selection: &'s FragmentSpreadSelection) { + self.visit_fragment_spread(&selection.spread); + self.visit_selection_set(&selection.selection_set); + } + + fn visit_selection(&mut self, selection: &'s Selection) { + match selection { + Selection::Field(field) => self.visit_field_selection(field), + Selection::InlineFragment(frag) => self.visit_inline_fragment_selection(frag), + Selection::FragmentSpread(frag) => self.visit_fragment_spread_selection(frag), } - for dir in self.fragment_directives.iter() { - collect_variables_from_directive(dir, variables); + } + + pub(crate) fn visit_selection_set(&mut self, selection_set: &'s SelectionSet) { + for selection in selection_set.iter() { + self.visit_selection(selection); } } -} -impl FragmentSpreadSelection { - fn collect_variables<'selection>(&'selection self, variables: &mut HashSet<&'selection Name>) { - self.spread.collect_variables(variables); - self.selection_set.collect_variables(variables); + /// Consume the collector and return the collected names. + pub(crate) fn into_inner(self) -> HashSet<&'s Name> { + self.variables } } -impl Selection { - fn collect_variables<'selection>(&'selection self, variables: &mut HashSet<&'selection Name>) { - match self { - Selection::Field(field) => field.collect_variables(variables), - Selection::InlineFragment(frag) => frag.collect_variables(variables), - Selection::FragmentSpread(frag) => frag.collect_variables(variables), - } +impl Fragment { + /// Returns the variable names that are used by this fragment. + pub(crate) fn used_variables(&self) -> HashSet<&'_ Name> { + let mut collector = VariableCollector::new(); + collector.visit_directive_list(&self.directives); + collector.visit_selection_set(&self.selection_set); + collector.into_inner() } } @@ -3769,15 +3778,9 @@ impl SelectionSet { /// Returns the variable names that are used by this selection set, including through fragment /// spreads. pub(crate) fn used_variables(&self) -> HashSet<&'_ Name> { - let mut variables = HashSet::new(); - self.collect_variables(&mut variables); - variables - } - - fn collect_variables<'selection>(&'selection self, variables: &mut HashSet<&'selection Name>) { - for selection in self.selections.values() { - selection.collect_variables(variables); - } + let mut collector = VariableCollector::new(); + collector.visit_selection_set(self); + collector.into_inner() } } diff --git a/apollo-federation/src/operation/optimize.rs b/apollo-federation/src/operation/optimize.rs index 3f21309a10..a4f95f9bb9 100644 --- a/apollo-federation/src/operation/optimize.rs +++ b/apollo-federation/src/operation/optimize.rs @@ -926,7 +926,7 @@ impl SelectionSet { // over fragment reuse, and so we do not want to invest a lot of time into improving // fragment reuse. We do the simple, less-than-ideal thing. if let Some(variable_definitions) = &context.operation_variables { - let fragment_variables = candidate.selection_set.used_variables(); + let fragment_variables = candidate.used_variables(); if fragment_variables .difference(variable_definitions) .next() diff --git a/apollo-federation/src/query_plan/fetch_dependency_graph.rs b/apollo-federation/src/query_plan/fetch_dependency_graph.rs index 754a3eec78..e76cf2ae78 100644 --- a/apollo-federation/src/query_plan/fetch_dependency_graph.rs +++ b/apollo-federation/src/query_plan/fetch_dependency_graph.rs @@ -45,6 +45,7 @@ use crate::operation::SelectionId; use crate::operation::SelectionMap; use crate::operation::SelectionSet; use crate::operation::TYPENAME_FIELD; +use crate::operation::VariableCollector; use crate::query_graph::extract_subgraphs_from_supergraph::FEDERATION_REPRESENTATIONS_ARGUMENTS_NAME; use crate::query_graph::extract_subgraphs_from_supergraph::FEDERATION_REPRESENTATIONS_VAR_NAME; use crate::query_graph::graph_path::concat_op_paths; @@ -2348,11 +2349,18 @@ impl FetchDependencyGraphNode { .transpose()?; let subgraph_schema = query_graph.schema_by_source(&self.subgraph_name)?; - let variable_usages = { - let set = selection.used_variables(); - let mut list = set.into_iter().cloned().collect::>(); - list.sort(); - list + // Narrow down the variable definitions to only the ones used in the subgraph operation. + let variable_definitions = { + let mut collector = VariableCollector::new(); + collector.visit_directive_list(operation_directives); + collector.visit_selection_set(&selection); + let used_variables = collector.into_inner(); + + variable_definitions + .iter() + .filter(|variable| used_variables.contains(&variable.name)) + .cloned() + .collect::>() }; let mut operation = if self.is_entity_fetch { @@ -2378,6 +2386,14 @@ impl FetchDependencyGraphNode { { operation.reuse_fragments(fragments)?; } + + let variable_usages = { + let mut list = operation.variables.iter() + .map(|variable| variable.name.clone()).collect::>(); + list.sort(); + list + }; + let operation_document = operation.try_into()?; let node = super::PlanNode::Fetch(Box::new(super::FetchNode { @@ -2539,20 +2555,11 @@ impl FetchDependencyGraphNode { fn operation_for_entities_fetch( subgraph_schema: &ValidFederationSchema, selection_set: SelectionSet, - all_variable_definitions: &[Node], + mut variable_definitions: Vec>, operation_directives: &Arc, operation_name: &Option, ) -> Result { - let mut variable_definitions: Vec> = - Vec::with_capacity(all_variable_definitions.len() + 1); - variable_definitions.push(representations_variable_definition(subgraph_schema)?); - let used_variables = selection_set.used_variables(); - variable_definitions.extend( - all_variable_definitions - .iter() - .filter(|definition| used_variables.contains(&definition.name)) - .cloned(), - ); + variable_definitions.insert(0, representations_variable_definition(subgraph_schema)?); let query_type_name = subgraph_schema.schema().root_operation(OperationType::Query).ok_or_else(|| SingleFederationError::InvalidSubgraph { @@ -2626,17 +2633,10 @@ fn operation_for_query_fetch( subgraph_schema: &ValidFederationSchema, root_kind: SchemaRootDefinitionKind, selection_set: SelectionSet, - variable_definitions: &[Node], + variable_definitions: Vec>, operation_directives: &Arc, operation_name: &Option, ) -> Result { - let used_variables = selection_set.used_variables(); - let variable_definitions = variable_definitions - .iter() - .filter(|definition| used_variables.contains(&definition.name)) - .cloned() - .collect(); - Ok(Operation { schema: subgraph_schema.clone(), root_kind, diff --git a/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs b/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs index a479a61863..11bc7ac655 100644 --- a/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs +++ b/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs @@ -169,7 +169,7 @@ fn test_if_directives_on_mutations_are_passed_down_to_subgraph_queries() { @r###" QueryPlan { Fetch(service: "subgraphA") { - mutation TestMutation__subgraphA__0 { + { updateFoo(bar: "something") @field { id @field bar @field From 6aeaa4e402f7ba90fda899ecfb52a9087f1a59db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9e=20Kooi?= Date: Wed, 31 Jul 2024 15:08:03 +0200 Subject: [PATCH 3/4] fmt --- apollo-federation/src/operation/mod.rs | 4 +++- .../src/query_plan/fetch_dependency_graph.rs | 9 ++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/apollo-federation/src/operation/mod.rs b/apollo-federation/src/operation/mod.rs index 0f5378575d..e6f4e688b7 100644 --- a/apollo-federation/src/operation/mod.rs +++ b/apollo-federation/src/operation/mod.rs @@ -3677,7 +3677,9 @@ pub(crate) struct VariableCollector<'s> { impl<'s> VariableCollector<'s> { pub(crate) fn new() -> Self { - Self { variables: Default::default() } + Self { + variables: Default::default(), + } } fn visit_value(&mut self, value: &'s executable::Value) { diff --git a/apollo-federation/src/query_plan/fetch_dependency_graph.rs b/apollo-federation/src/query_plan/fetch_dependency_graph.rs index e76cf2ae78..9eb638def7 100644 --- a/apollo-federation/src/query_plan/fetch_dependency_graph.rs +++ b/apollo-federation/src/query_plan/fetch_dependency_graph.rs @@ -44,8 +44,8 @@ use crate::operation::Selection; use crate::operation::SelectionId; use crate::operation::SelectionMap; use crate::operation::SelectionSet; -use crate::operation::TYPENAME_FIELD; use crate::operation::VariableCollector; +use crate::operation::TYPENAME_FIELD; use crate::query_graph::extract_subgraphs_from_supergraph::FEDERATION_REPRESENTATIONS_ARGUMENTS_NAME; use crate::query_graph::extract_subgraphs_from_supergraph::FEDERATION_REPRESENTATIONS_VAR_NAME; use crate::query_graph::graph_path::concat_op_paths; @@ -2388,8 +2388,11 @@ impl FetchDependencyGraphNode { } let variable_usages = { - let mut list = operation.variables.iter() - .map(|variable| variable.name.clone()).collect::>(); + let mut list = operation + .variables + .iter() + .map(|variable| variable.name.clone()) + .collect::>(); list.sort(); list }; From d1b515e8bab3f50e1af910b5963203eb84d79928 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ren=C3=A9e=20Kooi?= Date: Wed, 31 Jul 2024 17:30:06 +0200 Subject: [PATCH 4/4] fill in missing snapshot --- .../handles_operations_with_directives.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs b/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs index 11bc7ac655..7a79e150b2 100644 --- a/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs +++ b/apollo-federation/tests/query_plan/build_query_plan_tests/handles_operations_with_directives.rs @@ -259,7 +259,15 @@ fn subgraph_query_retains_the_query_variables_used_in_the_directives_applied_to_ test } "#, - @r#""# + @r###" + QueryPlan { + Fetch(service: "Subgraph1") { + { + test + } + }, + } + "### ); let fetch_nodes = find_fetch_nodes_for_subgraph("Subgraph1", &plan);