store: Asyncify VidBatcher.for_copy and for_prune

Author: lutter

store/postgres/src/copy.rs

@@ -336,7 +336,7 @@ impl TableState {
         dst: Arc<Table>,
         target_block: &BlockPtr,
     ) -> Result<Self, StoreError> {
-        let vid_range = VidRange::for_copy(conn, &src, target_block)?;
+        let vid_range = VidRange::for_copy(conn, &src, target_block).await?;
         let batcher =
             VidBatcher::load(conn, &src_layout.site.namespace, src.as_ref(), vid_range).await?;
         Ok(Self {

store/postgres/src/deployment_store.rs

@@ -1091,7 +1091,7 @@ impl DeploymentStore {
         let mut conn = self.get_conn().await?;
         let layout = self.layout(&mut conn, site).await?;
 
-        layout.find_many(&mut conn, ids_for_type, block)
+        layout.find_many(&mut conn, ids_for_type, block).await
     }
 
     pub(crate) async fn get_range(
@@ -1103,7 +1103,9 @@ impl DeploymentStore {
     ) -> Result<BTreeMap<BlockNumber, Vec<EntitySourceOperation>>, StoreError> {
         let mut conn = self.get_conn().await?;
         let layout = self.layout(&mut conn, site).await?;
-        layout.find_range(&mut conn, entity_types, causality_region, block_range)
+        layout
+            .find_range(&mut conn, entity_types, causality_region, block_range)
+            .await
     }
 
     pub(crate) async fn get_derived(

store/postgres/src/relational.rs

@@ -481,7 +481,7 @@ impl Layout {
     }
 
     // An optimization when looking up multiple entities, it will generate a single sql query using `UNION ALL`.
-    pub fn find_many(
+    pub async fn find_many(
         &self,
         conn: &mut PgConnection,
         ids_for_type: &BTreeMap<(EntityType, CausalityRegion), IdList>,
@@ -517,7 +517,7 @@ impl Layout {
         Ok(entities)
     }
 
-    pub fn find_range(
+    pub async fn find_range(
         &self,
         conn: &mut PgConnection,
         entity_types: Vec<EntityType>,

store/postgres/src/relational/prune.rs

@@ -105,7 +105,7 @@ impl TablePair {
         let column_list = Arc::new(self.column_list());
 
         // Determine the last vid that we need to copy
-        let range = VidRange::for_prune(conn, &self.src, earliest_block, final_block)?;
+        let range = VidRange::for_prune(conn, &self.src, earliest_block, final_block).await?;
         let mut batcher = VidBatcher::load(conn, &self.src_nsp, &self.src, range).await?;
         tracker.start_copy_final(conn, &self.src, range).await?;
 
@@ -173,7 +173,7 @@ impl TablePair {
         let column_list = Arc::new(self.column_list());
 
         // Determine the last vid that we need to copy
-        let range = VidRange::for_prune(conn, &self.src, final_block + 1, BLOCK_NUMBER_MAX)?;
+        let range = VidRange::for_prune(conn, &self.src, final_block + 1, BLOCK_NUMBER_MAX).await?;
         let mut batcher = VidBatcher::load(conn, &self.src.nsp, &self.src, range).await?;
         tracker.start_copy_nonfinal(conn, &self.src, range).await?;
 
@@ -478,7 +478,7 @@ impl Layout {
                 PruningStrategy::Delete => {
                     // Delete all entity versions whose range was closed
                     // before `req.earliest_block`
-                    let range = VidRange::for_prune(conn, &table, 0, req.earliest_block)?;
+                    let range = VidRange::for_prune(conn, &table, 0, req.earliest_block).await?;
                     let mut batcher =
                         VidBatcher::load(conn, &self.site.namespace, &table, range).await?;
 

store/postgres/src/vid_batcher.rs

@@ -274,7 +274,7 @@ impl VidRange {
     }
 
     /// Return the full range of `vid` values in the table `src`
-    pub fn for_copy(
+    pub async fn for_copy(
         conn: &mut PgConnection,
         src: &Table,
         target_block: &BlockPtr,
@@ -303,7 +303,7 @@ impl VidRange {
     /// Return the first and last vid of any entity that is visible in the
     /// block range from `first_block` (inclusive) to `last_block`
     /// (exclusive)
-    pub fn for_prune(
+    pub async fn for_prune(
         conn: &mut PgConnection,
         src: &Table,
         first_block: BlockNumber,

store/test-store/tests/postgres/relational_bytes.rs

@@ -300,6 +300,7 @@ async fn find_many() {
 
         let entities = layout
             .find_many(conn, &id_map, BLOCK_NUMBER_MAX)
+            .await
             .expect("Failed to read many things");
         assert_eq!(2, entities.len());