add types to the enities

Author: Zoran Cvetkov

graph/src/blockchain/block_stream.rs

@@ -437,33 +437,7 @@ async fn get_entities_for_range(
     schema: &InputSchema,
     from: BlockNumber,
     to: BlockNumber,
-    /*
-    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, Error> {
-        let mut entities_by_block = BTreeMap::new();
-
-        for entity_name in &filter.entities {
-            let entity_type = schema.entity_type(entity_name)?;
-
-            let entity_ranges = store.get_range(&entity_type, from..to)?;
-
-            for (block_number, entity_vec) in entity_ranges {
-                let mut entity_vec = entity_vec
-                    .into_iter()
-                    .map(|e| EntityWithType {
-                        entity_type: entity_type.clone(),
-                        entity: e,
-                    })
-                    .collect();
-
-                entities_by_block
-                    .entry(block_number)
-                    .and_modify(|existing_vec: &mut Vec<EntityWithType>| {
-                        existing_vec.append(&mut entity_vec);
-                    })
-                    .or_insert(entity_vec);
-            }
-    */
-) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, Error> {
+) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, Error> {
     let mut entity_types = vec![];
     for entity_name in &filter.entities {
         let entity_type = schema.entity_type(entity_name)?;

graph/src/components/store/mod.rs

@@ -26,6 +26,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 
+use crate::blockchain::block_stream::EntityWithType;
 use crate::blockchain::{Block, BlockHash, BlockPtr};
 use crate::cheap_clone::CheapClone;
 use crate::components::store::write::EntityModification;
@@ -1043,7 +1044,7 @@ impl ReadStore for EmptyStore {
         &self,
         _entity_types: Vec<EntityType>,
         _block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
         Ok(BTreeMap::new())
     }
 

graph/src/components/store/traits.rs

@@ -6,7 +6,7 @@ use async_trait::async_trait;
 use web3::types::{Address, H256};
 
 use super::*;
-use crate::blockchain::block_stream::FirehoseCursor;
+use crate::blockchain::block_stream::{EntityWithType, FirehoseCursor};
 use crate::blockchain::{BlockTime, ChainIdentifier};
 use crate::components::metrics::stopwatch::StopwatchMetrics;
 use crate::components::server::index_node::VersionInfo;
@@ -233,7 +233,7 @@ pub trait ReadStore: Send + Sync + 'static {
         &self,
         entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError>;
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError>;
 
     /// Reverse lookup
     fn get_derived(
@@ -261,7 +261,7 @@ impl<T: ?Sized + ReadStore> ReadStore for Arc<T> {
         &self,
         entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
         (**self).get_range(entity_types, block_range)
     }
 

store/postgres/src/deployment_store.rs

@@ -4,7 +4,7 @@ use diesel::pg::PgConnection;
 use diesel::r2d2::{ConnectionManager, PooledConnection};
 use diesel::{prelude::*, sql_query};
 use graph::anyhow::Context;
-use graph::blockchain::block_stream::FirehoseCursor;
+use graph::blockchain::block_stream::{EntityWithType, FirehoseCursor};
 use graph::blockchain::BlockTime;
 use graph::components::store::write::RowGroup;
 use graph::components::store::{
@@ -1061,7 +1061,7 @@ impl DeploymentStore {
         site: Arc<Site>,
         entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
         let mut conn = self.get_conn()?;
         let layout = self.layout(&mut conn, site)?;
         layout.find_range(&mut conn, entity_types, block_range)

store/postgres/src/relational.rs

@@ -24,6 +24,7 @@ use diesel::serialize::{Output, ToSql};
 use diesel::sql_types::Text;
 use diesel::{connection::SimpleConnection, Connection};
 use diesel::{debug_query, sql_query, OptionalExtension, PgConnection, QueryResult, RunQueryDsl};
+use graph::blockchain::block_stream::EntityWithType;
 use graph::blockchain::BlockTime;
 use graph::cheap_clone::CheapClone;
 use graph::components::store::write::{RowGroup, WriteChunk};
@@ -520,23 +521,30 @@ impl Layout {
         conn: &mut PgConnection,
         entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
         let mut tables = vec![];
+        let mut et_map: HashMap<String, Arc<EntityType>> = HashMap::new();
         for et in entity_types {
-            tables.push(self.table_for_entity(&et)?.as_ref())
+            tables.push(self.table_for_entity(&et)?.as_ref());
+            et_map.insert(et.to_string(), Arc::new(et));
         }
-        let mut entities: BTreeMap<BlockNumber, Vec<Entity>> = BTreeMap::new();
+        let mut entities: BTreeMap<BlockNumber, Vec<EntityWithType>> = BTreeMap::new();
         if let Some(vec) = FindRangeQuery::new(&tables, block_range)
             .get_results::<EntityData>(conn)
             .optional()?
         {
             for e in vec {
                 let block = e.clone().deserialize_block_number::<Entity>()?;
-                let en = e.deserialize_with_layout::<Entity>(self, None)?;
+                let entity_type = e.entity_type(&self.input_schema);
+                let entity = e.deserialize_with_layout::<Entity>(self, None)?;
+                let ewt = EntityWithType {
+                    entity_type,
+                    entity,
+                };
                 match entities.get_mut(&block) {
-                    Some(vec) => vec.push(en),
+                    Some(vec) => vec.push(ewt),
                     None => {
-                        let _ = entities.insert(block, vec![en]);
+                        let _ = entities.insert(block, vec![ewt]);
                     }
                 };
             }

store/postgres/src/writable.rs

@@ -5,7 +5,7 @@ use std::sync::{Mutex, RwLock, TryLockError as RwLockError};
 use std::time::Instant;
 use std::{collections::BTreeMap, sync::Arc};
 
-use graph::blockchain::block_stream::FirehoseCursor;
+use graph::blockchain::block_stream::{EntityWithType, FirehoseCursor};
 use graph::blockchain::BlockTime;
 use graph::components::store::{Batch, DeploymentCursorTracker, DerivedEntityQuery, ReadStore};
 use graph::constraint_violation;
@@ -356,7 +356,7 @@ impl SyncStore {
         &self,
         entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
         retry::forever(&self.logger, "get_range", || {
             self.writable.get_range(
                 self.site.cheap_clone(),
@@ -1235,7 +1235,7 @@ impl Queue {
         &self,
         entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
         self.store.get_range(entity_types, block_range)
     }
 
@@ -1455,7 +1455,7 @@ impl Writer {
         &self,
         entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
         match self {
             Writer::Sync(store) => store.get_range(entity_types, block_range),
             Writer::Async { queue, .. } => queue.get_range(entity_types, block_range),
@@ -1594,7 +1594,7 @@ impl ReadStore for WritableStore {
         &self,
         entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
         self.writer.get_range(entity_types, block_range)
     }
 

store/test-store/tests/graph/entity_cache.rs

@@ -1,4 +1,4 @@
-use graph::blockchain::block_stream::FirehoseCursor;
+use graph::blockchain::block_stream::{EntityWithType, FirehoseCursor};
 use graph::blockchain::BlockTime;
 use graph::components::store::{
     DeploymentCursorTracker, DerivedEntityQuery, GetScope, LoadRelatedRequest, ReadStore,
@@ -71,7 +71,7 @@ impl ReadStore for MockStore {
         &self,
         _entity_types: Vec<EntityType>,
         _block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
         todo!()
     }