Make VirtualSystemDocMapper trait async (#43536)

Makes `VirtualSystemDocMapper` trait async, which we will need to await `get_document` inside the implementation.

GitOrigin-RevId: 4238f6f8aa25bcd92fa3aa1704b282630e870b32

Author: emmaling27

crates/common/src/virtual_system_mapping.rs

@@ -4,9 +4,9 @@ use std::{
 };
 
 use anyhow::Context;
+use async_trait::async_trait;
 use imbl::OrdMap;
 use semver::Version;
-use tonic::async_trait;
 use value::{
     DeveloperDocumentId,
     NamespacedTableMapping,
@@ -25,8 +25,9 @@ use crate::{
     types::IndexName,
 };
 
+#[async_trait]
 pub trait VirtualSystemDocMapper: Send + Sync {
-    fn system_to_virtual_doc(
+    async fn system_to_virtual_doc(
         &self,
         tx: &mut dyn GetDocument,
         virtual_system_mapping: &VirtualSystemMapping,
@@ -52,6 +53,7 @@ pub struct NoopDocMapper;
 
 #[cfg(any(test, feature = "testing"))]
 pub mod test_virtual_system_mapping {
+    use async_trait::async_trait;
     use value::TableMapping;
 
     use super::NoopDocMapper;
@@ -68,8 +70,9 @@ pub mod test_virtual_system_mapping {
         },
     };
 
+    #[async_trait]
     impl VirtualSystemDocMapper for NoopDocMapper {
-        fn system_to_virtual_doc(
+        async fn system_to_virtual_doc(
             &self,
             _tx: &mut dyn GetDocument,
             _virtual_system_mapping: &VirtualSystemMapping,

crates/database/src/bootstrap_model/user_facing.rs

@@ -409,26 +409,29 @@ pub async fn index_range_batch<RT: Runtime>(
 
     for (&batch_key, fetch_result) in fetch_requests.keys().zip(fetch_results) {
         let virtual_table_version = virtual_table_versions.get(&batch_key).cloned();
-        let result = fetch_result.and_then(|IndexRangeResponse { page, cursor }| {
+        let result = try {
+            let IndexRangeResponse { page, cursor } = fetch_result?;
             let developer_results = match virtual_table_version {
-                Some(version) => page
-                    .into_iter()
-                    .map(|(key, doc, ts)| {
-                        let doc =
-                            VirtualTable::new(tx).system_to_virtual_doc(doc, version.clone())?;
-                        anyhow::Ok((key, doc, ts))
-                    })
-                    .try_collect()?,
+                Some(version) => {
+                    let mut converted_documents = vec![];
+                    for (key, doc, ts) in page {
+                        let doc = VirtualTable::new(tx)
+                            .system_to_virtual_doc(doc, version.clone())
+                            .await?;
+                        converted_documents.push((key, doc, ts));
+                    }
+                    converted_documents
+                },
                 None => page
                     .into_iter()
                     .map(|(key, doc, ts)| (key, doc.to_developer(), ts))
                     .collect(),
             };
-            anyhow::Ok(DeveloperIndexRangeResponse {
+            DeveloperIndexRangeResponse {
                 page: developer_results,
                 cursor,
-            })
-        });
+            }
+        };
         results.insert(batch_key, result);
     }
     assert_eq!(results.len(), batch_size);

crates/database/src/virtual_tables/mod.rs

@@ -62,14 +62,14 @@ impl<'a, RT: Runtime> VirtualTable<'a, RT> {
         let result = self.tx.get_inner(id_, system_table_name).await?;
         match result {
             Some((doc, ts)) => {
-                let doc = self.system_to_virtual_doc(doc, version)?;
+                let doc = self.system_to_virtual_doc(doc, version).await?;
                 Ok(Some((doc, ts)))
             },
             None => Ok(None),
         }
     }
 
-    pub fn system_to_virtual_doc(
+    pub async fn system_to_virtual_doc(
         &mut self,
         doc: ResolvedDocument,
         version: Option<Version>,
@@ -91,12 +91,14 @@ impl<'a, RT: Runtime> VirtualTable<'a, RT> {
         else {
             anyhow::bail!("System document cannot be converted to a virtual document")
         };
-        mapper.system_to_virtual_doc(
-            self.tx,
-            &virtual_system_mapping,
-            doc,
-            &table_mapping,
-            version,
-        )
+        mapper
+            .system_to_virtual_doc(
+                self.tx,
+                &virtual_system_mapping,
+                doc,
+                &table_mapping,
+                version,
+            )
+            .await
     }
 }

crates/model/src/file_storage/virtual_table.rs

@@ -3,6 +3,7 @@ use std::{
     sync::LazyLock,
 };
 
+use async_trait::async_trait;
 use common::{
     document::{
         DeveloperDocument,
@@ -42,8 +43,9 @@ static MIN_NPM_VERSION_FILE_STORAGE_V2: LazyLock<Version> =
 
 pub struct FileStorageDocMapper;
 
+#[async_trait]
 impl VirtualSystemDocMapper for FileStorageDocMapper {
-    fn system_to_virtual_doc(
+    async fn system_to_virtual_doc(
         &self,
         _tx: &mut dyn GetDocument,
         virtual_system_mapping: &VirtualSystemMapping,

crates/model/src/scheduled_jobs/virtual_table.rs

@@ -4,6 +4,7 @@ use std::{
     sync::LazyLock,
 };
 
+use async_trait::async_trait;
 use common::{
     document::{
         timestamp_to_ms,
@@ -43,8 +44,9 @@ static MIN_NPM_VERSION_SCHEDULED_JOBS_V1: LazyLock<Version> =
 
 pub struct ScheduledJobsDocMapper;
 
+#[async_trait]
 impl VirtualSystemDocMapper for ScheduledJobsDocMapper {
-    fn system_to_virtual_doc(
+    async fn system_to_virtual_doc(
         &self,
         _tx: &mut dyn GetDocument,
         virtual_system_mapping: &VirtualSystemMapping,