Skip to content

Commit

Permalink
graphql: Remove SelectedAttributes by add Field.selected_attrs
Browse files Browse the repository at this point in the history
  • Loading branch information
lutter committed Jan 23, 2024
1 parent 6a26a6a commit 4e28726
Show file tree
Hide file tree
Showing 3 changed files with 78 additions and 101 deletions.
76 changes: 70 additions & 6 deletions graphql/src/execution/ast.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand Down Expand Up @@ -101,18 +106,25 @@ impl SelectionSet {
pub fn fields_for(
&self,
obj_type: &ObjectType,
) -> Result<impl Iterator<Item = &Field>, QueryExecutionError> {
self.fields_for_name(&obj_type.name)
}

fn fields_for_name(
&self,
name: &str,
) -> Result<impl Iterator<Item = &Field>, 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())
Expand Down Expand Up @@ -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<AttributeNames, QueryExecutionError> {
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<String> = 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 {
Expand Down
77 changes: 4 additions & 73 deletions graphql/src/store/prefetch.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -682,15 +677,13 @@ impl<'a> Loader<'a> {
field: &a::Field,
) -> Result<(Vec<Node>, 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(),
field,
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(),
Expand Down Expand Up @@ -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<Self, QueryExecutionError> {
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)
}
}
26 changes: 4 additions & 22 deletions graphql/src/store/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -35,24 +33,23 @@ pub(crate) fn build_query<'a>(
types_for_interface: &'a BTreeMap<String, Vec<s::ObjectType>>,
max_first: u32,
max_skip: u32,
mut column_names: SelectedAttributes,
schema: &SchemaPair,
) -> Result<EntityQuery, QueryExecutionError> {
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::<Result<_, _>>()?,
});
let mut query = EntityQuery::new(parse_subgraph_id(entity)?, block, entity_types)
.range(build_range(field, max_first, max_skip)?);
Expand Down Expand Up @@ -933,7 +930,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema,
)
.unwrap()
Expand All @@ -951,7 +947,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema,
)
.unwrap()
Expand All @@ -974,7 +969,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema,
)
.unwrap()
Expand All @@ -995,7 +989,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema,
)
.unwrap()
Expand All @@ -1012,7 +1005,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema,
)
.unwrap()
Expand All @@ -1033,7 +1025,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema,
)
.unwrap()
Expand All @@ -1050,7 +1041,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema,
)
.unwrap()
Expand All @@ -1074,7 +1064,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema,
)
.unwrap()
Expand All @@ -1094,7 +1083,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema,
)
.unwrap()
Expand All @@ -1117,7 +1105,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema
)
.unwrap()
Expand All @@ -1138,7 +1125,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema
)
.unwrap()
Expand All @@ -1158,7 +1144,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema
)
.unwrap()
Expand All @@ -1181,7 +1166,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema
)
.unwrap()
Expand Down Expand Up @@ -1214,7 +1198,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema
)
.unwrap()
Expand Down Expand Up @@ -1250,7 +1233,6 @@ mod tests {
&BTreeMap::new(),
std::u32::MAX,
std::u32::MAX,
Default::default(),
&schema
)
.unwrap()
Expand Down

0 comments on commit 4e28726

Please sign in to comment.