all: Perform entity validation in EntityCache.set

That makes it possible to keep the SubgraphStore out of the runtime host
altogether and helps clarify that it doesn't need direct store access
except for ENS lookups

Author: lutter

core/src/subgraph/instance_manager.rs

@@ -404,7 +404,7 @@ where
         let host_builder = graph_runtime_wasm::RuntimeHostBuilder::new(
             chain.runtime_adapter(),
             self.link_resolver.cheap_clone(),
-            subgraph_store,
+            subgraph_store.ens_lookup(),
         );
 
         let features = manifest.features.clone();
@@ -1198,7 +1198,7 @@ async fn update_proof_of_indexing(
             digest: updated_proof_of_indexing,
         };
 
-        entity_cache.set(entity_key, new_poi_entity);
+        entity_cache.set(entity_key, new_poi_entity)?;
     }
 
     Ok(())

graph/src/components/store.rs

@@ -1485,8 +1485,33 @@ impl EntityCache {
         self.entity_op(key, EntityOp::Remove);
     }
 
-    pub fn set(&mut self, key: EntityKey, entity: Entity) {
-        self.entity_op(key, EntityOp::Update(entity))
+    /// Store the `entity` under the given `key`. The `entity` may be only a
+    /// partial entity; the cache will ensure partial updates get merged
+    /// with existing data. The entity will be validated against the
+    /// subgraph schema, and any errors will result in an `Err` being
+    /// returned.
+    pub fn set(&mut self, key: EntityKey, entity: Entity) -> Result<(), anyhow::Error> {
+        let is_valid = entity
+            .validate(&self.store.input_schema().document, &key)
+            .is_ok();
+
+        self.entity_op(key.clone(), EntityOp::Update(entity));
+
+        // The updates we were given are not valid by themselves; force a
+        // lookup in the database and check again with an entity that merges
+        // the existing entity with the changes
+        if !is_valid {
+            let entity = self.get(&key)?.ok_or_else(|| {
+                anyhow!(
+                    "Failed to read entity {}[{}] back from cache",
+                    key.entity_type,
+                    key.entity_id
+                )
+            })?;
+            entity.validate(&self.store.input_schema().document, &key)?;
+        }
+
+        Ok(())
     }
 
     pub fn append(&mut self, operations: Vec<EntityOperation>) {

graph/tests/entity_cache.rs

@@ -10,13 +10,27 @@ use std::sync::Arc;
 use graph::components::store::{EntityType, StoredDynamicDataSource, WritableStore};
 use graph::{
     components::store::{DeploymentId, DeploymentLocator},
-    prelude::{DeploymentHash, Entity, EntityCache, EntityKey, EntityModification, Value},
+    prelude::{anyhow, DeploymentHash, Entity, EntityCache, EntityKey, EntityModification, Value},
 };
 
 lazy_static! {
     static ref SUBGRAPH_ID: DeploymentHash = DeploymentHash::new("entity_cache").unwrap();
     static ref DEPLOYMENT: DeploymentLocator =
         DeploymentLocator::new(DeploymentId::new(-12), SUBGRAPH_ID.clone());
+    static ref SCHEMA: Arc<Schema> = Arc::new(
+        Schema::parse(
+            "
+            type Band @entity {
+                id: ID!
+                name: String!
+                founded: Int
+                label: String
+            }
+            ",
+            SUBGRAPH_ID.clone(),
+        )
+        .expect("Test schema invalid")
+    );
 }
 
 struct MockStore {
@@ -64,8 +78,17 @@ impl WritableStore for MockStore {
         unimplemented!()
     }
 
-    fn get(&self, _: &EntityKey) -> Result<Option<Entity>, StoreError> {
-        unimplemented!()
+    fn get(&self, key: &EntityKey) -> Result<Option<Entity>, StoreError> {
+        match self.get_many_res.get(&key.entity_type) {
+            Some(entities) => Ok(entities
+                .iter()
+                .find(|entity| entity.id().ok().as_ref() == Some(&key.entity_id))
+                .cloned()),
+            None => Err(StoreError::Unknown(anyhow!(
+                "nothing for type {}",
+                key.entity_type
+            ))),
+        }
     }
 
     fn transact_block_operations(
@@ -112,7 +135,7 @@ impl WritableStore for MockStore {
     }
 
     fn input_schema(&self) -> Arc<Schema> {
-        unimplemented!()
+        SCHEMA.clone()
     }
 }
 
@@ -149,13 +172,15 @@ fn insert_modifications() {
         "mogwai",
         vec![("id", "mogwai".into()), ("name", "Mogwai".into())],
     );
-    cache.set(mogwai_key.clone(), mogwai_data.clone());
+    cache.set(mogwai_key.clone(), mogwai_data.clone()).unwrap();
 
     let (sigurros_key, sigurros_data) = make_band(
         "sigurros",
         vec![("id", "sigurros".into()), ("name", "Sigur Ros".into())],
     );
-    cache.set(sigurros_key.clone(), sigurros_data.clone());
+    cache
+        .set(sigurros_key.clone(), sigurros_data.clone())
+        .unwrap();
 
     let result = cache.as_modifications();
     assert_eq!(
@@ -213,7 +238,7 @@ fn overwrite_modifications() {
             ("founded", 1995.into()),
         ],
     );
-    cache.set(mogwai_key.clone(), mogwai_data.clone());
+    cache.set(mogwai_key.clone(), mogwai_data.clone()).unwrap();
 
     let (sigurros_key, sigurros_data) = make_band(
         "sigurros",
@@ -223,7 +248,9 @@ fn overwrite_modifications() {
             ("founded", 1994.into()),
         ],
     );
-    cache.set(sigurros_key.clone(), sigurros_data.clone());
+    cache
+        .set(sigurros_key.clone(), sigurros_data.clone())
+        .unwrap();
 
     let result = cache.as_modifications();
     assert_eq!(
@@ -273,14 +300,14 @@ fn consecutive_modifications() {
             ("label", "Rock Action Records".into()),
         ],
     );
-    cache.set(update_key.clone(), update_data.clone());
+    cache.set(update_key.clone(), update_data.clone()).unwrap();
 
     // Then, just reset the "label".
     let (update_key, update_data) = make_band(
         "mogwai",
         vec![("id", "mogwai".into()), ("label", Value::Null)],
     );
-    cache.set(update_key.clone(), update_data.clone());
+    cache.set(update_key.clone(), update_data.clone()).unwrap();
 
     // We expect a single overwrite modification for the above that leaves "id"
     // and "name" untouched, sets "founded" and removes the "label" field.

runtime/test/src/common.rs

@@ -57,7 +57,6 @@ fn mock_host_exports(
         network,
         Arc::new(templates),
         Arc::new(graph_core::LinkResolver::from(IpfsClient::localhost())),
-        store,
         ens_lookup,
     )
 }

runtime/wasm/src/host.rs

@@ -8,7 +8,7 @@ use futures03::channel::oneshot::channel;
 use graph::blockchain::RuntimeAdapter;
 use graph::blockchain::{Blockchain, DataSource};
 use graph::blockchain::{HostFn, TriggerWithHandler};
-use graph::components::store::{EnsLookup, SubgraphStore};
+use graph::components::store::EnsLookup;
 use graph::components::subgraph::{MappingError, SharedProofOfIndexing};
 use graph::prelude::{
     RuntimeHost as RuntimeHostTrait, RuntimeHostBuilder as RuntimeHostBuilderTrait, *,
@@ -30,15 +30,15 @@ lazy_static! {
 pub struct RuntimeHostBuilder<C: Blockchain> {
     runtime_adapter: Arc<C::RuntimeAdapter>,
     link_resolver: Arc<dyn LinkResolver>,
-    store: Arc<dyn SubgraphStore>,
+    ens_lookup: Arc<dyn EnsLookup>,
 }
 
 impl<C: Blockchain> Clone for RuntimeHostBuilder<C> {
     fn clone(&self) -> Self {
         RuntimeHostBuilder {
             runtime_adapter: self.runtime_adapter.cheap_clone(),
             link_resolver: self.link_resolver.cheap_clone(),
-            store: self.store.cheap_clone(),
+            ens_lookup: self.ens_lookup.cheap_clone(),
         }
     }
 }
@@ -47,12 +47,12 @@ impl<C: Blockchain> RuntimeHostBuilder<C> {
     pub fn new(
         runtime_adapter: Arc<C::RuntimeAdapter>,
         link_resolver: Arc<dyn LinkResolver>,
-        store: Arc<dyn SubgraphStore>,
+        ens_lookup: Arc<dyn EnsLookup>,
     ) -> Self {
         RuntimeHostBuilder {
             runtime_adapter,
             link_resolver,
-            store,
+            ens_lookup,
         }
     }
 }
@@ -93,14 +93,13 @@ impl<C: Blockchain> RuntimeHostBuilderTrait<C> for RuntimeHostBuilder<C> {
         RuntimeHost::new(
             self.runtime_adapter.cheap_clone(),
             self.link_resolver.clone(),
-            self.store.clone(),
             network_name,
             subgraph_id,
             data_source,
             templates,
             mapping_request_sender,
             metrics,
-            self.store.ens_lookup(),
+            self.ens_lookup.cheap_clone(),
         )
     }
 }
@@ -120,7 +119,6 @@ where
     fn new(
         runtime_adapter: Arc<C::RuntimeAdapter>,
         link_resolver: Arc<dyn LinkResolver>,
-        store: Arc<dyn SubgraphStore>,
         network_name: String,
         subgraph_id: DeploymentHash,
         data_source: C::DataSource,
@@ -137,7 +135,6 @@ where
             network_name,
             templates,
             link_resolver,
-            store,
             ens_lookup,
         ));
 

runtime/wasm/src/host_exports.rs

@@ -67,7 +67,6 @@ pub struct HostExports<C: Blockchain> {
     causality_region: String,
     templates: Arc<Vec<C::DataSourceTemplate>>,
     pub(crate) link_resolver: Arc<dyn LinkResolver>,
-    store: Arc<dyn SubgraphStore>,
     ens_lookup: Arc<dyn EnsLookup>,
 }
 
@@ -78,7 +77,6 @@ impl<C: Blockchain> HostExports<C> {
         data_source_network: String,
         templates: Arc<Vec<C::DataSourceTemplate>>,
         link_resolver: Arc<dyn LinkResolver>,
-        store: Arc<dyn SubgraphStore>,
         ens_lookup: Arc<dyn EnsLookup>,
     ) -> Self {
         Self {
@@ -91,7 +89,6 @@ impl<C: Blockchain> HostExports<C> {
             data_source_network,
             templates,
             link_resolver,
-            store,
             ens_lookup,
         }
     }
@@ -168,7 +165,7 @@ impl<C: Blockchain> HostExports<C> {
         }
 
         id_insert_section.end();
-        let validation_section = stopwatch.start_section("host_export_store_set__validation");
+        let validation_section = stopwatch.start_section("host_export_store_set");
         let key = EntityKey {
             subgraph_id: self.subgraph_id.clone(),
             entity_type: EntityType::new(entity_type),
@@ -178,22 +175,9 @@ impl<C: Blockchain> HostExports<C> {
         gas.consume_host_fn(gas::STORE_SET.with_args(complexity::Linear, (&key, &data)))?;
 
         let entity = Entity::from(data);
-        let schema = self.store.input_schema(&self.subgraph_id)?;
-        let is_valid = entity.validate(&schema.document, &key).is_ok();
-        state.entity_cache.set(key.clone(), entity);
-
+        state.entity_cache.set(key.clone(), entity)?;
         validation_section.end();
-        // Validate the changes against the subgraph schema.
-        // If the set of fields we have is already valid, avoid hitting the DB.
-        if !is_valid {
-            stopwatch.start_section("host_export_store_set__post_validation");
-            let entity = state
-                .entity_cache
-                .get(&key)
-                .map_err(|e| HostExportError::Unknown(e.into()))?
-                .expect("we just stored this entity");
-            entity.validate(&schema.document, &key)?;
-        }
+
         Ok(())
     }