refactor: Separate a DataStorage and ETL Sink (#63)

* refactor: Separate a DataStorage and ETL Sink

* chore: fmt

Author: peasee

Cargo.lock

@@ -24,19 +24,19 @@ use tokio::task::JoinSet;
 use super::config::IngestorConfig;
 use super::dataset::Dataset;
 use super::metrics::{IngestResult, Metrics};
-use super::target::Target;
+use super::storage::DataStorage;
 
 pub struct DataGenerator {
     dataset: Arc<dyn Dataset>,
-    target: Arc<dyn Target>,
+    target: Arc<dyn DataStorage>,
     metrics: Metrics,
     semaphore: Arc<Semaphore>,
 }
 
 impl DataGenerator {
     pub fn new(
         dataset: Arc<dyn Dataset>,
-        target: Arc<dyn Target>,
+        target: Arc<dyn DataStorage>,
         config: &IngestorConfig,
         metrics: Metrics,
     ) -> Self {

crates/data-generation/src/generator.rs

@@ -18,6 +18,4 @@ pub mod config;
 pub mod dataset;
 pub mod generator;
 pub mod metrics;
-pub mod source;
 pub mod storage;
-pub mod target;

crates/data-generation/src/lib.rs

@@ -15,7 +15,7 @@ limitations under the License.
 */
 
 use clap::Parser;
-use data_generation::target::Target;
+use data_generation::storage::DataStorage;
 use tracing_subscriber::EnvFilter;
 
 use std::sync::Arc;
@@ -74,7 +74,7 @@ fn build(args: &CommonArgs) -> anyhow::Result<DataGenerator> {
 
     let ingestor = DataGenerator::new(
         dataset,
-        target as Arc<dyn Target>,
+        target as Arc<dyn DataStorage>,
         &ingestor_config,
         metrics,
     );

crates/data-generation/src/main.rs

@@ -18,7 +18,7 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, Instant};
 
-use crate::target::WriteResult;
+use crate::storage::WriteResult;
 
 #[derive(Clone)]
 pub struct Metrics {

crates/data-generation/src/metrics.rs

@@ -15,3 +15,54 @@ limitations under the License.
 */
 
 pub mod s3;
+
+use arrow::array::RecordBatch;
+use async_trait::async_trait;
+use std::collections::HashMap;
+
+pub struct ReadResult {
+    pub batches: Vec<RecordBatch>,
+    pub rows_read: u64,
+    pub bytes_read: u64,
+}
+
+pub struct WriteResult {
+    pub rows_written: u64,
+    pub bytes_written: u64,
+}
+
+#[async_trait]
+pub trait DataStorage: Send + Sync + 'static {
+    /// List available batch object paths for a given table.
+    async fn list_batches(&self, table_name: &str) -> anyhow::Result<Vec<String>>;
+
+    /// Read a single batch from the source by its batch ID and table name.
+    ///
+    /// Returns `Ok(None)` when the batch does not exist in the underlying
+    /// storage (e.g. the table has fewer batches than others). The caller
+    /// should treat this as the table having no more data.
+    ///
+    /// The concrete implementation is responsible for mapping `(table_name,
+    /// batch_id)` to the underlying storage path.
+    async fn read_batch(
+        &self,
+        table_name: &str,
+        batch_id: u64,
+    ) -> anyhow::Result<Option<ReadResult>>;
+
+    async fn write(
+        &self,
+        table_name: &str,
+        batch_id: u64,
+        batch: RecordBatch,
+    ) -> anyhow::Result<WriteResult>;
+
+    fn table_params(&self, table_name: &str) -> HashMap<String, serde_json::Value>;
+
+    /// Returns the list of file paths/URIs that would exist after a successful
+    /// generation for the given table and batch IDs.
+    ///
+    /// This is a planning method — no I/O is performed. Each implementation
+    /// maps `(table_name, batch_id)` to its own path scheme (e.g. an S3 URI).
+    fn expected_files(&self, table_name: &str, batch_ids: &[u64]) -> Vec<String>;
+}

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

@@ -21,6 +21,19 @@ use object_store::aws::AmazonS3Builder;
 use object_store::path::Path as ObjectPath;
 
 use crate::config::TargetConfig;
+use crate::storage::DataStorage;
+
+use arrow::array::RecordBatch;
+use async_trait::async_trait;
+use futures::TryStreamExt;
+use object_store::PutPayload;
+use parquet::arrow::ArrowWriter;
+use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+use parquet::basic::Compression;
+use parquet::file::properties::WriterProperties;
+use std::collections::HashMap;
+
+use super::{ReadResult, WriteResult};
 
 /// Unified S3 storage backend that implements both [`Source`] and [`Target`].
 ///
@@ -90,3 +103,125 @@ impl S3Storage {
         }
     }
 }
+
+#[async_trait]
+impl DataStorage for S3Storage {
+    fn expected_files(&self, table_name: &str, batch_ids: &[u64]) -> Vec<String> {
+        batch_ids
+            .iter()
+            .map(|id| {
+                if self.prefix.is_empty() {
+                    format!("s3://{}/{table_name}/batch-{id:06}.parquet", self.bucket)
+                } else {
+                    format!(
+                        "s3://{}/{}/{table_name}/batch-{id:06}.parquet",
+                        self.bucket, self.prefix
+                    )
+                }
+            })
+            .collect()
+    }
+
+    fn table_params(&self, table_name: &str) -> HashMap<String, serde_json::Value> {
+        let mut params = HashMap::new();
+        params.insert(
+            "connector".to_string(),
+            serde_json::Value::String("s3".to_string()),
+        );
+        params.insert(
+            "from".to_string(),
+            serde_json::Value::String(self.table_s3_path(table_name)),
+        );
+        params.insert(
+            "file_format".to_string(),
+            serde_json::Value::String("parquet".to_string()),
+        );
+
+        if let Some(region) = &self.region {
+            params.insert(
+                "s3_region".to_string(),
+                serde_json::Value::String(region.clone()),
+            );
+        }
+
+        params
+    }
+
+    async fn write(
+        &self,
+        table_name: &str,
+        batch_id: u64,
+        batch: RecordBatch,
+    ) -> anyhow::Result<WriteResult> {
+        let rows = batch.num_rows() as u64;
+        let schema = batch.schema();
+
+        // Serialize RecordBatch to Parquet bytes in memory
+        let props = WriterProperties::builder()
+            .set_compression(Compression::SNAPPY)
+            .build();
+
+        let mut buf = Vec::new();
+        let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props))?;
+        writer.write(&batch)?;
+        writer.close()?;
+
+        let bytes_written = buf.len() as u64;
+
+        // Upload to S3 with per-table directory structure
+        let path = self.batch_object_path(table_name, batch_id);
+
+        self.store.put(&path, PutPayload::from(buf)).await?;
+
+        Ok(WriteResult {
+            rows_written: rows,
+            bytes_written,
+        })
+    }
+
+    async fn list_batches(&self, table_name: &str) -> anyhow::Result<Vec<String>> {
+        let prefix = self.table_object_prefix(table_name);
+
+        let objects: Vec<_> = self.store.list(Some(&prefix)).try_collect().await?;
+
+        let paths: Vec<String> = objects
+            .into_iter()
+            .filter(|meta| meta.location.as_ref().ends_with(".parquet"))
+            .map(|meta| meta.location.to_string())
+            .collect();
+
+        Ok(paths)
+    }
+
+    async fn read_batch(
+        &self,
+        table_name: &str,
+        batch_id: u64,
+    ) -> anyhow::Result<Option<ReadResult>> {
+        let location = self.batch_object_path(table_name, batch_id);
+
+        let get_result = match self.store.get(&location).await {
+            Ok(r) => r,
+            Err(object_store::Error::NotFound { .. }) => return Ok(None),
+            Err(e) => return Err(e.into()),
+        };
+        let bytes = get_result.bytes().await?;
+        let bytes_read = bytes.len() as u64;
+
+        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes)?.build()?;
+
+        let mut batches = Vec::new();
+        let mut rows_read = 0u64;
+        for batch in reader {
+            let batch = batch?;
+            rows_read += batch.num_rows() as u64;
+            batches.push(batch);
+        }
+
+        Ok(Some(ReadResult {
+            batches,
+            rows_read,
+            bytes_read,
+        }))
+    }
+}