Subgraph composition sql more entities

graph/src/blockchain/block_stream.rs

@@ -1,5 +1,5 @@
 use crate::blockchain::SubgraphFilter;
-use crate::data_source::subgraph;
+use crate::data_source::{subgraph, CausalityRegion};
 use crate::substreams::Clock;
 use crate::substreams_rpc::response::Message as SubstreamsMessage;
 use crate::substreams_rpc::BlockScopedData;
@@ -426,9 +426,19 @@ async fn scan_subgraph_triggers<C: Blockchain>(
     }
 }
 
+#[derive(Debug)]
+pub enum EntitySubgraphOperation {
+    Create,
+    Modify,
+    Delete,
+}
+
+#[derive(Debug)]
 pub struct EntityWithType {
+    pub entity_op: EntitySubgraphOperation,
     pub entity_type: EntityType,
     pub entity: Entity,
+    pub vid: i64,
 }
 
 async fn get_entities_for_range(
@@ -438,32 +448,12 @@ async fn get_entities_for_range(
     from: BlockNumber,
     to: BlockNumber,
 ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, Error> {
-    let mut entities_by_block = BTreeMap::new();
-
+    let mut entity_types = vec![];
     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);
-        }
+        entity_types.push(entity_type);
     }
-
-    Ok(entities_by_block)
+    Ok(store.get_range(entity_types, CausalityRegion::ONCHAIN, from..to)?)
 }
 
 impl<C: Blockchain> TriggersAdapterWrapper<C> {

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;
@@ -1041,9 +1042,10 @@ impl ReadStore for EmptyStore {
 
     fn get_range(
         &self,
-        _entity_type: &EntityType,
+        _entity_types: Vec<EntityType>,
+        _causality_region: CausalityRegion,
         _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;
@@ -231,9 +231,10 @@ pub trait ReadStore: Send + Sync + 'static {
     /// Looks up entities using the given store key for a range of blocks.
     fn get_range(
         &self,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
+        causality_region: CausalityRegion,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError>;
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError>;
 
     /// Reverse lookup
     fn get_derived(
@@ -259,10 +260,11 @@ impl<T: ?Sized + ReadStore> ReadStore for Arc<T> {
 
     fn get_range(
         &self,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
+        causality_region: CausalityRegion,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
-        (**self).get_range(entity_type, block_range)
+    ) -> Result<BTreeMap<BlockNumber, Vec<EntityWithType>>, StoreError> {
+        (**self).get_range(entity_types, causality_region, block_range)
     }
 
     fn get_derived(

store/postgres/src/block_range.rs

@@ -135,33 +135,37 @@ impl<'a> QueryFragment<Pg> for BlockRangeUpperBoundClause<'a> {
 /// Helper for generating SQL fragments for selecting entities in a specific block range
 #[derive(Debug, Clone, Copy)]
 pub enum EntityBlockRange {
-    Mutable(BlockRange), // TODO: check if this is a proper type here (maybe Range<BlockNumber>?)
+    Mutable((BlockRange, bool)),
     Immutable(BlockRange),
 }
 
 impl EntityBlockRange {
-    pub fn new(immutable: bool, block_range: std::ops::Range<BlockNumber>) -> Self {
+    pub fn new(
+        immutable: bool,
+        block_range: std::ops::Range<BlockNumber>,
+        is_uppper_range: bool,
+    ) -> Self {
         let start: Bound<BlockNumber> = Bound::Included(block_range.start);
         let end: Bound<BlockNumber> = Bound::Excluded(block_range.end);
         let block_range: BlockRange = BlockRange(start, end);
         if immutable {
             Self::Immutable(block_range)
         } else {
-            Self::Mutable(block_range)
+            Self::Mutable((block_range, is_uppper_range))
         }
     }
 
     /// Output SQL that matches only rows whose block range contains `block`.
     pub fn contains<'b>(&'b self, out: &mut AstPass<'_, 'b, Pg>) -> QueryResult<()> {
         out.unsafe_to_cache_prepared();
         let block_range = match self {
-            EntityBlockRange::Mutable(br) => br,
+            EntityBlockRange::Mutable((br, _)) => br,
             EntityBlockRange::Immutable(br) => br,
         };
         let BlockRange(start, finish) = block_range;
 
         self.compare_column(out);
-        out.push_sql(" >= ");
+        out.push_sql(">= ");
         match start {
             Bound::Included(block) => out.push_bind_param::<Integer, _>(block)?,
             Bound::Excluded(block) => {
@@ -170,9 +174,9 @@ impl EntityBlockRange {
             }
             Bound::Unbounded => unimplemented!(),
         };
-        out.push_sql(" AND ");
+        out.push_sql(" and");
         self.compare_column(out);
-        out.push_sql(" <= ");
+        out.push_sql("<= ");
         match finish {
             Bound::Included(block) => {
                 out.push_bind_param::<Integer, _>(block)?;
@@ -186,7 +190,13 @@ impl EntityBlockRange {
 
     pub fn compare_column(&self, out: &mut AstPass<Pg>) {
         match self {
-            EntityBlockRange::Mutable(_) => out.push_sql(" lower(block_range) "),
+            EntityBlockRange::Mutable((_, is_upper_range)) => {
+                if *is_upper_range {
+                    out.push_sql(" upper(block_range) ")
+                } else {
+                    out.push_sql(" lower(block_range) ")
+                }
+            }
             EntityBlockRange::Immutable(_) => out.push_sql(" block$ "),
         }
     }

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::{
@@ -1059,12 +1059,13 @@ impl DeploymentStore {
     pub(crate) fn get_range(
         &self,
         site: Arc<Site>,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
+        causality_region: CausalityRegion,
         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_type, block_range)
+        layout.find_range(&mut conn, entity_types, causality_region, block_range)
     }
 
     pub(crate) fn get_derived(

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::{EntitySubgraphOperation, EntityWithType};
 use graph::blockchain::BlockTime;
 use graph::cheap_clone::CheapClone;
 use graph::components::store::write::{RowGroup, WriteChunk};
@@ -52,8 +53,8 @@ use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 
 use crate::relational_queries::{
-    ConflictingEntitiesData, ConflictingEntitiesQuery, FindChangesQuery, FindDerivedQuery,
-    FindPossibleDeletionsQuery, ReturnedEntityData,
+    ConflictingEntitiesData, ConflictingEntitiesQuery, EntityDataExt, FindChangesQuery,
+    FindDerivedQuery, FindPossibleDeletionsQuery, ReturnedEntityData,
 };
 use crate::{
     primary::{Namespace, Site},
@@ -518,26 +519,105 @@ impl Layout {
     pub fn find_range(
         &self,
         conn: &mut PgConnection,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
+        causality_region: CausalityRegion,
         block_range: Range<BlockNumber>,
-    ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
-        let table = self.table_for_entity(entity_type)?;
-        let mut entities: BTreeMap<BlockNumber, Vec<Entity>> = BTreeMap::new();
-        if let Some(vec) = FindRangeQuery::new(table.as_ref(), block_range)
-            .get_results::<EntityData>(conn)
+    ) -> 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());
+            et_map.insert(et.to_string(), Arc::new(et));
+        }
+        let mut entities: BTreeMap<BlockNumber, Vec<EntityWithType>> = BTreeMap::new();
+        let lower_vec = FindRangeQuery::new(&tables, causality_region, false, block_range.clone())
+            .get_results::<EntityDataExt>(conn)
             .optional()?
-        {
-            for e in vec {
-                let block = e.clone().deserialize_block_number::<Entity>()?;
-                let en = e.deserialize_with_layout::<Entity>(self, None)?;
-                match entities.get_mut(&block) {
-                    Some(vec) => vec.push(en),
-                    None => {
-                        let _ = entities.insert(block, vec![en]);
+            .unwrap_or_default();
+        let upper_vec = FindRangeQuery::new(&tables, causality_region, true, block_range)
+            .get_results::<EntityDataExt>(conn)
+            .optional()?
+            .unwrap_or_default();
+        let mut lower_iter = lower_vec.iter().fuse().peekable();
+        let mut upper_iter = upper_vec.iter().fuse().peekable();
+        let mut lower_now = lower_iter.next();
+        let mut upper_now = upper_iter.next();
+        let mut lower = lower_now.unwrap_or(&EntityDataExt::default()).clone();
+        let mut upper = upper_now.unwrap_or(&EntityDataExt::default()).clone();
+        let transform = |ede: EntityDataExt,
+                         entity_op: EntitySubgraphOperation|
+         -> Result<(EntityWithType, BlockNumber), StoreError> {
+            let e = EntityData::new(ede.entity, ede.data);
+            let block = ede.block_number;
+            let entity_type = e.entity_type(&self.input_schema);
+            let entity = e.deserialize_with_layout::<Entity>(self, None)?;
+            let vid = ede.vid;
+            let ewt = EntityWithType {
+                entity_op,
+                entity_type,
+                entity,
+                vid,
+            };
+            Ok((ewt, block))
+        };
+        while lower_now.is_some() || upper_now.is_some() {
+            let (ewt, block) = if lower_now.is_some() {
+                if upper_now.is_some() {
+                    if lower > upper {
+                        // we have upper bound at this block, but no lower bounds at the same block so it's deletion
+                        let (ewt, block) = transform(upper, EntitySubgraphOperation::Delete)?;
+                        // advance upper
+                        upper_now = upper_iter.next();
+                        upper = upper_now.unwrap_or(&EntityDataExt::default()).clone();
+                        (ewt, block)
+                    } else if lower < upper {
+                        // we have lower bound at this block but no upper bound at the same block so its creation
+                        let (ewt, block) = transform(lower, EntitySubgraphOperation::Create)?;
+                        // advance lower
+                        lower_now = lower_iter.next();
+                        lower = lower_now.unwrap_or(&EntityDataExt::default()).clone();
+                        (ewt, block)
+                    } else {
+                        assert!(upper == lower);
+                        let (ewt, block) = transform(lower, EntitySubgraphOperation::Modify)?;
+                        // advance both
+                        lower_now = lower_iter.next();
+                        lower = lower_now.unwrap_or(&EntityDataExt::default()).clone();
+                        upper_now = upper_iter.next();
+                        upper = upper_now.unwrap_or(&EntityDataExt::default()).clone();
+                        (ewt, block)
                     }
-                };
-            }
+                } else {
+                    // we have lower bound at this block but no upper bound at the same block so its creation
+                    let (ewt, block) = transform(lower, EntitySubgraphOperation::Create)?;
+                    // advance lower
+                    lower_now = lower_iter.next();
+                    lower = lower_now.unwrap_or(&EntityDataExt::default()).clone();
+                    (ewt, block)
+                }
+            } else {
+                // we have upper bound at this block, but no lower bounds at all so it's deletion
+                assert!(upper_now.is_some());
+                let (ewt, block) = transform(upper, EntitySubgraphOperation::Delete)?;
+                // advance upper
+                upper_now = upper_iter.next();
+                upper = upper_now.unwrap_or(&EntityDataExt::default()).clone();
+                (ewt, block)
+            };
+
+            match entities.get_mut(&block) {
+                Some(vec) => vec.push(ewt),
+                None => {
+                    let _ = entities.insert(block, vec![ewt]);
+                }
+            };
         }
+
+        // sort the elements in each blocks bucket by vid
+        for (_, vec) in &mut entities {
+            vec.sort_by(|a, b| a.vid.cmp(&b.vid));
+        }
+
         Ok(entities)
     }