From d4d4909a0fd695277c048ab016265f32d89cb843 Mon Sep 17 00:00:00 2001
From: incrypto32 <vpkrishnanand@gmail.com>
Date: Fri, 16 Aug 2024 16:57:36 +0530
Subject: [PATCH] graph: fix `get_entities_for_range` bug which didn't return
 all entity types

---
 graph/src/blockchain/block_stream.rs | 28 ++++++++++++++++------------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/graph/src/blockchain/block_stream.rs b/graph/src/blockchain/block_stream.rs
index 55e2da2da60..505353c15dd 100644
--- a/graph/src/blockchain/block_stream.rs
+++ b/graph/src/blockchain/block_stream.rs
@@ -21,7 +21,7 @@ use crate::components::store::{BlockNumber, DeploymentLocator, WritableStore};
 use crate::data::subgraph::UnifiedMappingApiVersion;
 use crate::firehose::{self, FirehoseEndpoint};
 use crate::futures03::stream::StreamExt as _;
-use crate::schema::{EntityType, InputSchema};
+use crate::schema::InputSchema;
 use crate::substreams_rpc::response::Message;
 use crate::{prelude::*, prometheus::labels};
 
@@ -433,18 +433,22 @@ async fn get_entities_for_range(
     from: BlockNumber,
     to: BlockNumber,
 ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, Error> {
-    let entity_types: Vec<EntityType> = filter
-        .entities
-        .iter()
-        .map(|e| schema.entity_type(e).unwrap())
-        .collect();
-    let mut entities = BTreeMap::new();
-    for entity_type in entity_types {
-        let range = from..to;
-        let mut entities_for_type = store.get_range(&entity_type, range)?;
-        entities.append(&mut entities_for_type);
+    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, mut entity_vec) in entity_ranges {
+            entities_by_block
+                .entry(block_number)
+                .and_modify(|existing_vec: &mut Vec<Entity>| existing_vec.append(&mut entity_vec))
+                .or_insert(entity_vec);
+        }
     }
-    Ok(entities)
+
+    Ok(entities_by_block)
 }
 
 impl<C: Blockchain> TriggersAdapterWrapper<C> {