graphql, store: Prefetch q::Value, not Entity, to avoid conversion

Author: leoyvens

graph/src/components/store.rs

@@ -826,6 +826,11 @@ pub trait Store: Send + Sync + 'static {
     /// Queries the store for entities that match the store query.
     fn find(&self, query: EntityQuery) -> Result<Vec<Entity>, QueryExecutionError>;
 
+    fn find_query_values(
+        &self,
+        query: EntityQuery,
+    ) -> Result<Vec<BTreeMap<String, graphql_parser::query::Value>>, QueryExecutionError>;
+
     /// Queries the store for a single entity matching the store query.
     fn find_one(&self, query: EntityQuery) -> Result<Option<Entity>, QueryExecutionError>;
 
@@ -1250,6 +1255,13 @@ impl Store for MockStore {
         unimplemented!()
     }
 
+    fn find_query_values(
+        &self,
+        _: EntityQuery,
+    ) -> Result<Vec<BTreeMap<String, graphql_parser::query::Value>>, QueryExecutionError> {
+        unimplemented!()
+    }
+
     fn find_one(&self, _query: EntityQuery) -> Result<Option<Entity>, QueryExecutionError> {
         unimplemented!()
     }

graphql/src/store/prefetch.rs

@@ -5,15 +5,14 @@ use graphql_parser::query as q;
 use graphql_parser::schema as s;
 use lazy_static::lazy_static;
 use std::collections::{BTreeMap, HashMap, HashSet};
-use std::ops::Deref;
 use std::rc::Rc;
 use std::time::Instant;
 
 use graph::data::graphql::ext::ObjectTypeExt;
 use graph::prelude::{
-    BlockNumber, ChildMultiplicity, Entity, EntityCollection, EntityFilter, EntityLink,
-    EntityOrder, EntityWindow, Logger, ParentLink, QueryExecutionError, Schema, Store,
-    Value as StoreValue, WindowAttribute,
+    BlockNumber, ChildMultiplicity, EntityCollection, EntityFilter, EntityLink, EntityOrder,
+    EntityWindow, Logger, ParentLink, QueryExecutionError, Schema, Store, Value as StoreValue,
+    WindowAttribute,
 };
 
 use crate::execution::{ExecutionContext, ObjectOrInterface, Resolver};
@@ -87,7 +86,7 @@ impl From<Option<q::TypeCondition>> for TypeCondition {
 /// has an entry mapping the response key to the list of nodes.
 #[derive(Debug, Clone)]
 struct Node {
-    entity: Entity,
+    entity: BTreeMap<String, q::Value>,
     /// We are using an `Rc` here for two reasons: it allows us to defer
     /// copying objects until the end, when converting to `q::Value` forces
     /// us to copy any child that is referenced by multiple parents. It also
@@ -97,8 +96,8 @@ struct Node {
     children: BTreeMap<String, Vec<Rc<Node>>>,
 }
 
-impl From<Entity> for Node {
-    fn from(entity: Entity) -> Self {
+impl From<BTreeMap<String, q::Value>> for Node {
+    fn from(entity: BTreeMap<String, q::Value>) -> Self {
         Node {
             entity,
             children: BTreeMap::default(),
@@ -137,7 +136,7 @@ fn is_root_node(nodes: &Vec<Node>) -> bool {
 
 fn make_root_node() -> Vec<Node> {
     vec![Node {
-        entity: Entity::new(),
+        entity: BTreeMap::new(),
         children: BTreeMap::default(),
     }]
 }
@@ -148,23 +147,44 @@ fn make_root_node() -> Vec<Node> {
 /// with any field of the entity.
 impl From<Node> for q::Value {
     fn from(node: Node) -> Self {
-        let mut map: BTreeMap<_, _> = node.entity.into();
+        let mut map = node.entity;
         for (key, nodes) in node.children.into_iter() {
             map.insert(format!("prefetch:{}", key), node_list_as_value(nodes));
         }
         q::Value::Object(map)
     }
 }
 
-impl Deref for Node {
-    type Target = Entity;
+trait ValueExt {
+    fn as_str(&self) -> Option<&str>;
+}
 
-    fn deref(&self) -> &Self::Target {
-        &self.entity
+impl ValueExt for q::Value {
+    fn as_str(&self) -> Option<&str> {
+        match self {
+            q::Value::String(s) => Some(s),
+            _ => None,
+        }
     }
 }
 
 impl Node {
+    fn id(&self) -> Result<String, graph::prelude::failure::Error> {
+        match self.get("id") {
+            None => Err(graph::prelude::failure::format_err!(
+                "Entity is missing an `id` attribute"
+            )),
+            Some(q::Value::String(s)) => Ok(s.to_owned()),
+            _ => Err(graph::prelude::failure::format_err!(
+                "Entity has non-string `id` attribute"
+            )),
+        }
+    }
+
+    fn get(&self, key: &str) -> Option<&q::Value> {
+        self.entity.get(key)
+    }
+
     fn typename(&self) -> &str {
         self.get("__typename")
             .expect("all entities have a __typename")
@@ -288,7 +308,7 @@ impl<'a> JoinCond<'a> {
                             .filter_map(|(id, node)| {
                                 node.get(*child_field)
                                     .and_then(|value| match value {
-                                        StoreValue::List(values) => {
+                                        q::Value::List(values) => {
                                             let values: Vec<_> = values
                                                 .into_iter()
                                                 .filter_map(|value| {
@@ -382,7 +402,7 @@ impl<'a> Join<'a> {
                 .get("g$parent_id")
                 .expect("the query that produces 'child' ensures there is always a g$parent_id")
             {
-                StoreValue::String(key) => grouped.entry(key).or_default().push(child.clone()),
+                q::Value::String(key) => grouped.entry(&key).or_default().push(child.clone()),
                 _ => unreachable!("the parent_id returned by the query is always a string"),
             }
         }
@@ -871,6 +891,6 @@ fn fetch<S: Store>(
     }
 
     store
-        .find(query)
+        .find_query_values(query)
         .map(|entities| entities.into_iter().map(|entity| entity.into()).collect())
 }

mock/src/store.rs

@@ -96,6 +96,13 @@ impl Store for MockStore {
         unimplemented!()
     }
 
+    fn find_query_values(
+        &self,
+        _: EntityQuery,
+    ) -> Result<Vec<BTreeMap<String, graphql_parser::query::Value>>, QueryExecutionError> {
+        unimplemented!()
+    }
+
     fn find_one(&self, _query: EntityQuery) -> Result<Option<Entity>, QueryExecutionError> {
         unimplemented!()
     }

store/postgres/src/relational_queries.rs

@@ -185,6 +185,14 @@ impl FromEntityData for Entity {
     }
 }
 
+impl FromEntityData for BTreeMap<String, graphql_parser::query::Value> {
+    type Value = graphql_parser::query::Value;
+
+    fn insert_entity_data(&mut self, key: String, v: Self::Value) {
+        self.insert(key, v);
+    }
+}
+
 pub trait FromColumnValue: Sized {
     fn is_null(&self) -> bool;
 
@@ -196,6 +204,79 @@ pub trait FromColumnValue: Sized {
     ) -> Result<Self, StoreError>;
 }
 
+impl FromColumnValue for graphql_parser::query::Value {
+    fn is_null(&self) -> bool {
+        self == &graphql_parser::query::Value::Null
+    }
+
+    fn from_string(s: String) -> Self {
+        graphql_parser::query::Value::String(s)
+    }
+
+    fn from_column_value(
+        column_type: &ColumnType,
+        json: serde_json::Value,
+    ) -> Result<graphql_parser::query::Value, StoreError> {
+        use graphql_parser::query::Value as q;
+        use serde_json::Value as j;
+        // Many possible conversion errors are already caught by how
+        // we define the schema; for example, we can only get a NULL for
+        // a column that is actually nullable
+        match (json, column_type) {
+            (j::Null, _) => Ok(q::Null),
+            (j::Bool(b), _) => Ok(q::Boolean(b)),
+            (j::Number(number), ColumnType::Int) => match number.as_i64() {
+                Some(i) => i32::try_from(i).map(|i| q::Int(i.into())).map_err(|e| {
+                    StoreError::Unknown(format_err!("failed to convert {} to Int: {}", number, e))
+                }),
+                None => Err(StoreError::Unknown(format_err!(
+                    "failed to convert {} to Int",
+                    number
+                ))),
+            },
+            (j::Number(number), ColumnType::BigDecimal) => {
+                let s = number.to_string();
+                scalar::BigDecimal::from_str(s.as_str())
+                    .map(|d| q::String(d.to_string()))
+                    .map_err(|e| {
+                        StoreError::Unknown(format_err!(
+                            "failed to convert {} to BigDecimal: {}",
+                            number,
+                            e
+                        ))
+                    })
+            }
+            (j::Number(number), ColumnType::BigInt) => Ok(q::String(number.to_string())),
+            (j::Number(number), column_type) => Err(StoreError::Unknown(format_err!(
+                "can not convert number {} to {:?}",
+                number,
+                column_type
+            ))),
+            (j::String(s), ColumnType::String) | (j::String(s), ColumnType::Enum(_)) => {
+                Ok(q::String(s))
+            }
+            (j::String(s), ColumnType::Bytes) => {
+                Ok(q::String(format!("0x{}", s.trim_start_matches("\\x"))))
+            }
+            (j::String(s), ColumnType::BytesId) => Ok(q::String(bytes_as_str(&s))),
+            (j::String(s), column_type) => Err(StoreError::Unknown(format_err!(
+                "can not convert string {} to {:?}",
+                s,
+                column_type
+            ))),
+            (j::Array(values), _) => Ok(q::List(
+                values
+                    .into_iter()
+                    .map(|v| Self::from_column_value(column_type, v))
+                    .collect::<Result<Vec<_>, _>>()?,
+            )),
+            (j::Object(_), _) => {
+                unimplemented!("objects as entity attributes are not needed/supported")
+            }
+        }
+    }
+}
+
 impl FromColumnValue for graph::prelude::Value {
     fn is_null(&self) -> bool {
         self == &Value::Null

store/postgres/src/store.rs

@@ -465,14 +465,16 @@ impl Store {
                 entity_ids: mut expected_entity_ids,
             } => {
                 // Execute query
+                let actual_entities =
                     self.execute_query::<Entity>(conn, query.clone())
-                    format_err!(
-                        "AbortUnless ({}): query execution error: {:?}, {}",
-                        description,
-                        query,
-                        e
-                    )
-                })?;
+                        .map_err(|e| {
+                            format_err!(
+                                "AbortUnless ({}): query execution error: {:?}, {}",
+                                description,
+                                query,
+                                e
+                            )
+                        })?;
 
                 // Extract IDs from entities
                 let mut actual_entity_ids: Vec<String> = actual_entities
@@ -1018,6 +1020,16 @@ impl StoreTrait for Store {
         self.execute_query(&conn, query)
     }
 
+    fn find_query_values(
+        &self,
+        query: EntityQuery,
+    ) -> Result<Vec<BTreeMap<String, graphql_parser::query::Value>>, QueryExecutionError> {
+        let conn = self
+            .get_entity_conn(&query.subgraph_id)
+            .map_err(|e| QueryExecutionError::StoreError(e.into()))?;
+        self.execute_query(&conn, query)
+    }
+
     fn find_one(&self, mut query: EntityQuery) -> Result<Option<Entity>, QueryExecutionError> {
         query.range = EntityRange::first(1);
 

store/postgres/tests/relational.rs

@@ -538,7 +538,7 @@ fn count_scalar_entities(conn: &PgConnection, layout: &Layout) -> usize {
     ]);
     let collection = EntityCollection::All(vec!["Scalar".to_owned()]);
     layout
-        .query(
+        .query::<Entity>(
             &*LOGGER,
             &conn,
             collection,
@@ -638,7 +638,7 @@ fn test_find(expected_entity_ids: Vec<&str>, query: EntityQuery) {
 
         let unordered = matches!(query.order, EntityOrder::Unordered);
         let entities = layout
-            .query(
+            .query::<Entity>(
                 &*LOGGER,
                 conn,
                 query.collection,
@@ -1374,7 +1374,7 @@ fn text_find(expected_entity_ids: Vec<&str>, filter: EntityFilter) {
         let query = query(vec!["Ferret"]).filter(filter).asc("id");
 
         let entities = layout
-            .query(
+            .query::<Entity>(
                 &*LOGGER,
                 conn,
                 query.collection,

store/postgres/tests/relational_bytes.rs

@@ -385,7 +385,7 @@ fn make_thing_tree(conn: &PgConnection, layout: &Layout) -> (Entity, Entity, Ent
 fn query() {
     fn fetch(conn: &PgConnection, layout: &Layout, coll: EntityCollection) -> Vec<String> {
         layout
-            .query(
+            .query::<Entity>(
                 &*LOGGER,
                 conn,
                 coll,