request all entity types in a single SQL query

Author: Zoran Cvetkov

graph/src/blockchain/block_stream.rs

@@ -437,33 +437,39 @@ 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();
-
+    /*
+    ) -> 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> {
+    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, from..to)?)
 }
 
 impl<C: Blockchain> TriggersAdapterWrapper<C> {

graph/src/components/store/mod.rs

@@ -1041,7 +1041,7 @@ impl ReadStore for EmptyStore {
 
     fn get_range(
         &self,
-        _entity_type: &EntityType,
+        _entity_types: Vec<EntityType>,
         _block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
         Ok(BTreeMap::new())

graph/src/components/store/traits.rs

@@ -231,7 +231,7 @@ 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>,
         block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError>;
 
@@ -259,10 +259,10 @@ impl<T: ?Sized + ReadStore> ReadStore for Arc<T> {
 
     fn get_range(
         &self,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
-        (**self).get_range(entity_type, block_range)
+        (**self).get_range(entity_types, block_range)
     }
 
     fn get_derived(

store/postgres/src/block_range.rs

@@ -161,7 +161,7 @@ impl EntityBlockRange {
         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 +170,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)?;

store/postgres/src/deployment_store.rs

@@ -1059,12 +1059,12 @@ impl DeploymentStore {
     pub(crate) fn get_range(
         &self,
         site: Arc<Site>,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, 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, block_range)
     }
 
     pub(crate) fn get_derived(

store/postgres/src/relational.rs

@@ -518,12 +518,15 @@ impl Layout {
     pub fn find_range(
         &self,
         conn: &mut PgConnection,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
-        let table = self.table_for_entity(entity_type)?;
+        let mut tables = vec![];
+        for et in entity_types {
+            tables.push(self.table_for_entity(&et)?.as_ref())
+        }
         let mut entities: BTreeMap<BlockNumber, Vec<Entity>> = BTreeMap::new();
-        if let Some(vec) = FindRangeQuery::new(table.as_ref(), block_range)
+        if let Some(vec) = FindRangeQuery::new(&tables, block_range)
             .get_results::<EntityData>(conn)
             .optional()?
         {

store/postgres/src/relational_queries.rs

@@ -2008,37 +2008,64 @@ impl<'a, Conn> RunQueryDsl<Conn> for FindQuery<'a> {}
 
 #[derive(Debug, Clone)]
 pub struct FindRangeQuery<'a> {
-    table: &'a Table,
-    eb_range: EntityBlockRange,
+    tables: &'a Vec<&'a Table>,
+    imm_range: EntityBlockRange,
+    mut_range: EntityBlockRange,
 }
 
 impl<'a> FindRangeQuery<'a> {
-    pub fn new(table: &'a Table, block_range: Range<BlockNumber>) -> Self {
-        let eb_range = EntityBlockRange::new(table.immutable, block_range);
-        Self { table, eb_range }
+    pub fn new(tables: &'a Vec<&Table>, block_range: Range<BlockNumber>) -> Self {
+        let imm_range = EntityBlockRange::new(true, block_range.clone());
+        let mut_range = EntityBlockRange::new(false, block_range);
+        Self {
+            tables,
+            imm_range,
+            mut_range,
+        }
     }
 }
 
 impl<'a> QueryFragment<Pg> for FindRangeQuery<'a> {
     fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> {
         out.unsafe_to_cache_prepared();
 
-        // Generate
-        //    select '..' as entity, to_jsonb(e.*) as data
-        //      from schema.table e where id = $1
-        out.push_sql("select ");
-        out.push_bind_param::<Text, _>(self.table.object.as_str())?;
-        out.push_sql(" as entity, to_jsonb(e.*) as data\n");
-        out.push_sql("  from ");
-        out.push_sql(self.table.qualified_name.as_str());
-        out.push_sql(" e\n where ");
-        // TODO: do we need to care about it?
-        // if self.table.has_causality_region {
-        //     out.push_sql("causality_region = ");
-        //     out.push_bind_param::<Integer, _>(&self.key.causality_region)?;
-        //     out.push_sql(" and ");
-        // }
-        self.eb_range.contains(&mut out)
+        let mut iter = self.tables.iter().peekable();
+        while let Some(table) = iter.next() {
+            // Generate
+            //    select '..' as entity, to_jsonb(e.*) as data, block$ as block_number
+            //      from schema.table e where id = $1
+            out.push_sql("select ");
+            out.push_bind_param::<Text, _>(table.object.as_str())?;
+            out.push_sql(" as entity, to_jsonb(e.*) as data,");
+            if table.immutable {
+                self.imm_range.compare_column(&mut out)
+            } else {
+                self.mut_range.compare_column(&mut out)
+            }
+            out.push_sql("as block_number\n");
+            out.push_sql("  from ");
+            out.push_sql(table.qualified_name.as_str());
+            out.push_sql(" e\n  where");
+            // TODO: do we need to care about it?
+            // if self.table.has_causality_region {
+            //     out.push_sql("causality_region = ");
+            //     out.push_bind_param::<Integer, _>(&self.key.causality_region)?;
+            //     out.push_sql(" and ");
+            // }
+            if table.immutable {
+                self.imm_range.contains(&mut out)?;
+            } else {
+                self.mut_range.contains(&mut out)?;
+            }
+            // more elements left?
+            if iter.peek().is_some() {
+                out.push_sql("\nunion all\n"); // with the next
+            } else {
+                out.push_sql("\norder by block_number"); // on the last
+            }
+        }
+
+        Ok(())
     }
 }
 

store/postgres/src/writable.rs

@@ -354,12 +354,15 @@ impl SyncStore {
 
     fn get_range(
         &self,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
         retry::forever(&self.logger, "get_range", || {
-            self.writable
-                .get_range(self.site.cheap_clone(), entity_type, block_range.clone())
+            self.writable.get_range(
+                self.site.cheap_clone(),
+                entity_types.clone(),
+                block_range.clone(),
+            )
         })
     }
 
@@ -1230,10 +1233,10 @@ impl Queue {
 
     fn get_range(
         &self,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
-        self.store.get_range(entity_type, block_range)
+        self.store.get_range(entity_types, block_range)
     }
 
     fn get_derived(
@@ -1450,12 +1453,12 @@ impl Writer {
 
     fn get_range(
         &self,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
         match self {
-            Writer::Sync(store) => store.get_range(entity_type, block_range),
-            Writer::Async { queue, .. } => queue.get_range(entity_type, block_range),
+            Writer::Sync(store) => store.get_range(entity_types, block_range),
+            Writer::Async { queue, .. } => queue.get_range(entity_types, block_range),
         }
     }
 
@@ -1589,10 +1592,10 @@ impl ReadStore for WritableStore {
 
     fn get_range(
         &self,
-        entity_type: &EntityType,
+        entity_types: Vec<EntityType>,
         block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
-        self.writer.get_range(entity_type, block_range)
+        self.writer.get_range(entity_types, block_range)
     }
 
     fn get_derived(

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

@@ -69,7 +69,7 @@ impl ReadStore for MockStore {
 
     fn get_range(
         &self,
-        _entity_type: &EntityType,
+        _entity_types: Vec<EntityType>,
         _block_range: Range<BlockNumber>,
     ) -> Result<BTreeMap<BlockNumber, Vec<Entity>>, StoreError> {
         todo!()