#![allow( clippy::assertions_on_constants, clippy::assertions_on_result_states, clippy::clone_on_copy, clippy::decimal_literal_representation, clippy::doc_markdown, clippy::empty_line_after_doc_comments, clippy::field_reassign_with_default, clippy::get_unwrap, clippy::identity_op, clippy::inconsistent_digit_grouping, clippy::indexing_slicing, clippy::integer_division, clippy::len_zero, clippy::let_underscore_must_use, clippy::manual_div_ceil, clippy::manual_let_else, clippy::manual_range_contains, clippy::modulo_arithmetic, clippy::needless_range_loop, clippy::non_ascii_literal, clippy::redundant_clone, clippy::shadow_reuse, clippy::shadow_same, clippy::shadow_unrelated, clippy::single_match_else, clippy::str_to_string, clippy::string_slice, clippy::tests_outside_test_module, clippy::too_many_lines, clippy::unnecessary_wraps, clippy::unseparated_literal_suffix, clippy::use_debug, clippy::useless_vec, clippy::wildcard_enum_match_arm, clippy::else_if_without_else, clippy::expect_used, clippy::missing_const_for_fn, clippy::similar_names, clippy::type_complexity, clippy::collapsible_else_if, clippy::doc_lazy_continuation, clippy::items_after_test_module, clippy::map_clone, clippy::multiple_unsafe_ops_per_block, clippy::unwrap_or_default, clippy::assign_op_pattern, clippy::needless_borrow, clippy::println_empty_string, clippy::unnecessary_cast, clippy::used_underscore_binding, clippy::create_dir, clippy::implicit_saturating_sub, clippy::exit, clippy::expect_fun_call, clippy::too_many_arguments, clippy::unnecessary_map_or, clippy::unwrap_used, dead_code, unused_imports, unused_variables, clippy::cloned_ref_to_slice_refs, clippy::neg_multiply, clippy::while_let_loop, clippy::bool_assert_comparison, clippy::excessive_precision, clippy::trivially_copy_pass_by_ref, clippy::op_ref, clippy::redundant_closure, clippy::unnecessary_lazy_evaluations, clippy::if_then_some_else_none, clippy::unnecessary_to_owned, clippy::single_component_path_imports, )] //! Download quarterly-chunked Databento data (any schema) with optional MinIO upload //! //! Downloads futures data for ML model training, split into ~90-day quarterly //! files for efficient caching and resume support. Schema (ohlcv-1m, mbp-10, //! etc.) is read from the universe TOML config. //! //! Configuration is read from a universe TOML file (default: //! `config/universe-futures-baseline.toml`). //! //! Usage: //! # Dry run (preview config and cost estimate) //! cargo run -p ml --example download_baseline --release -- --dry-run //! //! # Download OHLCV with confirmation prompt //! cargo run -p ml --example download_baseline --release //! //! # Download MBP-10 for ES.FUT and upload to MinIO //! cargo run -p ml --example download_baseline --release -- \ //! --universe-config config/universe-es-mbp10.toml \ //! --output-dir /tmp/mbp10 \ //! --rclone-dest "s3:foxhunt-training-data/futures-baseline-mbp10" \ //! --yes use anyhow::{Context, Result}; use chrono::{Datelike, NaiveDate}; use clap::Parser; use databento::historical::timeseries::GetRangeToFileParams; use databento::historical::DateTimeRange; use databento::HistoricalClient; use dbn::{Schema, SType}; use serde::Deserialize; use std::env; use std::fs; use std::path::PathBuf; use std::process::Command; use std::sync::Arc; use std::time::Instant; use tokio::sync::Semaphore; // --------------------------------------------------------------------------- // CLI // --------------------------------------------------------------------------- #[derive(Debug, Parser)] #[command( name = "download_baseline", about = "Download quarterly Databento data (any schema) with optional MinIO upload" )] struct Opts { /// Output directory for downloaded files (env: FOXHUNT_DATA_DIR) #[arg(long, env = "FOXHUNT_DATA_DIR", default_value = "test_data/futures-baseline")] output_dir: String, /// Path to universe configuration TOML #[arg(long, default_value = "config/universe-futures-baseline.toml")] universe_config: String, /// Preview only, do not download #[arg(long)] dry_run: bool, /// Skip confirmation prompt #[arg(long)] yes: bool, /// Upload each file to MinIO/S3 via rclone after download. /// Format: "s3:bucket/prefix" (e.g. "s3:foxhunt-training-data/futures-baseline-mbp10") #[arg(long)] rclone_dest: Option, /// Number of parallel downloads (default: 1 = sequential) #[arg(long, default_value = "1")] parallel: usize, } // --------------------------------------------------------------------------- // Universe config // --------------------------------------------------------------------------- #[derive(Debug, Deserialize)] struct UniverseConfig { universe: UniverseMeta, symbols: Vec, } #[derive(Debug, Deserialize)] struct UniverseMeta { name: String, date_range_start: String, date_range_end: String, databento_dataset: String, #[serde(default = "default_schema")] databento_schema: String, } fn default_schema() -> String { "ohlcv-1m".to_owned() } /// Map config schema string to dbn::Schema enum. fn parse_schema(s: &str) -> Result { match s { "ohlcv-1s" => Ok(Schema::Ohlcv1S), "ohlcv-1m" => Ok(Schema::Ohlcv1M), "ohlcv-1h" => Ok(Schema::Ohlcv1H), "ohlcv-1d" => Ok(Schema::Ohlcv1D), "trades" => Ok(Schema::Trades), "mbp-1" => Ok(Schema::Mbp1), "mbp-10" => Ok(Schema::Mbp10), "bbo-1s" => Ok(Schema::Bbo1S), "bbo-1m" => Ok(Schema::Bbo1M), other => Err(anyhow::anyhow!("Unknown Databento schema: {}", other)), } } #[derive(Debug, Deserialize)] struct SymbolEntry { symbol: String, } // --------------------------------------------------------------------------- // Quarter representation // --------------------------------------------------------------------------- #[derive(Debug, Clone)] struct Quarter { label: String, start: NaiveDate, end: NaiveDate, } /// Split a date range into calendar-quarter chunks. /// /// Each chunk begins at the later of `start` and the quarter boundary, and /// ends at the earlier of `end` and the next quarter boundary. fn generate_quarters(start: NaiveDate, end: NaiveDate) -> Vec { let mut quarters = Vec::new(); let mut cursor = start; while cursor < end { #[allow(clippy::integer_division)] let q = (cursor.month() - 1) / 3 + 1; // 1..4 let year = cursor.year(); let label = format!("{}-Q{}", year, q); // Next quarter boundary let next_q_start = if q == 4 { NaiveDate::from_ymd_opt(year + 1, 1, 1) } else { NaiveDate::from_ymd_opt(year, q * 3 + 1, 1) }; let chunk_end = match next_q_start { Some(nq) if nq < end => nq, _ => end, }; quarters.push(Quarter { label, start: cursor, end: chunk_end, }); cursor = chunk_end; } quarters } // --------------------------------------------------------------------------- // Conversion helpers // --------------------------------------------------------------------------- /// Convert a `chrono::NaiveDate` (interpreted as midnight UTC) to UNIX /// nanoseconds for the Databento `DateTimeRange` API. fn naive_date_to_unix_nanos(date: NaiveDate) -> u64 { let ts = date .and_hms_opt(0, 0, 0) .map(|dt| dt.and_utc().timestamp()) .unwrap_or(0); (ts as u64).saturating_mul(1_000_000_000) } /// Build a `DateTimeRange` from two `NaiveDate`s. fn date_range(start: NaiveDate, end: NaiveDate) -> Result { let start_ns = naive_date_to_unix_nanos(start); let end_ns = naive_date_to_unix_nanos(end); DateTimeRange::try_from((start_ns, end_ns)).map_err(|e| anyhow::anyhow!("{}", e)) } // --------------------------------------------------------------------------- // Download logic // --------------------------------------------------------------------------- struct DownloadStats { successful: usize, failed: usize, skipped: usize, total_bytes: u64, } impl DownloadStats { const fn new() -> Self { Self { successful: 0, failed: 0, skipped: 0, total_bytes: 0, } } } async fn download_quarter( client: &mut HistoricalClient, symbol: &str, quarter: &Quarter, dataset: &str, schema: Schema, output_dir: &std::path::Path, ) -> Result { let filename = format!("{}_{}.dbn.zst", symbol, quarter.label); let symbol_dir = output_dir.join(symbol); fs::create_dir_all(&symbol_dir).context("Failed to create symbol directory")?; let file_path = symbol_dir.join(&filename); // Resume support: skip if file exists and is non-empty if file_path.exists() { let meta = fs::metadata(&file_path).context("Failed to read file metadata")?; if meta.len() > 0 { return Ok(meta.len()); } } let dt_range = date_range(quarter.start, quarter.end)?; let params = GetRangeToFileParams::builder() .dataset(dataset.to_owned()) .symbols(vec![symbol.to_owned()]) .schema(schema) .stype_in(SType::Parent) .date_time_range(dt_range) .path(file_path.clone()) .build(); let _decoder = client .timeseries() .get_range_to_file(¶ms) .await .with_context(|| { format!( "Failed to download {} for {}", quarter.label, symbol ) })?; let meta = fs::metadata(&file_path).context("Failed to read downloaded file metadata")?; Ok(meta.len()) } /// Upload a single file to MinIO/S3 via rclone. /// /// Runs `rclone copyto ://`. /// Rclone flags for MinIO are sourced from RCLONE_* env vars (set in K8s job). fn rclone_upload(local_path: &std::path::Path, dest: &str, symbol: &str) -> Result<()> { let filename = local_path .file_name() .and_then(|n| n.to_str()) .ok_or_else(|| anyhow::anyhow!("Invalid filename: {:?}", local_path))?; let remote_path = format!(":{}/{}/{}", dest, symbol, filename); let status = Command::new("rclone") .args(["copyto", &local_path.to_string_lossy(), &remote_path]) .status() .context("Failed to execute rclone (is it installed?)")?; if !status.success() { anyhow::bail!("rclone upload failed for {}", remote_path); } println!(" [MINIO] uploaded → {}", remote_path); // Delete local file after successful upload to save scratch space if let Err(e) = fs::remove_file(local_path) { println!(" [WARN] failed to clean up local file: {}", e); } Ok(()) } // --------------------------------------------------------------------------- // Main // --------------------------------------------------------------------------- #[tokio::main] async fn main() -> Result<()> { let opts = Opts::parse(); // Load .env for API key dotenv::dotenv().ok(); // ------------------------------------------------------------------ // Parse universe config // ------------------------------------------------------------------ let config_text = fs::read_to_string(&opts.universe_config).with_context(|| { format!("Failed to read universe config: {}", opts.universe_config) })?; let config: UniverseConfig = toml::from_str(&config_text).context("Failed to parse universe config TOML")?; let start_date = NaiveDate::parse_from_str(&config.universe.date_range_start, "%Y-%m-%d") .context("Failed to parse date_range_start")?; let end_date = NaiveDate::parse_from_str(&config.universe.date_range_end, "%Y-%m-%d") .context("Failed to parse date_range_end")?; let schema = parse_schema(&config.universe.databento_schema)?; let symbols: Vec<&str> = config.symbols.iter().map(|s| s.symbol.as_str()).collect(); let quarters = generate_quarters(start_date, end_date); let total_days = (end_date - start_date).num_days(); // ------------------------------------------------------------------ // Print config summary // ------------------------------------------------------------------ println!("================================================================================"); println!("Futures Baseline Download - Databento (Quarterly Chunks)"); println!("================================================================================"); println!(); println!("Universe: {}", config.universe.name); println!("Dataset: {}", config.universe.databento_dataset); println!("Schema: {}", config.universe.databento_schema); println!("Date range: {} to {}", start_date, end_date); println!("Total days: {}", total_days); println!( "Symbols: {} ({})", symbols.len(), symbols.join(", ") ); println!("Quarters: {}", quarters.len()); println!("Output: {}", opts.output_dir); if let Some(ref dest) = opts.rclone_dest { println!("MinIO dest: :{}", dest); } println!(); // Print quarter breakdown println!("Quarter breakdown:"); for q in &quarters { let days = (q.end - q.start).num_days(); println!(" {} : {} to {} ({} days)", q.label, q.start, q.end, days); } println!(); // Cost estimate ($0.12 per symbol per day) let total_files = symbols.len() * quarters.len(); let estimated_cost = total_days as f64 * symbols.len() as f64 * 0.12; println!("Total files: {}", total_files); println!("Estimated cost: ${:.2} ({} symbols x {} days x $0.12/sym/day)", estimated_cost, symbols.len(), total_days); println!(); // ------------------------------------------------------------------ // Dry run exit // ------------------------------------------------------------------ if opts.dry_run { println!("[DRY RUN] Preview complete. Remove --dry-run to execute downloads."); return Ok(()); } // ------------------------------------------------------------------ // Confirmation // ------------------------------------------------------------------ if !opts.yes { println!( "This will download data and may incur costs (~${:.2}).", estimated_cost ); print!("Proceed? (yes/no): "); std::io::Write::flush(&mut std::io::stdout())?; let mut input = String::new(); std::io::stdin().read_line(&mut input)?; let trimmed = input.trim(); if !trimmed.eq_ignore_ascii_case("yes") && !trimmed.eq_ignore_ascii_case("y") { println!("Download cancelled."); return Ok(()); } println!(); } // ------------------------------------------------------------------ // Initialize client // ------------------------------------------------------------------ let api_key = env::var("DATABENTO_API_KEY") .context("DATABENTO_API_KEY not found in environment or .env file")?; let mut client = HistoricalClient::builder() .key(api_key.clone()) .map_err(|e| anyhow::anyhow!("Invalid API key: {}", e))? .build() .map_err(|e| anyhow::anyhow!("Failed to build Databento client: {}", e))?; let output_dir = PathBuf::from(&opts.output_dir); fs::create_dir_all(&output_dir)?; println!("Databento client initialized. Output: {}", opts.output_dir); println!(); // ------------------------------------------------------------------ // Download loop (parallel when --parallel > 1) // ------------------------------------------------------------------ let mut stats = DownloadStats::new(); let mut failed_items: Vec = Vec::new(); let parallel = opts.parallel.max(1); if parallel > 1 { println!("Downloading with {} parallel workers", parallel); println!(); } for symbol in &symbols { println!("[{}]", symbol); if parallel <= 1 { // Sequential path (original behavior) for quarter in &quarters { let filename = format!("{}_{}.dbn.zst", symbol, quarter.label); let symbol_dir = output_dir.join(symbol); let file_path = symbol_dir.join(&filename); if file_path.exists() { match fs::metadata(&file_path) { Ok(meta) if meta.len() > 0 => { println!(" [SKIP] {} already exists ({} bytes)", filename, meta.len()); stats.skipped += 1; stats.total_bytes += meta.len(); if let Some(ref dest) = opts.rclone_dest { if let Err(e) = rclone_upload(&file_path, dest, symbol) { println!(" [WARN] rclone upload failed: {}", e); } } continue; } _ => {} } } let t0 = Instant::now(); match download_quarter(&mut client, symbol, quarter, &config.universe.databento_dataset, schema, &output_dir).await { Ok(size) => { let elapsed = t0.elapsed().as_secs_f64(); println!(" [OK] {} -- {} bytes, {:.1}s", filename, size, elapsed); stats.successful += 1; stats.total_bytes += size; if let Some(ref dest) = opts.rclone_dest { if let Err(e) = rclone_upload(&file_path, dest, symbol) { println!(" [WARN] rclone upload failed: {}", e); } } } Err(e) => { println!(" [FAIL] {} -- {}", filename, e); stats.failed += 1; failed_items.push(format!("{}/{}", symbol, quarter.label)); } } } } else { // Parallel path — one HistoricalClient per task let sem = Arc::new(Semaphore::new(parallel)); let mut handles = Vec::new(); for quarter in quarters.clone() { let symbol = symbol.to_string(); let dataset = config.universe.databento_dataset.clone(); let out = output_dir.clone(); let rclone_dest = opts.rclone_dest.clone(); let api_key = api_key.clone(); let sem = Arc::clone(&sem); handles.push(tokio::spawn(async move { let _permit = sem.acquire().await.expect("semaphore closed"); let filename = format!("{}_{}.dbn.zst", symbol, quarter.label); let symbol_dir = out.join(&symbol); let file_path = symbol_dir.join(&filename); // Skip existing if file_path.exists() { if let Ok(meta) = fs::metadata(&file_path) { if meta.len() > 0 { println!(" [SKIP] {} already exists ({} bytes)", filename, meta.len()); if let Some(ref dest) = rclone_dest { if let Err(e) = rclone_upload(&file_path, dest, &symbol) { println!(" [WARN] rclone upload failed: {}", e); } } return (true, false, meta.len(), quarter.label.clone()); } } } // Each task gets its own client let mut task_client = match HistoricalClient::builder() .key(api_key) .map_err(|e| anyhow::anyhow!("{}", e)) .and_then(|b| b.build().map_err(|e| anyhow::anyhow!("{}", e))) { Ok(c) => c, Err(e) => { println!(" [FAIL] {} -- client error: {}", filename, e); return (false, true, 0u64, quarter.label.clone()); } }; let t0 = Instant::now(); match download_quarter(&mut task_client, &symbol, &quarter, &dataset, schema, &out).await { Ok(size) => { let elapsed = t0.elapsed().as_secs_f64(); println!(" [OK] {} -- {} bytes, {:.1}s", filename, size, elapsed); if let Some(ref dest) = rclone_dest { if let Err(e) = rclone_upload(&file_path, dest, &symbol) { println!(" [WARN] rclone upload failed: {}", e); } } (false, false, size, quarter.label.clone()) } Err(e) => { println!(" [FAIL] {} -- {}", filename, e); (false, true, 0, quarter.label.clone()) } } })); } for handle in handles { match handle.await { Ok((skipped, failed, bytes, label)) => { if skipped { stats.skipped += 1; } else if failed { stats.failed += 1; failed_items.push(format!("{}/{}", symbol, label)); } else { stats.successful += 1; } stats.total_bytes += bytes; } Err(e) => { println!(" [FAIL] task panicked: {}", e); stats.failed += 1; } } } } println!(); } // ------------------------------------------------------------------ // Summary // ------------------------------------------------------------------ println!("================================================================================"); println!("DOWNLOAD SUMMARY"); println!("================================================================================"); println!(); println!("Successful: {}", stats.successful); println!("Skipped: {}", stats.skipped); println!("Failed: {}", stats.failed); println!( "Total size: {:.1} MB", stats.total_bytes as f64 / 1_048_576.0 ); println!(); if !failed_items.is_empty() { println!("Failed downloads:"); for item in &failed_items { println!(" - {}", item); } println!(); println!("Re-run the command to retry failed downloads (existing files are skipped)."); } Ok(()) } #[cfg(test)] mod tests { use super::*; #[test] fn test_generate_quarters_full_range() { let start = NaiveDate::from_ymd_opt(2024, 3, 1).unwrap_or_default(); let end = NaiveDate::from_ymd_opt(2024, 12, 31).unwrap_or_default(); let qs = generate_quarters(start, end); assert_eq!(qs.len(), 4); // Q1(partial), Q2, Q3, Q4 assert_eq!(qs.first().map(|q| q.label.as_str()), Some("2024-Q1")); assert_eq!(qs.last().map(|q| q.label.as_str()), Some("2024-Q4")); } #[test] fn test_generate_quarters_cross_year() { let start = NaiveDate::from_ymd_opt(2024, 10, 1).unwrap_or_default(); let end = NaiveDate::from_ymd_opt(2025, 4, 1).unwrap_or_default(); let qs = generate_quarters(start, end); assert_eq!(qs.len(), 3); // 2024-Q4, 2025-Q1, 2025-Q2(partial) assert_eq!(qs.first().map(|q| q.label.as_str()), Some("2024-Q4")); assert_eq!(qs.last().map(|q| q.label.as_str()), Some("2025-Q2")); } #[test] fn test_generate_quarters_single_day() { let start = NaiveDate::from_ymd_opt(2024, 6, 15).unwrap_or_default(); let end = NaiveDate::from_ymd_opt(2024, 6, 16).unwrap_or_default(); let qs = generate_quarters(start, end); assert_eq!(qs.len(), 1); assert_eq!(qs.first().map(|q| q.label.as_str()), Some("2024-Q2")); } #[test] fn test_naive_date_to_unix_nanos() { let date = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap_or_default(); let nanos = naive_date_to_unix_nanos(date); // 2024-01-01 00:00:00 UTC = 1704067200 seconds assert_eq!(nanos, 1_704_067_200_000_000_000); } }