extract data fetch orchestration from table provider (#38)

Author: zfarrell

src/datafetch/mod.rs

@@ -1,9 +1,11 @@
 mod error;
 mod fetcher;
 pub mod native;
+mod orchestrator;
 mod types;
 
 pub use error::DataFetchError;
 pub use fetcher::DataFetcher;
 pub use native::{NativeFetcher, StreamingParquetWriter};
+pub use orchestrator::FetchOrchestrator;
 pub use types::{deserialize_arrow_schema, ColumnMetadata, TableMetadata};

src/datafetch/orchestrator.rs

@@ -0,0 +1,84 @@
+use anyhow::Result;
+use std::sync::Arc;
+
+use super::native::StreamingParquetWriter;
+use super::DataFetcher;
+use crate::catalog::CatalogManager;
+use crate::source::Source;
+use crate::storage::StorageManager;
+
+/// Orchestrates the full table fetch workflow: fetch from source → write to storage → update catalog.
+#[derive(Debug)]
+pub struct FetchOrchestrator {
+    fetcher: Arc<dyn DataFetcher>,
+    storage: Arc<dyn StorageManager>,
+    catalog: Arc<dyn CatalogManager>,
+}
+
+impl FetchOrchestrator {
+    pub fn new(
+        fetcher: Arc<dyn DataFetcher>,
+        storage: Arc<dyn StorageManager>,
+        catalog: Arc<dyn CatalogManager>,
+    ) -> Self {
+        Self {
+            fetcher,
+            storage,
+            catalog,
+        }
+    }
+
+    /// Fetch table data from source, write to cache storage, and update catalog metadata.
+    ///
+    /// Returns the URL of the cached parquet file.
+    pub async fn cache_table(
+        &self,
+        source: &Source,
+        connection_id: i32,
+        schema_name: &str,
+        table_name: &str,
+    ) -> Result<String> {
+        // Prepare cache write location
+        let write_path = self
+            .storage
+            .prepare_cache_write(connection_id, schema_name, table_name);
+
+        // Create writer
+        let mut writer = StreamingParquetWriter::new(write_path.clone());
+
+        // Fetch the table data into writer
+        self.fetcher
+            .fetch_table(
+                source,
+                None, // catalog
+                schema_name,
+                table_name,
+                &mut writer,
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("Failed to fetch table: {}", e))?;
+
+        // Close writer
+        writer
+            .close()
+            .map_err(|e| anyhow::anyhow!("Failed to close writer: {}", e))?;
+
+        // Finalize cache write (uploads to S3 if needed, returns URL)
+        let parquet_url = self
+            .storage
+            .finalize_cache_write(&write_path, connection_id, schema_name, table_name)
+            .await
+            .map_err(|e| anyhow::anyhow!("Failed to finalize cache write: {}", e))?;
+
+        // Update catalog with new path
+        if let Ok(Some(info)) = self
+            .catalog
+            .get_table(connection_id, schema_name, table_name)
+            .await
+        {
+            let _ = self.catalog.update_table_sync(info.id, &parquet_url).await;
+        }
+
+        Ok(parquet_url)
+    }
+}

src/datafusion/catalog_provider.rs

@@ -1,9 +1,8 @@
 use super::block_on;
 use super::schema_provider::RivetSchemaProvider;
 use crate::catalog::CatalogManager;
-use crate::datafetch::{DataFetcher, NativeFetcher};
+use crate::datafetch::FetchOrchestrator;
 use crate::source::Source;
-use crate::storage::StorageManager;
 use async_trait::async_trait;
 use datafusion::catalog::{CatalogProvider, SchemaProvider};
 use std::collections::HashMap;
@@ -17,9 +16,8 @@ pub struct RivetCatalogProvider {
     connection_name: String,
     source: Arc<Source>,
     catalog: Arc<dyn CatalogManager>,
+    orchestrator: Arc<FetchOrchestrator>,
     schemas: Arc<RwLock<HashMap<String, Arc<dyn SchemaProvider>>>>,
-    storage: Arc<dyn StorageManager>,
-    fetcher: Arc<dyn DataFetcher>,
 }
 
 impl RivetCatalogProvider {
@@ -28,16 +26,15 @@ impl RivetCatalogProvider {
         connection_name: String,
         source: Arc<Source>,
         catalog: Arc<dyn CatalogManager>,
-        storage: Arc<dyn StorageManager>,
+        orchestrator: Arc<FetchOrchestrator>,
     ) -> Self {
         Self {
             connection_id,
             connection_name,
             source,
             catalog,
+            orchestrator,
             schemas: Arc::new(RwLock::new(HashMap::new())),
-            storage,
-            fetcher: Arc::new(NativeFetcher::new()),
         }
     }
 
@@ -66,8 +63,7 @@ impl RivetCatalogProvider {
             schema_name.to_string(),
             self.source.clone(),
             self.catalog.clone(),
-            self.storage.clone(),
-            self.fetcher.clone(),
+            self.orchestrator.clone(),
         )) as Arc<dyn SchemaProvider>;
 
         schemas.insert(schema_name.to_string(), schema_provider.clone());

src/datafusion/lazy_table_provider.rs

@@ -1,4 +1,3 @@
-use anyhow::Result;
 use async_trait::async_trait;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::catalog::Session;
@@ -10,9 +9,8 @@ use std::any::Any;
 use std::sync::Arc;
 
 use crate::catalog::CatalogManager;
-use crate::datafetch::DataFetcher;
+use crate::datafetch::FetchOrchestrator;
 use crate::source::Source;
-use crate::storage::StorageManager;
 
 /// A lazy table provider that defers data fetching until scan() is called.
 ///
@@ -22,33 +20,29 @@ use crate::storage::StorageManager;
 #[derive(Debug)]
 pub struct LazyTableProvider {
     schema: SchemaRef,
-    fetcher: Arc<dyn DataFetcher>,
     source: Arc<Source>,
     catalog: Arc<dyn CatalogManager>,
-    storage: Arc<dyn StorageManager>,
+    orchestrator: Arc<FetchOrchestrator>,
     connection_id: i32,
     schema_name: String,
     table_name: String,
 }
 
 impl LazyTableProvider {
-    #[allow(clippy::too_many_arguments)]
     pub fn new(
         schema: SchemaRef,
-        fetcher: Arc<dyn DataFetcher>,
         source: Arc<Source>,
         catalog: Arc<dyn CatalogManager>,
-        storage: Arc<dyn StorageManager>,
+        orchestrator: Arc<FetchOrchestrator>,
         connection_id: i32,
         schema_name: String,
         table_name: String,
     ) -> Self {
         Self {
             schema,
-            fetcher,
             source,
             catalog,
-            storage,
+            orchestrator,
             connection_id,
             schema_name,
             table_name,
@@ -96,61 +90,15 @@ impl LazyTableProvider {
 
     /// Fetch the table data and update catalog
     async fn fetch_and_cache(&self) -> Result<String, DataFusionError> {
-        use crate::datafetch::native::StreamingParquetWriter;
-
-        // Prepare cache write location
-        let write_path = self.storage.prepare_cache_write(
-            self.connection_id,
-            &self.schema_name,
-            &self.table_name,
-        );
-
-        // Create writer
-        let mut writer = StreamingParquetWriter::new(write_path.clone());
-
-        // Fetch the table data into writer using the Source directly
-        self.fetcher
-            .fetch_table(
+        self.orchestrator
+            .cache_table(
                 &self.source,
-                None, // catalog
-                &self.schema_name,
-                &self.table_name,
-                &mut writer,
-            )
-            .await
-            .map_err(|e| {
-                DataFusionError::External(format!("Failed to fetch table: {}", e).into())
-            })?;
-
-        // Close writer
-        writer.close().map_err(|e| {
-            DataFusionError::External(format!("Failed to close writer: {}", e).into())
-        })?;
-
-        // Finalize cache write (uploads to S3 if needed, returns URL)
-        let parquet_url = self
-            .storage
-            .finalize_cache_write(
-                &write_path,
                 self.connection_id,
                 &self.schema_name,
                 &self.table_name,
             )
             .await
-            .map_err(|e| {
-                DataFusionError::External(format!("Failed to finalize cache write: {}", e).into())
-            })?;
-
-        // Update catalog with new path
-        if let Ok(Some(info)) = self
-            .catalog
-            .get_table(self.connection_id, &self.schema_name, &self.table_name)
-            .await
-        {
-            let _ = self.catalog.update_table_sync(info.id, &parquet_url).await;
-        }
-
-        Ok(parquet_url)
+            .map_err(|e| DataFusionError::External(format!("Failed to cache table: {}", e).into()))
     }
 }
 

src/datafusion/schema_provider.rs

@@ -4,12 +4,10 @@ use datafusion::datasource::TableProvider;
 use std::sync::Arc;
 
 use super::block_on;
+use super::lazy_table_provider::LazyTableProvider;
 use crate::catalog::CatalogManager;
-use crate::datafetch::{deserialize_arrow_schema, DataFetcher};
+use crate::datafetch::{deserialize_arrow_schema, FetchOrchestrator};
 use crate::source::Source;
-use crate::storage::StorageManager;
-
-use super::lazy_table_provider::LazyTableProvider;
 
 /// A schema provider that syncs tables on-demand from remote sources.
 /// Wraps MemorySchemaProvider for caching already-loaded tables.
@@ -21,31 +19,27 @@ pub struct RivetSchemaProvider {
     schema_name: String,
     source: Arc<Source>,
     catalog: Arc<dyn CatalogManager>,
+    orchestrator: Arc<FetchOrchestrator>,
     inner: Arc<MemorySchemaProvider>,
-    storage: Arc<dyn StorageManager>,
-    fetcher: Arc<dyn DataFetcher>,
 }
 
 impl RivetSchemaProvider {
-    #[allow(clippy::too_many_arguments)]
     pub fn new(
         connection_id: i32,
         connection_name: String,
         schema_name: String,
         source: Arc<Source>,
         catalog: Arc<dyn CatalogManager>,
-        storage: Arc<dyn StorageManager>,
-        fetcher: Arc<dyn DataFetcher>,
+        orchestrator: Arc<FetchOrchestrator>,
     ) -> Self {
         Self {
             connection_id,
             connection_name,
             schema_name,
             source,
             catalog,
+            orchestrator,
             inner: Arc::new(MemorySchemaProvider::new()),
-            storage,
-            fetcher,
         }
     }
 }
@@ -109,10 +103,9 @@ impl SchemaProvider for RivetSchemaProvider {
         // Create LazyTableProvider
         let provider = Arc::new(LazyTableProvider::new(
             schema,
-            self.fetcher.clone(),
             self.source.clone(),
             self.catalog.clone(),
-            self.storage.clone(),
+            self.orchestrator.clone(),
             self.connection_id,
             self.schema_name.clone(),
             name.to_string(),

src/engine.rs

@@ -1,5 +1,5 @@
 use crate::catalog::{CatalogManager, ConnectionInfo, SqliteCatalogManager, TableInfo};
-use crate::datafetch::DataFetcher;
+use crate::datafetch::{DataFetcher, FetchOrchestrator, NativeFetcher};
 use crate::datafusion::{block_on, RivetCatalogProvider};
 use crate::source::Source;
 use crate::storage::{FilesystemStorage, StorageManager};
@@ -23,6 +23,7 @@ pub struct RivetEngine {
     catalog: Arc<dyn CatalogManager>,
     df_ctx: SessionContext,
     storage: Arc<dyn StorageManager>,
+    orchestrator: Arc<FetchOrchestrator>,
 }
 
 impl RivetEngine {
@@ -200,7 +201,7 @@ impl RivetEngine {
                 conn.name.clone(),
                 Arc::new(source),
                 self.catalog.clone(),
-                self.storage.clone(),
+                self.orchestrator.clone(),
             )) as Arc<dyn CatalogProvider>;
 
             self.df_ctx.register_catalog(&conn.name, catalog_provider);
@@ -246,7 +247,7 @@ impl RivetEngine {
             name.to_string(),
             Arc::new(source),
             self.catalog.clone(),
-            self.storage.clone(),
+            self.orchestrator.clone(),
         )) as Arc<dyn CatalogProvider>;
 
         self.df_ctx.register_catalog(name, catalog_provider);
@@ -326,7 +327,7 @@ impl RivetEngine {
             conn.name.clone(),
             Arc::new(source),
             self.catalog.clone(),
-            self.storage.clone(),
+            self.orchestrator.clone(),
         )) as Arc<dyn CatalogProvider>;
 
         // register_catalog replaces existing catalog with same name
@@ -368,7 +369,7 @@ impl RivetEngine {
             conn.name.clone(),
             Arc::new(source),
             self.catalog.clone(),
-            self.storage.clone(),
+            self.orchestrator.clone(),
         )) as Arc<dyn CatalogProvider>;
 
         // register_catalog replaces existing catalog with same name
@@ -574,10 +575,19 @@ impl RivetEngineBuilder {
         let df_ctx = SessionContext::new();
         storage.register_with_datafusion(&df_ctx)?;
 
+        // Step 6: Create fetch orchestrator
+        let fetcher = Arc::new(NativeFetcher::new());
+        let orchestrator = Arc::new(FetchOrchestrator::new(
+            fetcher,
+            storage.clone(),
+            catalog.clone(),
+        ));
+
         let mut engine = RivetEngine {
             catalog,
             df_ctx,
             storage,
+            orchestrator,
         };
 
         // Register all existing connections as DataFusion catalogs