schwab_cli/agent/backtest/
prefetch.rs

1use std::path::Path;
2use std::sync::Arc;
3use std::time::Duration;
4
5use anyhow::{Context, Result};
6use chrono::{Datelike, NaiveDate, Utc};
7use schwab_market_data::MarketDataApi;
8use tokio::time::sleep;
9
10use crate::agent::paths::backtest_cache_path;
11use crate::rules::RulesConfig;
12
13use super::cache::BacktestCache;
14
15pub fn symbols_for_prefetch(rules: &RulesConfig) -> Vec<String> {
16    let mut out: Vec<String> = rules
17        .watchlist_items()
18        .into_iter()
19        .map(|w| w.symbol.trim().to_uppercase())
20        .filter(|s| !s.is_empty())
21        .collect();
22    let bench = rules.regime.benchmark_symbol.trim().to_uppercase();
23    if !bench.is_empty() && !out.contains(&bench) {
24        out.push(bench);
25    }
26    let vix = rules.regime.vix_symbol.trim().to_uppercase();
27    if !vix.is_empty() && !out.contains(&vix) {
28        out.push(vix);
29    }
30    // Schwab sometimes needs bare VIX without $.
31    if !out.iter().any(|s| s == "VIX") {
32        out.push("VIX".into());
33    }
34    if !out.iter().any(|s| s == "$VIX") {
35        out.push("$VIX".into());
36    }
37    out.sort();
38    out.dedup();
39    out
40}
41
42pub async fn prefetch_daily_bars(
43    market: &Arc<MarketDataApi>,
44    rules: &RulesConfig,
45    rules_path: &Path,
46    from: NaiveDate,
47    to: NaiveDate,
48    force: bool,
49) -> Result<BacktestCache> {
50    let cache_path = backtest_cache_path(rules_path);
51    if cache_path.is_file() && !force {
52        let existing = BacktestCache::load(&cache_path)?;
53        if existing.from <= from && existing.to >= to {
54            let needed = symbols_for_prefetch(rules);
55            let has_all = needed.iter().all(|s| existing.symbols.contains_key(s));
56            if has_all {
57                return Ok(existing);
58            }
59        }
60    }
61
62    let symbols = symbols_for_prefetch(rules);
63    let chunks = date_chunks(from, to);
64
65    let mut cache = if cache_path.is_file() {
66        BacktestCache::load(&cache_path).unwrap_or_else(|_| BacktestCache::new(from, to))
67    } else {
68        BacktestCache::new(from, to)
69    };
70    cache.from = cache.from.min(from);
71    cache.to = cache.to.max(to);
72    cache.fetched_at = Utc::now();
73
74    for (i, symbol) in symbols.iter().enumerate() {
75        if i > 0 {
76            sleep(Duration::from_millis(550)).await;
77        }
78        for (chunk_from, chunk_to) in &chunks {
79            let start_ms = naive_date_start_ms(*chunk_from)?;
80            let end_ms = naive_date_end_ms(*chunk_to)?;
81            match market
82                .price_history()
83                .get(
84                    symbol,
85                    Some("year"),
86                    Some(1),
87                    Some("daily"),
88                    Some(1),
89                    Some(start_ms),
90                    Some(end_ms),
91                    Some(false),
92                    Some(false),
93                )
94                .await
95            {
96                Ok(history) => {
97                    cache.ingest_schwab_history(symbol, &history);
98                    tracing::info!(
99                        "prefetch {symbol} {chunk_from}..{chunk_to}: {} bars",
100                        cache.symbols.get(symbol).map(|v| v.len()).unwrap_or(0)
101                    );
102                }
103                Err(e) => {
104                    tracing::warn!("prefetch {symbol} failed ({chunk_from}..{chunk_to}): {e:#}");
105                }
106            }
107            sleep(Duration::from_millis(350)).await;
108        }
109    }
110
111    cache.save(&cache_path)?;
112    Ok(cache)
113}
114
115fn date_chunks(from: NaiveDate, to: NaiveDate) -> Vec<(NaiveDate, NaiveDate)> {
116    let mut chunks = Vec::new();
117    let mut start = from;
118    while start <= to {
119        let year_end = NaiveDate::from_ymd_opt(start.year(), 12, 31).unwrap_or(to);
120        let end = year_end.min(to);
121        chunks.push((start, end));
122        if end >= to {
123            break;
124        }
125        start = end
126            .succ_opt()
127            .unwrap_or(to);
128    }
129    chunks
130}
131
132fn naive_date_start_ms(d: NaiveDate) -> Result<i64> {
133    let dt = d
134        .and_hms_opt(0, 0, 0)
135        .context("invalid date")?
136        .and_utc();
137    Ok(dt.timestamp_millis())
138}
139
140fn naive_date_end_ms(d: NaiveDate) -> Result<i64> {
141    let dt = d
142        .and_hms_opt(23, 59, 59)
143        .context("invalid date")?
144        .and_utc();
145    Ok(dt.timestamp_millis())
146}
147
148pub fn default_from_to(from: Option<NaiveDate>, to: Option<NaiveDate>) -> Result<(NaiveDate, NaiveDate)> {
149    let to = to.unwrap_or_else(|| {
150        let today = Utc::now().date_naive();
151        today.pred_opt().unwrap_or(today)
152    });
153    let from = from.unwrap_or_else(|| {
154        NaiveDate::from_ymd_opt(to.year() - 1, to.month(), to.day().min(28))
155            .unwrap_or(to - chrono::Duration::days(365))
156    });
157    anyhow::ensure!(from <= to, "from ({from}) must be <= to ({to})");
158    Ok((from, to))
159}
160
161pub fn parse_ymd(s: &str) -> Result<NaiveDate> {
162    NaiveDate::parse_from_str(s.trim(), "%Y-%m-%d")
163        .with_context(|| format!("invalid date {s:?}, expected YYYY-MM-DD"))
164}
schwab_cli/agent/backtest/prefetch.rs

schwab_cli/agent/backtest/
prefetch.rs
