always teardown after a setup (#86)

Entire-Checkpoint: bf89cf7ab0b7

Author: Jeadie

src/main.rs

@@ -36,6 +36,7 @@ mod commands;
 mod metrics;
 mod scenario;
 
+use crate::args::CommonArgs;
 use crate::commands::connect_system_adapter;
 use crate::scenario::Scenario;
 
@@ -76,87 +77,26 @@ pub enum SystemAdapterExecutionMode {
     DirectQuery,
 }
 
-#[tokio::main]
-async fn main() -> anyhow::Result<()> {
-    let _ = rustls::crypto::CryptoProvider::install_default(
-        rustls::crypto::aws_lc_rs::default_provider(),
-    );
-
-    tracing_subscriber::fmt()
-        .with_max_level(Level::INFO)
-        .with_env_filter(EnvFilter::from_default_env())
-        .init();
-
-    let cli = Cli::parse();
-
-    // --- Construct the ETL pipeline ---
-    let dataset_source = match &cli.common.scenario {
-        Scenario::TPCH => DatasetSource::Tpch,
-    };
-
-    let generation_config = GenerationDatasetConfig {
-        dataset_type: match &dataset_source {
-            DatasetSource::Tpch => "tpch".to_string(),
-            DatasetSource::SimpleSequence => "simple_sequence".to_string(),
-        },
-        scale_factor: cli.common.scale_factor,
-        num_steps: cli.common.etl_num_steps,
-    };
-
-    let source_config = TargetConfig {
-        bucket: cli.common.etl_bucket.clone(),
-        prefix: cli.common.etl_source_prefix.clone(),
-        region: cli.common.etl_region.clone(),
-        endpoint: cli.common.etl_endpoint.clone(),
-    };
-
-    let source = Arc::new(S3Storage::new(&source_config)?);
-
-    // --- Connect to the system adapter ---
-    let mut system_adapter_client = match connect_system_adapter(&cli.common).await {
-        Ok(system_adapter_client) => system_adapter_client,
-        Err(e) => {
-            return Err(anyhow::anyhow!("Failed to connect to system adapter: {e}"));
-        }
-    };
-
-    let run_id = uuid::Uuid::new_v4();
-    let mutations = MutationConfig::new(0.1, 0.1);
-
-    let setup_dataset = dataset_source.create(
-        &generation_config,
-        &mutations,
-        Arc::clone(&source) as Arc<dyn DataStorage>,
-    )?;
-    let datasets = create_tables_request_datasets(&setup_dataset);
-
-    let setup_metadata = std::collections::HashMap::from([
-        (
-            "executor_instance_type".to_string(),
-            serde_json::Value::String(cli.common.executor_instance_type.clone()),
-        ),
-        (
-            "table_format".to_string(),
-            serde_json::Value::String(cli.common.table_format.to_string()),
-        ),
-    ]);
-
-    let adbc_driver = match system_adapter_client.setup(run_id, setup_metadata).await {
-        Ok(response) => response,
-        Err(e) => {
-            return Err(anyhow::anyhow!("Failed to setup system adapter: {e}"));
-        }
-    };
-
+async fn run_benchmark(
+    common: &CommonArgs,
+    system_adapter_client: &mut system_adapter_protocol::Client,
+    run_id: uuid::Uuid,
+    adbc_driver: system_adapter_protocol::SetupResponse,
+    dataset_source: DatasetSource,
+    generation_config: &GenerationDatasetConfig,
+    mutations: &MutationConfig,
+    source: Arc<S3Storage>,
+    datasets: HashMap<String, system_adapter_protocol::DatasetConfig>,
+) -> anyhow::Result<()> {
     // --- Download checkpoints from S3 ---
-    let scenario_name = cli.common.scenario.to_string();
+    let scenario_name = common.scenario.to_string();
     let checkpoint_dir = tempfile::tempdir()?;
 
     let checkpoint_store = CheckpointStore::new(
-        &cli.common.etl_bucket,
-        &cli.common.etl_source_prefix,
-        cli.common.etl_region.as_deref(),
-        cli.common.etl_endpoint.as_deref(),
+        &common.etl_bucket,
+        &common.etl_source_prefix,
+        common.etl_region.as_deref(),
+        common.etl_endpoint.as_deref(),
     )?;
 
     let manifest = checkpoint_store.download_manifest().await.map_err(|e| {
@@ -199,38 +139,17 @@ async fn main() -> anyhow::Result<()> {
     let sink_kwargs = adbc_driver.db_kwargs.clone();
     let load_kwargs = adbc_driver.db_kwargs;
 
-    let adbc_conn: Option<AdbcConnection> = match AdbcConnection::create(&driver_name, sink_kwargs)
-    {
-        Ok(conn) => {
-            println!(
-                "ADBC connection established (driver: {})",
-                adbc_driver.driver
-            );
-            Some(conn)
-        }
-        Err(e) => {
-            eprintln!(
-                "Failed to create ADBC connection for driver {}: {e}",
-                adbc_driver.driver
-            );
-            None
-        }
-    };
-
-    let Some(adbc_conn) = adbc_conn else {
-        return Err(anyhow::anyhow!(
-            "ADBC connection is required to run benchmarks"
-        ));
-    };
+    let adbc_conn = AdbcConnection::create(&driver_name, sink_kwargs).map_err(|e| {
+        anyhow::anyhow!(
+            "Failed to create ADBC connection for driver {}: {e}",
+            driver_name
+        )
+    })?;
+    println!("ADBC connection established (driver: {})", driver_name);
 
     let target = Arc::new(AdbcSink::new_without_table_creation(adbc_conn, None));
-    let mut pipeline = ETLPipeline::new(
-        dataset_source,
-        &generation_config,
-        source,
-        target,
-        &mutations,
-    )?;
+    let mut pipeline =
+        ETLPipeline::new(dataset_source, generation_config, source, target, mutations)?;
 
     if let Err(e) = system_adapter_client.create_tables(run_id, datasets).await {
         pipeline.cancel();
@@ -250,14 +169,14 @@ async fn main() -> anyhow::Result<()> {
             pipeline.cancel();
             return Err(anyhow::anyhow!(
                 "Failed to create benchmark ADBC connection for driver {}: {e}",
-                adbc_driver.driver
+                driver_name
             ));
         }
     };
 
     commands::load::run(
-        &cli.common.scenario,
-        &cli.common,
+        &common.scenario,
+        common,
         load_conn,
         &mut pipeline,
         checkpoint_steps,
@@ -285,9 +204,98 @@ async fn main() -> anyhow::Result<()> {
         }
     }
 
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let _ = rustls::crypto::CryptoProvider::install_default(
+        rustls::crypto::aws_lc_rs::default_provider(),
+    );
+
+    tracing_subscriber::fmt()
+        .with_max_level(Level::INFO)
+        .with_env_filter(EnvFilter::from_default_env())
+        .init();
+
+    let cli = Cli::parse();
+
+    // --- Construct the ETL pipeline ---
+    let dataset_source = match &cli.common.scenario {
+        Scenario::TPCH => DatasetSource::Tpch,
+    };
+
+    let generation_config = GenerationDatasetConfig {
+        dataset_type: match &dataset_source {
+            DatasetSource::Tpch => "tpch".to_string(),
+            DatasetSource::SimpleSequence => "simple_sequence".to_string(),
+        },
+        scale_factor: cli.common.scale_factor,
+        num_steps: cli.common.etl_num_steps,
+    };
+
+    let source_config = TargetConfig {
+        bucket: cli.common.etl_bucket.clone(),
+        prefix: cli.common.etl_source_prefix.clone(),
+        region: cli.common.etl_region.clone(),
+        endpoint: cli.common.etl_endpoint.clone(),
+    };
+
+    let source = Arc::new(S3Storage::new(&source_config)?);
+
+    // --- Connect to the system adapter ---
+    let mut system_adapter_client = match connect_system_adapter(&cli.common).await {
+        Ok(system_adapter_client) => system_adapter_client,
+        Err(e) => {
+            return Err(anyhow::anyhow!("Failed to connect to system adapter: {e}"));
+        }
+    };
+
+    let run_id = uuid::Uuid::new_v4();
+    let mutations = MutationConfig::new(0.1, 0.1);
+
+    let setup_dataset = dataset_source.create(
+        &generation_config,
+        &mutations,
+        Arc::clone(&source) as Arc<dyn DataStorage>,
+    )?;
+    let datasets = create_tables_request_datasets(&setup_dataset);
+
+    let setup_metadata = std::collections::HashMap::from([
+        (
+            "executor_instance_type".to_string(),
+            serde_json::Value::String(cli.common.executor_instance_type.clone()),
+        ),
+        (
+            "table_format".to_string(),
+            serde_json::Value::String(cli.common.table_format.to_string()),
+        ),
+    ]);
+
+    let adbc_driver = match system_adapter_client.setup(run_id, setup_metadata).await {
+        Ok(response) => response,
+        Err(e) => {
+            return Err(anyhow::anyhow!("Failed to setup system adapter: {e}"));
+        }
+    };
+
+    let result = run_benchmark(
+        &cli.common,
+        &mut system_adapter_client,
+        run_id,
+        adbc_driver,
+        dataset_source,
+        &generation_config,
+        &mutations,
+        source,
+        datasets,
+    )
+    .await;
+
+    // After successful setup, always teardown even if there are errors in between.
     if let Err(e) = system_adapter_client.teardown(run_id).await {
-        return Err(anyhow::anyhow!("Failed to teardown system adapter: {e}"));
+        tracing::error!("Failed to teardown system adapter: {e}");
     }
 
-    Ok(())
+    result
 }