//! DBN File Uploader for MinIO //! //! Automatically monitors test_data/real/databento/ for new .dbn files, //! compresses them with gzip, checks for duplicates, and uploads to MinIO. //! //! # Features //! - File watching with configurable poll interval //! - Gzip compression before upload //! - Deduplication (checks if file already exists in MinIO) //! - Metadata tagging (symbol, schema, date range, file size) use flate2::write::GzEncoder; use flate2::Compression; use std::collections::HashMap; use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; use tokio::fs; use tokio::sync::RwLock; use tokio::time::interval; use tracing::{debug, error, info}; use crate::error::{DataError, Result as DataResult}; /// Configuration for DBN uploader #[derive(Debug, Clone)] pub struct DbnUploaderConfig { /// Directory to watch for new DBN files pub watch_path: PathBuf, /// MinIO bucket name pub bucket_name: String, /// Prefix for uploaded files (e.g., "training-data/") pub upload_prefix: String, /// Polling interval for new files pub poll_interval: Duration, /// Enable gzip compression before upload pub compression_enabled: bool, /// Enable deduplication check pub deduplication_enabled: bool, } impl Default for DbnUploaderConfig { fn default() -> Self { Self { watch_path: PathBuf::from("test_data/real/databento"), bucket_name: "ml-models".to_string(), upload_prefix: "training-data/".to_string(), poll_interval: Duration::from_secs(60), compression_enabled: true, deduplication_enabled: true, } } } /// Metadata extracted from DBN filename #[derive(Debug, Clone, PartialEq)] pub struct DbnMetadata { /// Trading symbol (e.g., "ES.FUT") pub symbol: String, /// Data schema (e.g., "ohlcv-1m") pub schema: String, /// Date range (e.g., "2024-01-02" or "2024-01-02_to_2024-01-31") pub date_range: String, /// Original file size in bytes pub file_size_bytes: u64, } impl DbnMetadata { /// Extract metadata from DBN filename /// /// Expected format: `{SYMBOL}_{SCHEMA}_{DATE_RANGE}.dbn` pub fn from_filename(filename: &str) -> DataResult { if !filename.ends_with(".dbn") { return Err(DataError::Validation { field: "filename".to_string(), message: format!("File must have .dbn extension: {}", filename), }); } let name = filename.trim_end_matches(".dbn"); let parts: Vec<&str> = name.split('_').collect(); if parts.len() < 3 { return Err(DataError::Validation { field: "filename".to_string(), message: format!( "Invalid DBN filename format (expected SYMBOL_SCHEMA_DATE): {}", filename ), }); } let symbol = parts[0].to_string(); let schema = parts[1].to_string(); let date_range = parts[2..].join("_"); Ok(Self { symbol, schema, date_range, file_size_bytes: 0, }) } /// Extract metadata from file (includes size) pub async fn from_file(path: &Path) -> DataResult { let filename = path .file_name() .ok_or_else(|| DataError::Validation { field: "path".to_string(), message: "Path has no filename".to_string(), })? .to_str() .ok_or_else(|| DataError::Validation { field: "filename".to_string(), message: "Filename is not valid UTF-8".to_string(), })?; let mut metadata = Self::from_filename(filename)?; let file_metadata = fs::metadata(path).await?; metadata.file_size_bytes = file_metadata.len(); Ok(metadata) } } /// DBN file uploader pub struct DbnUploader { config: DbnUploaderConfig, detected_files: Arc>>, /// Optional storage backend for MinIO/S3 uploads. /// When `None`, upload operations will return an error. storage_backend: Option>, } impl DbnUploader { /// Create new uploader without a storage backend. /// /// The uploader can scan and compress files, but `upload_file` will return /// an error unless a storage backend is attached via `with_storage`. pub async fn new(config: DbnUploaderConfig) -> DataResult { if !config.watch_path.exists() { return Err(DataError::Validation { field: "watch_path".to_string(), message: format!("Watch path does not exist: {:?}", config.watch_path), }); } info!( "Initializing DBN uploader: watch_path={:?}, bucket={}, prefix={}", config.watch_path, config.bucket_name, config.upload_prefix ); Ok(Self { config, detected_files: Arc::new(RwLock::new(Vec::new())), storage_backend: None, }) } /// Attach a storage backend (MinIO/S3) for actual uploads. /// /// Without a storage backend, `upload_file` and deduplication checks will /// return errors indicating no backend is configured. pub fn with_storage(mut self, backend: Arc) -> Self { self.storage_backend = Some(backend); self } /// Get list of detected files (for testing) pub async fn get_detected_files(&self) -> Vec { self.detected_files.read().await.clone() } /// Scan directory and return files immediately (for testing) /// /// This method is primarily for testing purposes, as `start_watching()` runs /// in a blocking loop. In production, use `start_watching()` which continuously /// monitors the directory. pub async fn scan_for_testing(&self) -> DataResult> { self.scan_directory().await } /// Scan directory for DBN files async fn scan_directory(&self) -> DataResult> { let mut dbn_files = Vec::new(); let mut entries = fs::read_dir(&self.config.watch_path).await?; while let Some(entry) = entries.next_entry().await? { let path = entry.path(); if path.extension().and_then(|s| s.to_str()) == Some("dbn") { debug!("Detected DBN file: {:?}", path); dbn_files.push(path); } } Ok(dbn_files) } /// Start watching for new files (blocking) pub async fn start_watching(&self) -> DataResult<()> { info!("Starting DBN file watcher..."); let mut ticker = interval(self.config.poll_interval); loop { ticker.tick().await; match self.scan_directory().await { Ok(files) => { let mut detected = self.detected_files.write().await; *detected = files; debug!("Scanned directory, found {} DBN files", detected.len()); }, Err(e) => { error!("Failed to scan directory: {}", e); }, } } } /// Check if file should be uploaded (deduplication). /// /// When deduplication is enabled and a storage backend is configured, /// this queries the backend to see if the upload key already exists. /// Returns `true` if the file should be uploaded, `false` if it already /// exists in the remote store. /// /// When no storage backend is configured, deduplication is skipped and /// the file is always considered uploadable. pub async fn should_upload_file(&self, path: &Path) -> DataResult { if !self.config.deduplication_enabled { return Ok(true); } let Some(backend) = &self.storage_backend else { // No backend configured -- cannot check remote, allow upload attempt debug!("No storage backend configured, skipping deduplication check"); return Ok(true); }; let key = Self::generate_upload_key(path, &self.config.upload_prefix)?; match backend.exists(&key).await { Ok(true) => { info!("File already exists in remote storage, skipping: {}", key); Ok(false) }, Ok(false) => Ok(true), Err(e) => { error!("Failed to check deduplication in remote storage: {}", e); // On error, allow the upload attempt rather than silently skipping Ok(true) }, } } /// Compress file with gzip pub async fn compress_file(path: &Path) -> DataResult> { let data = fs::read(path).await?; let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); encoder.write_all(&data)?; Ok(encoder.finish()?) } /// Generate MinIO upload key from a file path and prefix. /// /// Returns the key in the form `{prefix}{filename}.gz`. /// /// # Errors /// /// Returns `DataError::Validation` if the path has no filename or the /// filename is not valid UTF-8. pub fn generate_upload_key(path: &Path, prefix: &str) -> DataResult { let filename = path .file_name() .ok_or_else(|| DataError::Validation { field: "path".to_string(), message: format!("Path has no filename component: {:?}", path), })? .to_str() .ok_or_else(|| DataError::Validation { field: "filename".to_string(), message: format!("Filename is not valid UTF-8: {:?}", path), })?; Ok(format!("{}{}.gz", prefix, filename)) } /// Generate metadata tags for MinIO pub fn generate_metadata_tags(metadata: &DbnMetadata) -> HashMap { let mut tags = HashMap::new(); tags.insert("symbol".to_string(), metadata.symbol.clone()); tags.insert("schema".to_string(), metadata.schema.clone()); tags.insert("date_range".to_string(), metadata.date_range.clone()); tags.insert( "original_size".to_string(), metadata.file_size_bytes.to_string(), ); tags } /// Upload file to MinIO/S3 with metadata. /// /// Reads the file, optionally compresses it, checks for duplicates, and /// stores it via the configured storage backend. The upload key is derived /// from the file name and the configured prefix. /// /// # Errors /// /// Returns `DataError::Storage` if no storage backend is configured, or /// any storage/IO error that occurs during the upload. pub async fn upload_file(&self, path: &Path) -> DataResult<()> { info!("Uploading file: {:?}", path); let backend = self .storage_backend .as_ref() .ok_or_else(|| { DataError::Storage( "No storage backend configured. Call with_storage() to attach a \ MinIO/S3 backend before uploading." .to_string(), ) })? .clone(); let metadata = DbnMetadata::from_file(path).await?; // Check deduplication before doing expensive compression if !self.should_upload_file(path).await? { info!("Skipping duplicate file: {:?}", path); return Ok(()); } let data = if self.config.compression_enabled { Self::compress_file(path).await? } else { fs::read(path).await? }; let key = Self::generate_upload_key(path, &self.config.upload_prefix)?; let tags = Self::generate_metadata_tags(&metadata); info!( "Uploading to storage: key={}, size={} bytes, tags={:?}", key, data.len(), tags ); backend.store(&key, &data).await.map_err(|e| { DataError::Storage(format!("Failed to upload {} to storage backend: {}", key, e)) })?; info!( "Successfully uploaded {} ({} bytes, original {} bytes)", key, data.len(), metadata.file_size_bytes, ); Ok(()) } } #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn test_metadata_from_filename_simple() { let metadata = DbnMetadata::from_filename("ES.FUT_ohlcv-1m_2024-01-02.dbn").expect("Failed to parse"); assert_eq!(metadata.symbol, "ES.FUT"); assert_eq!(metadata.schema, "ohlcv-1m"); assert_eq!(metadata.date_range, "2024-01-02"); } #[tokio::test] async fn test_metadata_from_filename_date_range() { let metadata = DbnMetadata::from_filename("ZN.FUT_ohlcv-1m_2024-01-02_to_2024-01-31.dbn") .expect("Failed to parse"); assert_eq!(metadata.symbol, "ZN.FUT"); assert_eq!(metadata.schema, "ohlcv-1m"); assert_eq!(metadata.date_range, "2024-01-02_to_2024-01-31"); } #[tokio::test] async fn test_metadata_from_filename_invalid() { let result = DbnMetadata::from_filename("invalid.txt"); assert!(result.is_err()); } #[tokio::test] async fn test_generate_upload_key() { let path = PathBuf::from("ES.FUT_ohlcv-1m_2024-01-02.dbn"); let key = DbnUploader::generate_upload_key(&path, "training-data/") .expect("Failed to generate key"); assert_eq!(key, "training-data/ES.FUT_ohlcv-1m_2024-01-02.dbn.gz"); } #[tokio::test] async fn test_generate_upload_key_no_filename() { let path = PathBuf::from("/"); let result = DbnUploader::generate_upload_key(&path, "prefix/"); assert!(result.is_err(), "Should fail for path with no filename"); } #[tokio::test] async fn test_generate_metadata_tags() { let metadata = DbnMetadata { symbol: "ES.FUT".to_string(), schema: "ohlcv-1m".to_string(), date_range: "2024-01-02".to_string(), file_size_bytes: 1024, }; let tags = DbnUploader::generate_metadata_tags(&metadata); assert_eq!( tags.get("symbol").cloned().as_deref(), Some("ES.FUT") ); assert_eq!( tags.get("schema").cloned().as_deref(), Some("ohlcv-1m") ); assert_eq!( tags.get("date_range").cloned().as_deref(), Some("2024-01-02") ); assert_eq!( tags.get("original_size").cloned().as_deref(), Some("1024") ); } #[tokio::test] async fn test_upload_file_without_backend_returns_error() { let temp_dir = tempfile::TempDir::new().expect("Failed to create temp dir"); let test_file = temp_dir.path().join("ES.FUT_ohlcv-1m_2024-01-02.dbn"); fs::write(&test_file, b"test data") .await .expect("Failed to write"); let config = DbnUploaderConfig { watch_path: temp_dir.path().to_path_buf(), compression_enabled: false, deduplication_enabled: false, ..DbnUploaderConfig::default() }; let uploader = DbnUploader::new(config) .await .expect("Failed to create uploader"); let result = uploader.upload_file(&test_file).await; assert!(result.is_err(), "Should fail without storage backend"); let err_msg = format!("{}", result.unwrap_err()); assert!( err_msg.contains("No storage backend configured"), "Error should mention missing backend, got: {}", err_msg ); } #[tokio::test] async fn test_upload_file_with_local_storage_backend() { let watch_dir = tempfile::TempDir::new().expect("Failed to create watch dir"); let storage_dir = tempfile::TempDir::new().expect("Failed to create storage dir"); let test_file = watch_dir.path().join("ES.FUT_ohlcv-1m_2024-01-02.dbn"); fs::write(&test_file, b"fake dbn content for upload test") .await .expect("Failed to write"); // Use storage::local::LocalStorage as a real backend let local_config = storage::local::LocalStorageConfig { base_path: storage_dir.path().to_path_buf(), ..Default::default() }; let local_storage = storage::local::LocalStorage::new(local_config) .await .expect("Failed to create local storage"); let backend: Arc = Arc::new(local_storage); let config = DbnUploaderConfig { watch_path: watch_dir.path().to_path_buf(), compression_enabled: true, deduplication_enabled: false, ..DbnUploaderConfig::default() }; let uploader = DbnUploader::new(config) .await .expect("Failed to create uploader") .with_storage(backend.clone()); // First upload should succeed uploader .upload_file(&test_file) .await .expect("Upload should succeed"); // Verify the file was stored let key = "training-data/ES.FUT_ohlcv-1m_2024-01-02.dbn.gz"; let exists = backend.exists(key).await.expect("Exists check failed"); assert!(exists, "Uploaded file should exist in storage backend"); } #[tokio::test] async fn test_deduplication_with_storage_backend() { let watch_dir = tempfile::TempDir::new().expect("Failed to create watch dir"); let storage_dir = tempfile::TempDir::new().expect("Failed to create storage dir"); let test_file = watch_dir.path().join("ES.FUT_ohlcv-1m_2024-01-02.dbn"); fs::write(&test_file, b"dbn data") .await .expect("Failed to write"); let local_config = storage::local::LocalStorageConfig { base_path: storage_dir.path().to_path_buf(), ..Default::default() }; let local_storage = storage::local::LocalStorage::new(local_config) .await .expect("Failed to create local storage"); let backend: Arc = Arc::new(local_storage); let config = DbnUploaderConfig { watch_path: watch_dir.path().to_path_buf(), compression_enabled: false, deduplication_enabled: true, ..DbnUploaderConfig::default() }; let uploader = DbnUploader::new(config) .await .expect("Failed to create uploader") .with_storage(backend.clone()); // File does not exist yet, should_upload should be true let should = uploader .should_upload_file(&test_file) .await .expect("Check failed"); assert!(should, "Should upload new file"); // Upload once uploader .upload_file(&test_file) .await .expect("First upload should succeed"); // Now the file exists, should_upload should be false let should = uploader .should_upload_file(&test_file) .await .expect("Check failed"); assert!(!should, "Should NOT upload duplicate file"); } }