Put export implementation behind a trait object (#40094)

GitOrigin-RevId: f7dea429e43e6e21daf31b6dc112c0399db9d6f2

Author: goffrie

Cargo.lock

@@ -23,7 +23,7 @@ use database::{
     SystemMetadataModel,
 };
 use exports::{
-    export_inner,
+    interface::ExportProvider,
     ExportComponents,
 };
 use futures::{
@@ -63,6 +63,7 @@ pub struct ExportWorker<RT: Runtime> {
     pub(super) database: Database<RT>,
     pub(super) storage: Arc<dyn Storage>,
     pub(super) file_storage: Arc<dyn Storage>,
+    pub(super) export_provider: Arc<dyn ExportProvider<RT>>,
     pub(super) backoff: Backoff,
     pub(super) usage_tracking: UsageCounter,
     pub(super) instance_name: String,
@@ -75,6 +76,7 @@ impl<RT: Runtime> ExportWorker<RT> {
         database: Database<RT>,
         storage: Arc<dyn Storage>,
         file_storage: Arc<dyn Storage>,
+        export_provider: Arc<dyn ExportProvider<RT>>,
         usage_tracking: UsageCounter,
         instance_name: String,
     ) -> impl Future<Output = ()> + Send {
@@ -83,6 +85,7 @@ impl<RT: Runtime> ExportWorker<RT> {
             database,
             storage,
             file_storage,
+            export_provider,
             backoff: Backoff::new(INITIAL_BACKOFF, MAX_BACKOFF),
             usage_tracking,
             instance_name,
@@ -179,22 +182,22 @@ impl<RT: Runtime> ExportWorker<RT> {
         let database_snapshot = self.database.latest_database_snapshot()?;
         let snapshot_ts = *database_snapshot.timestamp();
         tracing::info!(%snapshot_ts, "Export {id} beginning...");
+        let components = ExportComponents {
+            runtime: self.runtime.clone(),
+            database: database_snapshot,
+            storage: self.storage.clone(),
+            file_storage: self.file_storage.clone(),
+            instance_name: self.instance_name.clone(),
+        };
         let (object_key, usage) = {
             let database_ = self.database.clone();
-            let export_future = async {
-                let database_ = self.database.clone();
-
-                export_inner(
-                    &ExportComponents {
-                        runtime: self.runtime.clone(),
-                        database: database_snapshot,
-                        storage: self.storage.clone(),
-                        file_storage: self.file_storage.clone(),
-                        instance_name: self.instance_name.clone(),
-                    },
-                    format,
-                    requestor,
-                    |msg| async {
+            let export_future = self.export_provider.export(
+                &components,
+                format,
+                requestor,
+                Box::new(move |msg| {
+                    let database_ = database_.clone();
+                    async move {
                         tracing::info!("Export {id} progress: {msg}");
                         database_
                             .execute_with_occ_retries(
@@ -221,12 +224,12 @@ impl<RT: Runtime> ExportWorker<RT> {
                             )
                             .await?;
                         Ok(())
-                    },
-                )
-                .await
-            };
-            tokio::pin!(export_future);
+                    }
+                    .boxed()
+                }),
+            );
 
+            let database_ = self.database.clone();
             // In parallel, monitor the export document to check for cancellation
             let monitor_export = async move {
                 loop {

crates/application/src/exports/worker.rs

@@ -21,6 +21,7 @@ use std::{
     },
 };
 
+use ::exports::interface::ExportProvider;
 use ::log_streaming::{
     LogManager,
     LogManagerClient,
@@ -687,6 +688,7 @@ impl<RT: Runtime> Application<RT> {
         fetch_client: Arc<dyn FetchClient>,
         local_log_sink: Option<String>,
         lease_lost_shutdown: ShutdownSignal,
+        export_provider: Arc<dyn ExportProvider<RT>>,
     ) -> anyhow::Result<Self> {
         let module_cache =
             ModuleCache::new(runtime.clone(), application_storage.modules_storage.clone()).await;
@@ -811,6 +813,7 @@ impl<RT: Runtime> Application<RT> {
             database.clone(),
             application_storage.exports_storage.clone(),
             application_storage.files_storage.clone(),
+            export_provider,
             database.usage_counter(),
             instance_name.clone(),
         );

crates/application/src/lib.rs

@@ -51,6 +51,7 @@ use events::usage::{
     NoOpUsageEventLogger,
     UsageEventLogger,
 };
+use exports::interface::InProcessExportProvider;
 use file_storage::{
     FileStorage,
     TransactionalFileStorage,
@@ -279,6 +280,7 @@ impl<RT: Runtime> ApplicationTestExt<RT> for Application<RT> {
             fetch_client,
             None, // local_log_sink
             ShutdownSignal::panic(),
+            Arc::new(InProcessExportProvider),
         )
         .await?;
 

crates/application/src/test_helpers.rs

@@ -23,6 +23,7 @@ testing = [
 
 [dependencies]
 anyhow = { workspace = true }
+async-trait = { workspace = true }
 async_zip = { workspace = true }
 bytes = { workspace = true }
 common = { workspace = true }

crates/exports/Cargo.toml

@@ -0,0 +1,42 @@
+use common::{
+    runtime::Runtime,
+    types::ObjectKey,
+};
+use futures::future::BoxFuture;
+use model::exports::types::{
+    ExportFormat,
+    ExportRequestor,
+};
+use usage_tracking::FunctionUsageTracker;
+
+use crate::ExportComponents;
+
+#[async_trait::async_trait]
+pub trait ExportProvider<RT: Runtime>: Send + Sync {
+    async fn export(
+        &self,
+        components: &ExportComponents<RT>,
+        format: ExportFormat,
+        requestor: ExportRequestor,
+        update_progress: Box<
+            dyn Fn(String) -> BoxFuture<'static, anyhow::Result<()>> + Send + Sync,
+        >,
+    ) -> anyhow::Result<(ObjectKey, FunctionUsageTracker)>;
+}
+
+pub struct InProcessExportProvider;
+
+#[async_trait::async_trait]
+impl<RT: Runtime> ExportProvider<RT> for InProcessExportProvider {
+    async fn export(
+        &self,
+        components: &ExportComponents<RT>,
+        format: ExportFormat,
+        requestor: ExportRequestor,
+        update_progress: Box<
+            dyn Fn(String) -> BoxFuture<'static, anyhow::Result<()>> + Send + Sync,
+        >,
+    ) -> anyhow::Result<(ObjectKey, FunctionUsageTracker)> {
+        crate::export_inner(components, format, requestor, &*update_progress).await
+    }
+}

crates/exports/src/interface.rs

@@ -79,6 +79,7 @@ use self::{
 };
 
 mod export_storage;
+pub mod interface;
 mod metrics;
 #[cfg(test)]
 mod tests;

crates/exports/src/lib.rs

@@ -34,6 +34,7 @@ testing = [
     "node_executor/testing",
     "sync/testing",
     "db_connection/testing",
+    "exports/testing",
 ]
 
 [dependencies]
@@ -56,6 +57,7 @@ db_connection = { workspace = true }
 either = { workspace = true }
 errors = { workspace = true }
 events = { workspace = true }
+exports = { workspace = true }
 fastrace = { workspace = true }
 file_storage = { workspace = true }
 function_runner = { workspace = true }
@@ -117,6 +119,7 @@ database = { workspace = true, features = ["testing"] }
 db_connection = { workspace = true, features = ["testing"] }
 errors = { workspace = true, features = ["testing"] }
 events = { workspace = true, features = ["testing"] }
+exports = { workspace = true, features = ["testing"] }
 function_runner = { workspace = true, features = ["testing"] }
 isolate = { workspace = true, features = ["testing"] }
 jsonschema = { workspace = true }

crates/local_backend/Cargo.toml

@@ -46,6 +46,7 @@ use common::{
 use config::LocalConfig;
 use database::Database;
 use events::usage::NoOpUsageEventLogger;
+use exports::interface::InProcessExportProvider;
 use file_storage::{
     FileStorage,
     TransactionalFileStorage,
@@ -241,6 +242,7 @@ pub async fn make_app(
         fetch_client,
         config.local_log_sink.clone(),
         preempt_tx.clone(),
+        Arc::new(InProcessExportProvider),
     )
     .await?;