diff --git a/apollo-router/src/query_planner/fetch.rs b/apollo-router/src/query_planner/fetch.rs index c7b4e7216f..598fe203f3 100644 --- a/apollo-router/src/query_planner/fetch.rs +++ b/apollo-router/src/query_planner/fetch.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::fmt::Display; use std::sync::Arc; @@ -98,13 +97,13 @@ pub(crate) struct FetchNode { pub(crate) struct Variables { pub(crate) variables: Object, - pub(crate) paths: HashMap, + pub(crate) inverted_paths: Vec>, } impl Variables { #[instrument(skip_all, level = "debug", name = "make_variables")] #[allow(clippy::too_many_arguments)] - pub(super) async fn new( + pub(super) fn new( requires: &[Selection], variable_usages: &[String], data: &Value, @@ -123,7 +122,7 @@ impl Variables { .map(|(variable_key, value)| (variable_key.clone(), value.clone())) })); - let mut paths: HashMap = HashMap::new(); + let mut inverted_paths: Vec> = Vec::new(); let mut values: IndexSet = IndexSet::new(); data.select_values_and_paths(schema, current_dir, |path, value| { @@ -132,11 +131,12 @@ impl Variables { rewrites::apply_rewrites(schema, &mut value, input_rewrites); match values.get_index_of(&value) { Some(index) => { - paths.insert(path.clone(), index); + inverted_paths[index].push(path.clone()); } None => { - paths.insert(path.clone(), values.len()); + inverted_paths.push(vec![path.clone()]); values.insert(value); + debug_assert!(inverted_paths.len() == values.len()); } } } @@ -151,7 +151,10 @@ impl Variables { variables.insert("representations", representations); - Some(Variables { variables, paths }) + Some(Variables { + variables, + inverted_paths, + }) } else { // with nested operations (Query or Mutation has an operation returning a Query or Mutation), // when the first fetch fails, the query plan will still execute up until the second fetch, @@ -176,7 +179,7 @@ impl Variables { .map(|(variable_key, value)| (variable_key.clone(), value.clone())) }) .collect::(), - paths: HashMap::new(), + inverted_paths: Vec::new(), }) } } @@ -198,7 +201,10 @@ impl FetchNode { .. } = self; - let Variables { variables, paths } = match Variables::new( + let Variables { + variables, + inverted_paths: paths, + } = match Variables::new( &self.requires, self.variable_usages.as_ref(), data, @@ -207,9 +213,7 @@ impl FetchNode { parameters.supergraph_request, parameters.schema, &self.input_rewrites, - ) - .await - { + ) { Some(variables) => variables, None => { return Ok((Value::Object(Object::default()), Vec::new())); @@ -304,15 +308,9 @@ impl FetchNode { &'a self, schema: &Schema, current_dir: &'a Path, - paths: HashMap, + inverted_paths: Vec>, response: graphql::Response, ) -> (Value, Vec) { - // for each entity in the response, find out the path where it must be inserted - let mut inverted_paths: HashMap> = HashMap::new(); - for (path, index) in paths.iter() { - (*inverted_paths.entry(*index).or_default()).push(path); - } - if !self.requires.is_empty() { let entities_path = Path(vec![json_ext::PathElement::Key("_entities".to_string())]); @@ -330,7 +328,7 @@ impl FetchNode { match path.0.get(1) { Some(json_ext::PathElement::Index(i)) => { for values_path in - inverted_paths.get(i).iter().flat_map(|v| v.iter()) + inverted_paths.get(*i).iter().flat_map(|v| v.iter()) { errors.push(Error { locations: error.locations.clone(), @@ -370,7 +368,7 @@ impl FetchNode { for (index, mut entity) in array.into_iter().enumerate() { rewrites::apply_rewrites(schema, &mut entity, &self.output_rewrites); - if let Some(paths) = inverted_paths.get(&index) { + if let Some(paths) = inverted_paths.get(index) { if paths.len() > 1 { for path in &paths[1..] { let _ = value.insert(path, entity.clone()); diff --git a/apollo-router/src/query_planner/selection.rs b/apollo-router/src/query_planner/selection.rs index 878ae72375..521882f40c 100644 --- a/apollo-router/src/query_planner/selection.rs +++ b/apollo-router/src/query_planner/selection.rs @@ -1,6 +1,6 @@ use serde::Deserialize; use serde::Serialize; -use serde_json_bytes::Entry; +use serde_json_bytes::ByteString; use crate::error::FetchError; use crate::json_ext::Object; @@ -53,16 +53,15 @@ pub(crate) fn select_object( selections: &[Selection], schema: &Schema, ) -> Result, FetchError> { - let mut output = Object::new(); + let mut output = Object::with_capacity(selections.len()); for selection in selections { match selection { Selection::Field(field) => { - if let Some(value) = select_field(content, field, schema)? { - match output.entry(field.name.to_owned()) { - Entry::Occupied(mut existing) => existing.get_mut().deep_merge(value), - Entry::Vacant(vacant) => { - vacant.insert(value); - } + if let Some((key, value)) = select_field(content, field, schema)? { + if let Some(o) = output.get_mut(field.name.as_str()) { + o.deep_merge(value); + } else { + output.insert(key.to_owned(), value); } } } @@ -81,13 +80,16 @@ pub(crate) fn select_object( Ok(Some(Value::Object(output))) } -fn select_field( - content: &Object, +fn select_field<'a>( + content: &'a Object, field: &Field, schema: &Schema, -) -> Result, FetchError> { - let res = match (content.get(field.name.as_str()), &field.selections) { - (Some(v), _) => select_value(v, field, schema), +) -> Result, FetchError> { + let res = match ( + content.get_key_value(field.name.as_str()), + &field.selections, + ) { + (Some((k, v)), _) => select_value(v, field, schema).map(|opt| opt.map(|v| (k, v))), (None, _) => Err(FetchError::ExecutionFieldNotFound { field: field.name.to_owned(), }), diff --git a/apollo-router/src/query_planner/subscription.rs b/apollo-router/src/query_planner/subscription.rs index d91014b83e..d409c5b136 100644 --- a/apollo-router/src/query_planner/subscription.rs +++ b/apollo-router/src/query_planner/subscription.rs @@ -209,9 +209,7 @@ impl SubscriptionNode { parameters.supergraph_request, parameters.schema, &self.input_rewrites, - ) - .await - { + ) { Some(variables) => variables, None => { return Ok(Vec::new());