From 4e28726d0e7d1dea2b372f003fc8a4a1f9ba6474 Mon Sep 17 00:00:00 2001 From: David Lutterkort Date: Fri, 15 Dec 2023 18:23:17 -0800 Subject: [PATCH] graphql: Remove SelectedAttributes by add Field.selected_attrs --- graphql/src/execution/ast.rs | 76 +++++++++++++++++++++++++++++++--- graphql/src/store/prefetch.rs | 77 ++--------------------------------- graphql/src/store/query.rs | 26 ++---------- 3 files changed, 78 insertions(+), 101 deletions(-) diff --git a/graphql/src/execution/ast.rs b/graphql/src/execution/ast.rs index 26f1e076e97..df1a9b3de5f 100644 --- a/graphql/src/execution/ast.rs +++ b/graphql/src/execution/ast.rs @@ -1,10 +1,15 @@ -use std::collections::HashSet; +use std::collections::{BTreeSet, HashSet}; use graph::{ - components::store::ChildMultiplicity, - data::graphql::ObjectOrInterface, + components::store::{AttributeNames, ChildMultiplicity}, + constraint_violation, + data::graphql::{ext::DirectiveFinder as _, ObjectOrInterface}, + env::ENV_VARS, prelude::{anyhow, q, r, s, QueryExecutionError, ValueMap}, - schema::{ast::ObjectType, ApiSchema}, + schema::{ + ast::{self as sast, ObjectType}, + ApiSchema, + }, }; use graphql_parser::Pos; @@ -101,18 +106,25 @@ impl SelectionSet { pub fn fields_for( &self, obj_type: &ObjectType, + ) -> Result, QueryExecutionError> { + self.fields_for_name(&obj_type.name) + } + + fn fields_for_name( + &self, + name: &str, ) -> Result, QueryExecutionError> { let item = self .items .iter() - .find(|(our_type, _)| our_type == obj_type) + .find(|(our_type, _)| our_type.name == name) .ok_or_else(|| { // see: graphql-bug-compat // Once queries are validated, this can become a panic since // users won't be able to trigger this any more QueryExecutionError::ValidationError( None, - format!("invalid query: no fields for type `{}`", obj_type.name), + format!("invalid query: no fields for type `{}`", name), ) })?; Ok(item.1.iter()) @@ -261,6 +273,58 @@ impl Field { fn is_leaf(&self) -> bool { self.selection_set.is_empty() } + + /// Return the set of attributes that should be selected for this field. + /// If `ENV_VARS.enable_select_by_specific_attributes` is `false`, + /// return `AttributeNames::All + pub fn selected_attrs( + &self, + object_type: &s::ObjectType, + ) -> Result { + if !ENV_VARS.enable_select_by_specific_attributes { + return Ok(AttributeNames::All); + } + + let fields = self.selection_set.fields_for_name(&object_type.name)?; + + // Extract the attributes we should select from `selection_set`. In + // particular, disregard derived fields since they are not stored + let mut column_names: BTreeSet = fields + .filter(|field| { + // Keep fields that are not derived and for which we + // can find the field type + sast::get_field(object_type, &field.name) + .map(|field_type| !field_type.is_derived()) + .unwrap_or(false) + }) + .filter_map(|field| { + if field.name.starts_with("__") { + None + } else { + Some(field.name.clone()) + } + }) + .collect(); + + // We need to also select the `orderBy` field if there is one. + // Because of how the API Schema is set up, `orderBy` can only have + // an enum value + match self.argument_value("orderBy") { + None => { /* nothing to do */ } + Some(r::Value::Enum(e)) => { + column_names.insert(e.clone()); + } + Some(v) => { + return Err(constraint_violation!( + "'orderBy' attribute must be an enum but is {:?}", + v + ) + .into()); + } + } + + Ok(AttributeNames::Select(column_names)) + } } impl ValueMap for Field { diff --git a/graphql/src/store/prefetch.rs b/graphql/src/store/prefetch.rs index 80f46bb3d28..c8c3e34dafe 100644 --- a/graphql/src/store/prefetch.rs +++ b/graphql/src/store/prefetch.rs @@ -1,7 +1,6 @@ //! Run a GraphQL query and fetch all the entitied needed to build the //! final result -use graph::constraint_violation; use graph::data::query::Trace; use graph::data::store::Id; use graph::data::store::IdList; @@ -16,15 +15,11 @@ use std::rc::Rc; use std::time::Instant; use graph::data::graphql::*; -use graph::schema::{ast as sast, EntityType, InputSchema}; -use graph::{ - data::graphql::ext::DirectiveFinder, - prelude::{ - s, AttributeNames, ChildMultiplicity, EntityCollection, EntityFilter, EntityLink, - EntityOrder, EntityWindow, ParentLink, QueryExecutionError, Value as StoreValue, - WindowAttribute, ENV_VARS, - }, +use graph::prelude::{ + s, AttributeNames, ChildMultiplicity, EntityCollection, EntityFilter, EntityLink, EntityOrder, + EntityWindow, ParentLink, QueryExecutionError, Value as StoreValue, WindowAttribute, ENV_VARS, }; +use graph::schema::{ast as sast, EntityType, InputSchema}; use crate::execution::ast as a; use crate::metrics::GraphQLMetrics; @@ -682,7 +677,6 @@ impl<'a> Loader<'a> { field: &a::Field, ) -> Result<(Vec, Trace), QueryExecutionError> { let input_schema = self.resolver.store.input_schema()?; - let selected_attrs = SelectedAttributes::new(field)?; let mut query = build_query( join.child_type(), self.resolver.block_number(), @@ -690,7 +684,6 @@ impl<'a> Loader<'a> { self.ctx.query.schema.types_for_interface(), self.ctx.max_first, self.ctx.max_skip, - selected_attrs, &super::query::SchemaPair { api: self.ctx.query.schema.clone(), input: input_schema.cheap_clone(), @@ -748,65 +741,3 @@ impl<'a> Loader<'a> { Ok(()) } } - -#[derive(Debug, Default, Clone)] -pub(crate) struct SelectedAttributes<'a>(BTreeMap<&'a String, AttributeNames>); - -impl<'a> SelectedAttributes<'a> { - // "Select by Specific Attribute Names" is an experimental feature and can be disabled completely. - // If this environment variable is set, the program will use an empty collection that, - // effectively, causes the `AttributeNames::All` variant to be used as a fallback value for all - // queries. - fn new(field: &'a a::Field) -> Result { - if !ENV_VARS.enable_select_by_specific_attributes { - return Ok(SelectedAttributes(BTreeMap::new())); - } - - // Extract the attributes we should select from `selection_set`. In - // particular, disregard derived fields since they are not stored - let mut map = BTreeMap::new(); - for (object_type, fields) in field.selection_set.fields() { - let column_names = fields - .filter(|field| { - // Keep fields that are not derived and for which we - // can find the field type - sast::get_field(object_type, &field.name) - .map(|field_type| !field_type.is_derived()) - .unwrap_or(false) - }) - .filter_map(|field| { - if field.name.starts_with("__") { - None - } else { - Some(field.name.clone()) - } - }) - .collect(); - map.insert(&object_type.name, AttributeNames::Select(column_names)); - } - // We need to also select the `orderBy` field if there is one. - // Because of how the API Schema is set up, `orderBy` can only have - // an enum value - match field.argument_value("orderBy") { - None => { /* nothing to do */ } - Some(r::Value::Enum(e)) => { - for columns in map.values_mut() { - columns.add_str(e); - } - } - Some(v) => { - return Err(constraint_violation!( - "'orderBy' attribute must be an enum but is {:?}", - v - ) - .into()); - } - } - - Ok(SelectedAttributes(map)) - } - - pub fn get(&mut self, obj_type: &s::ObjectType) -> AttributeNames { - self.0.remove(&obj_type.name).unwrap_or(AttributeNames::All) - } -} diff --git a/graphql/src/store/query.rs b/graphql/src/store/query.rs index 5e58b3cdd17..23d62020097 100644 --- a/graphql/src/store/query.rs +++ b/graphql/src/store/query.rs @@ -12,8 +12,6 @@ use graph::schema::{ApiSchema, EntityType, InputSchema}; use crate::execution::ast as a; -use super::prefetch::SelectedAttributes; - #[derive(Debug)] enum OrderDirection { Ascending, @@ -35,24 +33,23 @@ pub(crate) fn build_query<'a>( types_for_interface: &'a BTreeMap>, max_first: u32, max_skip: u32, - mut column_names: SelectedAttributes, schema: &SchemaPair, ) -> Result { let entity = entity.into(); let entity_types = EntityCollection::All(match &entity { ObjectOrInterface::Object(object) => { - let selected_columns = column_names.get(object); + let selected_columns = field.selected_attrs(object)?; let entity_type = schema.input.entity_type(*object).unwrap(); vec![(entity_type, selected_columns)] } ObjectOrInterface::Interface(interface) => types_for_interface[&interface.name] .iter() .map(|o| { - let selected_columns = column_names.get(o); + let selected_columns = field.selected_attrs(o); let entity_type = schema.input.entity_type(o).unwrap(); - (entity_type, selected_columns) + selected_columns.map(|selected_columns| (entity_type, selected_columns)) }) - .collect(), + .collect::>()?, }); let mut query = EntityQuery::new(parse_subgraph_id(entity)?, block, entity_types) .range(build_range(field, max_first, max_skip)?); @@ -933,7 +930,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema, ) .unwrap() @@ -951,7 +947,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema, ) .unwrap() @@ -974,7 +969,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema, ) .unwrap() @@ -995,7 +989,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema, ) .unwrap() @@ -1012,7 +1005,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema, ) .unwrap() @@ -1033,7 +1025,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema, ) .unwrap() @@ -1050,7 +1041,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema, ) .unwrap() @@ -1074,7 +1064,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema, ) .unwrap() @@ -1094,7 +1083,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema, ) .unwrap() @@ -1117,7 +1105,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema ) .unwrap() @@ -1138,7 +1125,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema ) .unwrap() @@ -1158,7 +1144,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema ) .unwrap() @@ -1181,7 +1166,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema ) .unwrap() @@ -1214,7 +1198,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema ) .unwrap() @@ -1250,7 +1233,6 @@ mod tests { &BTreeMap::new(), std::u32::MAX, std::u32::MAX, - Default::default(), &schema ) .unwrap()