6befda18a5
Adds `ml/examples/download_baseline.rs` that downloads 730 days of Databento OHLCV-1m data in quarterly chunks for 4 CME futures symbols. Features: - Reads universe config from TOML (symbols, date range, dataset) - Splits date range into calendar-quarter chunks (~90 days each) - Resume support: skips existing non-empty files - Uses `get_range_to_file` for streaming writes to .dbn.zst - Dry-run mode with cost estimate ($0.12/symbol/day) - Confirmation prompt (skippable with --yes) - Per-file progress with timing and byte counts - Failure-tolerant: logs errors and continues Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
474 lines
16 KiB
Rust
474 lines
16 KiB
Rust
//! Download 730 days of Databento OHLCV-1m data in quarterly chunks
|
|
//!
|
|
//! Downloads futures baseline data for ML model training, split into
|
|
//! ~90-day quarterly files for efficient caching and resume support.
|
|
//!
|
|
//! 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 with confirmation prompt
|
|
//! cargo run -p ml --example download_baseline --release
|
|
//!
|
|
//! # Skip confirmation prompt
|
|
//! cargo run -p ml --example download_baseline --release -- --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;
|
|
use serde::Deserialize;
|
|
use std::env;
|
|
use std::fs;
|
|
use std::path::PathBuf;
|
|
use std::time::Instant;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// CLI
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[derive(Debug, Parser)]
|
|
#[command(
|
|
name = "download_baseline",
|
|
about = "Download quarterly OHLCV-1m futures data from Databento"
|
|
)]
|
|
struct Opts {
|
|
/// Output directory for downloaded files
|
|
#[arg(long, default_value = "data/cache/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,
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Universe config
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct UniverseConfig {
|
|
universe: UniverseMeta,
|
|
symbols: Vec<SymbolEntry>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct UniverseMeta {
|
|
name: String,
|
|
#[allow(dead_code)]
|
|
description: String,
|
|
date_range_start: String,
|
|
date_range_end: String,
|
|
#[allow(dead_code)]
|
|
bar_size: String,
|
|
databento_dataset: String,
|
|
#[allow(dead_code)]
|
|
databento_schema: String,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct SymbolEntry {
|
|
symbol: String,
|
|
#[allow(dead_code)]
|
|
exchange: String,
|
|
#[allow(dead_code)]
|
|
asset_class: String,
|
|
#[allow(dead_code)]
|
|
trading_symbol: Option<String>,
|
|
#[allow(dead_code)]
|
|
min_daily_volume: Option<f64>,
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 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<Quarter> {
|
|
let mut quarters = Vec::new();
|
|
let mut cursor = start;
|
|
|
|
while cursor < end {
|
|
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<DateTimeRange> {
|
|
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 {
|
|
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,
|
|
output_dir: &PathBuf,
|
|
) -> Result<u64> {
|
|
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_string())
|
|
.symbols(vec![symbol.to_string()])
|
|
.schema(Schema::Ohlcv1M)
|
|
.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())
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 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 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: ohlcv-1m");
|
|
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);
|
|
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)
|
|
.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
|
|
// ------------------------------------------------------------------
|
|
let mut stats = DownloadStats::new();
|
|
let mut failed_items: Vec<String> = Vec::new();
|
|
|
|
for symbol in &symbols {
|
|
println!("[{}]", symbol);
|
|
|
|
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);
|
|
|
|
// Check for existing file before calling download
|
|
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();
|
|
continue;
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
let t0 = Instant::now();
|
|
match download_quarter(
|
|
&mut client,
|
|
symbol,
|
|
quarter,
|
|
&config.universe.databento_dataset,
|
|
&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;
|
|
}
|
|
Err(e) => {
|
|
println!(" [FAIL] {} -- {}", filename, e);
|
|
stats.failed += 1;
|
|
failed_items.push(format!("{}/{}", symbol, quarter.label));
|
|
}
|
|
}
|
|
}
|
|
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);
|
|
}
|
|
}
|