refactor module fetcher (#26197)

make names more explicit and allow more complex module fetcher code inline

GitOrigin-RevId: d1db3471f7856482cecfad9953ec8420d6fa4387

Author: ldanilek

crates/application/src/module_cache/mod.rs

@@ -15,6 +15,7 @@ use database::{
     Transaction,
 };
 use futures::FutureExt;
+use isolate::environment::helpers::module_loader::get_module;
 use keybroker::Identity;
 use model::{
     config::module_loader::ModuleLoader,
@@ -36,24 +37,6 @@ use value::{
 
 mod metrics;
 
-#[derive(Clone)]
-pub struct ModuleVersionFetcher<RT: Runtime> {
-    database: Database<RT>,
-    // TODO(lee) read module source from storage.
-    #[allow(unused)]
-    modules_storage: Arc<dyn Storage>,
-}
-
-impl<RT: Runtime> ModuleVersionFetcher<RT> {
-    async fn generate_value(
-        self,
-        key: (ResolvedDocumentId, ModuleVersion),
-    ) -> anyhow::Result<FullModuleSource> {
-        let mut tx = self.database.begin(Identity::system()).await?;
-        ModuleModel::new(&mut tx).get_source(key.0, key.1).await
-    }
-}
-
 #[derive(Clone)]
 pub struct ModuleCache<RT: Runtime> {
     database: Database<RT>,
@@ -97,19 +80,21 @@ impl<RT: Runtime> ModuleLoader<RT> for ModuleCache<RT> {
             .id(&MODULE_VERSIONS_TABLE)?;
         if tx.writes().has_written_to(&module_versions_table_id) {
             let source = ModuleModel::new(tx)
-                .get_source(module_metadata.id(), module_metadata.latest_version)
+                .get_source_from_db(module_metadata.id(), module_metadata.latest_version)
                 .await?;
             return Ok(Arc::new(source));
         }
 
         let key = (module_metadata.id(), module_metadata.latest_version);
-        let fetcher = ModuleVersionFetcher {
-            database: self.database.clone(),
-            modules_storage: self.modules_storage.clone(),
-        };
+        let mut cache_tx = self.database.begin(Identity::system()).await?;
+        let modules_storage = self.modules_storage.clone();
         let result = self
             .cache
-            .get(key, fetcher.generate_value(key).boxed())
+            .get(
+                key,
+                async move { get_module(&mut cache_tx, modules_storage, module_metadata).await }
+                    .boxed(),
+            )
             .await?;
         // Record read dependency on the module version so the transactions
         // read same is the same regardless if we hit the cache or not.

crates/function_runner/src/module_cache.rs

@@ -26,10 +26,7 @@ use model::{
     },
 };
 use storage::Storage;
-use value::{
-    ResolvedDocumentId,
-    TableNamespace,
-};
+use value::ResolvedDocumentId;
 
 use crate::in_memory_indexes::TransactionIngredients;
 
@@ -72,15 +69,18 @@ impl<RT: Runtime> ModuleLoader<RT> for FunctionRunnerModuleLoader<RT> {
         // this module loader was created.
         assert_eq!(tx.begin_timestamp(), self.transaction_ingredients.ts);
 
+        let namespace = tx
+            .table_mapping()
+            .tablet_namespace(module_metadata.id().table().tablet_id)?;
         // If this transaction wrote to module_versions (true for REPLs), we cannot use
         // the cache, load the module directly.
         let module_versions_table_id = tx
             .table_mapping()
-            .namespace(TableNamespace::Global)
+            .namespace(namespace)
             .id(&MODULE_VERSIONS_TABLE)?;
         if tx.writes().has_written_to(&module_versions_table_id) {
             let source = ModuleModel::new(tx)
-                .get_source(module_metadata.id(), module_metadata.latest_version)
+                .get_source_from_db(module_metadata.id(), module_metadata.latest_version)
                 .await?;
             return Ok(Arc::new(source));
         }
@@ -90,13 +90,15 @@ impl<RT: Runtime> ModuleLoader<RT> for FunctionRunnerModuleLoader<RT> {
             module_id: module_metadata.id(),
             module_version: module_metadata.latest_version,
         };
-        let transaction = self.transaction_ingredients.clone().try_into()?;
+        let mut transaction = self.transaction_ingredients.clone().try_into()?;
+        let modules_storage = self.modules_storage.clone();
         let result = self
             .cache
             .0
             .get(
                 key.clone(),
-                get_module(transaction, self.modules_storage.clone(), module_metadata).boxed(),
+                async move { get_module(&mut transaction, modules_storage, module_metadata).await }
+                    .boxed(),
             )
             .await?;
         // Record read dependency on the module version so the transactions

crates/isolate/src/environment/helpers/module_loader.rs

@@ -21,14 +21,14 @@ use crate::{
 };
 
 pub async fn get_module<RT: Runtime>(
-    mut tx: Transaction<RT>,
+    tx: &mut Transaction<RT>,
     // TODO(lee) fetch from module storage
     _modules_storage: Arc<dyn Storage>,
     module_metadata: ParsedDocument<ModuleMetadata>,
 ) -> anyhow::Result<FullModuleSource> {
     let _timer = module_load_timer();
-    let source = ModuleModel::new(&mut tx)
-        .get_source(module_metadata.id(), module_metadata.latest_version)
+    let source = ModuleModel::new(tx)
+        .get_source_from_db(module_metadata.id(), module_metadata.latest_version)
         .await?;
     Ok(source)
 }

crates/model/src/modules/mod.rs

@@ -9,7 +9,6 @@ use common::{
         CanonicalizedComponentFunctionPath,
         CanonicalizedComponentModulePath,
         ComponentDefinitionId,
-        COMPONENTS_ENABLED,
     },
     document::{
         ParsedDocument,
@@ -184,10 +183,6 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
         &mut self,
         component: ComponentDefinitionId,
     ) -> anyhow::Result<Vec<ParsedDocument<ModuleMetadata>>> {
-        // TODO(CX-6379): Remove this branch once we've made modules component-aware.
-        if !*COMPONENTS_ENABLED {
-            anyhow::ensure!(component.is_root());
-        }
         let index_query = Query::full_table_scan(MODULES_TABLE.clone(), Order::Asc);
         let mut query_stream = ResolvedQuery::new(self.tx, component.into(), index_query)?;
 
@@ -256,7 +251,11 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
             order: Order::Asc,
         };
         let module_query = Query::index_range(index_range);
-        let mut query_stream = ResolvedQuery::new(self.tx, TableNamespace::Global, module_query)?;
+        let namespace = self
+            .tx
+            .table_mapping()
+            .tablet_namespace(module_id.table().tablet_id)?;
+        let mut query_stream = ResolvedQuery::new(self.tx, namespace, module_query)?;
         let module_version: ParsedDocument<ModuleVersionMetadata> = query_stream
             .expect_at_most_one(self.tx)
             .await?
@@ -269,7 +268,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
         Ok(module_version)
     }
 
-    pub async fn get_source(
+    pub async fn get_source_from_db(
         &mut self,
         module_id: ResolvedDocumentId,
         version: ModuleVersion,
@@ -299,12 +298,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
     ) -> anyhow::Result<Option<ParsedDocument<ModuleMetadata>>> {
         let timer = get_module_metadata_timer();
 
-        // TODO(CX-6379): Remove this branch once we've made modules component-aware.
-        let is_system = if !*COMPONENTS_ENABLED {
-            path.as_root_module_path()?.is_system()
-        } else {
-            path.module_path.is_system()
-        };
+        let is_system = path.module_path.is_system();
         if is_system && !(self.tx.identity().is_admin() || self.tx.identity().is_system()) {
             anyhow::bail!(unauthorized_error("get_module"))
         }
@@ -435,14 +429,8 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
         &mut self,
         path: CanonicalizedComponentModulePath,
     ) -> anyhow::Result<Option<ParsedDocument<ModuleMetadata>>> {
-        // TODO(CX-6379): Remove this branch once we've made modules component-aware.
-        let module_path = if !*COMPONENTS_ENABLED {
-            path.as_root_module_path()?
-        } else {
-            &path.module_path
-        };
         let namespace = path.component.into();
-        let module_path = ConvexValue::try_from(module_path.as_str())?;
+        let module_path = ConvexValue::try_from(path.module_path.as_str())?;
         let index_range = IndexRange {
             index_name: MODULE_INDEX_BY_PATH.clone(),
             range: vec![IndexRangeExpression::Eq(
@@ -468,12 +456,7 @@ impl<'a, RT: Runtime> ModuleModel<'a, RT> {
         &mut self,
         path: &CanonicalizedComponentFunctionPath,
     ) -> anyhow::Result<anyhow::Result<AnalyzedFunction>> {
-        // TODO(CX-6379): Remove this branch once we've made modules component-aware.
-        let udf_path = if !*COMPONENTS_ENABLED {
-            path.as_root_udf_path()?
-        } else {
-            &path.udf_path
-        };
+        let udf_path = &path.udf_path;
         let Some(module) = self.get_metadata_for_function(path.clone()).await? else {
             let err = ModuleNotFoundError::new(udf_path.module().as_str());
             return Ok(Err(ErrorMetadata::bad_request(