graph, graphql, server: Encapsulate maps in Value in a new type

Author: lutter

graph/src/data/graphql/object_macro.rs

@@ -1,13 +1,13 @@
+use crate::data::value::Object;
 use crate::prelude::q;
 use crate::prelude::r;
-use std::collections::BTreeMap;
 use std::iter::FromIterator;
 
 /// Creates a `graphql_parser::query::Value::Object` from key/value pairs.
 /// If you don't need to determine which keys are included dynamically at runtime
 /// consider using the `object! {}` macro instead.
 pub fn object_value(data: Vec<(&str, r::Value)>) -> r::Value {
-    r::Value::Object(BTreeMap::from_iter(
+    r::Value::Object(Object::from_iter(
         data.into_iter().map(|(k, v)| (k.to_string(), v)),
     ))
 }
@@ -93,7 +93,7 @@ macro_rules! object {
                 let value = $crate::data::graphql::object_macro::IntoValue::into_value($value);
                 result.insert(stringify!($name).to_string(), value);
             )*
-            $crate::prelude::r::Value::Object(result)
+            $crate::prelude::r::Value::object(result)
         }
     };
     ($($name:ident: $value:expr),*) => {

graph/src/data/graphql/values.rs

@@ -1,8 +1,9 @@
 use anyhow::{anyhow, Error};
-use std::collections::{BTreeMap, HashMap};
+use std::collections::HashMap;
 use std::convert::TryFrom;
 use std::str::FromStr;
 
+use crate::data::value::Object;
 use crate::prelude::{r, BigInt, Entity};
 use web3::types::{H160, H256};
 
@@ -166,7 +167,7 @@ impl ValueMap for r::Value {
     }
 }
 
-impl ValueMap for &BTreeMap<String, r::Value> {
+impl ValueMap for &Object {
     fn get_required<T>(&self, key: &str) -> Result<T, Error>
     where
         T: TryFromValue,

graph/src/data/query/result.rs

@@ -1,12 +1,12 @@
 use super::error::{QueryError, QueryExecutionError};
+use crate::data::value::Object;
 use crate::prelude::{r, CacheWeight, DeploymentHash};
 use http::header::{
     ACCESS_CONTROL_ALLOW_HEADERS, ACCESS_CONTROL_ALLOW_METHODS, ACCESS_CONTROL_ALLOW_ORIGIN,
     CONTENT_TYPE,
 };
 use serde::ser::*;
 use serde::Serialize;
-use std::collections::BTreeMap;
 use std::convert::TryFrom;
 use std::sync::Arc;
 
@@ -39,7 +39,7 @@ where
     ser.end()
 }
 
-pub type Data = BTreeMap<String, r::Value>;
+pub type Data = Object;
 
 #[derive(Debug)]
 /// A collection of query results that is serialized as a single result.
@@ -270,8 +270,8 @@ impl From<Vec<QueryExecutionError>> for QueryResult {
     }
 }
 
-impl From<Data> for QueryResult {
-    fn from(val: Data) -> Self {
+impl From<Object> for QueryResult {
+    fn from(val: Object) -> Self {
         QueryResult::new(val)
     }
 }
@@ -309,7 +309,7 @@ fn multiple_data_items() {
     use serde_json::json;
 
     fn make_obj(key: &str, value: &str) -> Arc<QueryResult> {
-        let mut map = BTreeMap::new();
+        let mut map = Object::new();
         map.insert(key.to_owned(), r::Value::String(value.to_owned()));
         Arc::new(map.into())
     }

graph/src/data/value.rs

@@ -3,6 +3,78 @@ use serde::ser::{SerializeMap, SerializeSeq, Serializer};
 use serde::Serialize;
 use std::collections::BTreeMap;
 use std::convert::TryFrom;
+use std::iter::FromIterator;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct Object(BTreeMap<String, Value>);
+
+impl Object {
+    pub fn new() -> Self {
+        Self(BTreeMap::default())
+    }
+
+    pub fn get(&self, key: &str) -> Option<&Value> {
+        self.0.get(key)
+    }
+
+    pub fn remove(&mut self, key: &str) -> Option<Value> {
+        self.0.remove(key)
+    }
+
+    pub fn iter(&self) -> std::collections::btree_map::Iter<String, Value> {
+        self.into_iter()
+    }
+
+    fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    pub fn extend(&mut self, other: Object) {
+        self.0.extend(other.0)
+    }
+
+    pub fn insert(&mut self, key: String, value: Value) -> Option<Value> {
+        self.0.insert(key, value)
+    }
+}
+
+impl FromIterator<(String, Value)> for Object {
+    fn from_iter<T: IntoIterator<Item = (String, Value)>>(iter: T) -> Self {
+        Object(BTreeMap::from_iter(iter))
+    }
+}
+
+impl IntoIterator for Object {
+    type Item = (String, Value);
+
+    type IntoIter = std::collections::btree_map::IntoIter<String, Value>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
+}
+
+impl<'a> IntoIterator for &'a Object {
+    type Item = (&'a String, &'a Value);
+
+    type IntoIter = std::collections::btree_map::Iter<'a, String, Value>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.iter()
+    }
+}
+
+impl CacheWeight for Object {
+    fn indirect_weight(&self) -> usize {
+        self.0.indirect_weight()
+    }
+}
+
+impl Default for Object {
+    fn default() -> Self {
+        Self(BTreeMap::default())
+    }
+}
 
 #[derive(Debug, Clone, PartialEq)]
 pub enum Value {
@@ -13,10 +85,14 @@ pub enum Value {
     Null,
     Enum(String),
     List(Vec<Value>),
-    Object(BTreeMap<String, Value>),
+    Object(Object),
 }
 
 impl Value {
+    pub fn object(map: BTreeMap<String, Value>) -> Self {
+        Value::Object(Object(map))
+    }
+
     pub fn is_null(&self) -> bool {
         matches!(self, Value::Null)
     }
@@ -164,7 +240,7 @@ impl TryFrom<q::Value> for Value {
                     let value = Value::try_from(value)?;
                     rmap.insert(key, value);
                 }
-                Ok(Value::Object(rmap))
+                Ok(Value::object(rmap))
             }
         }
     }

graphql/src/execution/execution.rs

@@ -1,7 +1,7 @@
 use super::cache::{QueryBlockCache, QueryCache};
 use crossbeam::atomic::AtomicCell;
 use graph::{
-    data::schema::META_FIELD_NAME,
+    data::{schema::META_FIELD_NAME, value::Object},
     prelude::{s, CheapClone},
     util::timed_rw_lock::TimedMutex,
 };
@@ -11,7 +11,7 @@ use stable_hash::crypto::SetHasher;
 use stable_hash::prelude::*;
 use stable_hash::utils::stable_hash;
 use std::borrow::ToOwned;
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{HashMap, HashSet};
 use std::iter;
 use std::time::Instant;
 
@@ -128,7 +128,7 @@ impl CacheWeight for WeightedResult {
 impl Default for WeightedResult {
     fn default() -> Self {
         WeightedResult {
-            result: Arc::new(QueryResult::new(BTreeMap::default())),
+            result: Arc::new(QueryResult::new(Object::default())),
             weight: 0,
         }
     }
@@ -290,7 +290,7 @@ pub fn execute_root_selection_set_uncached(
     ctx: &ExecutionContext<impl Resolver>,
     selection_set: &q::SelectionSet,
     root_type: &s::ObjectType,
-) -> Result<BTreeMap<String, r::Value>, Vec<QueryExecutionError>> {
+) -> Result<Object, Vec<QueryExecutionError>> {
     // Split the top-level fields into introspection fields and
     // regular data fields
     let mut data_set = q::SelectionSet {
@@ -320,7 +320,7 @@ pub fn execute_root_selection_set_uncached(
 
     // If we are getting regular data, prefetch it from the database
     let mut values = if data_set.items.is_empty() && meta_items.is_empty() {
-        BTreeMap::default()
+        Object::default()
     } else {
         let initial_data = ctx.resolver.prefetch(&ctx, &data_set)?;
         data_set.items.extend(meta_items);
@@ -506,14 +506,14 @@ fn execute_selection_set_to_map<'a>(
     selection_sets: impl Iterator<Item = &'a q::SelectionSet>,
     object_type: &s::ObjectType,
     prefetched_value: Option<r::Value>,
-) -> Result<BTreeMap<String, r::Value>, Vec<QueryExecutionError>> {
+) -> Result<Object, Vec<QueryExecutionError>> {
     let mut prefetched_object = match prefetched_value {
         Some(r::Value::Object(object)) => Some(object),
         Some(_) => unreachable!(),
         None => None,
     };
     let mut errors: Vec<QueryExecutionError> = Vec::new();
-    let mut result_map: BTreeMap<String, r::Value> = BTreeMap::new();
+    let mut result_map = Object::new();
 
     // Group fields with the same response key, so we can execute them together
     let grouped_field_set = collect_fields(ctx, object_type, selection_sets);
@@ -1063,7 +1063,7 @@ pub fn coerce_argument_values<'a>(
                 if argument_def.name == "text".to_string() {
                     coerced_values.insert(
                         argument_def.name.as_str(),
-                        r::Value::Object(BTreeMap::from_iter(vec![(field.name.clone(), value)])),
+                        r::Value::Object(Object::from_iter(vec![(field.name.clone(), value)])),
                     );
                 } else {
                     coerced_values.insert(&argument_def.name, value);

graphql/src/execution/resolver.rs

@@ -95,7 +95,7 @@ pub trait Resolver: Sized + Send + Sync + 'static {
     ) -> Option<&'a s::ObjectType> {
         let concrete_type_name = match object_value {
             // All objects contain `__typename`
-            r::Value::Object(data) => match &data["__typename"] {
+            r::Value::Object(data) => match &data.get("__typename").unwrap() {
                 r::Value::String(name) => name.clone(),
                 _ => unreachable!("__typename must be a string"),
             },

graphql/src/store/prefetch.rs

@@ -3,6 +3,7 @@
 
 use anyhow::{anyhow, Error};
 use graph::constraint_violation;
+use graph::data::value::Object;
 use graph::prelude::{r, CacheWeight};
 use graph::slog::warn;
 use graph::util::cache_weight;
@@ -208,7 +209,7 @@ impl From<Node> for r::Value {
         for (key, nodes) in node.children.into_iter() {
             map.insert(format!("prefetch:{}", key), node_list_as_value(nodes));
         }
-        r::Value::Object(map)
+        r::Value::object(map)
     }
 }
 
@@ -541,8 +542,7 @@ pub fn run(
 ) -> Result<r::Value, Vec<QueryExecutionError>> {
     execute_root_selection_set(resolver, ctx, selection_set).map(|nodes| {
         result_size.observe(nodes.weight());
-        let map = BTreeMap::default();
-        r::Value::Object(nodes.into_iter().fold(map, |mut map, node| {
+        r::Value::Object(nodes.into_iter().fold(Object::default(), |mut map, node| {
             // For root nodes, we only care about the children
             for (key, nodes) in node.children.into_iter() {
                 map.insert(format!("prefetch:{}", key), node_list_as_value(nodes));

graphql/src/store/query.rs

@@ -1,6 +1,7 @@
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque};
 use std::mem::discriminant;
 
+use graph::data::value::Object;
 use graph::prelude::*;
 use graph::{components::store::EntityType, data::graphql::ObjectOrInterface};
 
@@ -125,7 +126,7 @@ fn build_filter(
 }
 
 fn build_fulltext_filter_from_object(
-    object: &BTreeMap<String, r::Value>,
+    object: &Object,
 ) -> Result<Option<EntityFilter>, QueryExecutionError> {
     object.iter().next().map_or(
         Err(QueryExecutionError::FulltextQueryRequiresFilter),
@@ -145,7 +146,7 @@ fn build_fulltext_filter_from_object(
 /// Parses a GraphQL input object into an EntityFilter, if present.
 fn build_filter_from_object(
     entity: ObjectOrInterface,
-    object: &BTreeMap<String, r::Value>,
+    object: &Object,
 ) -> Result<Option<EntityFilter>, QueryExecutionError> {
     Ok(Some(EntityFilter::And({
         object
@@ -242,7 +243,7 @@ fn build_order_by(
 }
 
 fn build_fulltext_order_by_from_object(
-    object: &BTreeMap<String, r::Value>,
+    object: &Object,
 ) -> Result<Option<(String, ValueType)>, QueryExecutionError> {
     object.iter().next().map_or(
         Err(QueryExecutionError::FulltextQueryRequiresFilter),
@@ -345,6 +346,7 @@ pub fn collect_entities_from_query_field(
 mod tests {
     use graph::{
         components::store::EntityType,
+        data::value::Object,
         prelude::s::{Directive, Field, InputValue, ObjectType, Type, Value as SchemaValue},
     };
     use graphql_parser::Pos;
@@ -724,7 +726,7 @@ mod tests {
         let mut args = default_arguments();
         args.insert(
             &whre,
-            r::Value::Object(BTreeMap::from_iter(vec![(
+            r::Value::Object(Object::from_iter(vec![(
                 "name_ends_with".to_string(),
                 r::Value::String("ello".to_string()),
             )])),

graphql/src/store/resolver.rs

@@ -2,6 +2,7 @@ use std::collections::{BTreeMap, HashMap};
 use std::result;
 use std::sync::Arc;
 
+use graph::data::value::Object;
 use graph::data::{
     graphql::{object, ObjectOrInterface},
     schema::META_FIELD_TYPE,
@@ -220,7 +221,7 @@ impl StoreResolver {
                 "__typename".to_string(),
                 r::Value::String(META_FIELD_TYPE.to_string()),
             );
-            return Ok((None, Some(r::Value::Object(map))));
+            return Ok((None, Some(r::Value::object(map))));
         }
         Ok((prefetched_object, None))
     }
@@ -328,8 +329,9 @@ impl Resolver for StoreResolver {
             // or a different field queried under the response key `_meta`.
             ErrorPolicy::Deny => {
                 let data = result.take_data();
-                let meta = data.and_then(|mut d| d.remove_entry("_meta"));
-                result.set_data(meta.map(|m| BTreeMap::from_iter(Some(m))));
+                let meta =
+                    data.and_then(|d| d.get("_meta").map(|m| ("_meta".to_string(), m.clone())));
+                result.set_data(meta.map(|m| Object::from_iter(Some(m))));
             }
             ErrorPolicy::Allow => (),
         }

graphql/src/values/coercion.rs

@@ -89,7 +89,7 @@ fn coerce_to_definition<'a>(
                         },
                     );
                 }
-                Ok(r::Value::Object(coerced_object))
+                Ok(r::Value::object(coerced_object))
             }
             _ => Err(value),
         },