refactor: Move duplicated S3 source and target structs into a single struct (#56)

Author: peasee

crates/data-generation/src/lib.rs

@@ -19,4 +19,5 @@ pub mod dataset;
 pub mod generator;
 pub mod metrics;
 pub mod source;
+pub mod storage;
 pub mod target;

crates/data-generation/src/main.rs

@@ -25,7 +25,7 @@ use data_generation::dataset;
 use data_generation::dataset::tpch::TpchDataset;
 use data_generation::generator::DataGenerator;
 use data_generation::metrics::{IngestResult, Metrics};
-use data_generation::target::s3::S3Target;
+use data_generation::storage::s3::S3Storage;
 
 fn print_summary(result: &IngestResult) {
     println!("  Duration:          {:?}", result.elapsed);
@@ -69,7 +69,7 @@ fn build(args: &CommonArgs) -> anyhow::Result<DataGenerator> {
         other => anyhow::bail!("Unknown dataset type: {other}. Supported: tpch"),
     };
 
-    let target = Arc::new(S3Target::new(&target_config)?);
+    let target = Arc::new(S3Storage::new(&target_config)?);
     let metrics = Metrics::new();
 
     let ingestor = DataGenerator::new(

crates/data-generation/src/source/s3.rs

@@ -14,62 +14,18 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use std::sync::Arc;
-
 use async_trait::async_trait;
 use futures::TryStreamExt;
-use object_store::ObjectStore;
-use object_store::aws::AmazonS3Builder;
-use object_store::path::Path as ObjectPath;
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 
-use crate::config::TargetConfig;
+use crate::storage::s3::S3Storage;
 
 use super::{ReadResult, Source};
 
-/// Reads Parquet data from S3 (or S3-compatible storage).
-#[derive(Clone)]
-pub struct S3Source {
-    store: Arc<dyn ObjectStore>,
-    prefix: String,
-}
-
-impl S3Source {
-    /// Create a new [`S3Source`] from the same config used for [`S3Target`].
-    ///
-    /// The source and target typically share the same bucket/prefix, so they
-    /// reuse [`TargetConfig`] to avoid duplicating configuration structs.
-    pub fn new(config: &TargetConfig) -> anyhow::Result<Self> {
-        let mut builder = AmazonS3Builder::from_env().with_bucket_name(&config.bucket);
-
-        if let Some(region) = &config.region {
-            builder = builder.with_region(region);
-        }
-        if let Some(endpoint) = &config.endpoint
-            && !endpoint.is_empty()
-        {
-            builder = builder.with_endpoint(endpoint);
-            if endpoint.starts_with("http://") {
-                builder = builder.with_allow_http(true);
-            }
-        }
-
-        let store = Arc::new(builder.build()?);
-        Ok(Self {
-            store,
-            prefix: config.prefix.clone(),
-        })
-    }
-}
-
 #[async_trait]
-impl Source for S3Source {
+impl Source for S3Storage {
     async fn list_batches(&self, table_name: &str) -> anyhow::Result<Vec<String>> {
-        let prefix = if self.prefix.is_empty() {
-            ObjectPath::from(format!("{table_name}/"))
-        } else {
-            ObjectPath::from(format!("{}/{table_name}/", self.prefix))
-        };
+        let prefix = self.table_object_prefix(table_name);
 
         let objects: Vec<_> = self.store.list(Some(&prefix)).try_collect().await?;
 
@@ -87,14 +43,7 @@ impl Source for S3Source {
         table_name: &str,
         batch_id: u64,
     ) -> anyhow::Result<Option<ReadResult>> {
-        let location = if self.prefix.is_empty() {
-            ObjectPath::from(format!("{table_name}/batch-{batch_id:06}.parquet"))
-        } else {
-            ObjectPath::from(format!(
-                "{}/{table_name}/batch-{batch_id:06}.parquet",
-                self.prefix
-            ))
-        };
+        let location = self.batch_object_path(table_name, batch_id);
 
         let get_result = match self.store.get(&location).await {
             Ok(r) => r,

crates/data-generation/src/storage/mod.rs

@@ -0,0 +1,17 @@
+/*
+Copyright 2024-2025 The Spice.ai OSS Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+pub mod s3;

crates/data-generation/src/storage/s3.rs

@@ -0,0 +1,92 @@
+/*
+Copyright 2024-2025 The Spice.ai OSS Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+use std::sync::Arc;
+
+use object_store::ObjectStore;
+use object_store::aws::AmazonS3Builder;
+use object_store::path::Path as ObjectPath;
+
+use crate::config::TargetConfig;
+
+/// Unified S3 storage backend that implements both [`Source`] and [`Target`].
+///
+/// The trait implementations live in their respective modules
+/// (`source/s3.rs` and `target/s3.rs`) to keep concerns separated.
+#[derive(Clone)]
+pub struct S3Storage {
+    pub(crate) store: Arc<dyn ObjectStore>,
+    pub(crate) bucket: String,
+    pub(crate) prefix: String,
+    pub(crate) region: Option<String>,
+}
+
+impl S3Storage {
+    pub fn new(config: &TargetConfig) -> anyhow::Result<Self> {
+        let mut builder = AmazonS3Builder::from_env().with_bucket_name(&config.bucket);
+
+        if let Some(region) = &config.region {
+            tracing::info!("S3 storage with region: {region}");
+            builder = builder.with_region(region);
+        }
+        if let Some(endpoint) = &config.endpoint
+            && !endpoint.is_empty()
+        {
+            builder = builder.with_endpoint(endpoint);
+            if endpoint.starts_with("http://") {
+                builder = builder.with_allow_http(true);
+            }
+        }
+
+        let store = Arc::new(builder.build()?);
+        Ok(Self {
+            store,
+            bucket: config.bucket.clone(),
+            prefix: config.prefix.clone(),
+            region: config.region.clone(),
+        })
+    }
+
+    /// Returns the S3 URI for a given table name (e.g. `s3://bucket/prefix/customer/`).
+    pub fn table_s3_path(&self, table_name: &str) -> String {
+        if self.prefix.is_empty() {
+            format!("s3://{}/{table_name}/", self.bucket)
+        } else {
+            format!("s3://{}/{}/{table_name}/", self.bucket, self.prefix)
+        }
+    }
+
+    /// Returns the [`ObjectPath`] for a batch file within a table directory.
+    pub(crate) fn batch_object_path(&self, table_name: &str, batch_id: u64) -> ObjectPath {
+        if self.prefix.is_empty() {
+            ObjectPath::from(format!("{table_name}/batch-{batch_id:06}.parquet"))
+        } else {
+            ObjectPath::from(format!(
+                "{}/{table_name}/batch-{batch_id:06}.parquet",
+                self.prefix
+            ))
+        }
+    }
+
+    /// Returns the [`ObjectPath`] prefix for listing objects in a table directory.
+    pub(crate) fn table_object_prefix(&self, table_name: &str) -> ObjectPath {
+        if self.prefix.is_empty() {
+            ObjectPath::from(format!("{table_name}/"))
+        } else {
+            ObjectPath::from(format!("{}/{table_name}/", self.prefix))
+        }
+    }
+}

crates/data-generation/src/target/s3.rs

@@ -15,71 +15,20 @@ limitations under the License.
 */
 
 use std::collections::HashMap;
-use std::sync::Arc;
 
 use arrow::array::RecordBatch;
 use async_trait::async_trait;
-use object_store::aws::AmazonS3Builder;
-use object_store::path::Path as ObjectPath;
-use object_store::{ObjectStore, PutPayload};
+use object_store::PutPayload;
 use parquet::arrow::ArrowWriter;
 use parquet::basic::Compression;
 use parquet::file::properties::WriterProperties;
 
-use crate::config::TargetConfig;
+use crate::storage::s3::S3Storage;
 
 use super::{Target, WriteResult};
 
-#[derive(Clone)]
-pub struct S3Target {
-    store: Arc<dyn ObjectStore>,
-    bucket: String,
-    prefix: String,
-    table_format: String,
-    executor_instance_type: String,
-    region: Option<String>,
-}
-
-impl S3Target {
-    pub fn new(config: &TargetConfig) -> anyhow::Result<Self> {
-        let mut builder = AmazonS3Builder::from_env().with_bucket_name(&config.bucket);
-
-        if let Some(region) = &config.region {
-            tracing::info!("S3 Target with region: {region}");
-            builder = builder.with_region(region);
-        }
-        if let Some(endpoint) = &config.endpoint
-            && !endpoint.is_empty()
-        {
-            builder = builder.with_endpoint(endpoint);
-            if endpoint.starts_with("http://") {
-                builder = builder.with_allow_http(true);
-            }
-        }
-
-        let store = Arc::new(builder.build()?);
-        Ok(Self {
-            store,
-            bucket: config.bucket.clone(),
-            prefix: config.prefix.clone(),
-            table_format: config.table_format.to_string(),
-            executor_instance_type: config.executor_instance_type.clone(),
-            region: config.region.clone(),
-        })
-    }
-
-    /// Returns the S3 URI for a given table name (e.g. `s3://bucket/prefix/customer/`).
-    pub fn table_s3_path(&self, table_name: &str) -> String {
-        if self.prefix.is_empty() {
-            format!("s3://{}/{table_name}/", self.bucket)
-        } else {
-            format!("s3://{}/{}/{table_name}/", self.bucket, self.prefix)
-        }
-    }
-}
-
 #[async_trait]
-impl Target for S3Target {
+impl Target for S3Storage {
     fn expected_files(&self, table_name: &str, batch_ids: &[u64]) -> Vec<String> {
         batch_ids
             .iter()
@@ -110,14 +59,6 @@ impl Target for S3Target {
             "file_format".to_string(),
             serde_json::Value::String("parquet".to_string()),
         );
-        params.insert(
-            "table_format".to_string(),
-            serde_json::Value::String(self.table_format.clone()),
-        );
-        params.insert(
-            "executor_instance_type".to_string(),
-            serde_json::Value::String(self.executor_instance_type.clone()),
-        );
 
         if let Some(region) = &self.region {
             params.insert(
@@ -151,14 +92,7 @@ impl Target for S3Target {
         let bytes_written = buf.len() as u64;
 
         // Upload to S3 with per-table directory structure
-        let path = if self.prefix.is_empty() {
-            ObjectPath::from(format!("{table_name}/batch-{batch_id:06}.parquet"))
-        } else {
-            ObjectPath::from(format!(
-                "{}/{table_name}/batch-{batch_id:06}.parquet",
-                self.prefix
-            ))
-        };
+        let path = self.batch_object_path(table_name, batch_id);
 
         self.store.put(&path, PutPayload::from(buf)).await?;
 

crates/etl/src/main.rs

@@ -18,8 +18,7 @@ use std::sync::Arc;
 
 use clap::Parser;
 use data_generation::config::{DatasetConfig, TableFormat, TargetConfig};
-use data_generation::source::s3::S3Source;
-use data_generation::target::s3::S3Target;
+use data_generation::storage::s3::S3Storage;
 use etl::{DatasetSource, ETLPipeline, PipelineState, StopReason};
 use tracing_subscriber::EnvFilter;
 
@@ -127,8 +126,8 @@ async fn main() -> anyhow::Result<()> {
     let dataset_source = cli.dataset_source()?;
     let dataset_config = cli.dataset_config();
 
-    let source = Arc::new(S3Source::new(&cli.source_config())?);
-    let target = Arc::new(S3Target::new(&cli.target_config())?);
+    let source = Arc::new(S3Storage::new(&cli.source_config())?);
+    let target = Arc::new(S3Storage::new(&cli.target_config())?);
 
     let mut pipeline = ETLPipeline::new(dataset_source, &dataset_config, source, target)?;
 

src/main.rs

@@ -19,8 +19,7 @@ use std::sync::Arc;
 use adbc_client::AdbcConnection;
 use clap::Parser;
 use data_generation::config::{DatasetConfig as GenerationDatasetConfig, TargetConfig};
-use data_generation::source::s3::S3Source;
-use data_generation::target::s3::S3Target;
+use data_generation::storage::s3::S3Storage;
 use etl::{DatasetSource, ETLPipeline, PipelineState, StopReason};
 use test_framework::{anyhow, rustls};
 use tracing::Level;
@@ -101,8 +100,8 @@ async fn main() -> anyhow::Result<()> {
         executor_instance_type: cli.common.executor_instance_type.clone(),
     };
 
-    let source = Arc::new(S3Source::new(&source_config)?);
-    let target = Arc::new(S3Target::new(&target_config)?);
+    let source = Arc::new(S3Storage::new(&source_config)?);
+    let target = Arc::new(S3Storage::new(&target_config)?);
 
     let mut pipeline = ETLPipeline::new(dataset_source, &generation_config, source, target)?;