store, graph, graphql: Fix prefetching bug when selecting by specific attributes: allow orderBy

tilacog · web-flow · commit eb5e78fa569a · 2021-11-22T14:33:04.000-03:00
This fixes errors in prefetching phase for queries that orders the
selection set with the `orderBy` attribute.

The changes rely on complementing the selection set with the fields
used for ordering. The problematic part is that each visited selection
set is unaware of those ordering fields due to them not being in scope
during the recursive invocations of the `execute_selection_set`
function.

This led to malformed SQL queries that tried to `ORDER BY`
columns that were not being selected in sub-queries.

The solution presented here was to pass those complementary
fields as arguments along the recursive traversal and injecting them
into their respective selection sets.
diff --git a/graph/src/components/store.rs b/graph/src/components/store.rs
@@ -1785,7 +1785,7 @@ pub enum AttributeNames {
 }
 
 impl AttributeNames {
-    pub fn insert(&mut self, column_name: &str) {
+    fn insert(&mut self, column_name: &str) {
         match self {
             AttributeNames::All => {
                 let mut set = BTreeSet::new();
@@ -1798,14 +1798,27 @@ impl AttributeNames {
         }
     }
 
-    pub fn update(&mut self, field: &q::Field) {
-        // ignore "meta" field names
-        if field.name.starts_with("__") {
+    /// Adds a attribute name. Ignores meta fields.
+    pub fn add(&mut self, field: &q::Field) {
+        if Self::is_meta_field(&field.name) {
             return;
         }
         self.insert(&field.name)
     }
 
+    /// Adds a attribute name. Ignores meta fields.
+    pub fn add_str(&mut self, field_name: &str) {
+        if Self::is_meta_field(field_name) {
+            return;
+        }
+        self.insert(field_name);
+    }
+
+    /// Returns `true` for meta field names, `false` otherwise.
+    fn is_meta_field(field_name: &str) -> bool {
+        field_name.starts_with("__")
+    }
+
     pub fn extend(&mut self, other: Self) {
         use AttributeNames::*;
         match (self, other) {
diff --git a/graph/src/data/graphql/object_or_interface.rs b/graph/src/data/graphql/object_or_interface.rs
@@ -1,5 +1,6 @@
 use crate::prelude::Schema;
 use crate::{components::store::EntityType, prelude::s};
+use std::cmp::Ordering;
 use std::collections::BTreeMap;
 use std::hash::{Hash, Hasher};
 use std::mem;
@@ -32,6 +33,24 @@ impl<'a> Hash for ObjectOrInterface<'a> {
     }
 }
 
+impl<'a> PartialOrd for ObjectOrInterface<'a> {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl<'a> Ord for ObjectOrInterface<'a> {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        use ObjectOrInterface::*;
+        match (self, other) {
+            (Object(a), Object(b)) => a.name.cmp(&b.name),
+            (Interface(a), Interface(b)) => a.name.cmp(&b.name),
+            (Interface(_), Object(_)) => Ordering::Less,
+            (Object(_), Interface(_)) => Ordering::Greater,
+        }
+    }
+}
+
 impl<'a> From<&'a s::ObjectType> for ObjectOrInterface<'a> {
     fn from(object: &'a s::ObjectType) -> Self {
         ObjectOrInterface::Object(object)
diff --git a/graphql/src/store/prefetch.rs b/graphql/src/store/prefetch.rs
@@ -2,6 +2,7 @@
 //! final result
 
 use anyhow::{anyhow, Error};
+use graph::constraint_violation;
 use graph::prelude::{r, CacheWeight};
 use graph::slog::warn;
 use graph::util::cache_weight::btree_node_size;
@@ -18,12 +19,12 @@ use graph::{
     prelude::{
         q, s, ApiSchema, AttributeNames, BlockNumber, ChildMultiplicity, EntityCollection,
         EntityFilter, EntityLink, EntityOrder, EntityWindow, Logger, ParentLink,
-        QueryExecutionError, QueryStore, Value as StoreValue, WindowAttribute,
+        QueryExecutionError, QueryStore, StoreError, Value as StoreValue, WindowAttribute,
     },
 };
 
 use crate::execution::{ExecutionContext, Resolver};
-use crate::query::ast as qast;
+use crate::query::ast::{self as qast, get_argument_value};
 use crate::runner::ResultSizeMetrics;
 use crate::schema::ast as sast;
 use crate::store::{build_query, StoreResolver};
@@ -49,6 +50,10 @@ lazy_static! {
 
 type GroupedFieldSet<'a> = IndexMap<&'a str, CollectedResponseKey<'a>>;
 
+/// Used for associating objects or interfaces and the field names used in `orderBy` query field
+/// attributes.
+type ComplementaryFields<'a> = BTreeMap<ObjectOrInterface<'a>, String>;
+
 /// An `ObjectType` with `Hash` and `Eq` derived from the name.
 #[derive(Clone, Debug)]
 pub struct ObjectCondition<'a>(&'a s::ObjectType);
@@ -555,10 +560,17 @@ fn execute_root_selection_set(
 ) -> Result<Vec<Node>, Vec<QueryExecutionError>> {
     // Obtain the root Query type and fail if there isn't one
     let query_type = ctx.query.schema.query_type.as_ref().into();
-    let grouped_field_set = collect_fields(ctx, query_type, once(selection_set));
+    let (grouped_field_set, _complementary_fields) =
+        collect_fields(ctx, query_type, once(selection_set));
 
     // Execute the root selection set against the root query type
-    execute_selection_set(resolver, ctx, make_root_node(), grouped_field_set)
+    execute_selection_set(
+        resolver,
+        ctx,
+        make_root_node(),
+        grouped_field_set,
+        ComplementaryFields::new(),
+    )
 }
 
 fn check_result_size(logger: &Logger, size: usize) -> Result<(), QueryExecutionError> {
@@ -576,6 +588,7 @@ fn execute_selection_set<'a>(
     ctx: &'a ExecutionContext<impl Resolver>,
     mut parents: Vec<Node>,
     grouped_field_set: GroupedFieldSet<'a>,
+    mut complementary_fields: ComplementaryFields<'a>,
 ) -> Result<Vec<Node>, Vec<QueryExecutionError>> {
     let schema = &ctx.query.schema;
     let mut errors: Vec<QueryExecutionError> = Vec::new();
@@ -619,7 +632,7 @@ fn execute_selection_set<'a>(
                 &field.name,
             );
             // Group fields with the same response key, so we can execute them together
-            let mut grouped_field_set =
+            let (mut grouped_field_set, new_complementary_fields) =
                 collect_fields(ctx, child_type, fields.iter().map(|f| &f.selection_set));
 
             // "Select by Specific Attribute Names" is an experimental feature and can be disabled completely.
@@ -630,8 +643,10 @@ fn execute_selection_set<'a>(
                 if *DISABLE_EXPERIMENTAL_FEATURE_SELECT_BY_SPECIFIC_ATTRIBUTE_NAMES {
                     BTreeMap::new()
                 } else {
-                    CollectedAttributeNames::consolidate_column_names(&mut grouped_field_set)
-                        .resolve_interfaces(&ctx.query.schema.types_for_interface())
+                    let mut collected =
+                        CollectedAttributeNames::consolidate_column_names(&mut grouped_field_set);
+                    collected.populate_complementary_fields(&mut complementary_fields);
+                    collected.resolve_interfaces(&ctx.query.schema.types_for_interface())
                 };
 
             match execute_field(
@@ -645,7 +660,13 @@ fn execute_selection_set<'a>(
                 collected_columns,
             ) {
                 Ok(children) => {
-                    match execute_selection_set(resolver, ctx, children, grouped_field_set) {
+                    match execute_selection_set(
+                        resolver,
+                        ctx,
+                        children,
+                        grouped_field_set,
+                        new_complementary_fields,
+                    ) {
                         Ok(children) => {
                             Join::perform(&mut parents, children, response_key);
                             let weight =
@@ -662,6 +683,23 @@ fn execute_selection_set<'a>(
         }
     }
 
+    // Confidence check: all complementary fields must be consumed, otherwise constructed SQL
+    // queries will be malformed.
+    if !*DISABLE_EXPERIMENTAL_FEATURE_SELECT_BY_SPECIFIC_ATTRIBUTE_NAMES {
+        complementary_fields
+            .into_iter()
+            .for_each(|(parent, complementary_field)| {
+                errors.push(
+                    constraint_violation!(
+                        "Complementary field \"{}\" was not prefetched by its parent: {}",
+                        complementary_field,
+                        parent.name().to_string(),
+                    )
+                    .into(),
+                )
+            });
+    }
+
     if errors.is_empty() {
         Ok(parents)
     } else {
@@ -755,8 +793,9 @@ fn collect_fields<'a>(
     ctx: &'a ExecutionContext<impl Resolver>,
     parent_ty: ObjectOrInterface<'a>,
     selection_sets: impl Iterator<Item = &'a q::SelectionSet>,
-) -> GroupedFieldSet<'a> {
+) -> (GroupedFieldSet<'a>, ComplementaryFields<'a>) {
     let mut grouped_fields = IndexMap::new();
+    let mut complementary_fields = ComplementaryFields::new();
 
     for selection_set in selection_sets {
         collect_fields_inner(
@@ -765,6 +804,7 @@ fn collect_fields<'a>(
             selection_set,
             &mut HashSet::new(),
             &mut grouped_fields,
+            &mut complementary_fields,
         );
     }
 
@@ -778,7 +818,7 @@ fn collect_fields<'a>(
         }
     }
 
-    grouped_fields
+    (grouped_fields, complementary_fields)
 }
 
 // When querying an object type, `type_condition` will always be that object type, even if it passes
@@ -792,6 +832,7 @@ fn collect_fields_inner<'a>(
     selection_set: &'a q::SelectionSet,
     visited_fragments: &mut HashSet<&'a str>,
     output: &mut GroupedFieldSet<'a>,
+    complementary_fields: &mut ComplementaryFields<'a>,
 ) {
     fn collect_fragment<'a>(
         ctx: &'a ExecutionContext<impl Resolver>,
@@ -800,6 +841,7 @@ fn collect_fields_inner<'a>(
         frag_selection_set: &'a q::SelectionSet,
         visited_fragments: &mut HashSet<&'a str>,
         output: &mut GroupedFieldSet<'a>,
+        complementary_fields: &mut ComplementaryFields<'a>,
     ) {
         let schema = &ctx.query.schema.document();
         let fragment_ty = match frag_ty_condition {
@@ -820,6 +862,7 @@ fn collect_fields_inner<'a>(
                 &frag_selection_set,
                 visited_fragments,
                 output,
+                complementary_fields,
             );
         } else {
             // This is an interface fragment in the root selection for an interface.
@@ -842,6 +885,7 @@ fn collect_fields_inner<'a>(
                     &frag_selection_set,
                     visited_fragments,
                     output,
+                    complementary_fields,
                 );
             }
         }
@@ -863,6 +907,34 @@ fn collect_fields_inner<'a>(
                     type_condition,
                     field,
                 );
+
+                // Collect complementary fields used in the `orderBy` query attribute, if present.
+                if let Some(arguments) = get_argument_value(&field.arguments, "orderBy") {
+                    let schema_field = type_condition.field(&field.name).expect(&format!(
+                        "the field {:?} to exist in {:?}",
+                        &field.name,
+                        &type_condition.name()
+                    ));
+                    let field_name = sast::get_field_name(&schema_field.field_type);
+                    let object_or_interface_for_field = ctx
+                        .query
+                        .schema
+                        .document()
+                        .object_or_interface(&field_name)
+                        .expect(&format!(
+                            "The field {:?} to exist in the Document",
+                            field_name
+                        ));
+                    match arguments {
+                        graphql_parser::schema::Value::Enum(complementary_field_name) => {
+                            complementary_fields.insert(
+                                object_or_interface_for_field,
+                                complementary_field_name.clone(),
+                            );
+                        }
+                        _ => unimplemented!("unsure on what to do about other variants"),
+                    }
+                }
             }
 
             q::Selection::FragmentSpread(spread) => {
@@ -878,6 +950,7 @@ fn collect_fields_inner<'a>(
                         &fragment.selection_set,
                         visited_fragments,
                         output,
+                        complementary_fields,
                     );
                 }
             }
@@ -890,6 +963,7 @@ fn collect_fields_inner<'a>(
                     &fragment.selection_set,
                     visited_fragments,
                     output,
+                    complementary_fields,
                 );
             }
         }
@@ -999,7 +1073,22 @@ impl<'a> CollectedAttributeNames<'a> {
         self.0
             .entry(object_or_interface)
             .or_insert(AttributeNames::All)
-            .update(field);
+            .add(field);
+    }
+
+    /// Injects complementary fields that were collected priviously in upper hierarchical levels of
+    /// the query into `self`.
+    fn populate_complementary_fields(
+        &mut self,
+        complementary_fields: &mut ComplementaryFields<'a>,
+    ) {
+        for (object_or_interface, selected_attributes) in self.0.iter_mut() {
+            if let Some(complementary_field_name) =
+                complementary_fields.remove(&object_or_interface)
+            {
+                selected_attributes.add_str(&complementary_field_name)
+            }
+        }
     }
 
     /// Consume this instance and transform it into a mapping from
@@ -1090,7 +1179,7 @@ fn filter_derived_fields(
                         None // field does not exist
                     }
                 })
-                .for_each(|col| filtered.insert(&col));
+                .for_each(|col| filtered.add_str(&col));
             filtered
         }
     }
diff --git a/store/postgres/src/relational_queries.rs b/store/postgres/src/relational_queries.rs
