3d63f73c08
3 retries with exponential backoff (2s, 4s, 8s) for transient connection failures (IncompleteBody/UnexpectedEof). Partial files are deleted before retry so skip-existing logic works correctly. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
726 lines
26 KiB
Rust
726 lines
26 KiB
Rust
#![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<String>,
|
|
|
|
/// 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<SymbolEntry>,
|
|
}
|
|
|
|
#[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<Schema> {
|
|
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<Quarter> {
|
|
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<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 {
|
|
const fn new() -> Self {
|
|
Self {
|
|
successful: 0,
|
|
failed: 0,
|
|
skipped: 0,
|
|
total_bytes: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Maximum retries for transient download failures (connection resets, EOF).
|
|
const MAX_RETRIES: u32 = 3;
|
|
|
|
async fn download_quarter(
|
|
client: &mut HistoricalClient,
|
|
symbol: &str,
|
|
quarter: &Quarter,
|
|
dataset: &str,
|
|
schema: Schema,
|
|
output_dir: &std::path::Path,
|
|
) -> 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 mut last_err = None;
|
|
for attempt in 0..=MAX_RETRIES {
|
|
if attempt > 0 {
|
|
// Remove partial file from previous failed attempt
|
|
if file_path.exists() {
|
|
let _ = fs::remove_file(&file_path);
|
|
}
|
|
let delay = std::time::Duration::from_secs(2u64.pow(attempt));
|
|
println!(
|
|
" [RETRY] {} attempt {}/{} in {}s",
|
|
filename,
|
|
attempt + 1,
|
|
MAX_RETRIES + 1,
|
|
delay.as_secs()
|
|
);
|
|
tokio::time::sleep(delay).await;
|
|
}
|
|
|
|
let params = GetRangeToFileParams::builder()
|
|
.dataset(dataset.to_owned())
|
|
.symbols(vec![symbol.to_owned()])
|
|
.schema(schema)
|
|
.stype_in(SType::Parent)
|
|
.date_time_range(dt_range.clone())
|
|
.path(file_path.clone())
|
|
.build();
|
|
|
|
match client.timeseries().get_range_to_file(¶ms).await {
|
|
Ok(_decoder) => {
|
|
let meta =
|
|
fs::metadata(&file_path).context("Failed to read downloaded file metadata")?;
|
|
return Ok(meta.len());
|
|
}
|
|
Err(e) => {
|
|
last_err = Some(anyhow::anyhow!(
|
|
"Failed to download {} for {}: {}",
|
|
quarter.label,
|
|
symbol,
|
|
e
|
|
));
|
|
}
|
|
}
|
|
}
|
|
|
|
Err(last_err.unwrap())
|
|
}
|
|
|
|
/// Upload a single file to MinIO/S3 via rclone.
|
|
///
|
|
/// Runs `rclone copyto <local_path> :<dest>/<symbol>/<filename>`.
|
|
/// 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<String> = 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);
|
|
}
|
|
}
|