From f7c9498b31e1597677ed9b6de41b55532480b4e3 Mon Sep 17 00:00:00 2001 From: Dotan Simha Date: Thu, 10 Jul 2025 14:20:48 +0300 Subject: [PATCH 1/4] feat(qp): introduce single/multi state for fetch-step and selections, replace `output` --- bin/dev-cli/src/main.rs | 6 +- .../src/ast/minification/error.rs | 4 +- .../src/ast/minification/stats.rs | 1 + .../src/ast/minification/transform.rs | 1 + lib/query-planner/src/ast/mismatch_finder.rs | 52 +-- .../src/ast/type_aware_selection.rs | 141 +------ lib/query-planner/src/planner/fetch/error.rs | 5 +- .../src/planner/fetch/fetch_graph.rs | 345 ++++++++---------- .../src/planner/fetch/fetch_step_data.rs | 55 ++- lib/query-planner/src/planner/fetch/mod.rs | 2 + .../apply_internal_aliases_patching.rs | 130 +++---- .../deduplicate_and_prune_fetch_steps.rs | 8 +- .../optimize/merge_children_with_parents.rs | 3 +- .../src/planner/fetch/optimize/merge_leafs.rs | 8 +- .../fetch/optimize/merge_passthrough_child.rs | 48 ++- .../planner/fetch/optimize/merge_siblings.rs | 8 +- .../src/planner/fetch/optimize/mod.rs | 4 +- .../optimize/turn_mutations_into_sequence.rs | 10 +- .../planner/fetch/optimize/type_mismatches.rs | 14 +- .../src/planner/fetch/optimize/utils.rs | 62 ++-- .../src/planner/fetch/selections.rs | 315 ++++++++++++++++ lib/query-planner/src/planner/fetch/state.rs | 17 + lib/query-planner/src/planner/mod.rs | 7 +- lib/query-planner/src/planner/plan_nodes.rs | 45 +-- lib/query-planner/src/planner/query_plan.rs | 11 +- lib/query-planner/src/tests/mod.rs | 17 + 26 files changed, 775 insertions(+), 544 deletions(-) create mode 100644 lib/query-planner/src/planner/fetch/selections.rs create mode 100644 lib/query-planner/src/planner/fetch/state.rs diff --git a/bin/dev-cli/src/main.rs b/bin/dev-cli/src/main.rs index 6d4aaaa1..f802c877 100644 --- a/bin/dev-cli/src/main.rs +++ b/bin/dev-cli/src/main.rs @@ -9,6 +9,7 @@ use query_planner::graph::PlannerOverrideContext; use query_planner::planner::best::find_best_combination; use query_planner::planner::fetch::fetch_graph::build_fetch_graph_from_query_tree; use query_planner::planner::fetch::fetch_graph::FetchGraph; +use query_planner::planner::fetch::state::MultiTypeFetchStep; use query_planner::planner::plan_nodes::QueryPlan; use query_planner::planner::query_plan::build_query_plan_from_fetch_graph; use query_planner::planner::tree::query_tree::QueryTree; @@ -124,7 +125,10 @@ fn process_consumer_schema(path: &str) { println!("{}", consumer_schema.document); } -fn process_fetch_graph(supergraph_path: &str, operation_path: &str) -> FetchGraph { +fn process_fetch_graph( + supergraph_path: &str, + operation_path: &str, +) -> FetchGraph { let (graph, query_tree, supergraph_state) = process_merged_tree(supergraph_path, operation_path); diff --git a/lib/query-planner/src/ast/minification/error.rs b/lib/query-planner/src/ast/minification/error.rs index dd50441e..6bda2dba 100644 --- a/lib/query-planner/src/ast/minification/error.rs +++ b/lib/query-planner/src/ast/minification/error.rs @@ -6,6 +6,6 @@ pub enum MinificationError { FieldNotFound(String, String), #[error("Unsupported fragment spread")] UnsupportedFragmentSpread, - #[error("Unsupported field in `_entities`: {0}")] - UnsupportedFieldInEntities(String), + #[error("Unsupported field in `_entities`: {0}.{1}")] + UnsupportedFieldInEntities(String, String), } diff --git a/lib/query-planner/src/ast/minification/stats.rs b/lib/query-planner/src/ast/minification/stats.rs index c8236b05..d668d9ec 100644 --- a/lib/query-planner/src/ast/minification/stats.rs +++ b/lib/query-planner/src/ast/minification/stats.rs @@ -90,6 +90,7 @@ fn walk_and_collect_stats( } return Err(MinificationError::UnsupportedFieldInEntities( + type_def.name().to_string(), field.name.clone(), )); } diff --git a/lib/query-planner/src/ast/minification/transform.rs b/lib/query-planner/src/ast/minification/transform.rs index 0be97927..6baa7f3c 100644 --- a/lib/query-planner/src/ast/minification/transform.rs +++ b/lib/query-planner/src/ast/minification/transform.rs @@ -185,6 +185,7 @@ fn transform_field( } return Err(MinificationError::UnsupportedFieldInEntities( + type_name.to_string(), field.name.clone(), )); } diff --git a/lib/query-planner/src/ast/mismatch_finder.rs b/lib/query-planner/src/ast/mismatch_finder.rs index cbd5192e..7239232c 100644 --- a/lib/query-planner/src/ast/mismatch_finder.rs +++ b/lib/query-planner/src/ast/mismatch_finder.rs @@ -7,8 +7,8 @@ use crate::{ merge_path::{MergePath, Segment}, selection_item::SelectionItem, selection_set::{FieldSelection, SelectionSet}, - type_aware_selection::TypeAwareSelection, }, + planner::fetch::{selections::FetchStepSelections, state::MultiTypeFetchStep}, state::{ subgraph_state::{SubgraphDefinition, SubgraphState}, supergraph_state::{SubgraphName, SupergraphState, TypeNode}, @@ -20,46 +20,42 @@ pub struct SelectionMismatchFinder<'a> { supergraph_state: &'a SupergraphState, } -type MismatchesFound = Vec; +type MismatchesFound = Vec<(String, MergePath)>; impl<'a> SelectionMismatchFinder<'a> { pub fn new(supergraph_state: &'a SupergraphState) -> Self { Self { supergraph_state } } - #[instrument(level = "trace", skip_all, fields( - subgraph_name, - selection = format!("{}", selection_set) - ))] + #[instrument(level = "trace", skip_all, fields(subgraph_name,))] pub fn find_mismatches_in_node( &self, subgraph_name: &SubgraphName, - selection_set: &TypeAwareSelection, + selections: &FetchStepSelections, ) -> MismatchesFound { + let mut mismtaches_found = MismatchesFound::new(); let subgraph_state = self .supergraph_state .subgraphs_state .get(subgraph_name) .unwrap(); - let entrypoint_type = subgraph_state - .definitions - .get(&selection_set.type_name) - .unwrap(); - - let mut mismtaches_found = MismatchesFound::new(); - let start_path = MergePath::default(); - - handle_selection_set( - self.supergraph_state, - subgraph_state, - entrypoint_type, - &selection_set.selection_set, - start_path, - &mut mismtaches_found, - ); - - trace!("found total of {} mismatches", mismtaches_found.len(),); + for (definition_name, selection_set) in selections.iter_selections() { + let entrypoint_type = subgraph_state.definitions.get(definition_name).unwrap(); + let start_path = MergePath::default(); + + handle_selection_set( + definition_name, + self.supergraph_state, + subgraph_state, + entrypoint_type, + selection_set, + start_path, + &mut mismtaches_found, + ); + + trace!("found total of {} mismatches", mismtaches_found.len()); + } mismtaches_found } @@ -73,6 +69,7 @@ impl<'a> SelectionMismatchFinder<'a> { selection = format!("{}", selection_set) ))] fn handle_selection_set<'field, 'schema>( + root_def_type_name: &str, supergraph_state: &'schema SupergraphState, subgraph_state: &'schema SubgraphState, parent_def: &'schema SubgraphDefinition, @@ -105,6 +102,7 @@ fn handle_selection_set<'field, 'schema>( )); let next_parent_type_name = handle_field( + root_def_type_name, supergraph_state, type_def, field, @@ -117,6 +115,7 @@ fn handle_selection_set<'field, 'schema>( next_parent_type_name.and_then(|n| subgraph_state.definitions.get(n)) { handle_selection_set( + root_def_type_name, supergraph_state, subgraph_state, next_parent_def, @@ -156,6 +155,7 @@ fn handle_selection_set<'field, 'schema>( /// /// Returns the return type of the selection, if the inner selection needs to be processed (in case nested selections are defined). fn handle_field<'field, 'schema>( + root_def_type_name: &str, state: &'schema SupergraphState, parent_def: &'schema SubgraphDefinition, field: &'field FieldSelection, @@ -201,7 +201,7 @@ fn handle_field<'field, 'schema>( field_path, ); - mismatches_found.push(field_path.clone()); + mismatches_found.push((root_def_type_name.to_string(), field_path.clone())); } } } else { diff --git a/lib/query-planner/src/ast/type_aware_selection.rs b/lib/query-planner/src/ast/type_aware_selection.rs index 0f0c0cdc..729d4f5a 100644 --- a/lib/query-planner/src/ast/type_aware_selection.rs +++ b/lib/query-planner/src/ast/type_aware_selection.rs @@ -2,18 +2,11 @@ use std::{fmt::Display, hash::Hash}; use crate::ast::{ merge_path::{Condition, Segment}, - safe_merge::{AliasesRecords, SafeSelectionSetMerger}, selection_set::{FieldSelection, InlineFragmentSelection}, }; use super::{merge_path::MergePath, selection_item::SelectionItem, selection_set::SelectionSet}; -#[derive(Debug, Clone, thiserror::Error)] -pub enum TypeAwareSelectionError { - #[error("Failed to locate path '{0}' in selection set '{1}'")] - PathNotFound(String, String), -} - #[derive(Debug, Clone)] pub struct TypeAwareSelection { pub type_name: String, @@ -75,62 +68,6 @@ impl TypeAwareSelection { pub fn add(&mut self, to_add: &Self) { merge_selection_set(&mut self.selection_set, &to_add.selection_set, false); } - - pub fn add_at_path( - &mut self, - to_add: &Self, - add_at_fetch_path: MergePath, - as_first: bool, - ) -> Result<(), TypeAwareSelectionError> { - if let Some(source) = - find_selection_set_by_path_mut(&mut self.selection_set, &add_at_fetch_path) - { - merge_selection_set(source, &to_add.selection_set, as_first); - - Ok(()) - } else { - Err(TypeAwareSelectionError::PathNotFound( - add_at_fetch_path.to_string(), - self.selection_set.to_string(), - )) - } - } - - pub fn add_at_path_and_solve_conflicts( - &mut self, - to_add: &Self, - add_at_fetch_path: MergePath, - (self_used_for_requires, other_used_for_requires): (bool, bool), - as_first: bool, - ) -> Option { - if let Some(source) = - find_selection_set_by_path_mut(&mut self.selection_set, &add_at_fetch_path) - { - let mut merger = SafeSelectionSetMerger::default(); - let aliases_made = merger.merge_selection_set( - source, - &to_add.selection_set, - (self_used_for_requires, other_used_for_requires), - as_first, - ); - - if !aliases_made.is_empty() { - return Some(aliases_made); - } - } - - None - } - - pub fn has_typename_at_path(&self, lookup_path: &MergePath) -> bool { - find_selection_set_by_path( - &self.selection_set, - // We pass None for the condition, because we are checking - // the presence of __typename at the given path, not type __typename with a condition. - &lookup_path.push(Segment::Field("__typename".to_string(), 0, None)), - ) - .is_some() - } } fn selection_item_is_subset_of(source: &SelectionItem, target: &SelectionItem) -> bool { @@ -154,7 +91,7 @@ fn selection_item_is_subset_of(source: &SelectionItem, target: &SelectionItem) - } } -fn selection_items_are_subset_of(source: &[SelectionItem], target: &[SelectionItem]) -> bool { +pub fn selection_items_are_subset_of(source: &[SelectionItem], target: &[SelectionItem]) -> bool { target.iter().all(|target_node| { source .iter() @@ -162,7 +99,7 @@ fn selection_items_are_subset_of(source: &[SelectionItem], target: &[SelectionIt }) } -fn merge_selection_set(target: &mut SelectionSet, source: &SelectionSet, as_first: bool) { +pub fn merge_selection_set(target: &mut SelectionSet, source: &SelectionSet, as_first: bool) { if source.items.is_empty() { return; } @@ -217,80 +154,6 @@ fn merge_selection_set(target: &mut SelectionSet, source: &SelectionSet, as_firs } } -pub fn find_selection_set_by_path<'a>( - root_selection_set: &'a SelectionSet, - path: &MergePath, -) -> Option<&'a SelectionSet> { - let mut current_selection_set = root_selection_set; - - for path_element in path.inner.iter() { - match path_element { - Segment::List => { - continue; - } - Segment::Cast(type_name, condition) => { - let next_selection_set_option = - current_selection_set - .items - .iter() - .find_map(|item| match item { - SelectionItem::Field(_) => None, - SelectionItem::InlineFragment(f) => { - if f.type_condition.eq(type_name) - && fragment_condition_equal(condition, f) - { - Some(&f.selections) - } else { - None - } - } - SelectionItem::FragmentSpread(_) => None, - }); - - match next_selection_set_option { - Some(next_set) => { - current_selection_set = next_set; - } - None => { - return None; - } - } - } - Segment::Field(field_name, args_hash, condition) => { - let next_selection_set_option = - current_selection_set - .items - .iter() - .find_map(|item| match item { - SelectionItem::Field(field) => { - if &field.name == field_name - && field.arguments_hash() == *args_hash - && field_condition_equal(condition, field) - { - Some(&field.selections) - } else { - None - } - } - SelectionItem::InlineFragment(..) => None, - SelectionItem::FragmentSpread(_) => None, - }); - - match next_selection_set_option { - Some(next_set) => { - current_selection_set = next_set; - } - None => { - return None; - } - } - } - } - } - - Some(current_selection_set) -} - pub fn find_selection_set_by_path_mut<'a>( root_selection_set: &'a mut SelectionSet, path: &MergePath, diff --git a/lib/query-planner/src/planner/fetch/error.rs b/lib/query-planner/src/planner/fetch/error.rs index 4f24047f..48400c94 100644 --- a/lib/query-planner/src/planner/fetch/error.rs +++ b/lib/query-planner/src/planner/fetch/error.rs @@ -1,7 +1,6 @@ use crate::{ - ast::type_aware_selection::TypeAwareSelectionError, graph::{error::GraphError, node::Node}, - planner::walker::error::WalkOperationError, + planner::{fetch::selections::FetchStepSelectionsError, walker::error::WalkOperationError}, }; #[derive(Debug, thiserror::Error)] @@ -47,7 +46,7 @@ pub enum FetchGraphError { #[error("Expected @requires")] MissingRequires, #[error(transparent)] - SelectionSetManipulationError(#[from] TypeAwareSelectionError), + SelectionSetManipulationError(#[from] FetchStepSelectionsError), } impl From for FetchGraphError { diff --git a/lib/query-planner/src/planner/fetch/fetch_graph.rs b/lib/query-planner/src/planner/fetch/fetch_graph.rs index 01353478..0654562f 100644 --- a/lib/query-planner/src/planner/fetch/fetch_graph.rs +++ b/lib/query-planner/src/planner/fetch/fetch_graph.rs @@ -6,6 +6,8 @@ use crate::graph::edge::{Edge, FieldMove, InterfaceObjectTypeMove, PlannerOverri use crate::graph::node::Node; use crate::graph::Graph; use crate::planner::fetch::fetch_step_data::{FetchStepData, FetchStepFlags, FetchStepKind}; +use crate::planner::fetch::selections::FetchStepSelections; +use crate::planner::fetch::state::{MultiTypeFetchStep, SingleTypeFetchStep}; use crate::planner::plan_nodes::{FetchNodePathSegment, FetchRewrite, ValueSetter}; use crate::planner::tree::query_tree::QueryTree; use crate::planner::tree::query_tree_node::{MutationFieldPosition, QueryTreeNode}; @@ -26,31 +28,64 @@ use tracing::{instrument, trace}; use super::error::FetchGraphError; #[derive(Debug, Clone)] -pub struct FetchGraph { - pub(crate) graph: StableDiGraph, +pub struct FetchGraph { + pub(crate) graph: StableDiGraph, ()>, pub root_index: Option, } -impl Default for FetchGraph { +impl Default for FetchGraph { fn default() -> Self { Self::new() } } -impl FetchGraph { +impl FetchGraph { + pub fn to_multi_type(self) -> FetchGraph { + let new_graph = self + .graph + .map(|_, w| w.clone().into_multi_type(), |_, _| ()); + + FetchGraph { + graph: new_graph, + root_index: self.root_index, + } + } + pub fn new() -> Self { Self { graph: StableDiGraph::new(), root_index: None, } } +} - pub fn all_nodes(&self) -> NodeReferences<'_, FetchStepData> { +impl FetchGraph { + pub fn all_nodes(&self) -> NodeReferences<'_, FetchStepData> { self.graph.node_references() } + + pub fn get_step_data_mut( + &mut self, + index: NodeIndex, + ) -> Result<&mut FetchStepData, FetchGraphError> { + self.graph + .node_weight_mut(index) + .ok_or(FetchGraphError::MissingStep( + index.index(), + String::from("when getting mutable step data"), + )) + } } -impl FetchGraph { +impl FetchGraph { + pub fn is_ancestor_or_descendant(&self, a: NodeIndex, b: NodeIndex) -> bool { + self.is_descendant_of(a, b) || self.is_descendant_of(b, a) + } + + pub fn step_indices(&self) -> NodeIndices> { + self.graph.node_indices() + } + pub fn parents_of(&self, index: NodeIndex) -> petgraph::stable_graph::Edges<'_, (), Directed> { self.graph.edges_directed(index, Direction::Incoming) } @@ -63,16 +98,10 @@ impl FetchGraph { has_path_connecting(&self.graph, ancestor, descendant, None) } - /// Checks if one is ancestor of the other and vice versa - pub fn is_ancestor_or_descendant(&self, a: NodeIndex, b: NodeIndex) -> bool { - self.is_descendant_of(a, b) || self.is_descendant_of(b, a) - } - - pub fn step_indices(&self) -> NodeIndices { - self.graph.node_indices() - } - - pub fn get_step_data(&self, index: NodeIndex) -> Result<&FetchStepData, FetchGraphError> { + pub fn get_step_data( + &self, + index: NodeIndex, + ) -> Result<&FetchStepData, FetchGraphError> { self.graph .node_weight(index) .ok_or(FetchGraphError::MissingStep( @@ -81,49 +110,6 @@ impl FetchGraph { )) } - pub fn get_step_data_mut( - &mut self, - index: NodeIndex, - ) -> Result<&mut FetchStepData, FetchGraphError> { - self.graph - .node_weight_mut(index) - .ok_or(FetchGraphError::MissingStep( - index.index(), - String::from("when getting mutable step data"), - )) - } - - pub fn get_pair_of_steps_mut( - &mut self, - index1: NodeIndex, - index2: NodeIndex, - ) -> Result<(&mut FetchStepData, &mut FetchStepData), FetchGraphError> { - // `index_twice_mut` panics when indexes are equal - if index1 == index2 { - return Err(FetchGraphError::SameNodeIndex(index1.index())); - } - - // `index_twice_mut` panics when nodes do not exist - if self.graph.node_weight(index1).is_none() { - return Err(FetchGraphError::MissingStep( - index1.index(), - String::from("when checking existence"), - )); - } - if self.graph.node_weight(index2).is_none() { - return Err(FetchGraphError::MissingStep( - index2.index(), - String::from("when checking existence"), - )); - } - - Ok(self.graph.index_twice_mut(index1, index2)) - } - - #[instrument(level = "trace",skip_all, fields( - parent = parent_index.index(), - child = child_index.index(), - ))] pub fn connect(&mut self, parent_index: NodeIndex, child_index: NodeIndex) -> EdgeIndex { self.graph.update_edge(parent_index, child_index, ()) } @@ -139,13 +125,13 @@ impl FetchGraph { self.graph.remove_node(index).is_some_and(|_| true) } - pub fn add_step(&mut self, data: FetchStepData) -> NodeIndex { + pub fn add_step(&mut self, data: FetchStepData) -> NodeIndex { self.graph.add_node(data) } pub fn bfs(&self, root_index: NodeIndex, mut visitor: F) -> Option where - F: FnMut(&NodeIndex, &FetchStepData) -> bool, + F: FnMut(&NodeIndex, &FetchStepData) -> bool, { self.graph.node_weight(root_index)?; @@ -165,14 +151,16 @@ impl FetchGraph { None } +} +impl FetchGraph { #[instrument(level = "trace", skip_all)] pub fn collect_variable_usages(&mut self) -> Result<(), FetchGraphError> { let nodes_idx = self.graph.node_indices().collect::>(); for node_idx in nodes_idx { let step_data = self.get_step_data_mut(node_idx)?; - let usage = step_data.output.selection_set.variable_usages(); + let usage = step_data.output.variable_usages(); if !usage.is_empty() { step_data.variable_usages = Some(usage); @@ -181,9 +169,42 @@ impl FetchGraph { Ok(()) } + + pub fn get_pair_of_steps_mut( + &mut self, + index1: NodeIndex, + index2: NodeIndex, + ) -> Result< + ( + &mut FetchStepData, + &mut FetchStepData, + ), + FetchGraphError, + > { + // `index_twice_mut` panics when indexes are equal + if index1 == index2 { + return Err(FetchGraphError::SameNodeIndex(index1.index())); + } + + // `index_twice_mut` panics when nodes do not exist + if self.graph.node_weight(index1).is_none() { + return Err(FetchGraphError::MissingStep( + index1.index(), + String::from("when checking existence"), + )); + } + if self.graph.node_weight(index2).is_none() { + return Err(FetchGraphError::MissingStep( + index2.index(), + String::from("when checking existence"), + )); + } + + Ok(self.graph.index_twice_mut(index1, index2)) + } } -impl Display for FetchGraph { +impl Display for FetchGraph { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { writeln!(f, "Nodes:")?; for node_index in self.graph.node_indices() { @@ -240,12 +261,16 @@ impl Display for FetchGraph { } } -fn create_noop_fetch_step(fetch_graph: &mut FetchGraph, created_from_requires: bool) -> NodeIndex { +fn create_noop_fetch_step( + fetch_graph: &mut FetchGraph, + created_from_requires: bool, +) -> NodeIndex { let flags = if created_from_requires { FetchStepFlags::USED_FOR_REQUIRES } else { FetchStepFlags::empty() }; + fetch_graph.add_step(FetchStepData { service_name: SubgraphName::any(), response_path: MergePath::default(), @@ -253,10 +278,7 @@ fn create_noop_fetch_step(fetch_graph: &mut FetchGraph, created_from_requires: b selection_set: SelectionSet::default(), type_name: "*".to_string(), }, - output: TypeAwareSelection { - selection_set: SelectionSet::default(), - type_name: "*".to_string(), - }, + output: FetchStepSelections::new_empty(), flags, condition: None, kind: FetchStepKind::Root, @@ -270,7 +292,7 @@ fn create_noop_fetch_step(fetch_graph: &mut FetchGraph, created_from_requires: b } fn create_fetch_step_for_entity_call( - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, subgraph_name: &SubgraphName, input_type_name: &str, output_type_name: &str, @@ -292,10 +314,7 @@ fn create_fetch_step_for_entity_call( }, type_name: input_type_name.to_string(), }, - output: TypeAwareSelection { - selection_set: SelectionSet::default(), - type_name: output_type_name.to_string(), - }, + output: FetchStepSelections::new(output_type_name), flags, condition: condition.cloned(), kind: FetchStepKind::Entity, @@ -309,7 +328,7 @@ fn create_fetch_step_for_entity_call( } fn create_fetch_step_for_root_move( - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, root_step_index: NodeIndex, subgraph_name: &SubgraphName, type_name: &str, @@ -322,10 +341,7 @@ fn create_fetch_step_for_root_move( selection_set: SelectionSet::default(), type_name: type_name.to_string(), }, - output: TypeAwareSelection { - selection_set: SelectionSet::default(), - type_name: type_name.to_string(), - }, + output: FetchStepSelections::new(type_name), flags: FetchStepFlags::empty(), condition: None, kind: FetchStepKind::Root, @@ -345,7 +361,7 @@ fn create_fetch_step_for_root_move( // TODO: simplfy args #[allow(clippy::too_many_arguments)] fn ensure_fetch_step_for_subgraph( - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, parent_fetch_step_index: NodeIndex, subgraph_name: &SubgraphName, input_type_name: &str, @@ -443,7 +459,7 @@ fn ensure_fetch_step_for_subgraph( } fn ensure_fetch_step_for_requirement( - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, parent_fetch_step_index: NodeIndex, subgraph_name: &SubgraphName, type_name: &String, @@ -521,11 +537,13 @@ fn ensure_fetch_step_for_requirement( #[instrument(level = "trace", skip_all, fields( count = query_node.children.len(), parent_fetch_step_index = parent_fetch_step_index.index(), - requiring_fetch_step_index = requiring_fetch_step_index.map(|s| s.index()) + requiring_fetch_step_index = requiring_fetch_step_index.map(|s| s.index()), + fetch_path = fetch_path.to_string(), + response_path = response_path.to_string() ))] fn process_children_for_fetch_steps( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: NodeIndex, @@ -561,11 +579,13 @@ fn process_children_for_fetch_steps( // TODO: simplfy args #[allow(clippy::too_many_arguments)] #[instrument(level = "trace",skip_all, fields( - count = query_node.requirements.len() + count = query_node.requirements.len(), + fetch_path = fetch_path.to_string(), + response_path = response_path.to_string() ))] fn process_requirements_for_fetch_steps( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: NodeIndex, @@ -598,10 +618,9 @@ fn process_requirements_for_fetch_steps( // TODO: simplfy args #[allow(clippy::too_many_arguments)] -#[instrument(level = "trace", skip_all)] fn process_noop_edge( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: Option, @@ -629,27 +648,6 @@ fn process_noop_edge( ) } -fn add_typename_field_to_output( - fetch_step: &mut FetchStepData, - type_name: &str, - add_at: &MergePath, -) -> Result<(), FetchGraphError> { - trace!("adding __typename field to output for type '{}'", type_name); - - fetch_step.output.add_at_path( - &TypeAwareSelection { - selection_set: SelectionSet { - items: vec![SelectionItem::Field(FieldSelection::new_typename())], - }, - type_name: type_name.to_string(), - }, - add_at.clone(), - true, - )?; - - Ok(()) -} - // TODO: simplfy args #[allow(clippy::too_many_arguments)] #[instrument(level = "trace",skip_all, fields( @@ -659,7 +657,7 @@ fn add_typename_field_to_output( ))] fn process_entity_move_edge( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: NodeIndex, @@ -745,7 +743,9 @@ fn process_entity_move_edge( } let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?; - add_typename_field_to_output(parent_fetch_step, output_type_name, fetch_path)?; + parent_fetch_step + .output + .add_selection_typename(fetch_path)?; // Make the fetch step a child of the parent fetch step trace!( @@ -790,7 +790,7 @@ fn process_entity_move_edge( ))] fn process_interface_object_type_move_edge( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: NodeIndex, @@ -883,7 +883,9 @@ fn process_interface_object_type_move_edge( })); let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?; - add_typename_field_to_output(parent_fetch_step, interface_type_name, fetch_path)?; + parent_fetch_step + .output + .add_selection_typename(fetch_path)?; // In all cases it's `__typename` that needs to be resolved by another subgraph. trace!("Creating a fetch step for requirement of @interfaceObject"); @@ -964,7 +966,7 @@ fn process_interface_object_type_move_edge( ))] fn process_subgraph_entrypoint_edge( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: NodeIndex, @@ -1003,11 +1005,11 @@ fn process_subgraph_entrypoint_edge( 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() + fetch_path = fetch_path.to_string(), ))] fn process_abstract_edge( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: NodeIndex, @@ -1015,39 +1017,28 @@ fn process_abstract_edge( response_path: &MergePath, fetch_path: &MergePath, target_type_name: &String, - edge_index: &EdgeIndex, condition: Option<&Condition>, ) -> Result, FetchGraphError> { - let head_index = graph.get_edge_head(edge_index)?; - let head = graph.node(head_index)?; - let head_type_name = match head { - Node::SubgraphType(t) => &t.name, - _ => return Err(FetchGraphError::ExpectedSubgraphType), - }; - let parent_fetch_step = fetch_graph.get_step_data_mut(parent_fetch_step_index)?; trace!( "adding output field '__typename' and 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::Field(FieldSelection::new_typename()), - SelectionItem::InlineFragment(InlineFragmentSelection { - type_condition: target_type_name.clone(), - selections: SelectionSet::default(), - skip_if: None, - include_if: None, - }), - ], - }, - type_name: head_type_name.clone(), + fetch_path, + SelectionSet { + items: vec![ + SelectionItem::Field(FieldSelection::new_typename()), + SelectionItem::InlineFragment(InlineFragmentSelection { + type_condition: target_type_name.clone(), + selections: SelectionSet::default(), + skip_if: None, + include_if: None, + }), + ], }, - fetch_path.clone(), - false, )?; let child_response_path = response_path.push(Segment::Cast(target_type_name.clone(), None)); @@ -1079,11 +1070,11 @@ fn process_abstract_edge( leaf = field_move.is_leaf, list = field_move.is_list, response_path = response_path.to_string(), - fetch_path = fetch_path.to_string() + fetch_path = fetch_path.to_string(), ))] fn process_plain_field_edge( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: NodeIndex, @@ -1111,21 +1102,17 @@ fn process_plain_field_edge( ); parent_fetch_step.output.add_at_path( - &TypeAwareSelection { - selection_set: SelectionSet { - items: vec![SelectionItem::Field(FieldSelection { - name: field_move.name.to_string(), - 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, - })], - }, - type_name: field_move.type_name.to_string(), + fetch_path, + SelectionSet { + items: vec![SelectionItem::Field(FieldSelection { + name: field_move.name.to_string(), + 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, + })], }, - fetch_path.clone(), - false, )?; let child_segment = query_node.selection_alias().unwrap_or(&field_move.name); @@ -1168,10 +1155,11 @@ fn process_plain_field_edge( #[instrument(level = "trace",skip_all, fields( parent_fetch_step_index = parent_fetch_step_index.index(), requiring_fetch_step_index = requiring_fetch_step_index.map(|s| s.index()), + response_path = response_path.to_string(), ))] fn process_requires_field_edge( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: NodeIndex, @@ -1199,12 +1187,6 @@ fn process_requires_field_edge( .as_ref() .ok_or(FetchGraphError::MissingRequires)?; - let tail_node_index = graph.get_edge_tail(&edge_index)?; - let tail_node = graph.node(tail_node_index)?; - let tail_type_name = match tail_node { - Node::SubgraphType(t) => &t.name, - _ => return Err(FetchGraphError::ExpectedSubgraphType), - }; let head_node_index = graph.get_edge_head(&edge_index)?; let head_node = graph.node(head_node_index)?; let (head_type_name, head_subgraph_name, is_interface_object) = match head_node { @@ -1217,11 +1199,13 @@ fn process_requires_field_edge( override_context, query_node.requirements.first().unwrap(), )?; - trace!("Key to re-enter: {}", key_to_reenter_subgraph); let parent_fetch_step = fetch_graph.get_step_data(parent_fetch_step_index)?; // In case of a field with `@requires`, the parent will be the current subgraph we're in. - let real_parent_fetch_step_index = match parent_fetch_step.output.type_name != *head_type_name { + let real_parent_fetch_step_index = match !parent_fetch_step + .output + .is_selecting_definition(head_type_name) + { // If the parent's output resolves a different type, then it's a root type. // We can use that as a parent. true => parent_fetch_step_index, @@ -1259,21 +1243,17 @@ fn process_requires_field_edge( let step_for_children = fetch_graph.get_step_data_mut(step_for_children_index)?; step_for_children.output.add_at_path( - &TypeAwareSelection { - selection_set: SelectionSet { - items: vec![SelectionItem::Field(FieldSelection { - name: field_move.name.to_string(), - 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, - })], - }, - type_name: tail_type_name.clone(), + &MergePath::default(), + SelectionSet { + items: vec![SelectionItem::Field(FieldSelection { + name: field_move.name.to_string(), + 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, + })], }, - MergePath::default(), - false, )?; if *is_interface_object { @@ -1341,9 +1321,8 @@ fn process_requires_field_edge( ); real_parent_fetch_step.output.add_at_path( - key_to_reenter_subgraph, - key_to_reenter_at.clone(), - false, + key_to_reenter_at, + key_to_reenter_subgraph.clone().selection_set, )?; fetch_graph.connect(real_parent_fetch_step_index, step_for_requirements_index); @@ -1505,7 +1484,7 @@ fn find_satisfiable_key<'a>( #[allow(clippy::too_many_arguments)] fn process_query_node( graph: &Graph, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, override_context: &PlannerOverrideContext, query_node: &QueryTreeNode, parent_fetch_step_index: Option, @@ -1593,7 +1572,6 @@ fn process_query_node( response_path, fetch_path, type_name, - &edge_index, condition, ), Edge::InterfaceObjectTypeMove(InterfaceObjectTypeMove { @@ -1629,7 +1607,7 @@ fn process_query_node( } } -pub fn find_graph_roots(graph: &FetchGraph) -> Vec { +pub fn find_graph_roots(graph: &FetchGraph) -> Vec { let mut roots = Vec::new(); // Iterate over all nodes in the graph @@ -1653,7 +1631,7 @@ pub fn build_fetch_graph_from_query_tree( supergraph: &SupergraphState, override_context: &PlannerOverrideContext, query_tree: QueryTree, -) -> Result { +) -> Result, FetchGraphError> { let mut fetch_graph = FetchGraph::new(); process_query_node( @@ -1688,6 +1666,7 @@ pub fn build_fetch_graph_from_query_tree( // fine to unwrap as we have already checked the length fetch_graph.root_index = Some(*root_indexes.first().unwrap()); + let mut fetch_graph = fetch_graph.to_multi_type(); fetch_graph.optimize(supergraph)?; fetch_graph.collect_variable_usages()?; diff --git a/lib/query-planner/src/planner/fetch/fetch_step_data.rs b/lib/query-planner/src/planner/fetch/fetch_step_data.rs index ac73b1a8..1cd381d7 100644 --- a/lib/query-planner/src/planner/fetch/fetch_step_data.rs +++ b/lib/query-planner/src/planner/fetch/fetch_step_data.rs @@ -12,7 +12,12 @@ use crate::{ type_aware_selection::{find_arguments_conflicts, TypeAwareSelection}, }, planner::{ - fetch::fetch_graph::FetchGraph, plan_nodes::FetchRewrite, + fetch::{ + fetch_graph::FetchGraph, + selections::FetchStepSelections, + state::{MultiTypeFetchStep, SingleTypeFetchStep}, + }, + plan_nodes::FetchRewrite, tree::query_tree_node::MutationFieldPosition, }, state::supergraph_state::SubgraphName, @@ -29,11 +34,11 @@ bitflags! { } #[derive(Debug, Clone)] -pub struct FetchStepData { +pub struct FetchStepData { pub service_name: SubgraphName, pub response_path: MergePath, pub input: TypeAwareSelection, - pub output: TypeAwareSelection, + pub output: FetchStepSelections, pub kind: FetchStepKind, pub flags: FetchStepFlags, pub condition: Option, @@ -42,7 +47,7 @@ pub struct FetchStepData { pub mutation_field_position: MutationFieldPosition, pub input_rewrites: Option>, pub output_rewrites: Option>, - pub internal_aliases_locations: AliasesRecords, + pub internal_aliases_locations: Vec<(String, AliasesRecords)>, } #[derive(Debug, Clone, PartialEq)] @@ -51,18 +56,20 @@ pub enum FetchStepKind { Root, } -impl Display for FetchStepData { +impl Display for FetchStepData { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "{}/{} {} → {} at $.{}", - self.input.type_name, - self.service_name, - self.input, - self.output, - self.response_path.join("."), + "{}/{} {} → ", + self.input.type_name, self.service_name, self.input, )?; + for (def_name, selections) in self.output.iter() { + write!(f, "{}/{} ", def_name, selections)?; + } + + write!(f, " at $.{}", self.response_path.join("."))?; + if self.flags.contains(FetchStepFlags::USED_FOR_REQUIRES) { write!(f, " [@requires]")?; } @@ -82,7 +89,7 @@ impl Display for FetchStepData { } } -impl FetchStepData { +impl FetchStepData { pub fn pretty_write( &self, writer: &mut std::fmt::Formatter<'_>, @@ -112,13 +119,15 @@ impl FetchStepData { rewrites.push(rewrite); } } +} +impl FetchStepData { pub fn can_merge( &self, self_index: NodeIndex, other_index: NodeIndex, other: &Self, - fetch_graph: &FetchGraph, + fetch_graph: &FetchGraph, ) -> bool { if self_index == other_index { return false; @@ -179,3 +188,23 @@ impl FetchStepData { true } } + +impl FetchStepData { + pub fn into_multi_type(self) -> FetchStepData { + FetchStepData:: { + service_name: self.service_name, + response_path: self.response_path, + input: self.input, + output: self.output.into_multi_type(), + kind: self.kind, + flags: self.flags, + condition: self.condition, + variable_usages: self.variable_usages, + variable_definitions: self.variable_definitions, + mutation_field_position: self.mutation_field_position, + input_rewrites: self.input_rewrites, + output_rewrites: self.output_rewrites, + internal_aliases_locations: self.internal_aliases_locations, + } + } +} diff --git a/lib/query-planner/src/planner/fetch/mod.rs b/lib/query-planner/src/planner/fetch/mod.rs index 148dc4b5..767b57c2 100644 --- a/lib/query-planner/src/planner/fetch/mod.rs +++ b/lib/query-planner/src/planner/fetch/mod.rs @@ -2,3 +2,5 @@ pub(crate) mod error; pub mod fetch_graph; pub(crate) mod fetch_step_data; pub(crate) mod optimize; +pub(crate) mod selections; +pub mod state; diff --git a/lib/query-planner/src/planner/fetch/optimize/apply_internal_aliases_patching.rs b/lib/query-planner/src/planner/fetch/optimize/apply_internal_aliases_patching.rs index 03efa7e8..56813132 100644 --- a/lib/query-planner/src/planner/fetch/optimize/apply_internal_aliases_patching.rs +++ b/lib/query-planner/src/planner/fetch/optimize/apply_internal_aliases_patching.rs @@ -7,10 +7,10 @@ use crate::{ selection_item::SelectionItem, type_aware_selection::find_selection_set_by_path_mut, }, - planner::fetch::{error::FetchGraphError, fetch_graph::FetchGraph}, + planner::fetch::{error::FetchGraphError, fetch_graph::FetchGraph, state::MultiTypeFetchStep}, }; -impl FetchGraph { +impl FetchGraph { /// This method applies internal aliasing for fields in the fetch graph. /// In case a fetch step contains a record of alias made to an output field, it needs to be propagated to all descendants steps that depends on this /// output field, in multiple locations: @@ -36,70 +36,75 @@ impl FetchGraph { nodes_with_aliases.len(), ); - while let Some((aliased_node_index, aliases_locations)) = nodes_with_aliases.pop() { - let mut bfs = Bfs::new(&self.graph, aliased_node_index); - - trace!( - "Iterating step [{}], total of {} aliased fields in output selections", - aliased_node_index.index(), - aliases_locations.len() - ); - - // Iterate and find all possible children of a node that needed aliasing. - // We can't really tell which nodes are affected, as they might be at any level of the hierarchy, so we travel the graph. - while let Some(decendent_idx) = bfs.next(&self.graph) { - if decendent_idx != aliased_node_index { - let decendent = self.get_step_data_mut(decendent_idx)?; - - trace!( - "Checking if decendent [{}] is relevant for aliasing patching...", - decendent_idx.index() - ); - - for (alias_path, new_name) in aliases_locations.iter() { - // Last segment is the field that was aliased - let maybe_patched_field = alias_path.last(); - // Build a path without the alias path, to make sure we don't patch the wrong field - let relative_path = decendent.response_path.slice_from(alias_path.len()); - - if let Some(Segment::Field(field_name, args_hash, condition)) = - maybe_patched_field - { - trace!( + while let Some((aliased_node_index, scoped_aliases_locations)) = nodes_with_aliases.pop() { + for (root_type_name, aliases_locations) in scoped_aliases_locations { + let mut bfs = Bfs::new(&self.graph, aliased_node_index); + + trace!( + "Iterating step [{}], total of {} aliased fields in output selections of type {}", + aliased_node_index.index(), + aliases_locations.len(), + root_type_name + ); + + // Iterate and find all possible children of a node that needed aliasing. + // We can't really tell which nodes are affected, as they might be at any level of the hierarchy, so we travel the graph. + while let Some(decendent_idx) = bfs.next(&self.graph) { + if decendent_idx != aliased_node_index { + let decendent = self.get_step_data_mut(decendent_idx)?; + + trace!( + "Checking if decendent [{}] is relevant for aliasing patching...", + decendent_idx.index() + ); + + for (alias_path, new_name) in aliases_locations.iter() { + // Last segment is the field that was aliased + let maybe_patched_field = alias_path.last(); + // Build a path without the alias path, to make sure we don't patch the wrong field + let relative_path = + decendent.response_path.slice_from(alias_path.len()); + + if let Some(Segment::Field(field_name, args_hash, condition)) = + maybe_patched_field + { + trace!( "field '{}' was aliased, relative selection path: '{}', checking if need to patch selection '{}'", field_name, relative_path, decendent.input.selection_set ); - // First, check if the node's input selection set contains the field that was aliased - if let Some(selection) = find_selection_set_by_path_mut( - &mut decendent.input.selection_set, - &relative_path, - ) { - trace!("found selection to patch: {}", selection); - let item_to_patch = selection.items.iter_mut().find(|item| matches!(item, SelectionItem::Field(field) if field.name == *field_name && field.arguments_hash() == *args_hash)); - - if let Some(SelectionItem::Field(field_to_patch)) = item_to_patch { - field_to_patch.alias = Some(field_to_patch.name.clone()); - field_to_patch.name = new_name.clone(); - - trace!( + // First, check if the node's input selection set contains the field that was aliased + if let Some(selection) = find_selection_set_by_path_mut( + &mut decendent.input.selection_set, + &relative_path, + ) { + trace!("found selection to patch: {}", selection); + let item_to_patch = selection.items.iter_mut().find(|item| matches!(item, SelectionItem::Field(field) if field.name == *field_name && field.arguments_hash() == *args_hash)); + + if let Some(SelectionItem::Field(field_to_patch)) = + item_to_patch + { + field_to_patch.alias = Some(field_to_patch.name.clone()); + field_to_patch.name = new_name.clone(); + + trace!( "path '{}' found in selection, patched applied, new selection: {}", relative_path, field_to_patch ); + } + } else { + trace!( + "path '{}' was not found in selection '{}', skipping...", + relative_path, + decendent.input.selection_set + ); } - } else { - trace!( - "path '{}' was not found in selection '{}', skipping...", - relative_path, - decendent.input.selection_set - ); - } - // Then, check if the node's response_path is using the part that was aliased - let segment_idx_to_patch = decendent + // Then, check if the node's response_path is using the part that was aliased + let segment_idx_to_patch = decendent .response_path .inner .iter() @@ -112,8 +117,8 @@ impl FetchGraph { } }); - if let Some(segment_idx_to_patch) = segment_idx_to_patch { - trace!( + if let Some(segment_idx_to_patch) = segment_idx_to_patch { + trace!( "Node [{}] is using aliased field {} in response_path (segment idx: {}, alias: {:?})", decendent_idx.index(), field_name, @@ -121,13 +126,14 @@ impl FetchGraph { alias_path ); - let mut new_path = (*decendent.response_path.inner).to_vec(); + let mut new_path = (*decendent.response_path.inner).to_vec(); - if let Some(Segment::Field(name, _, _)) = - new_path.get_mut(segment_idx_to_patch) - { - *name = new_name.clone(); - decendent.response_path = MergePath::new(new_path); + if let Some(Segment::Field(name, _, _)) = + new_path.get_mut(segment_idx_to_patch) + { + *name = new_name.clone(); + decendent.response_path = MergePath::new(new_path); + } } } } diff --git a/lib/query-planner/src/planner/fetch/optimize/deduplicate_and_prune_fetch_steps.rs b/lib/query-planner/src/planner/fetch/optimize/deduplicate_and_prune_fetch_steps.rs index 416f6a2b..ebba8f4b 100644 --- a/lib/query-planner/src/planner/fetch/optimize/deduplicate_and_prune_fetch_steps.rs +++ b/lib/query-planner/src/planner/fetch/optimize/deduplicate_and_prune_fetch_steps.rs @@ -6,10 +6,10 @@ use tracing::{instrument, trace}; use crate::planner::fetch::{ error::FetchGraphError, fetch_graph::FetchGraph, - optimize::utils::is_reachable_via_alternative_upstream_path, + optimize::utils::is_reachable_via_alternative_upstream_path, state::MultiTypeFetchStep, }; -impl FetchGraph { +impl FetchGraph { /// Removes redundant direct dependencies from a FetchStep graph. /// /// ```text @@ -29,9 +29,7 @@ impl FetchGraph { Err(_) => return false, }; - if !step.output.selection_set.items.is_empty() - && self.parents_of(step_index).next().is_some() - { + if !step.output.is_empty() && self.parents_of(step_index).next().is_some() { return false; } diff --git a/lib/query-planner/src/planner/fetch/optimize/merge_children_with_parents.rs b/lib/query-planner/src/planner/fetch/optimize/merge_children_with_parents.rs index 28f7839c..29172614 100644 --- a/lib/query-planner/src/planner/fetch/optimize/merge_children_with_parents.rs +++ b/lib/query-planner/src/planner/fetch/optimize/merge_children_with_parents.rs @@ -5,9 +5,10 @@ use tracing::{instrument, trace}; use crate::planner::fetch::{ error::FetchGraphError, fetch_graph::FetchGraph, optimize::utils::perform_fetch_step_merge, + state::MultiTypeFetchStep, }; -impl FetchGraph { +impl FetchGraph { #[instrument(level = "trace", skip_all)] pub(crate) fn merge_children_with_parents(&mut self) -> Result<(), FetchGraphError> { let root_index = self diff --git a/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs b/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs index 653f3985..da35effa 100644 --- a/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs +++ b/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs @@ -5,17 +5,17 @@ use crate::{ ast::type_aware_selection::find_arguments_conflicts, planner::fetch::{ error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData, - optimize::utils::perform_fetch_step_merge, + optimize::utils::perform_fetch_step_merge, state::MultiTypeFetchStep, }, }; -impl FetchStepData { +impl FetchStepData { pub fn can_merge_leafs( &self, self_index: NodeIndex, other_index: NodeIndex, other: &Self, - fetch_graph: &FetchGraph, + fetch_graph: &FetchGraph, ) -> bool { if self_index == other_index { return false; @@ -66,7 +66,7 @@ impl FetchStepData { } } -impl FetchGraph { +impl FetchGraph { #[instrument(level = "trace", skip_all)] /// This optimization is about merging leaf nodes in the fetch nodes with other nodes. /// It reduces the number of fetch steps, without degrading the query performance. diff --git a/lib/query-planner/src/planner/fetch/optimize/merge_passthrough_child.rs b/lib/query-planner/src/planner/fetch/optimize/merge_passthrough_child.rs index 0ea0d859..335326c9 100644 --- a/lib/query-planner/src/planner/fetch/optimize/merge_passthrough_child.rs +++ b/lib/query-planner/src/planner/fetch/optimize/merge_passthrough_child.rs @@ -7,14 +7,18 @@ use petgraph::{ }; use tracing::{instrument, trace}; -use crate::planner::fetch::{ - error::FetchGraphError, - fetch_graph::FetchGraph, - fetch_step_data::{FetchStepData, FetchStepFlags}, +use crate::{ + ast::type_aware_selection::selection_items_are_subset_of, + planner::fetch::{ + error::FetchGraphError, + fetch_graph::FetchGraph, + fetch_step_data::{FetchStepData, FetchStepFlags}, + state::MultiTypeFetchStep, + }, }; -impl FetchGraph { - /// When a child's input contains its output, +impl FetchGraph { + /// When a child has the input identical as the output, /// it gets squashed into its parent. /// Its children becomes children of the parent. #[instrument(level = "trace", skip_all)] @@ -84,13 +88,13 @@ impl FetchGraph { } } -impl FetchStepData { +impl FetchStepData { pub(crate) fn can_merge_passthrough_child( &self, self_index: NodeIndex, other_index: NodeIndex, other: &Self, - fetch_graph: &FetchGraph, + fetch_graph: &FetchGraph, ) -> bool { if self_index == other_index { return false; @@ -112,7 +116,19 @@ impl FetchStepData { return false; } - other.input.contains(&other.output) + // TODO: Use selection_items_are_subset_of + for (output_def_name, output_selections) in other.output.iter_selections() { + if &other.input.type_name == output_def_name + && selection_items_are_subset_of( + &other.input.selection_set.items, + &output_selections.items, + ) + { + return true; + } + } + + false } } @@ -120,21 +136,19 @@ impl FetchStepData { fn perform_passthrough_child_merge( self_index: NodeIndex, other_index: NodeIndex, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, ) -> Result<(), FetchGraphError> { let (me, other) = fetch_graph.get_pair_of_steps_mut(self_index, other_index)?; + let path = other.response_path.slice_from(me.response_path.len()); trace!( - "merging fetch steps [{}] + [{}]", + "merging fetch steps [{}] + [{}] at path {}", self_index.index(), - other_index.index() + other_index.index(), + path ); - me.output.add_at_path( - &other.output, - other.response_path.slice_from(me.response_path.len()), - false, - )?; + me.output.migrate_from_another(&other.output, &path)?; let mut children_indexes: Vec = vec![]; let mut parents_indexes: Vec = vec![]; diff --git a/lib/query-planner/src/planner/fetch/optimize/merge_siblings.rs b/lib/query-planner/src/planner/fetch/optimize/merge_siblings.rs index d22b480a..d2722233 100644 --- a/lib/query-planner/src/planner/fetch/optimize/merge_siblings.rs +++ b/lib/query-planner/src/planner/fetch/optimize/merge_siblings.rs @@ -5,10 +5,10 @@ use tracing::{instrument, trace}; use crate::planner::fetch::{ error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData, - optimize::utils::perform_fetch_step_merge, + optimize::utils::perform_fetch_step_merge, state::MultiTypeFetchStep, }; -impl FetchGraph { +impl FetchGraph { #[instrument(level = "trace", skip_all)] pub(crate) fn merge_siblings(&mut self) -> Result<(), FetchGraphError> { let root_index = self @@ -104,13 +104,13 @@ impl FetchGraph { } } -impl FetchStepData { +impl FetchStepData { pub(crate) fn can_merge_siblings( &self, self_index: NodeIndex, other_index: NodeIndex, other: &Self, - fetch_graph: &FetchGraph, + fetch_graph: &FetchGraph, ) -> bool { // First, check if the base conditions for merging are met. let can_merge_base = self.can_merge(self_index, other_index, other, fetch_graph); diff --git a/lib/query-planner/src/planner/fetch/optimize/mod.rs b/lib/query-planner/src/planner/fetch/optimize/mod.rs index 3b732ea9..95f4c9c4 100644 --- a/lib/query-planner/src/planner/fetch/optimize/mod.rs +++ b/lib/query-planner/src/planner/fetch/optimize/mod.rs @@ -11,11 +11,11 @@ mod utils; use tracing::instrument; use crate::{ - planner::fetch::{error::FetchGraphError, fetch_graph::FetchGraph}, + planner::fetch::{error::FetchGraphError, fetch_graph::FetchGraph, state::MultiTypeFetchStep}, state::supergraph_state::SupergraphState, }; -impl FetchGraph { +impl FetchGraph { #[instrument(level = "trace", skip_all)] pub fn optimize(&mut self, supergraph_state: &SupergraphState) -> Result<(), FetchGraphError> { // Run optimization passes repeatedly until the graph stabilizes, as one optimization can create diff --git a/lib/query-planner/src/planner/fetch/optimize/turn_mutations_into_sequence.rs b/lib/query-planner/src/planner/fetch/optimize/turn_mutations_into_sequence.rs index 0a4ac9cb..704f82b5 100644 --- a/lib/query-planner/src/planner/fetch/optimize/turn_mutations_into_sequence.rs +++ b/lib/query-planner/src/planner/fetch/optimize/turn_mutations_into_sequence.rs @@ -4,9 +4,11 @@ use petgraph::{ }; use tracing::instrument; -use crate::planner::fetch::{error::FetchGraphError, fetch_graph::FetchGraph}; +use crate::planner::fetch::{ + error::FetchGraphError, fetch_graph::FetchGraph, state::MultiTypeFetchStep, +}; -impl FetchGraph { +impl FetchGraph { #[instrument(level = "trace", skip_all)] pub(crate) fn turn_mutations_into_sequence(&mut self) -> Result<(), FetchGraphError> { let root_index = self @@ -63,13 +65,13 @@ impl FetchGraph { } fn is_mutation_fetch_step( - fetch_graph: &FetchGraph, + fetch_graph: &FetchGraph, fetch_step_index: NodeIndex, ) -> Result { for edge_ref in fetch_graph.children_of(fetch_step_index) { let child = fetch_graph.get_step_data(edge_ref.target().id())?; - if child.output.type_name.ne("Mutation") { + if !child.output.is_selecting_definition("Mutation") { return Ok(false); } } diff --git a/lib/query-planner/src/planner/fetch/optimize/type_mismatches.rs b/lib/query-planner/src/planner/fetch/optimize/type_mismatches.rs index b8cfd1dd..2f31c660 100644 --- a/lib/query-planner/src/planner/fetch/optimize/type_mismatches.rs +++ b/lib/query-planner/src/planner/fetch/optimize/type_mismatches.rs @@ -10,19 +10,19 @@ use crate::{ type_aware_selection::{field_condition_equal, find_selection_set_by_path_mut}, }, planner::{ - fetch::{error::FetchGraphError, fetch_graph::FetchGraph}, + fetch::{error::FetchGraphError, fetch_graph::FetchGraph, state::MultiTypeFetchStep}, plan_nodes::{FetchNodePathSegment, FetchRewrite, KeyRenamer}, }, state::supergraph_state::SupergraphState, }; -impl FetchGraph { +impl FetchGraph { #[instrument(level = "trace", skip_all)] pub(crate) fn fix_conflicting_type_mismatches( &mut self, supergraph: &SupergraphState, ) -> Result<(), FetchGraphError> { - let mut pending_patches = Vec::<(NodeIndex, Vec)>::new(); + let mut pending_patches = Vec::<(NodeIndex, Vec<(String, MergePath)>)>::new(); for (node_index, node) in self.all_nodes() { if self.root_index.is_some_and(|v| v == node_index) { @@ -53,7 +53,7 @@ impl FetchGraph { node_index.index() ); - for mismatch_path in mismatches_paths { + for (root_def_name, mismatch_path) in mismatches_paths { let mut merger = SafeSelectionSetMerger::default(); if let Some(Segment::Field(field_lookup, args_hash_lookup, condition)) = @@ -61,9 +61,13 @@ impl FetchGraph { { // TODO: We can avoid this cut and slice thing, if we return "SelectionItem" instead of "SelectionSet" inside "find_selection_set_by_path_mut". let lookup_path = &mismatch_path.without_last(); + let root_def_selections = node + .output + .selections_for_definition_mut(&root_def_name) + .expect("missing definition in step"); if let Some(selection_set) = - find_selection_set_by_path_mut(&mut node.output.selection_set, lookup_path) + find_selection_set_by_path_mut(root_def_selections, lookup_path) { let next_alias = merger.safe_next_alias_name(&selection_set.items); let item = selection_set diff --git a/lib/query-planner/src/planner/fetch/optimize/utils.rs b/lib/query-planner/src/planner/fetch/optimize/utils.rs index 07ddb85d..7d5ad38b 100644 --- a/lib/query-planner/src/planner/fetch/optimize/utils.rs +++ b/lib/query-planner/src/planner/fetch/optimize/utils.rs @@ -6,14 +6,9 @@ use petgraph::{ }; use tracing::{instrument, trace}; -use crate::{ - ast::{ - merge_path::Condition, selection_item::SelectionItem, - selection_set::InlineFragmentSelection, - }, - planner::fetch::{ - error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepFlags, - }, +use crate::planner::fetch::{ + error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepFlags, + state::MultiTypeFetchStep, }; // Return true in case an alias was applied during the merge process. @@ -21,7 +16,7 @@ use crate::{ pub(crate) fn perform_fetch_step_merge( self_index: NodeIndex, other_index: NodeIndex, - fetch_graph: &mut FetchGraph, + fetch_graph: &mut FetchGraph, ) -> Result<(), FetchGraphError> { let (me, other) = fetch_graph.get_pair_of_steps_mut(self_index, other_index)?; @@ -44,29 +39,11 @@ pub(crate) fn perform_fetch_step_merge( // { products { ... on Product @skip(if: $bool) { __typename id price } } } // // We can't do it when both are entity calls fetch steps. - other.output.add(&other.input); - - let old_selection_set = other.output.selection_set.clone(); - match condition { - Condition::Include(var_name) => { - other.output.selection_set.items = - vec![SelectionItem::InlineFragment(InlineFragmentSelection { - type_condition: other.output.type_name.clone(), - selections: old_selection_set, - skip_if: None, - include_if: Some(var_name), - })]; - } - Condition::Skip(var_name) => { - other.output.selection_set.items = - vec![SelectionItem::InlineFragment(InlineFragmentSelection { - type_condition: other.output.type_name.clone(), - selections: old_selection_set, - skip_if: Some(var_name), - include_if: None, - })]; - } - } + other + .output + .add(&other.input.type_name, &other.input.selection_set)?; + + other.output.wrap_with_condition(condition); } else if me.is_entity_call() && other.condition.is_some() { // We don't want to pass a condition // to a regular (non-entity call) fetch step, @@ -74,19 +51,22 @@ pub(crate) fn perform_fetch_step_merge( me.condition = other.condition.take(); } - let maybe_aliases = me.output.add_at_path_and_solve_conflicts( + let scoped_aliases = me.output.safe_migrate_from_another( &other.output, - other.response_path.slice_from(me.response_path.len()), + &other.response_path.slice_from(me.response_path.len()), ( me.flags.contains(FetchStepFlags::USED_FOR_REQUIRES), other.flags.contains(FetchStepFlags::USED_FOR_REQUIRES), ), - false, - ); - - // In cases where merging a step resulted in internal aliasing, keep a record of the aliases. - if let Some(aliases) = maybe_aliases { - me.internal_aliases_locations.extend(aliases); + )?; + + if !scoped_aliases.is_empty() { + trace!( + "Total of {} alises applied during safe merge of selections", + scoped_aliases.len() + ); + // In cases where merging a step resulted in internal aliasing, keep a record of the aliases. + me.internal_aliases_locations.extend(scoped_aliases); } if let Some(input_rewrites) = other.input_rewrites.take() { @@ -145,7 +125,7 @@ pub(crate) fn perform_fetch_step_merge( /// The search starts from all direct parents of `child_index` *except* /// `target_ancestor_index`, and follows incoming edges from there. pub fn is_reachable_via_alternative_upstream_path( - graph: &FetchGraph, + graph: &FetchGraph, child_index: NodeIndex, target_ancestor_index: NodeIndex, ) -> Result { diff --git a/lib/query-planner/src/planner/fetch/selections.rs b/lib/query-planner/src/planner/fetch/selections.rs new file mode 100644 index 00000000..149e0efc --- /dev/null +++ b/lib/query-planner/src/planner/fetch/selections.rs @@ -0,0 +1,315 @@ +use std::{ + collections::{BTreeSet, HashMap}, + fmt::Display, + marker::PhantomData, +}; + +use crate::{ + ast::{ + merge_path::{Condition, MergePath}, + safe_merge::{AliasesRecords, SafeSelectionSetMerger}, + selection_item::SelectionItem, + selection_set::{FieldSelection, InlineFragmentSelection, SelectionSet}, + type_aware_selection::{find_selection_set_by_path_mut, merge_selection_set}, + }, + planner::fetch::state::{MultiTypeFetchStep, SingleTypeFetchStep}, +}; + +#[derive(Debug, thiserror::Error)] +pub enum FetchStepSelectionsError { + #[error("Unexpected missing definition: {0}")] + UnexpectedMissingDefinition(String), + #[error("Path '{0}' cannot be found in selection set of type {1}")] + MissingPathInSelection(String, String), +} + +#[derive(Debug, Clone)] +pub struct FetchStepSelections { + selections: HashMap, + _state: PhantomData, +} + +impl Display for FetchStepSelections { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let as_selection_set: SelectionSet = self.into(); + + write!(f, "{}", as_selection_set) + } +} + +impl From<&FetchStepSelections> for SelectionSet { + fn from(value: &FetchStepSelections) -> Self { + if value.selections.len() == 1 { + let (type_name, selections) = value.selections.iter().next().unwrap(); + + if type_name == "Query" || type_name == "Mutation" || type_name == "Subscription" { + return selections.clone(); + } + } + + SelectionSet { + items: value + .selections + .iter() + .map(|(def_name, selections)| { + SelectionItem::InlineFragment(InlineFragmentSelection { + type_condition: def_name.clone(), + include_if: None, + skip_if: None, + selections: selections.clone(), + }) + }) + .collect(), + } + } +} + +impl From<&FetchStepSelections> for SelectionSet { + fn from(value: &FetchStepSelections) -> Self { + let (_type_name, selections) = value.selections.iter().next().unwrap(); + + selections.clone() + } +} + +impl FetchStepSelections { + pub fn into_multi_type(self) -> FetchStepSelections { + FetchStepSelections { + _state: Default::default(), + selections: self.selections, + } + } + + pub fn definition_name(&self) -> &str { + self.selections + .keys() + .next() + .expect("SingleTypeFetchStep should have exactly one selection") + } + + pub fn selection_set(&self) -> &SelectionSet { + self.selections + .iter() + .next() + .expect("SingleTypeFetchStep should have exactly one selection") + .1 + } + + pub fn add_at_path( + &mut self, + fetch_path: &MergePath, + selection_set: SelectionSet, + ) -> Result<(), FetchStepSelectionsError> { + let def_name = self.definition_name().to_string(); + + self.add_at_path_inner(&def_name, fetch_path, selection_set, false) + } + + pub fn add_selection_typename( + &mut self, + fetch_path: &MergePath, + ) -> Result<(), FetchStepSelectionsError> { + let def_name = self.definition_name().to_string(); + + self.add_at_path_inner( + &def_name, + fetch_path, + SelectionSet { + items: vec![SelectionItem::Field(FieldSelection::new_typename())], + }, + true, + ) + } +} + +impl FetchStepSelections { + pub fn is_selecting_definition(&self, definition_name: &str) -> bool { + self.selections.contains_key(definition_name) + } + + pub fn is_empty(&self) -> bool { + self.selections.is_empty() + || self + .selections + .values() + .all(|selection_set| selection_set.is_empty()) + } + + pub fn iter(&self) -> impl Iterator { + self.selections.iter() + } + + pub fn variable_usages(&self) -> BTreeSet { + let mut usages = BTreeSet::new(); + + for selection_set in self.selections.values() { + usages.extend(selection_set.variable_usages()); + } + + usages + } + + pub fn selections_for_definition_mut( + &mut self, + definition_name: &str, + ) -> Option<&mut SelectionSet> { + self.selections.get_mut(definition_name) + } + + fn add_at_path_inner( + &mut self, + definition_name: &str, + fetch_path: &MergePath, + selection_set: SelectionSet, + as_first: bool, + ) -> Result<(), FetchStepSelectionsError> { + let current = self + .selections_for_definition_mut(definition_name) + .ok_or_else(|| { + FetchStepSelectionsError::UnexpectedMissingDefinition(definition_name.to_string()) + })?; + + let selection_set_at_path = find_selection_set_by_path_mut(current, fetch_path) + .ok_or_else(|| { + FetchStepSelectionsError::MissingPathInSelection( + fetch_path.to_string(), + definition_name.to_string(), + ) + })?; + + merge_selection_set(selection_set_at_path, &selection_set, as_first); + + Ok(()) + } +} + +impl FetchStepSelections { + pub fn try_as_single(&self) -> Option<&str> { + if self.selections.len() == 1 { + self.selections.keys().next().map(|key| key.as_str()) + } else { + None + } + } + + pub fn iter_selections(&self) -> impl Iterator { + self.selections.iter() + } + + pub fn migrate_from_another( + &mut self, + other: &Self, + fetch_path: &MergePath, + ) -> Result<(), FetchStepSelectionsError> { + let maybe_merge_into = self.try_as_single().map(|str| str.to_string()); + + for (definition_name, selection_set) in other.iter_selections() { + let target_type = maybe_merge_into.as_ref().unwrap_or(definition_name); + self.add_at_path_inner(target_type, fetch_path, selection_set.clone(), false)?; + } + + Ok(()) + } + + pub fn add( + &mut self, + definition_name: &str, + selection_set: &SelectionSet, + ) -> Result<(), FetchStepSelectionsError> { + let selections = self + .selections_for_definition_mut(definition_name) + .ok_or_else(|| { + FetchStepSelectionsError::UnexpectedMissingDefinition(definition_name.to_string()) + })?; + + merge_selection_set(selections, selection_set, false); + + Ok(()) + } + + pub fn safe_migrate_from_another( + &mut self, + other: &Self, + fetch_path: &MergePath, + (self_used_for_requires, other_used_for_requires): (bool, bool), + ) -> Result, FetchStepSelectionsError> { + let mut aliases_made: Vec<(String, AliasesRecords)> = Vec::new(); + let maybe_merge_into = self.try_as_single().map(|str| str.to_string()); + + for (definition_name, selection_set) in other.iter_selections() { + let target_type = maybe_merge_into.as_ref().unwrap_or(definition_name); + let current = self + .selections_for_definition_mut(target_type) + .ok_or_else(|| { + FetchStepSelectionsError::UnexpectedMissingDefinition(target_type.to_string()) + })?; + + let selection_at_path = find_selection_set_by_path_mut(current, fetch_path) + .ok_or_else(|| { + FetchStepSelectionsError::MissingPathInSelection( + fetch_path.to_string(), + target_type.to_string(), + ) + })?; + + let mut merger = SafeSelectionSetMerger::default(); + let current_aliases_made = merger.merge_selection_set( + selection_at_path, + selection_set, + (self_used_for_requires, other_used_for_requires), + false, + ); + + if !current_aliases_made.is_empty() { + aliases_made.push((target_type.to_string(), current_aliases_made)); + } + } + + Ok(aliases_made) + } + + pub fn wrap_with_condition(&mut self, condition: Condition) { + for (def_name, selection_set) in self.selections.iter_mut() { + let prev = selection_set.clone(); + match &condition { + Condition::Include(var_name) => { + selection_set.items = + vec![SelectionItem::InlineFragment(InlineFragmentSelection { + type_condition: def_name.to_string(), + selections: prev, + skip_if: None, + include_if: Some(var_name.clone()), + })]; + } + Condition::Skip(var_name) => { + selection_set.items = + vec![SelectionItem::InlineFragment(InlineFragmentSelection { + type_condition: def_name.to_string(), + selections: prev, + skip_if: Some(var_name.clone()), + include_if: None, + })]; + } + } + } + } +} + +impl FetchStepSelections { + pub fn new(definition_name: &str) -> Self { + let mut map = HashMap::new(); + map.insert(definition_name.to_string(), SelectionSet::default()); + + Self { + _state: Default::default(), + selections: map, + } + } + + pub fn new_empty() -> Self { + Self { + _state: Default::default(), + selections: Default::default(), + } + } +} diff --git a/lib/query-planner/src/planner/fetch/state.rs b/lib/query-planner/src/planner/fetch/state.rs new file mode 100644 index 00000000..bd9ecf2f --- /dev/null +++ b/lib/query-planner/src/planner/fetch/state.rs @@ -0,0 +1,17 @@ +/// The "state" representation of a fetch step and fetch graph. +/// +/// We split the work we do on the FetchGraph to to distinct parts: +/// +/// 1. Graph building: mainly happens in `fetch_step.rs` and is using `SingleTypeFetchStep` as state. +/// This ensures that the step is built with a single type in mind. +/// 2. Graph optimization: mainly happens in `fetch/optimize/` files, and using `MultiTypeFetchStep` as state. +/// This ensures that the step is optimized with multiple types in mind. +/// +/// With this setup, we can ensure that some parts of the logic/capabilities that can be performed on a selection set or a step are limited, +/// and either scoped to a single type or a multi-type context. + +#[derive(Debug, Clone)] +pub struct SingleTypeFetchStep; + +#[derive(Debug, Clone)] +pub struct MultiTypeFetchStep; diff --git a/lib/query-planner/src/planner/mod.rs b/lib/query-planner/src/planner/mod.rs index 34c6b422..797f5f72 100644 --- a/lib/query-planner/src/planner/mod.rs +++ b/lib/query-planner/src/planner/mod.rs @@ -10,7 +10,10 @@ use crate::{ ast::operation::{OperationDefinition, VariableDefinition}, consumer_schema::ConsumerSchema, graph::{edge::PlannerOverrideContext, error::GraphError, Graph}, - planner::{best::find_best_combination, fetch::fetch_graph::FetchGraph}, + planner::{ + best::find_best_combination, + fetch::{fetch_graph::FetchGraph, state::MultiTypeFetchStep}, + }, state::supergraph_state::SupergraphState, }; @@ -114,7 +117,7 @@ impl Planner { } pub fn add_variables_to_fetch_steps( - graph: &mut FetchGraph, + graph: &mut FetchGraph, variables: &Option>, ) -> Result<(), PlannerError> { if let Some(variables) = variables { diff --git a/lib/query-planner/src/planner/plan_nodes.rs b/lib/query-planner/src/planner/plan_nodes.rs index f68d3592..8ed9d148 100644 --- a/lib/query-planner/src/planner/plan_nodes.rs +++ b/lib/query-planner/src/planner/plan_nodes.rs @@ -8,7 +8,7 @@ use crate::{ type_aware_selection::TypeAwareSelection, value::Value, }, - planner::fetch::fetch_step_data::FetchStepData, + planner::fetch::{fetch_step_data::FetchStepData, state::MultiTypeFetchStep}, state::supergraph_state::{OperationKind, SupergraphState, TypeNode}, utils::pretty_display::{get_indent, PrettyDisplay}, }; @@ -268,10 +268,9 @@ fn create_input_selection_set(input_selections: &TypeAwareSelection) -> Selectio } fn create_output_operation( - step: &FetchStepData, + step: &FetchStepData, supergraph: &SupergraphState, ) -> SubgraphFetchOperation { - let type_aware_selection = &step.output; let mut variables = vec![VariableDefinition { name: "representations".to_string(), variable_type: TypeNode::NonNull(Box::new(TypeNode::List(Box::new(TypeNode::NonNull( @@ -291,14 +290,7 @@ fn create_output_operation( selection_set: SelectionSet { items: vec![SelectionItem::Field(FieldSelection { name: "_entities".to_string(), - selections: SelectionSet { - items: vec![SelectionItem::InlineFragment(InlineFragmentSelection { - selections: type_aware_selection.selection_set.clone(), - type_condition: type_aware_selection.type_name.clone(), - skip_if: None, - include_if: None, - })], - }, + selections: (&step.output).into(), alias: None, arguments: Some( ( @@ -322,24 +314,22 @@ fn create_output_operation( } } -impl From<&FetchStepData> for OperationKind { - fn from(step: &FetchStepData) -> Self { - let type_name = step.output.type_name.as_str(); - - if type_name == "Query" { - OperationKind::Query - } else if type_name == "Mutation" { - OperationKind::Mutation - } else if type_name == "Subscription" { - OperationKind::Subscription - } else { - OperationKind::Query +impl From<&FetchStepData> for OperationKind { + fn from(step: &FetchStepData) -> Self { + match &step.input.type_name { + str if str == "Query" => OperationKind::Query, + str if str == "Mutation" => OperationKind::Mutation, + str if str == "Subscription" => OperationKind::Subscription, + _ => OperationKind::Query, } } } impl FetchNode { - pub fn from_fetch_step(step: &FetchStepData, supergraph: &SupergraphState) -> Self { + pub fn from_fetch_step( + step: &FetchStepData, + supergraph: &SupergraphState, + ) -> Self { match step.is_entity_call() { true => FetchNode { service_name: step.service_name.0.clone(), @@ -355,7 +345,7 @@ impl FetchNode { let operation_def = OperationDefinition { name: None, operation_kind: Some(step.into()), - selection_set: step.output.selection_set.clone(), + selection_set: (&step.output).into(), variable_definitions: step.variable_definitions.clone(), }; let document = @@ -381,7 +371,10 @@ impl FetchNode { } impl PlanNode { - pub fn from_fetch_step(step: &FetchStepData, supergraph: &SupergraphState) -> Self { + pub fn from_fetch_step( + step: &FetchStepData, + supergraph: &SupergraphState, + ) -> Self { let node = if step.response_path.is_empty() { PlanNode::Fetch(FetchNode::from_fetch_step(step, supergraph)) } else { diff --git a/lib/query-planner/src/planner/query_plan.rs b/lib/query-planner/src/planner/query_plan.rs index d0c15993..f39d6022 100644 --- a/lib/query-planner/src/planner/query_plan.rs +++ b/lib/query-planner/src/planner/query_plan.rs @@ -2,7 +2,10 @@ use std::collections::{HashMap, VecDeque}; use petgraph::{graph::NodeIndex, visit::EdgeRef}; -use crate::{planner::plan_nodes::ConditionNode, state::supergraph_state::SupergraphState}; +use crate::{ + planner::{fetch::state::MultiTypeFetchStep, plan_nodes::ConditionNode}, + state::supergraph_state::SupergraphState, +}; use super::{ error::QueryPlanError, @@ -16,11 +19,11 @@ use super::{ /// when its in-degree becomes zero. pub struct InDegree<'a> { state: HashMap, - fetch_graph: &'a FetchGraph, + fetch_graph: &'a FetchGraph, } impl<'a> InDegree<'a> { - pub fn new(fetch_graph: &'a FetchGraph) -> Result { + pub fn new(fetch_graph: &'a FetchGraph) -> Result { let mut state: HashMap = HashMap::new(); let root_index = fetch_graph.root_index.ok_or(QueryPlanError::NoRoot)?; @@ -71,7 +74,7 @@ impl<'a> InDegree<'a> { #[tracing::instrument(level = "trace", skip_all)] pub fn build_query_plan_from_fetch_graph( - fetch_graph: FetchGraph, + fetch_graph: FetchGraph, supergraph: &SupergraphState, ) -> Result { let root_index = fetch_graph.root_index.ok_or(QueryPlanError::NoRoot)?; diff --git a/lib/query-planner/src/tests/mod.rs b/lib/query-planner/src/tests/mod.rs index 0d869af4..db3021da 100644 --- a/lib/query-planner/src/tests/mod.rs +++ b/lib/query-planner/src/tests/mod.rs @@ -19,3 +19,20 @@ mod requires_requires; mod root_types; mod testkit; mod union; + +use crate::{ + tests::testkit::{build_query_plan, init_logger}, + utils::parsing::parse_operation, +}; + +#[test] +fn test_bench_operation() -> Result<(), Box> { + init_logger(); + let document = parse_operation( + &std::fs::read_to_string("../../bench/operation.graphql") + .expect("Unable to read input file"), + ); + let _query_plan = build_query_plan("../../bench/supergraph.graphql", document)?; + + Ok(()) +} From 280164fca19a270eed3916b1631c83c6488c09f9 Mon Sep 17 00:00:00 2001 From: Dotan Simha Date: Wed, 16 Jul 2025 23:48:55 +0300 Subject: [PATCH 2/4] feat(qp): apply new single/multi typestate and replace `input` --- lib/query-planner/src/ast/selection_set.rs | 247 +++++++++++++++- .../src/ast/type_aware_selection.rs | 264 +----------------- lib/query-planner/src/graph/edge.rs | 2 +- .../src/planner/fetch/fetch_graph.rs | 65 +++-- .../src/planner/fetch/fetch_step_data.rs | 35 ++- .../apply_internal_aliases_patching.rs | 25 +- .../src/planner/fetch/optimize/merge_leafs.rs | 21 +- .../fetch/optimize/merge_passthrough_child.rs | 15 +- .../planner/fetch/optimize/type_mismatches.rs | 2 +- .../src/planner/fetch/optimize/utils.rs | 20 +- .../src/planner/fetch/selections.rs | 79 ++++-- lib/query-planner/src/planner/plan_nodes.rs | 30 +- 12 files changed, 436 insertions(+), 369 deletions(-) diff --git a/lib/query-planner/src/ast/selection_set.rs b/lib/query-planner/src/ast/selection_set.rs index 9e233b24..cfb60302 100644 --- a/lib/query-planner/src/ast/selection_set.rs +++ b/lib/query-planner/src/ast/selection_set.rs @@ -6,7 +6,10 @@ use std::{ hash::Hash, }; -use crate::utils::pretty_display::{get_indent, PrettyDisplay}; +use crate::{ + ast::merge_path::{Condition, MergePath, Segment}, + utils::pretty_display::{get_indent, PrettyDisplay}, +}; use super::{arguments::ArgumentsMap, selection_item::SelectionItem}; @@ -55,10 +58,24 @@ impl Display for SelectionSet { } impl SelectionSet { + pub fn cost(&self) -> u64 { + let mut cost = 1; + + for node in &self.items { + cost += node.cost(); + } + + cost + } + pub fn is_empty(&self) -> bool { self.items.is_empty() } + pub fn contains(&self, other: &Self) -> bool { + selection_items_are_subset_of(&self.items, &other.items) + } + pub fn variable_usages(&self) -> BTreeSet { self.items .iter() @@ -343,6 +360,234 @@ impl PrettyDisplay for InlineFragmentSelection { } } +pub fn selection_items_are_subset_of(source: &[SelectionItem], target: &[SelectionItem]) -> bool { + target.iter().all(|target_node| { + source + .iter() + .any(|source_node| selection_item_is_subset_of(source_node, target_node)) + }) +} + +fn selection_item_is_subset_of(source: &SelectionItem, target: &SelectionItem) -> bool { + match (source, target) { + (SelectionItem::Field(source_field), SelectionItem::Field(target_field)) => { + if source_field.name != target_field.name { + return false; + } + + if source_field.is_leaf() != target_field.is_leaf() { + return false; + } + + selection_items_are_subset_of( + &source_field.selections.items, + &target_field.selections.items, + ) + } + // TODO: support fragments + _ => false, + } +} + +pub fn merge_selection_set(target: &mut SelectionSet, source: &SelectionSet, as_first: bool) { + if source.items.is_empty() { + return; + } + + let mut pending_items = Vec::with_capacity(source.items.len()); + for source_item in source.items.iter() { + let mut found = false; + for target_item in target.items.iter_mut() { + match (source_item, target_item) { + (SelectionItem::Field(source_field), SelectionItem::Field(target_field)) => { + if source_field == target_field { + found = true; + merge_selection_set( + &mut target_field.selections, + &source_field.selections, + as_first, + ); + break; + } + } + ( + SelectionItem::InlineFragment(source_fragment), + SelectionItem::InlineFragment(target_fragment), + ) => { + if source_fragment.type_condition == target_fragment.type_condition { + found = true; + merge_selection_set( + &mut target_fragment.selections, + &source_fragment.selections, + as_first, + ); + break; + } + } + _ => {} + } + } + + if !found { + pending_items.push(source_item.clone()) + } + } + + if !pending_items.is_empty() { + if as_first { + let mut new_items = pending_items; + new_items.append(&mut target.items); + target.items = new_items; + } else { + target.items.extend(pending_items); + } + } +} + +pub fn find_selection_set_by_path_mut<'a>( + root_selection_set: &'a mut SelectionSet, + path: &MergePath, +) -> Option<&'a mut SelectionSet> { + let mut current_selection_set = root_selection_set; + + for path_element in path.inner.iter() { + match path_element { + Segment::List => { + continue; + } + Segment::Cast(type_name, condition) => { + let next_selection_set_option = + current_selection_set + .items + .iter_mut() + .find_map(|item| match item { + SelectionItem::Field(_) => None, + SelectionItem::InlineFragment(f) => { + if f.type_condition.eq(type_name) + && fragment_condition_equal(condition, f) + { + Some(&mut f.selections) + } else { + None + } + } + SelectionItem::FragmentSpread(_) => None, + }); + + match next_selection_set_option { + Some(next_set) => { + current_selection_set = next_set; + } + None => { + return None; + } + } + } + Segment::Field(field_name, args_hash, condition) => { + let next_selection_set_option = + current_selection_set + .items + .iter_mut() + .find_map(|item| match item { + SelectionItem::Field(field) => { + if field.selection_identifier() == field_name + && field.arguments_hash() == *args_hash + && field_condition_equal(condition, field) + { + Some(&mut field.selections) + } else { + None + } + } + SelectionItem::InlineFragment(..) => None, + SelectionItem::FragmentSpread(_) => None, + }); + + match next_selection_set_option { + Some(next_set) => { + current_selection_set = next_set; + } + None => { + return None; + } + } + } + } + } + Some(current_selection_set) +} + +pub fn field_condition_equal(cond: &Option, field: &FieldSelection) -> bool { + match cond { + Some(cond) => match cond { + Condition::Include(var_name) => { + field.include_if.as_ref().is_some_and(|v| v == var_name) + } + Condition::Skip(var_name) => field.skip_if.as_ref().is_some_and(|v| v == var_name), + }, + None => field.include_if.is_none() && field.skip_if.is_none(), + } +} + +fn fragment_condition_equal(cond: &Option, fragment: &InlineFragmentSelection) -> bool { + match cond { + Some(cond) => match cond { + Condition::Include(var_name) => { + fragment.include_if.as_ref().is_some_and(|v| v == var_name) + } + Condition::Skip(var_name) => fragment.skip_if.as_ref().is_some_and(|v| v == var_name), + }, + None => fragment.include_if.is_none() && fragment.skip_if.is_none(), + } +} + +/// Find the arguments conflicts between two selections. +/// Returns a vector of tuples containing the indices of conflicting fields in both "source" and "other" +/// Both indices are returned in order to allow for easy resolution of conflicts later, in either side. +pub fn find_arguments_conflicts( + source: &SelectionSet, + other: &SelectionSet, +) -> Vec<(usize, usize)> { + other + .items + .iter() + .enumerate() + .filter_map(|(index, other_selection)| { + if let SelectionItem::Field(other_field) = other_selection { + let other_identifier = other_field.selection_identifier(); + let other_args_hash = other_field.arguments_hash(); + + let existing_in_self = + source + .items + .iter() + .enumerate() + .find_map(|(self_index, self_selection)| { + if let SelectionItem::Field(self_field) = self_selection { + // If the field selection identifier matches and the arguments hash is different, + // then it means that we can't merge the two input siblings + if self_field.selection_identifier() == other_identifier + && self_field.arguments_hash() != other_args_hash + { + return Some(self_index); + } + } + + None + }); + + if let Some(existing_index) = existing_in_self { + return Some((existing_index, index)); + } + + return None; + } + + None + }) + .collect() +} + #[cfg(test)] mod tests { use crate::ast::value::Value; diff --git a/lib/query-planner/src/ast/type_aware_selection.rs b/lib/query-planner/src/ast/type_aware_selection.rs index 729d4f5a..49ccb150 100644 --- a/lib/query-planner/src/ast/type_aware_selection.rs +++ b/lib/query-planner/src/ast/type_aware_selection.rs @@ -1,11 +1,6 @@ use std::{fmt::Display, hash::Hash}; -use crate::ast::{ - merge_path::{Condition, Segment}, - selection_set::{FieldSelection, InlineFragmentSelection}, -}; - -use super::{merge_path::MergePath, selection_item::SelectionItem, selection_set::SelectionSet}; +use super::selection_set::SelectionSet; #[derive(Debug, Clone)] pub struct TypeAwareSelection { @@ -38,260 +33,3 @@ impl PartialEq for TypeAwareSelection { } impl Eq for TypeAwareSelection {} - -impl TypeAwareSelection { - pub fn new(type_name: String, selection_set: SelectionSet) -> Self { - Self { - type_name, - selection_set, - } - } - - pub fn cost(&self) -> u64 { - let mut cost = 1; - - for node in &self.selection_set.items { - cost += node.cost(); - } - - cost - } - - pub fn contains(&self, other: &Self) -> bool { - if self.type_name != other.type_name { - return false; - } - - selection_items_are_subset_of(&self.selection_set.items, &other.selection_set.items) - } - - pub fn add(&mut self, to_add: &Self) { - merge_selection_set(&mut self.selection_set, &to_add.selection_set, false); - } -} - -fn selection_item_is_subset_of(source: &SelectionItem, target: &SelectionItem) -> bool { - match (source, target) { - (SelectionItem::Field(source_field), SelectionItem::Field(target_field)) => { - if source_field.name != target_field.name { - return false; - } - - if source_field.is_leaf() != target_field.is_leaf() { - return false; - } - - selection_items_are_subset_of( - &source_field.selections.items, - &target_field.selections.items, - ) - } - // TODO: support fragments - _ => false, - } -} - -pub fn selection_items_are_subset_of(source: &[SelectionItem], target: &[SelectionItem]) -> bool { - target.iter().all(|target_node| { - source - .iter() - .any(|source_node| selection_item_is_subset_of(source_node, target_node)) - }) -} - -pub fn merge_selection_set(target: &mut SelectionSet, source: &SelectionSet, as_first: bool) { - if source.items.is_empty() { - return; - } - - let mut pending_items = Vec::with_capacity(source.items.len()); - for source_item in source.items.iter() { - let mut found = false; - for target_item in target.items.iter_mut() { - match (source_item, target_item) { - (SelectionItem::Field(source_field), SelectionItem::Field(target_field)) => { - if source_field == target_field { - found = true; - merge_selection_set( - &mut target_field.selections, - &source_field.selections, - as_first, - ); - break; - } - } - ( - SelectionItem::InlineFragment(source_fragment), - SelectionItem::InlineFragment(target_fragment), - ) => { - if source_fragment.type_condition == target_fragment.type_condition { - found = true; - merge_selection_set( - &mut target_fragment.selections, - &source_fragment.selections, - as_first, - ); - break; - } - } - _ => {} - } - } - - if !found { - pending_items.push(source_item.clone()) - } - } - - if !pending_items.is_empty() { - if as_first { - let mut new_items = pending_items; - new_items.append(&mut target.items); - target.items = new_items; - } else { - target.items.extend(pending_items); - } - } -} - -pub fn find_selection_set_by_path_mut<'a>( - root_selection_set: &'a mut SelectionSet, - path: &MergePath, -) -> Option<&'a mut SelectionSet> { - let mut current_selection_set = root_selection_set; - - for path_element in path.inner.iter() { - match path_element { - Segment::List => { - continue; - } - Segment::Cast(type_name, condition) => { - let next_selection_set_option = - current_selection_set - .items - .iter_mut() - .find_map(|item| match item { - SelectionItem::Field(_) => None, - SelectionItem::InlineFragment(f) => { - if f.type_condition.eq(type_name) - && fragment_condition_equal(condition, f) - { - Some(&mut f.selections) - } else { - None - } - } - SelectionItem::FragmentSpread(_) => None, - }); - - match next_selection_set_option { - Some(next_set) => { - current_selection_set = next_set; - } - None => { - return None; - } - } - } - Segment::Field(field_name, args_hash, condition) => { - let next_selection_set_option = - current_selection_set - .items - .iter_mut() - .find_map(|item| match item { - SelectionItem::Field(field) => { - if field.selection_identifier() == field_name - && field.arguments_hash() == *args_hash - && field_condition_equal(condition, field) - { - Some(&mut field.selections) - } else { - None - } - } - SelectionItem::InlineFragment(..) => None, - SelectionItem::FragmentSpread(_) => None, - }); - - match next_selection_set_option { - Some(next_set) => { - current_selection_set = next_set; - } - None => { - return None; - } - } - } - } - } - Some(current_selection_set) -} - -/// Find the arguments conflicts between two selections. -/// Returns a vector of tuples containing the indices of conflicting fields in both "source" and "other" -/// Both indices are returned in order to allow for easy resolution of conflicts later, in either side. -pub fn find_arguments_conflicts( - source: &TypeAwareSelection, - other: &TypeAwareSelection, -) -> Vec<(usize, usize)> { - other - .selection_set - .items - .iter() - .enumerate() - .filter_map(|(index, other_selection)| { - if let SelectionItem::Field(other_field) = other_selection { - let other_identifier = other_field.selection_identifier(); - let other_args_hash = other_field.arguments_hash(); - - let existing_in_self = source.selection_set.items.iter().enumerate().find_map( - |(self_index, self_selection)| { - if let SelectionItem::Field(self_field) = self_selection { - // If the field selection identifier matches and the arguments hash is different, - // then it means that we can't merge the two input siblings - if self_field.selection_identifier() == other_identifier - && self_field.arguments_hash() != other_args_hash - { - return Some(self_index); - } - } - - None - }, - ); - - if let Some(existing_index) = existing_in_self { - return Some((existing_index, index)); - } - - return None; - } - - None - }) - .collect() -} - -pub fn field_condition_equal(cond: &Option, field: &FieldSelection) -> bool { - match cond { - Some(cond) => match cond { - Condition::Include(var_name) => { - field.include_if.as_ref().is_some_and(|v| v == var_name) - } - Condition::Skip(var_name) => field.skip_if.as_ref().is_some_and(|v| v == var_name), - }, - None => field.include_if.is_none() && field.skip_if.is_none(), - } -} - -fn fragment_condition_equal(cond: &Option, fragment: &InlineFragmentSelection) -> bool { - match cond { - Some(cond) => match cond { - Condition::Include(var_name) => { - fragment.include_if.as_ref().is_some_and(|v| v == var_name) - } - Condition::Skip(var_name) => fragment.skip_if.as_ref().is_some_and(|v| v == var_name), - }, - None => fragment.include_if.is_none() && fragment.skip_if.is_none(), - } -} diff --git a/lib/query-planner/src/graph/edge.rs b/lib/query-planner/src/graph/edge.rs index 78cd83be..a0deaa04 100644 --- a/lib/query-planner/src/graph/edge.rs +++ b/lib/query-planner/src/graph/edge.rs @@ -250,7 +250,7 @@ impl Edge { }; let requirement_cost = match self.requirements() { - Some(selection) => selection.cost(), + Some(selection) => selection.selection_set.cost(), None => 0, }; diff --git a/lib/query-planner/src/planner/fetch/fetch_graph.rs b/lib/query-planner/src/planner/fetch/fetch_graph.rs index 0654562f..55c2b650 100644 --- a/lib/query-planner/src/planner/fetch/fetch_graph.rs +++ b/lib/query-planner/src/planner/fetch/fetch_graph.rs @@ -274,10 +274,7 @@ fn create_noop_fetch_step( fetch_graph.add_step(FetchStepData { service_name: SubgraphName::any(), response_path: MergePath::default(), - input: TypeAwareSelection { - selection_set: SelectionSet::default(), - type_name: "*".to_string(), - }, + input: FetchStepSelections::new_empty(), output: FetchStepSelections::new_empty(), flags, condition: None, @@ -305,15 +302,17 @@ fn create_fetch_step_for_entity_call( } else { FetchStepFlags::empty() }; + let mut input = FetchStepSelections::new(input_type_name); + input + .add(&SelectionSet { + items: vec![SelectionItem::Field(FieldSelection::new_typename())], + }) + .unwrap(); + fetch_graph.add_step(FetchStepData { service_name: subgraph_name.clone(), response_path: response_path.clone(), - input: TypeAwareSelection { - selection_set: SelectionSet { - items: vec![SelectionItem::Field(FieldSelection::new_typename())], - }, - type_name: input_type_name.to_string(), - }, + input, output: FetchStepSelections::new(output_type_name), flags, condition: condition.cloned(), @@ -337,10 +336,7 @@ fn create_fetch_step_for_root_move( let idx = fetch_graph.add_step(FetchStepData { service_name: subgraph_name.clone(), response_path: MergePath::default(), - input: TypeAwareSelection { - selection_set: SelectionSet::default(), - type_name: type_name.to_string(), - }, + input: FetchStepSelections::new(type_name), output: FetchStepSelections::new(type_name), flags: FetchStepFlags::empty(), condition: None, @@ -383,7 +379,7 @@ fn ensure_fetch_step_for_subgraph( return None; } - if fetch_step.input.type_name != *input_type_name { + if fetch_step.input.definition_name() != input_type_name { return None; } @@ -396,7 +392,12 @@ fn ensure_fetch_step_for_subgraph( } if let Some(key) = &key { - if !fetch_step.input.contains(key) { + if fetch_step.input.definition_name() != key.type_name + || !fetch_step + .input + .selection_set() + .contains(&key.selection_set) + { // requested key fields are not part of the input return None; } @@ -440,7 +441,7 @@ fn ensure_fetch_step_for_subgraph( ); if let Some(selection) = key { let step = fetch_graph.get_step_data_mut(step_index)?; - step.input.add(selection) + step.input.add(&selection.selection_set)? } trace!( @@ -480,7 +481,7 @@ fn ensure_fetch_step_for_requirement( return None; } - if fetch_step.input.type_name != *type_name { + if fetch_step.input.definition_name() != *type_name { return None; } @@ -488,7 +489,10 @@ fn ensure_fetch_step_for_requirement( return None; } - if !fetch_step.input.contains(requirement) { + if !fetch_step + .input + .contains(&requirement.type_name, &requirement.selection_set) + { return None; } @@ -717,7 +721,7 @@ fn process_entity_move_edge( requirement, fetch_step_index.index() ); - fetch_step.input.add(&requirement); + fetch_step.input.add(&requirement.selection_set)?; if is_interface { // We use `output_type_name` as there's no connection from `Interface` to `Object`, @@ -854,19 +858,20 @@ fn process_interface_object_type_move_edge( requirement, step_for_children_index.index() ); - step_for_children.input.add(&requirement); + step_for_children.input.add(&requirement.selection_set)?; let key_to_reenter_subgraph = find_satisfiable_key( graph, override_context, query_node.requirements.first().unwrap(), )?; - step_for_children.input.add(&requirement); trace!( "adding key '{}' to fetch step [{}]", key_to_reenter_subgraph, step_for_children_index.index() ); - step_for_children.input.add(key_to_reenter_subgraph); + step_for_children + .input + .add(&key_to_reenter_subgraph.selection_set)?; trace!( "adding input rewrite '... on {} {{ __typename }}' to '{}' of [{}]", @@ -904,7 +909,9 @@ fn process_interface_object_type_move_edge( key_to_reenter_subgraph, step_for_requirements_index.index() ); - step_for_requirements.input.add(key_to_reenter_subgraph); + step_for_requirements + .input + .add(&key_to_reenter_subgraph.selection_set)?; // // Given `f0 { ... on User { f1 } }` where f1 is a field contributed by @interfaceObject, @@ -1277,13 +1284,15 @@ fn process_requires_field_edge( requires, step_for_children_index.index() ); - step_for_children.input.add(requires); + step_for_children.input.add(&requires.selection_set)?; trace!( "Adding {} to fetch([{}]).input (key re-enter)", key_to_reenter_subgraph, step_for_children_index.index() ); - step_for_children.input.add(key_to_reenter_subgraph); + step_for_children + .input + .add(&key_to_reenter_subgraph.selection_set)?; trace!("Creating a fetch step for requirement of @requires"); let step_for_requirements_index = create_fetch_step_for_entity_call( @@ -1301,7 +1310,9 @@ fn process_requires_field_edge( key_to_reenter_subgraph, step_for_requirements_index.index() ); - step_for_requirements.input.add(key_to_reenter_subgraph); + step_for_requirements + .input + .add(&key_to_reenter_subgraph.selection_set)?; let real_parent_fetch_step = fetch_graph.get_step_data_mut(real_parent_fetch_step_index)?; diff --git a/lib/query-planner/src/planner/fetch/fetch_step_data.rs b/lib/query-planner/src/planner/fetch/fetch_step_data.rs index 1cd381d7..e7bd3530 100644 --- a/lib/query-planner/src/planner/fetch/fetch_step_data.rs +++ b/lib/query-planner/src/planner/fetch/fetch_step_data.rs @@ -9,7 +9,7 @@ use crate::{ merge_path::{Condition, MergePath}, operation::VariableDefinition, safe_merge::AliasesRecords, - type_aware_selection::{find_arguments_conflicts, TypeAwareSelection}, + selection_set::find_arguments_conflicts, }, planner::{ fetch::{ @@ -37,7 +37,7 @@ bitflags! { pub struct FetchStepData { pub service_name: SubgraphName, pub response_path: MergePath, - pub input: TypeAwareSelection, + pub input: FetchStepSelections, pub output: FetchStepSelections, pub kind: FetchStepKind, pub flags: FetchStepFlags, @@ -58,17 +58,19 @@ pub enum FetchStepKind { impl Display for FetchStepData { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "{}/{} {} → ", - self.input.type_name, self.service_name, self.input, - )?; + write!(f, "[{}]: ", self.service_name)?; + + for (def_name, selections) in self.input.iter() { + write!(f, "{}/{} ", def_name, selections)?; + } + + write!(f, "-> ")?; for (def_name, selections) in self.output.iter() { write!(f, "{}/{} ", def_name, selections)?; } - write!(f, " at $.{}", self.response_path.join("."))?; + write!(f, "at $.{}", self.response_path.join("."))?; if self.flags.contains(FetchStepFlags::USED_FOR_REQUIRES) { write!(f, " [@requires]")?; @@ -99,9 +101,7 @@ impl FetchStepData { } pub fn is_entity_call(&self) -> bool { - self.input.type_name != "Query" - && self.input.type_name != "Mutation" - && self.input.type_name != "Subscription" + self.kind == FetchStepKind::Entity } pub fn add_input_rewrite(&mut self, rewrite: FetchRewrite) { @@ -155,13 +155,18 @@ impl FetchStepData { } } - let input_conflicts = find_arguments_conflicts(&self.input, &other.input); + let has_input_conflicts = + FetchStepSelections::iter_matching_types(&self.input, &other.input, |_, s1, s2| { + find_arguments_conflicts(s1, s2) + }) + .iter() + .any(|(_, conflicts)| !conflicts.is_empty()); - if !input_conflicts.is_empty() { + if has_input_conflicts { trace!( "preventing merge of [{}]+[{}] due to input conflicts", self_index.index(), - other_index.index() + other_index.index(), ); return false; @@ -194,7 +199,7 @@ impl FetchStepData { FetchStepData:: { service_name: self.service_name, response_path: self.response_path, - input: self.input, + input: self.input.into_multi_type(), output: self.output.into_multi_type(), kind: self.kind, flags: self.flags, diff --git a/lib/query-planner/src/planner/fetch/optimize/apply_internal_aliases_patching.rs b/lib/query-planner/src/planner/fetch/optimize/apply_internal_aliases_patching.rs index 56813132..e7a427cc 100644 --- a/lib/query-planner/src/planner/fetch/optimize/apply_internal_aliases_patching.rs +++ b/lib/query-planner/src/planner/fetch/optimize/apply_internal_aliases_patching.rs @@ -5,7 +5,7 @@ use crate::{ ast::{ merge_path::{MergePath, Segment}, selection_item::SelectionItem, - type_aware_selection::find_selection_set_by_path_mut, + selection_set::find_selection_set_by_path_mut, }, planner::fetch::{error::FetchGraphError, fetch_graph::FetchGraph, state::MultiTypeFetchStep}, }; @@ -68,18 +68,29 @@ impl FetchGraph { if let Some(Segment::Field(field_name, args_hash, condition)) = maybe_patched_field { + // TODO: Avoid "except" here of course. + let decendent_type_name = decendent + .input + .try_as_single() + .expect("to be a single input") + .to_string(); + + let selection = decendent + .input + .selections_for_definition_mut(&decendent_type_name) + .expect("selection set is missing"); + trace!( "field '{}' was aliased, relative selection path: '{}', checking if need to patch selection '{}'", field_name, relative_path, - decendent.input.selection_set + selection ); // First, check if the node's input selection set contains the field that was aliased - if let Some(selection) = find_selection_set_by_path_mut( - &mut decendent.input.selection_set, - &relative_path, - ) { + if let Some(selection) = + find_selection_set_by_path_mut(selection, &relative_path) + { trace!("found selection to patch: {}", selection); let item_to_patch = selection.items.iter_mut().find(|item| matches!(item, SelectionItem::Field(field) if field.name == *field_name && field.arguments_hash() == *args_hash)); @@ -99,7 +110,7 @@ impl FetchGraph { trace!( "path '{}' was not found in selection '{}', skipping...", relative_path, - decendent.input.selection_set + selection ); } diff --git a/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs b/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs index da35effa..7e790c64 100644 --- a/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs +++ b/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs @@ -2,10 +2,11 @@ use petgraph::graph::NodeIndex; use tracing::{instrument, trace}; use crate::{ - ast::type_aware_selection::find_arguments_conflicts, + ast::selection_set::find_arguments_conflicts, planner::fetch::{ error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData, - optimize::utils::perform_fetch_step_merge, state::MultiTypeFetchStep, + optimize::utils::perform_fetch_step_merge, selections::FetchStepSelections, + state::MultiTypeFetchStep, }, }; @@ -29,7 +30,7 @@ impl FetchStepData { return false; } - if self.input.type_name != other.input.type_name { + if !self.input.selecting_same_types(&other.input) { return false; } @@ -56,9 +57,19 @@ impl FetchStepData { return false; } - let input_conflicts = find_arguments_conflicts(&self.input, &other.input); + let input_conflicts = FetchStepSelections::::iter_matching_types( + &self.input, + &other.input, + |_, self_selections, other_selections| { + find_arguments_conflicts(self_selections, other_selections) + }, + ); - if !input_conflicts.is_empty() { + let has_conflicts = input_conflicts + .iter() + .any(|(_, conflicts)| !conflicts.is_empty()); + + if has_conflicts { return false; } diff --git a/lib/query-planner/src/planner/fetch/optimize/merge_passthrough_child.rs b/lib/query-planner/src/planner/fetch/optimize/merge_passthrough_child.rs index 335326c9..252c75c0 100644 --- a/lib/query-planner/src/planner/fetch/optimize/merge_passthrough_child.rs +++ b/lib/query-planner/src/planner/fetch/optimize/merge_passthrough_child.rs @@ -8,7 +8,7 @@ use petgraph::{ use tracing::{instrument, trace}; use crate::{ - ast::type_aware_selection::selection_items_are_subset_of, + ast::selection_set::selection_items_are_subset_of, planner::fetch::{ error::FetchGraphError, fetch_graph::FetchGraph, @@ -116,15 +116,12 @@ impl FetchStepData { return false; } - // TODO: Use selection_items_are_subset_of for (output_def_name, output_selections) in other.output.iter_selections() { - if &other.input.type_name == output_def_name - && selection_items_are_subset_of( - &other.input.selection_set.items, - &output_selections.items, - ) - { - return true; + if let Some(input_selections) = other.input.selections_for_definition(output_def_name) { + if selection_items_are_subset_of(&input_selections.items, &output_selections.items) + { + return true; + } } } diff --git a/lib/query-planner/src/planner/fetch/optimize/type_mismatches.rs b/lib/query-planner/src/planner/fetch/optimize/type_mismatches.rs index 2f31c660..1e111e7e 100644 --- a/lib/query-planner/src/planner/fetch/optimize/type_mismatches.rs +++ b/lib/query-planner/src/planner/fetch/optimize/type_mismatches.rs @@ -7,7 +7,7 @@ use crate::{ mismatch_finder::SelectionMismatchFinder, safe_merge::SafeSelectionSetMerger, selection_item::SelectionItem, - type_aware_selection::{field_condition_equal, find_selection_set_by_path_mut}, + selection_set::{field_condition_equal, find_selection_set_by_path_mut}, }, planner::{ fetch::{error::FetchGraphError, fetch_graph::FetchGraph, state::MultiTypeFetchStep}, diff --git a/lib/query-planner/src/planner/fetch/optimize/utils.rs b/lib/query-planner/src/planner/fetch/optimize/utils.rs index 7d5ad38b..8bdacb45 100644 --- a/lib/query-planner/src/planner/fetch/optimize/utils.rs +++ b/lib/query-planner/src/planner/fetch/optimize/utils.rs @@ -6,9 +6,12 @@ use petgraph::{ }; use tracing::{instrument, trace}; -use crate::planner::fetch::{ - error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepFlags, - state::MultiTypeFetchStep, +use crate::{ + ast::merge_path::MergePath, + planner::fetch::{ + error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepFlags, + state::MultiTypeFetchStep, + }, }; // Return true in case an alias was applied during the merge process. @@ -41,7 +44,11 @@ pub(crate) fn perform_fetch_step_merge( // We can't do it when both are entity calls fetch steps. other .output - .add(&other.input.type_name, &other.input.selection_set)?; + .migrate_from_another(&other.input, &MergePath::default())?; + // TODO: Validate + // other + // .output + // .add(&other.input.type_name, &other.input.selection_set)?; other.output.wrap_with_condition(condition); } else if me.is_entity_call() && other.condition.is_some() { @@ -80,12 +87,13 @@ pub(crate) fn perform_fetch_step_merge( // It's safe to not check if a condition was turned into an inline fragment, // because if a condition is present and "me" is a non-entity fetch step, // then the type_name values of the inputs are different. - if me.input.type_name == other.input.type_name { + if me.input.selecting_same_types(&other.input) { if me.response_path != other.response_path { return Err(FetchGraphError::MismatchedResponsePath); } - me.input.add(&other.input); + me.input + .migrate_from_another(&other.input, &MergePath::default())?; } let mut children_indexes: Vec = vec![]; diff --git a/lib/query-planner/src/planner/fetch/selections.rs b/lib/query-planner/src/planner/fetch/selections.rs index 149e0efc..a450e7be 100644 --- a/lib/query-planner/src/planner/fetch/selections.rs +++ b/lib/query-planner/src/planner/fetch/selections.rs @@ -9,8 +9,10 @@ use crate::{ merge_path::{Condition, MergePath}, safe_merge::{AliasesRecords, SafeSelectionSetMerger}, selection_item::SelectionItem, - selection_set::{FieldSelection, InlineFragmentSelection, SelectionSet}, - type_aware_selection::{find_selection_set_by_path_mut, merge_selection_set}, + selection_set::{ + find_selection_set_by_path_mut, merge_selection_set, selection_items_are_subset_of, + FieldSelection, InlineFragmentSelection, SelectionSet, + }, }, planner::fetch::state::{MultiTypeFetchStep, SingleTypeFetchStep}, }; @@ -95,6 +97,14 @@ impl FetchStepSelections { .1 } + pub fn selection_set_mut(&mut self) -> &mut SelectionSet { + self.selections + .iter_mut() + .next() + .expect("SingleTypeFetchStep should have exactly one selection") + .1 + } + pub fn add_at_path( &mut self, fetch_path: &MergePath, @@ -123,6 +133,14 @@ impl FetchStepSelections { } impl FetchStepSelections { + pub fn contains(&self, definition_name: &str, selection_set: &SelectionSet) -> bool { + if let Some(self_selections) = self.selections.get(definition_name) { + return selection_items_are_subset_of(&self_selections.items, &selection_set.items); + } + + false + } + pub fn is_selecting_definition(&self, definition_name: &str) -> bool { self.selections.contains_key(definition_name) } @@ -156,6 +174,10 @@ impl FetchStepSelections { self.selections.get_mut(definition_name) } + pub fn selections_for_definition(&self, definition_name: &str) -> Option<&SelectionSet> { + self.selections.get(definition_name) + } + fn add_at_path_inner( &mut self, definition_name: &str, @@ -184,6 +206,37 @@ impl FetchStepSelections { } impl FetchStepSelections { + pub fn selecting_same_types(&self, other: &Self) -> bool { + if self.selections.len() != other.selections.len() { + return false; + } + + for key in self.selections.keys() { + if !other.selections.contains_key(key) { + return false; + } + } + + true + } + + pub fn iter_matching_types<'a, 'b, R>( + input: &'a FetchStepSelections, + other: &'b FetchStepSelections, + mut callback: impl FnMut(&str, &SelectionSet, &SelectionSet) -> R, + ) -> Vec<(&'a str, R)> { + let mut result: Vec<(&'a str, R)> = Vec::new(); + + for (definition_name, input_selections) in input.iter_selections() { + if let Some(other_selections) = other.selections.get(definition_name) { + let r = callback(definition_name, input_selections, other_selections); + result.push((definition_name, r)); + } + } + + result + } + pub fn try_as_single(&self) -> Option<&str> { if self.selections.len() == 1 { self.selections.keys().next().map(|key| key.as_str()) @@ -211,22 +264,6 @@ impl FetchStepSelections { Ok(()) } - pub fn add( - &mut self, - definition_name: &str, - selection_set: &SelectionSet, - ) -> Result<(), FetchStepSelectionsError> { - let selections = self - .selections_for_definition_mut(definition_name) - .ok_or_else(|| { - FetchStepSelectionsError::UnexpectedMissingDefinition(definition_name.to_string()) - })?; - - merge_selection_set(selections, selection_set, false); - - Ok(()) - } - pub fn safe_migrate_from_another( &mut self, other: &Self, @@ -296,6 +333,12 @@ impl FetchStepSelections { } impl FetchStepSelections { + pub fn add(&mut self, selection_set: &SelectionSet) -> Result<(), FetchStepSelectionsError> { + merge_selection_set(self.selection_set_mut(), selection_set, false); + + Ok(()) + } + pub fn new(definition_name: &str) -> Self { let mut map = HashMap::new(); map.insert(definition_name.to_string(), SelectionSet::default()); diff --git a/lib/query-planner/src/planner/plan_nodes.rs b/lib/query-planner/src/planner/plan_nodes.rs index 8ed9d148..0a334fa2 100644 --- a/lib/query-planner/src/planner/plan_nodes.rs +++ b/lib/query-planner/src/planner/plan_nodes.rs @@ -4,11 +4,12 @@ use crate::{ minification::minify_operation, operation::{OperationDefinition, SubgraphFetchOperation, VariableDefinition}, selection_item::SelectionItem, - selection_set::{FieldSelection, InlineFragmentSelection, SelectionSet}, - type_aware_selection::TypeAwareSelection, + selection_set::{FieldSelection, SelectionSet}, value::Value, }, - planner::fetch::{fetch_step_data::FetchStepData, state::MultiTypeFetchStep}, + planner::fetch::{ + fetch_step_data::FetchStepData, selections::FetchStepSelections, state::MultiTypeFetchStep, + }, state::supergraph_state::{OperationKind, SupergraphState, TypeNode}, utils::pretty_display::{get_indent, PrettyDisplay}, }; @@ -256,15 +257,12 @@ impl PlanNode { } } -fn create_input_selection_set(input_selections: &TypeAwareSelection) -> SelectionSet { - SelectionSet { - items: vec![SelectionItem::InlineFragment(InlineFragmentSelection { - selections: input_selections.selection_set.strip_for_plan_input(), - type_condition: input_selections.type_name.clone(), - skip_if: None, - include_if: None, - })], - } +fn create_input_selection_set( + input_selections: &FetchStepSelections, +) -> SelectionSet { + let selection_set: SelectionSet = input_selections.into(); + + selection_set.strip_for_plan_input() } fn create_output_operation( @@ -316,10 +314,10 @@ fn create_output_operation( impl From<&FetchStepData> for OperationKind { fn from(step: &FetchStepData) -> Self { - match &step.input.type_name { - str if str == "Query" => OperationKind::Query, - str if str == "Mutation" => OperationKind::Mutation, - str if str == "Subscription" => OperationKind::Subscription, + match step.input.iter().next().unwrap().0.as_str() { + "Query" => OperationKind::Query, + "Mutation" => OperationKind::Mutation, + "Subscription" => OperationKind::Subscription, _ => OperationKind::Query, } } From 4fc5ad6013c7f899183cd3d1797e2507ff373f9a Mon Sep 17 00:00:00 2001 From: Dotan Simha Date: Thu, 17 Jul 2025 20:56:14 +0300 Subject: [PATCH 3/4] feat(qp): introduce new optimization to batch siblings using multi-type steps --- lib/query-planner/src/ast/merge_path.rs | 11 + .../src/planner/fetch/fetch_step_data.rs | 78 +--- .../fetch/optimize/batch_multi_type.rs | 129 ++++++ .../optimize/merge_children_with_parents.rs | 2 +- .../src/planner/fetch/optimize/merge_leafs.rs | 26 +- .../planner/fetch/optimize/merge_siblings.rs | 7 +- .../src/planner/fetch/optimize/mod.rs | 2 + .../src/planner/fetch/optimize/utils.rs | 97 ++++- .../src/planner/fetch/selections.rs | 17 +- lib/query-planner/src/planner/plan_nodes.rs | 6 +- lib/query-planner/src/tests/interface.rs | 395 ++++++------------ lib/query-planner/src/tests/issues.rs | 44 +- lib/query-planner/src/tests/mod.rs | 261 +++++++++++- 13 files changed, 655 insertions(+), 420 deletions(-) create mode 100644 lib/query-planner/src/planner/fetch/optimize/batch_multi_type.rs diff --git a/lib/query-planner/src/ast/merge_path.rs b/lib/query-planner/src/ast/merge_path.rs index bc439201..ea522572 100644 --- a/lib/query-planner/src/ast/merge_path.rs +++ b/lib/query-planner/src/ast/merge_path.rs @@ -179,6 +179,17 @@ impl MergePath { } self.common_prefix_len(other) == other.len() } + + pub fn without_type_castings(&self) -> Self { + let new_segments = self + .inner + .iter() + .filter(|segment| !matches!(segment, Segment::Cast(_, _))) + .cloned() + .collect::>(); + + Self::new(new_segments) + } } impl Display for MergePath { diff --git a/lib/query-planner/src/planner/fetch/fetch_step_data.rs b/lib/query-planner/src/planner/fetch/fetch_step_data.rs index e7bd3530..1ffd5f2a 100644 --- a/lib/query-planner/src/planner/fetch/fetch_step_data.rs +++ b/lib/query-planner/src/planner/fetch/fetch_step_data.rs @@ -1,19 +1,16 @@ use std::{collections::BTreeSet, fmt::Display}; use bitflags::bitflags; -use petgraph::{graph::NodeIndex, visit::EdgeRef}; -use tracing::trace; +use petgraph::graph::NodeIndex; use crate::{ ast::{ merge_path::{Condition, MergePath}, operation::VariableDefinition, safe_merge::AliasesRecords, - selection_set::find_arguments_conflicts, }, planner::{ fetch::{ - fetch_graph::FetchGraph, selections::FetchStepSelections, state::{MultiTypeFetchStep, SingleTypeFetchStep}, }, @@ -121,76 +118,9 @@ impl FetchStepData { } } -impl FetchStepData { - pub fn can_merge( - &self, - self_index: NodeIndex, - other_index: NodeIndex, - other: &Self, - fetch_graph: &FetchGraph, - ) -> bool { - if self_index == other_index { - return false; - } - - if self.service_name != other.service_name { - return false; - } - - // We allow to merge root with entity calls by adding an inline fragment with the @include/@skip - if self.is_entity_call() && other.is_entity_call() && self.condition != other.condition { - return false; - } - - // If both are entities, their response_paths should match, - // as we can't merge entity calls resolving different entities - if matches!(self.kind, FetchStepKind::Entity) && self.kind == other.kind { - if !self.response_path.eq(&other.response_path) { - return false; - } - } else { - // otherwise we can merge - if !other.response_path.starts_with(&self.response_path) { - return false; - } - } - - let has_input_conflicts = - FetchStepSelections::iter_matching_types(&self.input, &other.input, |_, s1, s2| { - find_arguments_conflicts(s1, s2) - }) - .iter() - .any(|(_, conflicts)| !conflicts.is_empty()); - - if has_input_conflicts { - trace!( - "preventing merge of [{}]+[{}] due to input conflicts", - self_index.index(), - other_index.index(), - ); - - return false; - } - - // if the `other` FetchStep has a single parent and it's `this` FetchStep - if fetch_graph.parents_of(other_index).count() == 1 - && fetch_graph - .parents_of(other_index) - .all(|edge| edge.source() == self_index) - { - return true; - } - - // if they do not share parents, they can't be merged - if !fetch_graph.parents_of(self_index).all(|self_edge| { - fetch_graph - .parents_of(other_index) - .any(|other_edge| other_edge.source() == self_edge.source()) - }) { - return false; - } - - true +impl FetchStepData { + pub fn is_fetching_multiple_types(&self) -> bool { + self.input.is_fetching_multiple_types() || self.output.is_fetching_multiple_types() } } diff --git a/lib/query-planner/src/planner/fetch/optimize/batch_multi_type.rs b/lib/query-planner/src/planner/fetch/optimize/batch_multi_type.rs new file mode 100644 index 00000000..51ac7980 --- /dev/null +++ b/lib/query-planner/src/planner/fetch/optimize/batch_multi_type.rs @@ -0,0 +1,129 @@ +use std::collections::{HashMap, VecDeque}; + +use petgraph::{graph::NodeIndex, Direction}; +use tracing::{instrument, trace}; + +use crate::planner::fetch::{ + error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData, + optimize::utils::perform_fetch_step_merge, state::MultiTypeFetchStep, +}; + +impl FetchGraph { + #[instrument(level = "trace", skip_all)] + pub(crate) fn batch_multi_type(&mut self) -> Result<(), FetchGraphError> { + let root_index = self + .root_index + .ok_or(FetchGraphError::NonSingleRootStep(0))?; + // Breadth-First Search (BFS) starting from the root node. + let mut queue = VecDeque::from([root_index]); + + while let Some(parent_index) = queue.pop_front() { + let mut merges_to_perform = Vec::<(NodeIndex, NodeIndex)>::new(); + let mut node_indexes: HashMap = HashMap::new(); + let siblings_indices = self + .graph + .neighbors_directed(parent_index, Direction::Outgoing) + .collect::>(); + + for (i, sibling_index) in siblings_indices.iter().enumerate() { + queue.push_back(*sibling_index); + let current = self.get_step_data(*sibling_index)?; + + for other_sibling_index in siblings_indices.iter().skip(i + 1) { + trace!( + "checking if [{}] and [{}] can be batched", + sibling_index.index(), + other_sibling_index.index() + ); + + let other_sibling = self.get_step_data(*other_sibling_index)?; + + if current.can_be_batched_with(other_sibling) { + trace!( + "Found multi-type batching optimization: [{}] <- [{}]", + sibling_index.index(), + other_sibling_index.index() + ); + // Register their original indexes in the map. + node_indexes.insert(*sibling_index, *sibling_index); + node_indexes.insert(*other_sibling_index, *other_sibling_index); + + merges_to_perform.push((*sibling_index, *other_sibling_index)); + } + } + } + + for (child_index, other_child_index) in merges_to_perform { + // Get the latest indexes for the nodes, accounting for previous merges. + let child_index_latest = node_indexes + .get(&child_index) + .ok_or(FetchGraphError::IndexMappingLost)?; + let other_child_index_latest = node_indexes + .get(&other_child_index) + .ok_or(FetchGraphError::IndexMappingLost)?; + + if child_index_latest == other_child_index_latest { + continue; + } + + if self.is_ancestor_or_descendant(*child_index_latest, *other_child_index_latest) { + continue; + } + + let (me, other) = + self.get_pair_of_steps_mut(*child_index_latest, *other_child_index_latest)?; + + // We "declare" the known type for the step, so later merging will be possible into that type instead of failing with an error. + for (input_type_name, _) in other.input.iter_selections() { + me.input.declare_known_type(input_type_name); + } + + // We "declare" the known type for the step, so later merging will be possible into that type instead of failing with an error. + for (output_type_name, _) in other.output.iter_selections() { + me.output.declare_known_type(output_type_name); + } + + perform_fetch_step_merge( + *child_index_latest, + *other_child_index_latest, + self, + true, + )?; + + // Because `other_child` was merged into `child`, + // then everything that was pointing to `other_child` + // has to point to the `child`. + node_indexes.insert(*other_child_index_latest, *child_index_latest); + } + } + + Ok(()) + } +} + +impl FetchStepData { + pub fn can_be_batched_with(&self, other: &Self) -> bool { + if self.kind != other.kind { + return false; + } + + if self.service_name != other.service_name { + return false; + } + + if !self.is_entity_call() || !other.is_entity_call() { + return false; + } + + if self.response_path.without_type_castings() != other.response_path.without_type_castings() + { + return false; + } + + if self.has_arguments_conflicts_with(other) { + return false; + } + + true + } +} diff --git a/lib/query-planner/src/planner/fetch/optimize/merge_children_with_parents.rs b/lib/query-planner/src/planner/fetch/optimize/merge_children_with_parents.rs index 29172614..daf57385 100644 --- a/lib/query-planner/src/planner/fetch/optimize/merge_children_with_parents.rs +++ b/lib/query-planner/src/planner/fetch/optimize/merge_children_with_parents.rs @@ -63,7 +63,7 @@ impl FetchGraph { .get(&child_index) .ok_or(FetchGraphError::IndexMappingLost)?; - perform_fetch_step_merge(*parent_index_latest, *child_index_latest, self)?; + perform_fetch_step_merge(*parent_index_latest, *child_index_latest, self, false)?; // Because `child` was merged into `parent`, // then everything that was pointing to `child` diff --git a/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs b/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs index 7e790c64..82431c28 100644 --- a/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs +++ b/lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs @@ -1,13 +1,9 @@ use petgraph::graph::NodeIndex; use tracing::{instrument, trace}; -use crate::{ - ast::selection_set::find_arguments_conflicts, - planner::fetch::{ - error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData, - optimize::utils::perform_fetch_step_merge, selections::FetchStepSelections, - state::MultiTypeFetchStep, - }, +use crate::planner::fetch::{ + error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData, + optimize::utils::perform_fetch_step_merge, state::MultiTypeFetchStep, }; impl FetchStepData { @@ -57,19 +53,7 @@ impl FetchStepData { return false; } - let input_conflicts = FetchStepSelections::::iter_matching_types( - &self.input, - &other.input, - |_, self_selections, other_selections| { - find_arguments_conflicts(self_selections, other_selections) - }, - ); - - let has_conflicts = input_conflicts - .iter() - .any(|(_, conflicts)| !conflicts.is_empty()); - - if has_conflicts { + if self.has_arguments_conflicts_with(other) { return false; } @@ -85,7 +69,7 @@ impl FetchGraph { /// meaning the overall depth (amount of parallel layers) is not increased. pub(crate) fn merge_leafs(&mut self) -> Result<(), FetchGraphError> { while let Some((target_idx, leaf_idx)) = self.find_merge_candidate()? { - perform_fetch_step_merge(target_idx, leaf_idx, self)?; + perform_fetch_step_merge(target_idx, leaf_idx, self, false)?; } Ok(()) diff --git a/lib/query-planner/src/planner/fetch/optimize/merge_siblings.rs b/lib/query-planner/src/planner/fetch/optimize/merge_siblings.rs index d2722233..1ff908e5 100644 --- a/lib/query-planner/src/planner/fetch/optimize/merge_siblings.rs +++ b/lib/query-planner/src/planner/fetch/optimize/merge_siblings.rs @@ -92,7 +92,12 @@ impl FetchGraph { .get(&other_child_index) .ok_or(FetchGraphError::IndexMappingLost)?; - perform_fetch_step_merge(*child_index_latest, *other_child_index_latest, self)?; + perform_fetch_step_merge( + *child_index_latest, + *other_child_index_latest, + self, + false, + )?; // Because `other_child` was merged into `child`, // then everything that was pointing to `other_child` diff --git a/lib/query-planner/src/planner/fetch/optimize/mod.rs b/lib/query-planner/src/planner/fetch/optimize/mod.rs index 95f4c9c4..68c5bd51 100644 --- a/lib/query-planner/src/planner/fetch/optimize/mod.rs +++ b/lib/query-planner/src/planner/fetch/optimize/mod.rs @@ -1,4 +1,5 @@ mod apply_internal_aliases_patching; +mod batch_multi_type; mod deduplicate_and_prune_fetch_steps; mod merge_children_with_parents; mod merge_leafs; @@ -29,6 +30,7 @@ impl FetchGraph { self.merge_siblings()?; self.merge_leafs()?; self.deduplicate_and_prune_fetch_steps()?; + self.batch_multi_type()?; let node_count_after = self.graph.node_count(); let edge_count_after = self.graph.edge_count(); diff --git a/lib/query-planner/src/planner/fetch/optimize/utils.rs b/lib/query-planner/src/planner/fetch/optimize/utils.rs index 8bdacb45..a6a219e0 100644 --- a/lib/query-planner/src/planner/fetch/optimize/utils.rs +++ b/lib/query-planner/src/planner/fetch/optimize/utils.rs @@ -7,9 +7,12 @@ use petgraph::{ use tracing::{instrument, trace}; use crate::{ - ast::merge_path::MergePath, + ast::{merge_path::MergePath, selection_set::find_arguments_conflicts}, planner::fetch::{ - error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepFlags, + error::FetchGraphError, + fetch_graph::FetchGraph, + fetch_step_data::{FetchStepData, FetchStepFlags, FetchStepKind}, + selections::FetchStepSelections, state::MultiTypeFetchStep, }, }; @@ -20,6 +23,7 @@ pub(crate) fn perform_fetch_step_merge( self_index: NodeIndex, other_index: NodeIndex, fetch_graph: &mut FetchGraph, + force_merge_inputs: bool, ) -> Result<(), FetchGraphError> { let (me, other) = fetch_graph.get_pair_of_steps_mut(self_index, other_index)?; @@ -45,10 +49,6 @@ pub(crate) fn perform_fetch_step_merge( other .output .migrate_from_another(&other.input, &MergePath::default())?; - // TODO: Validate - // other - // .output - // .add(&other.input.type_name, &other.input.selection_set)?; other.output.wrap_with_condition(condition); } else if me.is_entity_call() && other.condition.is_some() { @@ -84,10 +84,13 @@ pub(crate) fn perform_fetch_step_merge( } } - // It's safe to not check if a condition was turned into an inline fragment, - // because if a condition is present and "me" is a non-entity fetch step, - // then the type_name values of the inputs are different. - if me.input.selecting_same_types(&other.input) { + if force_merge_inputs { + me.input + .migrate_from_another(&other.input, &MergePath::default())?; + } else if me.input.selecting_same_types(&other.input) { + // It's safe to not check if a condition was turned into an inline fragment, + // because if a condition is present and "me" is a non-entity fetch step, + // then the type_name values of the inputs are different. if me.response_path != other.response_path { return Err(FetchGraphError::MismatchedResponsePath); } @@ -174,3 +177,77 @@ pub fn is_reachable_via_alternative_upstream_path( // no indirect path exists Ok(false) } + +impl FetchStepData { + pub fn can_merge( + &self, + self_index: NodeIndex, + other_index: NodeIndex, + other: &Self, + fetch_graph: &FetchGraph, + ) -> bool { + if self_index == other_index { + return false; + } + + if self.service_name != other.service_name { + return false; + } + + // We allow to merge root with entity calls by adding an inline fragment with the @include/@skip + if self.is_entity_call() && other.is_entity_call() && self.condition != other.condition { + return false; + } + + // If both are entities, their response_paths should match, + // as we can't merge entity calls resolving different entities + if matches!(self.kind, FetchStepKind::Entity) && self.kind == other.kind { + if !self.response_path.eq(&other.response_path) { + return false; + } + } else { + // otherwise we can merge + if !other.response_path.starts_with(&self.response_path) { + return false; + } + } + + if self.has_arguments_conflicts_with(other) { + return false; + } + + // if the `other` FetchStep has a single parent and it's `this` FetchStep + if fetch_graph.parents_of(other_index).count() == 1 + && fetch_graph + .parents_of(other_index) + .all(|edge| edge.source() == self_index) + { + return true; + } + + // if they do not share parents, they can't be merged + if !fetch_graph.parents_of(self_index).all(|self_edge| { + fetch_graph + .parents_of(other_index) + .any(|other_edge| other_edge.source() == self_edge.source()) + }) { + return false; + } + + true + } + + pub fn has_arguments_conflicts_with(&self, other: &Self) -> bool { + let input_conflicts = FetchStepSelections::::iter_matching_types( + &self.input, + &other.input, + |_, self_selections, other_selections| { + find_arguments_conflicts(self_selections, other_selections) + }, + ); + + input_conflicts + .iter() + .any(|(_, conflicts)| !conflicts.is_empty()) + } +} diff --git a/lib/query-planner/src/planner/fetch/selections.rs b/lib/query-planner/src/planner/fetch/selections.rs index a450e7be..2321aebd 100644 --- a/lib/query-planner/src/planner/fetch/selections.rs +++ b/lib/query-planner/src/planner/fetch/selections.rs @@ -1,5 +1,5 @@ use std::{ - collections::{BTreeSet, HashMap}, + collections::{BTreeMap, BTreeSet}, fmt::Display, marker::PhantomData, }; @@ -27,7 +27,7 @@ pub enum FetchStepSelectionsError { #[derive(Debug, Clone)] pub struct FetchStepSelections { - selections: HashMap, + selections: BTreeMap, _state: PhantomData, } @@ -133,6 +133,10 @@ impl FetchStepSelections { } impl FetchStepSelections { + pub fn is_fetching_multiple_types(&self) -> bool { + self.selections.len() > 1 + } + pub fn contains(&self, definition_name: &str, selection_set: &SelectionSet) -> bool { if let Some(self_selections) = self.selections.get(definition_name) { return selection_items_are_subset_of(&self_selections.items, &selection_set.items); @@ -249,6 +253,13 @@ impl FetchStepSelections { self.selections.iter() } + /// Creates a slot in the internal HashMap and will allow to add selections for the given definition name. + /// Without that, trying to add selections using any function will either fail or result in trying to force-add to a root type. + /// Calling this method is crucial if you wish to create multi-type steps. + pub fn declare_known_type(&mut self, def_name: &str) { + self.selections.entry(def_name.to_string()).or_default(); + } + pub fn migrate_from_another( &mut self, other: &Self, @@ -340,7 +351,7 @@ impl FetchStepSelections { } pub fn new(definition_name: &str) -> Self { - let mut map = HashMap::new(); + let mut map = BTreeMap::new(); map.insert(definition_name.to_string(), SelectionSet::default()); Self { diff --git a/lib/query-planner/src/planner/plan_nodes.rs b/lib/query-planner/src/planner/plan_nodes.rs index 0a334fa2..71f67047 100644 --- a/lib/query-planner/src/planner/plan_nodes.rs +++ b/lib/query-planner/src/planner/plan_nodes.rs @@ -378,7 +378,11 @@ impl PlanNode { } else { PlanNode::Flatten(FlattenNode { // it's cheaper to clone response_path (Arc etc), rather then cloning the step - path: step.response_path.clone().into(), + path: if step.is_fetching_multiple_types() { + step.response_path.without_type_castings().into() + } else { + step.response_path.clone().into() + }, node: Box::new(PlanNode::Fetch(FetchNode::from_fetch_step( step, supergraph, ))), diff --git a/lib/query-planner/src/tests/interface.rs b/lib/query-planner/src/tests/interface.rs index 215639d8..78c2a28b 100644 --- a/lib/query-planner/src/tests/interface.rs +++ b/lib/query-planner/src/tests/interface.rs @@ -513,8 +513,7 @@ fn type_expand_interface_field() -> Result<(), Box> { ); let query_plan = build_query_plan("fixture/tests/abstract-types.supergraph.graphql", document)?; - // TODO: this should be batched (turned into a sequence of two fetch steps) - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "products") { @@ -533,45 +532,36 @@ fn type_expand_interface_field() -> Result<(), Box> { } } }, - Parallel { - Flatten(path: "products.@|[Magazine]") { - Fetch(service: "reviews") { - { - ... on Magazine { - __typename - id - } - } => - { - ... on Magazine { - reviews { - id - } - } + Flatten(path: "products.@") { + Fetch(service: "reviews") { + { + ... on Book { + __typename + id } - }, - }, - Flatten(path: "products.@|[Book]") { - Fetch(service: "reviews") { - { - ... on Book { - __typename - id - } - } => - { - ... on Book { - reviews { - id - } - } + ... on Magazine { + __typename + id } - }, + } => + { + ... on Book { + reviews { + ...a } + } + ... on Magazine { + reviews { + ...a } + } + } + fragment a on Review { + id + } }, }, }, }, - "#); + "###); Ok(()) } @@ -604,9 +594,7 @@ fn requires_on_field_with_args_test() -> Result<(), Box> { ); let query_plan = build_query_plan("fixture/tests/abstract-types.supergraph.graphql", document)?; - // TODO: there are duplicated calls in the plan, - // but it's because of the #206 - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "products") { @@ -640,10 +628,10 @@ fn requires_on_field_with_args_test() -> Result<(), Box> { } }, Parallel { - Flatten(path: "magazine.@|[Magazine]") { + Flatten(path: "magazine.@") { Fetch(service: "inventory") { { - ... on Magazine { + ... on Book { __typename dimensions { size @@ -651,21 +639,7 @@ fn requires_on_field_with_args_test() -> Result<(), Box> { } id } - } => - { ... on Magazine { - delivery(zip: "1234") { - fastestDelivery - estimatedDelivery - } - } - } - }, - }, - Flatten(path: "magazine.@|[Book]") { - Fetch(service: "inventory") { - { - ... on Book { __typename dimensions { size @@ -677,17 +651,23 @@ fn requires_on_field_with_args_test() -> Result<(), Box> { { ... on Book { delivery(zip: "1234") { - fastestDelivery - estimatedDelivery - } + ...a } + } + ... on Magazine { + delivery(zip: "1234") { + ...a } } } + fragment a on DeliveryEstimates { + fastestDelivery + estimatedDelivery + } }, }, - Flatten(path: "book.@|[Magazine]") { + Flatten(path: "book.@") { Fetch(service: "inventory") { { - ... on Magazine { + ... on Book { __typename dimensions { size @@ -695,21 +675,7 @@ fn requires_on_field_with_args_test() -> Result<(), Box> { } id } - } => - { ... on Magazine { - delivery(zip: "1234") { - fastestDelivery - estimatedDelivery - } - } - } - }, - }, - Flatten(path: "book.@|[Book]") { - Fetch(service: "inventory") { - { - ... on Book { __typename dimensions { size @@ -721,17 +687,23 @@ fn requires_on_field_with_args_test() -> Result<(), Box> { { ... on Book { delivery(zip: "1234") { - fastestDelivery - estimatedDelivery - } + ...a } + } + ... on Magazine { + delivery(zip: "1234") { + ...a } } } + fragment a on DeliveryEstimates { + fastestDelivery + estimatedDelivery + } }, }, }, }, }, - "#); + "###); Ok(()) } @@ -762,9 +734,7 @@ fn nested_interface_field_with_inline_fragments() -> Result<(), Box> ); let query_plan = build_query_plan("fixture/tests/abstract-types.supergraph.graphql", document)?; - // TODO: there are duplicated calls in the plan, - // but it's because of the #206 - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "products") { @@ -783,82 +753,47 @@ fn nested_interface_field_with_inline_fragments() -> Result<(), Box> } } }, - Parallel { - Flatten(path: "products.@|[Magazine]") { - Fetch(service: "reviews") { - { - ... on Magazine { - __typename - id - } - } => - { - ... on Magazine { - reviews { - product { - id - __typename - ... on Book { - __typename - id - } - ... on Magazine { - __typename - id - } - } - } - } + Flatten(path: "products.@") { + Fetch(service: "reviews") { + { + ... on Book { + __typename + id } - }, - }, - Flatten(path: "products.@|[Book]") { - Fetch(service: "reviews") { - { + ... on Magazine { + __typename + id + } + } => + { + ... on Book { + reviews { + ...a } + } + ... on Magazine { + reviews { + ...a } + } + } + fragment a on Review { + product { + id + __typename ... on Book { __typename id } - } => - { - ... on Book { - reviews { - product { - id - __typename - ... on Book { - __typename - id - } - ... on Magazine { - __typename - id - } - } - } - } - } - }, - }, - }, - Parallel { - Flatten(path: "products.@|[Magazine].reviews.@.product|[Magazine]") { - Fetch(service: "products") { - { ... on Magazine { __typename id } - } => - { - ... on Magazine { - sku - } } - }, + } }, + }, + Parallel { Include(if: $title) { - Flatten(path: "products.@|[Magazine].reviews.@.product|[Book]") { + Flatten(path: "products.@|[Book].reviews.@.product|[Book]") { Fetch(service: "books") { { ... on Book { @@ -889,27 +824,10 @@ fn nested_interface_field_with_inline_fragments() -> Result<(), Box> } }, }, - Include(if: $title) { - Flatten(path: "products.@|[Book].reviews.@.product|[Book]") { - Fetch(service: "books") { - { - ... on Book { - __typename - id - } - } => - { - ... on Book { - title - } - } - }, - }, - }, }, }, }, - "#); + "###); Ok(()) } @@ -943,9 +861,7 @@ fn nested_interface_field_with_redundant_inline_fragments() -> Result<(), Box Result<(), Box - { - ... on Magazine { - reviews { - product { - id - __typename - ... on Book { - __typename - id - } - ... on Magazine { - __typename - id - } - } - } - } + Flatten(path: "products.@") { + Fetch(service: "reviews") { + { + ... on Book { + __typename + id } - }, - }, - Flatten(path: "products.@|[Book]") { - Fetch(service: "reviews") { - { + ... on Magazine { + __typename + id + } + } => + { + ... on Book { + reviews { + ...a } + } + ... on Magazine { + reviews { + ...a } + } + } + fragment a on Review { + product { + id + __typename ... on Book { __typename id } - } => - { - ... on Book { - reviews { - product { - id - __typename - ... on Book { - __typename - id - } - ... on Magazine { - __typename - id - } - } - } - } - } - }, - }, - }, - Parallel { - Flatten(path: "products.@|[Magazine].reviews.@.product|[Magazine]") { - Fetch(service: "products") { - { ... on Magazine { __typename id } - } => - { - ... on Magazine { - sku - } } - }, - }, - Include(if: $title) { - Flatten(path: "products.@|[Magazine].reviews.@.product|[Book]") { - Fetch(service: "products") { - { - ... on Book { - __typename - id - } - } => - { - ... on Book { - sku - } - } - }, - }, + } }, + }, + Parallel { Include(if: $title) { - Flatten(path: "products.@|[Magazine].reviews.@.product|[Book]") { + Flatten(path: "products.@|[Book].reviews.@.product|[Book]") { Fetch(service: "books") { { ... on Book { @@ -1072,50 +936,25 @@ fn nested_interface_field_with_redundant_inline_fragments() -> Result<(), Box - { - ... on Magazine { - sku - } - } - }, - }, Include(if: $title) { - Flatten(path: "products.@|[Book].reviews.@.product|[Book]") { + Flatten(path: "products.@.reviews.@.product") { Fetch(service: "products") { { ... on Book { __typename id } - } => - { - ... on Book { - sku - } - } - }, - }, - }, - Include(if: $title) { - Flatten(path: "products.@|[Book].reviews.@.product|[Book]") { - Fetch(service: "books") { - { - ... on Book { + ... on Magazine { __typename id } } => { ... on Book { - title + sku + } + ... on Magazine { + sku } } }, @@ -1124,7 +963,7 @@ fn nested_interface_field_with_redundant_inline_fragments() -> Result<(), Box Result<(), Box> { ); let query_plan = build_query_plan("fixture/issues/281.supergraph.graphql", document)?; - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "a") { @@ -59,36 +59,20 @@ fn issue_281_test() -> Result<(), Box> { id } }, - Parallel { - Flatten(path: "viewer.review|[UserReview].product") { - Fetch(service: "b") { - { - ... on Product { - __typename - id - } - } => - { - ... on Product { - pid - } + Flatten(path: "viewer.review|[UserReview].product") { + Fetch(service: "b") { + { + ... on Product { + __typename + id } - }, - }, - Flatten(path: "viewer.review|[AnonymousReview].product") { - Fetch(service: "b") { - { - ... on Product { - __typename - id - } - } => - { - ... on Product { - b - } + } => + { + ... on Product { + pid + b } - }, + } }, }, Flatten(path: "viewer.review|[UserReview].product") { @@ -124,7 +108,7 @@ fn issue_281_test() -> Result<(), Box> { }, }, }, - "#); + "###); Ok(()) } diff --git a/lib/query-planner/src/tests/mod.rs b/lib/query-planner/src/tests/mod.rs index db3021da..32401319 100644 --- a/lib/query-planner/src/tests/mod.rs +++ b/lib/query-planner/src/tests/mod.rs @@ -32,7 +32,266 @@ fn test_bench_operation() -> Result<(), Box> { &std::fs::read_to_string("../../bench/operation.graphql") .expect("Unable to read input file"), ); - let _query_plan = build_query_plan("../../bench/supergraph.graphql", document)?; + let query_plan = build_query_plan("../../bench/supergraph.graphql", document)?; + + insta::assert_snapshot!(format!("{}", query_plan), @r###" + QueryPlan { + Sequence { + Parallel { + Fetch(service: "products") { + { + topProducts { + __typename + upc + name + price + weight + } + } + }, + Fetch(service: "accounts") { + { + users { + __typename + id + username + name + } + } + }, + }, + Parallel { + Flatten(path: "topProducts.@") { + Fetch(service: "inventory") { + { + ... on Product { + __typename + price + weight + upc + } + } => + { + ... on Product { + shippingEstimate + inStock + } + } + }, + }, + Flatten(path: "topProducts.@") { + Fetch(service: "reviews") { + { + ... on Product { + __typename + upc + } + } => + { + ... on Product { + reviews { + id + body + author { + __typename + id + reviews { + id + body + product { + __typename + upc + } + } + username + } + } + } + } + }, + }, + Flatten(path: "users.@") { + Fetch(service: "reviews") { + { + ... on User { + __typename + id + } + } => + { + ... on User { + reviews { + id + body + product { + __typename + upc + reviews { + id + body + author { + __typename + id + reviews { + id + body + product { + __typename + upc + } + } + username + } + } + } + } + } + } + }, + }, + }, + Parallel { + Flatten(path: "topProducts.@.reviews.@.author.reviews.@.product") { + Fetch(service: "products") { + { + ... on Product { + __typename + upc + } + } => + { + ... on Product { + price + weight + name + } + } + }, + }, + Flatten(path: "topProducts.@.reviews.@.author") { + Fetch(service: "accounts") { + { + ... on User { + __typename + id + } + } => + { + ... on User { + name + } + } + }, + }, + Flatten(path: "users.@.reviews.@.product") { + Fetch(service: "products") { + { + ... on Product { + __typename + upc + } + } => + { + ... on Product { + price + weight + name + } + } + }, + }, + Flatten(path: "users.@.reviews.@.product.reviews.@.author.reviews.@.product") { + Fetch(service: "products") { + { + ... on Product { + __typename + upc + } + } => + { + ... on Product { + price + weight + name + } + } + }, + }, + Flatten(path: "users.@.reviews.@.product.reviews.@.author") { + Fetch(service: "accounts") { + { + ... on User { + __typename + id + } + } => + { + ... on User { + name + } + } + }, + }, + }, + Parallel { + Flatten(path: "topProducts.@.reviews.@.author.reviews.@.product") { + Fetch(service: "inventory") { + { + ... on Product { + __typename + upc + price + weight + } + } => + { + ... on Product { + inStock + shippingEstimate + } + } + }, + }, + Flatten(path: "users.@.reviews.@.product") { + Fetch(service: "inventory") { + { + ... on Product { + __typename + upc + price + weight + } + } => + { + ... on Product { + inStock + shippingEstimate + } + } + }, + }, + Flatten(path: "users.@.reviews.@.product.reviews.@.author.reviews.@.product") { + Fetch(service: "inventory") { + { + ... on Product { + __typename + upc + price + weight + } + } => + { + ... on Product { + inStock + shippingEstimate + } + } + }, + }, + }, + }, + }, + "###); Ok(()) } From dd05ef473c03f2972e6c3c0256062fb5a8daba88 Mon Sep 17 00:00:00 2001 From: Dotan Simha Date: Tue, 29 Jul 2025 09:57:38 +0300 Subject: [PATCH 4/4] try --- .../audit-abstract-types.supergraph.graphql | 256 ++++++++++++++++++ lib/query-planner/src/planner/plan_nodes.rs | 6 +- lib/query-planner/src/tests/alias.rs | 6 +- lib/query-planner/src/tests/interface.rs | 137 +++++++++- .../src/tests/interface_object.rs | 6 +- .../tests/interface_object_with_requires.rs | 15 +- lib/query-planner/src/tests/issues.rs | 6 +- lib/query-planner/src/tests/overrides.rs | 6 +- lib/query-planner/src/tests/provides.rs | 19 +- lib/query-planner/src/tests/union.rs | 28 +- 10 files changed, 431 insertions(+), 54 deletions(-) create mode 100644 lib/query-planner/fixture/tests/audit-abstract-types.supergraph.graphql diff --git a/lib/query-planner/fixture/tests/audit-abstract-types.supergraph.graphql b/lib/query-planner/fixture/tests/audit-abstract-types.supergraph.graphql new file mode 100644 index 00000000..7d333535 --- /dev/null +++ b/lib/query-planner/fixture/tests/audit-abstract-types.supergraph.graphql @@ -0,0 +1,256 @@ +schema + @link(url: "https://specs.apollo.dev/link/v1.0") + @link(url: "https://specs.apollo.dev/join/v0.3", for: EXECUTION) + @link(url: "https://specs.apollo.dev/inaccessible/v0.2", for: SECURITY) { + query: Query +} + +directive @inaccessible on FIELD_DEFINITION | OBJECT | INTERFACE | UNION | ARGUMENT_DEFINITION | SCALAR | ENUM | ENUM_VALUE | INPUT_OBJECT | INPUT_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 +) 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 Agency @join__type(graph: AGENCY, key: "id") @join__type(graph: PRODUCTS) { + id: ID! + companyName: String @join__field(graph: AGENCY) + email: Email @join__field(graph: AGENCY) +} + +type Book implements Product & Similar + @join__implements(graph: INVENTORY, interface: "Product") + @join__implements(graph: PRODUCTS, interface: "Product") + @join__implements(graph: PRODUCTS, interface: "Similar") + @join__implements(graph: REVIEWS, interface: "Product") + @join__implements(graph: REVIEWS, interface: "Similar") + @join__type(graph: BOOKS, key: "id") + @join__type(graph: INVENTORY, key: "id") + @join__type(graph: PRODUCTS, key: "id") + @join__type(graph: REVIEWS, key: "id") { + id: ID! + title: String @join__field(graph: BOOKS) + dimensions: ProductDimension + @join__field(graph: INVENTORY, external: true) + @join__field(graph: PRODUCTS) + delivery(zip: String): DeliveryEstimates + @join__field(graph: INVENTORY, requires: "dimensions { size weight }") + sku: String @join__field(graph: PRODUCTS) + createdBy: User @join__field(graph: PRODUCTS) + similar: [Book] + @join__field(graph: PRODUCTS) + @join__field(graph: REVIEWS, external: true) + hidden: Boolean @join__field(graph: PRODUCTS) + publisherType: PublisherType @join__field(graph: PRODUCTS) + reviewsCount: Int! @join__field(graph: REVIEWS) + reviewsScore: Float! @join__field(graph: REVIEWS) + reviews: [Review!]! @join__field(graph: REVIEWS) + reviewsOfSimilar: [Review!]! + @join__field(graph: REVIEWS, requires: "similar { id }") +} + +type DeliveryEstimates @join__type(graph: INVENTORY) { + estimatedDelivery: String + fastestDelivery: String +} + +type Email @join__type(graph: AGENCY) { + address: String +} + +type Group @join__type(graph: AGENCY, key: "id") { + id: ID! + name: String + email: String +} + +scalar join__FieldSet + +enum join__Graph { + AGENCY + @join__graph( + name: "agency" + url: "http://localhost:4200/abstract-types/agency" + ) + BOOKS + @join__graph( + name: "books" + url: "http://localhost:4200/abstract-types/books" + ) + INVENTORY + @join__graph( + name: "inventory" + url: "http://localhost:4200/abstract-types/inventory" + ) + MAGAZINES + @join__graph( + name: "magazines" + url: "http://localhost:4200/abstract-types/magazines" + ) + PRODUCTS + @join__graph( + name: "products" + url: "http://localhost:4200/abstract-types/products" + ) + REVIEWS + @join__graph( + name: "reviews" + url: "http://localhost:4200/abstract-types/reviews" + ) + USERS + @join__graph( + name: "users" + url: "http://localhost:4200/abstract-types/users" + ) +} + +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 Magazine implements Product & Similar + @join__implements(graph: INVENTORY, interface: "Product") + @join__implements(graph: PRODUCTS, interface: "Product") + @join__implements(graph: PRODUCTS, interface: "Similar") + @join__implements(graph: REVIEWS, interface: "Product") + @join__implements(graph: REVIEWS, interface: "Similar") + @join__type(graph: INVENTORY, key: "id") + @join__type(graph: MAGAZINES, key: "id") + @join__type(graph: PRODUCTS, key: "id") + @join__type(graph: REVIEWS, key: "id") { + id: ID! + dimensions: ProductDimension + @join__field(graph: INVENTORY, external: true) + @join__field(graph: PRODUCTS) + delivery(zip: String): DeliveryEstimates + @join__field(graph: INVENTORY, requires: "dimensions { size weight }") + title: String @join__field(graph: MAGAZINES) + sku: String @join__field(graph: PRODUCTS) + createdBy: User @join__field(graph: PRODUCTS) + similar: [Magazine] + @join__field(graph: PRODUCTS) + @join__field(graph: REVIEWS, external: true) + hidden: Boolean @join__field(graph: PRODUCTS) + publisherType: PublisherType @join__field(graph: PRODUCTS) + reviewsCount: Int! @join__field(graph: REVIEWS) + reviewsScore: Float! @join__field(graph: REVIEWS) + reviews: [Review!]! @join__field(graph: REVIEWS) + reviewsOfSimilar: [Review!]! + @join__field(graph: REVIEWS, requires: "similar { id }") +} + +interface Product + @join__type(graph: INVENTORY) + @join__type(graph: PRODUCTS) + @join__type(graph: REVIEWS) { + id: ID! + dimensions: ProductDimension + @join__field(graph: INVENTORY) + @join__field(graph: PRODUCTS) + delivery(zip: String): DeliveryEstimates @join__field(graph: INVENTORY) + sku: String @join__field(graph: PRODUCTS) + createdBy: User @join__field(graph: PRODUCTS) + hidden: Boolean @inaccessible @join__field(graph: PRODUCTS) + reviewsCount: Int! @join__field(graph: REVIEWS) + reviewsScore: Float! @join__field(graph: REVIEWS) + reviews: [Review!]! @join__field(graph: REVIEWS) +} + +type ProductDimension + @join__type(graph: INVENTORY) + @join__type(graph: PRODUCTS) { + size: String + weight: Float +} + +union PublisherType + @join__type(graph: AGENCY) + @join__type(graph: PRODUCTS) + @join__unionMember(graph: AGENCY, member: "Agency") + @join__unionMember(graph: PRODUCTS, member: "Agency") + @join__unionMember(graph: AGENCY, member: "Group") + @join__unionMember(graph: PRODUCTS, member: "Self") = + | Agency + | Group + | Self + +type Query + @join__type(graph: AGENCY) + @join__type(graph: BOOKS) + @join__type(graph: INVENTORY) + @join__type(graph: MAGAZINES) + @join__type(graph: PRODUCTS) + @join__type(graph: REVIEWS) + @join__type(graph: USERS) { + books: [Book] @join__field(graph: BOOKS) + magazines: [Magazine] @join__field(graph: MAGAZINES) + products: [Product] @join__field(graph: PRODUCTS) + similar(id: ID!): [Product] @join__field(graph: PRODUCTS) + review(id: Int!): Review @join__field(graph: REVIEWS) +} + +type Review @join__type(graph: REVIEWS) { + id: Int! + body: String! + product: Product +} + +type Self @join__type(graph: PRODUCTS) { + email: String +} + +interface Similar @join__type(graph: PRODUCTS) @join__type(graph: REVIEWS) { + similar: [Product] +} + +type User + @join__type(graph: PRODUCTS, key: "email") + @join__type(graph: USERS, key: "email") { + email: ID! + totalProductsCreated: Int + name: String @join__field(graph: USERS) +} diff --git a/lib/query-planner/src/planner/plan_nodes.rs b/lib/query-planner/src/planner/plan_nodes.rs index 71f67047..4348b55e 100644 --- a/lib/query-planner/src/planner/plan_nodes.rs +++ b/lib/query-planner/src/planner/plan_nodes.rs @@ -378,11 +378,7 @@ impl PlanNode { } else { PlanNode::Flatten(FlattenNode { // it's cheaper to clone response_path (Arc etc), rather then cloning the step - path: if step.is_fetching_multiple_types() { - step.response_path.without_type_castings().into() - } else { - step.response_path.clone().into() - }, + path: step.response_path.without_type_castings().into(), node: Box::new(PlanNode::Fetch(FetchNode::from_fetch_step( step, supergraph, ))), diff --git a/lib/query-planner/src/tests/alias.rs b/lib/query-planner/src/tests/alias.rs index 4a259e98..3e569f27 100644 --- a/lib/query-planner/src/tests/alias.rs +++ b/lib/query-planner/src/tests/alias.rs @@ -45,7 +45,7 @@ fn circular_reference_interface() -> Result<(), Box> { document, )?; - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "a") { @@ -76,7 +76,7 @@ fn circular_reference_interface() -> Result<(), Box> { id } }, - Flatten(path: "product|[Book]") { + Flatten(path: "product") { Fetch(service: "b") { { ... on Book { @@ -93,7 +93,7 @@ fn circular_reference_interface() -> Result<(), Box> { }, }, }, - "#); + "###); Ok(()) } diff --git a/lib/query-planner/src/tests/interface.rs b/lib/query-planner/src/tests/interface.rs index 78c2a28b..7f3ab854 100644 --- a/lib/query-planner/src/tests/interface.rs +++ b/lib/query-planner/src/tests/interface.rs @@ -4,6 +4,137 @@ use crate::{ }; use std::error::Error; +#[test] +fn audit_test_case_5() -> Result<(), Box> { + init_logger(); + let document = parse_operation( + r#" + { + products { + id + reviews { + product { + sku + ... on Magazine { + title + } + ... on Book { + reviewsCount + } + } + } + } + } + "#, + ); + let query_plan = build_query_plan( + "fixture/tests/audit-abstract-types.supergraph.graphql", + document, + )?; + + insta::assert_snapshot!(format!("{}", query_plan), @r###" + QueryPlan { + Sequence { + Fetch(service: "products") { + { + products { + id + __typename + ... on Book { + __typename + id + } + ... on Magazine { + __typename + id + } + } + } + }, + Flatten(path: "products.@") { + Fetch(service: "reviews") { + { + ... on Book { + __typename + id + } + ... on Magazine { + __typename + id + } + } => + { + ... on Book { + reviews { + ...a } + } + ... on Magazine { + reviews { + ...a } + } + } + fragment a on Review { + product { + __typename + ... on Book { + __typename + id + reviewsCount + } + ... on Magazine { + __typename + id + } + } + } + }, + }, + Parallel { + Flatten(path: "products.@.reviews.@.product") { + Fetch(service: "products") { + { + ... on Book { + __typename + id + } + ... on Magazine { + __typename + id + } + } => + { + ... on Book { + sku + } + ... on Magazine { + sku + } + } + }, + }, + Flatten(path: "products.@.reviews.@.product") { + Fetch(service: "magazines") { + { + ... on Magazine { + __typename + id + } + } => + { + ... on Magazine { + title + } + } + }, + }, + }, + }, + }, + "###); + + Ok(()) +} + /// Tests querying the `node` interface field using two different aliases (`account` and `chat`). /// Verifies that aliases work correctly when querying the same interface field with different IDs. #[test] @@ -793,7 +924,7 @@ fn nested_interface_field_with_inline_fragments() -> Result<(), Box> }, Parallel { Include(if: $title) { - Flatten(path: "products.@|[Book].reviews.@.product|[Book]") { + Flatten(path: "products.@.reviews.@.product") { Fetch(service: "books") { { ... on Book { @@ -809,7 +940,7 @@ fn nested_interface_field_with_inline_fragments() -> Result<(), Box> }, }, }, - Flatten(path: "products.@|[Book].reviews.@.product|[Magazine]") { + Flatten(path: "products.@.reviews.@.product") { Fetch(service: "products") { { ... on Magazine { @@ -920,7 +1051,7 @@ fn nested_interface_field_with_redundant_inline_fragments() -> Result<(), Box Result<(), Box> { id } }, - Flatten(path: "viewer.review|[UserReview].product") { + Flatten(path: "viewer.review.product") { Fetch(service: "b") { { ... on Product { @@ -75,7 +75,7 @@ fn issue_281_test() -> Result<(), Box> { } }, }, - Flatten(path: "viewer.review|[UserReview].product") { + Flatten(path: "viewer.review.product") { Fetch(service: "c") { { ... on Product { @@ -91,7 +91,7 @@ fn issue_281_test() -> Result<(), Box> { } }, }, - Flatten(path: "viewer.review|[UserReview].product") { + Flatten(path: "viewer.review.product") { Fetch(service: "d") { { ... on Product { diff --git a/lib/query-planner/src/tests/overrides.rs b/lib/query-planner/src/tests/overrides.rs index 0838a25d..bf6becf4 100644 --- a/lib/query-planner/src/tests/overrides.rs +++ b/lib/query-planner/src/tests/overrides.rs @@ -182,7 +182,7 @@ fn override_object_field_but_interface_is_requested() -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box> { document, )?; - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "b") { @@ -245,7 +245,7 @@ fn provides_on_union() -> Result<(), Box> { } } }, - Flatten(path: "media|[Movie]") { + Flatten(path: "media") { Fetch(service: "c") { { ... on Movie { @@ -262,8 +262,8 @@ fn provides_on_union() -> Result<(), Box> { }, }, }, - "#); - insta::assert_snapshot!(format!("{}", serde_json::to_string_pretty(&query_plan).unwrap_or_default()), @r#" + "###); + insta::assert_snapshot!(format!("{}", serde_json::to_string_pretty(&query_plan).unwrap_or_default()), @r###" { "kind": "QueryPlan", "node": { @@ -280,9 +280,6 @@ fn provides_on_union() -> Result<(), Box> { "path": [ { "Field": "media" - }, - { - "Cast": "Movie" } ], "node": { @@ -311,7 +308,7 @@ fn provides_on_union() -> Result<(), Box> { ] } } - "#); + "###); Ok(()) } @@ -390,7 +387,7 @@ fn provides_on_interface_2_test() -> Result<(), Box> { document, )?; - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "b") { @@ -415,7 +412,7 @@ fn provides_on_interface_2_test() -> Result<(), Box> { } } }, - Flatten(path: "media|[Book].animals.@|[Cat]") { + Flatten(path: "media.animals.@") { Fetch(service: "c") { { ... on Cat { @@ -432,7 +429,7 @@ fn provides_on_interface_2_test() -> Result<(), Box> { }, }, }, - "#); + "###); Ok(()) } diff --git a/lib/query-planner/src/tests/union.rs b/lib/query-planner/src/tests/union.rs index 34ec7bea..0336cfb5 100644 --- a/lib/query-planner/src/tests/union.rs +++ b/lib/query-planner/src/tests/union.rs @@ -141,7 +141,7 @@ fn union_member_entity_call() -> Result<(), Box> { document, )?; - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "a") { @@ -157,7 +157,7 @@ fn union_member_entity_call() -> Result<(), Box> { } } }, - Flatten(path: "aMedia|[Book]") { + Flatten(path: "aMedia") { Fetch(service: "b") { { ... on Book { @@ -174,7 +174,7 @@ fn union_member_entity_call() -> Result<(), Box> { }, }, }, - "#); + "###); Ok(()) } @@ -212,7 +212,7 @@ fn union_member_entity_call_many_local() -> Result<(), Box> { document, )?; - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "a") { @@ -234,7 +234,7 @@ fn union_member_entity_call_many_local() -> Result<(), Box> { } } }, - Flatten(path: "viewer.song|[Book]") { + Flatten(path: "viewer.song") { Fetch(service: "b") { { ... on Book { @@ -251,7 +251,7 @@ fn union_member_entity_call_many_local() -> Result<(), Box> { }, }, }, - "#); + "###); Ok(()) } @@ -321,7 +321,7 @@ fn union_member_entity_call_many() -> Result<(), Box> { document, )?; - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Parallel { @@ -372,7 +372,7 @@ fn union_member_entity_call_many() -> Result<(), Box> { } }, }, - Flatten(path: "viewer.song|[Book]") { + Flatten(path: "viewer.song") { Fetch(service: "b") { { ... on Book { @@ -389,7 +389,7 @@ fn union_member_entity_call_many() -> Result<(), Box> { }, }, }, - "#); + "###); Ok(()) } @@ -422,7 +422,7 @@ fn union_overfetching_test() -> Result<(), Box> { document, )?; - insta::assert_snapshot!(format!("{}", query_plan), @r#" + insta::assert_snapshot!(format!("{}", query_plan), @r###" QueryPlan { Sequence { Fetch(service: "a") { @@ -445,7 +445,7 @@ fn union_overfetching_test() -> Result<(), Box> { } }, Parallel { - Flatten(path: "review|[UserReview].product") { + Flatten(path: "review.product") { Fetch(service: "d") { { ... on Product { @@ -460,7 +460,7 @@ fn union_overfetching_test() -> Result<(), Box> { } }, }, - Flatten(path: "review|[UserReview].product") { + Flatten(path: "review.product") { Fetch(service: "c") { { ... on Product { @@ -475,7 +475,7 @@ fn union_overfetching_test() -> Result<(), Box> { } }, }, - Flatten(path: "review|[AnonymousReview].product") { + Flatten(path: "review.product") { Fetch(service: "b") { { ... on Product { @@ -493,7 +493,7 @@ fn union_overfetching_test() -> Result<(), Box> { }, }, }, - "#); + "###); Ok(()) }