1use std::cell::Cell;
21use std::time::{Duration, Instant};
22
23use serde_json::Value;
24
25#[derive(Clone, Copy, PartialEq, Eq, Debug)]
27pub enum WriteMode {
28 Overwrite,
30 Append,
32 Resume,
34}
35
36#[derive(Debug, Clone)]
38pub enum HttpError {
39 Status(u16),
41 Transport(String),
43}
44
45impl std::fmt::Display for HttpError {
46 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47 match self {
48 HttpError::Status(code) => write!(f, "HTTP {code}"),
49 HttpError::Transport(msg) => write!(f, "transport error: {msg}"),
50 }
51 }
52}
53
54impl HttpError {
55 pub fn retryable(&self) -> bool {
57 match self {
58 HttpError::Transport(_) => true,
59 HttpError::Status(code) => *code == 429 || (500..600).contains(code),
60 }
61 }
62}
63
64pub trait HttpClient {
66 fn get(&self, url: &str) -> Result<Vec<u8>, HttpError>;
68}
69
70pub trait RetrySettings {
72 fn rate_limit_per_min(&self) -> u32;
74 fn max_retries(&self) -> u32;
76 fn backoff_base(&self) -> Duration;
78}
79
80#[cfg(feature = "ureq")]
82pub struct UreqClient;
83
84#[cfg(feature = "ureq")]
85impl UreqClient {
86 pub fn new() -> Self {
87 UreqClient
88 }
89}
90
91#[cfg(feature = "ureq")]
92impl Default for UreqClient {
93 fn default() -> Self {
94 Self::new()
95 }
96}
97
98#[cfg(feature = "ureq")]
99impl HttpClient for UreqClient {
100 fn get(&self, url: &str) -> Result<Vec<u8>, HttpError> {
101 match ureq::get(url).call() {
102 Ok(resp) => resp
103 .into_body()
104 .with_config()
105 .limit(256 * 1024 * 1024)
106 .read_to_vec()
107 .map_err(|e| HttpError::Transport(e.to_string())),
108 Err(ureq::Error::StatusCode(code)) => Err(HttpError::Status(code)),
109 Err(e) => Err(HttpError::Transport(e.to_string())),
110 }
111 }
112}
113
114pub fn redact(url: &str) -> String {
117 for key in ["api_token=", "token=", "apikey=", "api_key="] {
118 if let Some(i) = url.find(key) {
119 let start = i + key.len();
120 let end = url[start..]
121 .find('&')
122 .map(|j| start + j)
123 .unwrap_or(url.len());
124 return format!("{}***{}", &url[..start], &url[end..]);
125 }
126 }
127 url.to_string()
128}
129
130pub struct Fetcher<'a, H: HttpClient, C: RetrySettings> {
134 http: &'a H,
135 cfg: &'a C,
136 last_request: Cell<Option<Instant>>,
137}
138
139impl<'a, H: HttpClient, C: RetrySettings> Fetcher<'a, H, C> {
140 pub fn new(http: &'a H, cfg: &'a C) -> Self {
141 Fetcher {
142 http,
143 cfg,
144 last_request: Cell::new(None),
145 }
146 }
147
148 fn throttle(&self) {
149 let rpm = self.cfg.rate_limit_per_min();
150 if rpm == 0 {
151 return;
152 }
153 let min_interval = Duration::from_secs_f64(60.0 / rpm as f64);
154 if let Some(prev) = self.last_request.get() {
155 let elapsed = prev.elapsed();
156 if elapsed < min_interval {
157 std::thread::sleep(min_interval - elapsed);
158 }
159 }
160 self.last_request.set(Some(Instant::now()));
161 }
162
163 pub fn get(&self, url: &str) -> Result<Vec<u8>, String> {
166 let max_retries = self.cfg.max_retries();
167 let mut attempt = 0u32;
168 loop {
169 self.throttle();
170 match self.http.get(url) {
171 Ok(body) => return Ok(body),
172 Err(e) if e.retryable() && attempt < max_retries => {
173 let wait = self.cfg.backoff_base() * 2u32.pow(attempt.min(16));
174 eprintln!(
175 " retry {}/{} after {}: {} ({:?})",
176 attempt + 1,
177 max_retries,
178 e,
179 redact(url),
180 wait
181 );
182 if !wait.is_zero() {
183 std::thread::sleep(wait);
184 }
185 attempt += 1;
186 }
187 Err(e) => return Err(format!("{e} for {}", redact(url))),
188 }
189 }
190 }
191
192 pub fn get_json(&self, url: &str) -> Result<Value, String> {
194 let body = self.get(url)?;
195 serde_json::from_slice(&body).map_err(|e| format!("bad JSON from {}: {e}", redact(url)))
196 }
197
198 pub fn get_rows(&self, url: &str) -> Result<Vec<Value>, String> {
205 match self.get_json(url)? {
206 Value::Array(rows) => Ok(rows),
207 Value::Object(map) => {
208 if let Some(msg) = ["message", "error", "Error Message"]
209 .iter()
210 .find_map(|k| map.get(*k).and_then(Value::as_str))
211 {
212 Err(format!("API error: {msg}"))
213 } else {
214 Err(format!("expected a JSON array from {}", redact(url)))
215 }
216 }
217 _ => Err(format!("expected a JSON array from {}", redact(url))),
218 }
219 }
220}
221
222#[cfg(test)]
223mod tests {
224 use super::*;
225 use std::cell::RefCell;
226
227 struct Policy {
229 rpm: u32,
230 retries: u32,
231 backoff: Duration,
232 }
233 impl RetrySettings for Policy {
234 fn rate_limit_per_min(&self) -> u32 {
235 self.rpm
236 }
237 fn max_retries(&self) -> u32 {
238 self.retries
239 }
240 fn backoff_base(&self) -> Duration {
241 self.backoff
242 }
243 }
244 fn policy(retries: u32) -> Policy {
245 Policy {
246 rpm: 0,
247 retries,
248 backoff: Duration::ZERO,
249 }
250 }
251
252 #[test]
253 fn redacts_all_token_shapes() {
254 assert_eq!(
255 redact("https://x/candle?symbol=AAPL&token=SECRET"),
256 "https://x/candle?symbol=AAPL&token=***"
257 );
258 assert_eq!(
259 redact("https://x?apikey=SECRET&r=D"),
260 "https://x?apikey=***&r=D"
261 );
262 assert_eq!(
263 redact("https://eodhd.com/eod/AAPL?api_token=SECRET&fmt=json"),
264 "https://eodhd.com/eod/AAPL?api_token=***&fmt=json"
265 );
266 assert_eq!(redact("https://x/no-token"), "https://x/no-token");
267 }
268
269 #[test]
270 fn retryable_classification() {
271 assert!(HttpError::Status(429).retryable());
272 assert!(HttpError::Status(503).retryable());
273 assert!(!HttpError::Status(404).retryable());
274 assert!(HttpError::Transport("x".into()).retryable());
275 }
276
277 #[test]
278 fn http_error_display() {
279 assert_eq!(HttpError::Status(429).to_string(), "HTTP 429");
280 assert!(HttpError::Transport("boom".into())
281 .to_string()
282 .contains("boom"));
283 }
284
285 struct SeqHttp {
286 calls: RefCell<Vec<Result<Vec<u8>, HttpError>>>,
287 }
288 impl HttpClient for SeqHttp {
289 fn get(&self, _url: &str) -> Result<Vec<u8>, HttpError> {
290 let mut q = self.calls.borrow_mut();
291 if q.is_empty() {
292 Err(HttpError::Status(500))
293 } else {
294 q.remove(0)
295 }
296 }
297 }
298
299 #[test]
300 fn fetcher_retries_then_succeeds() {
301 let http = SeqHttp {
302 calls: RefCell::new(vec![
303 Err(HttpError::Status(503)),
304 Ok(br#"{"ok":true}"#.to_vec()),
305 ]),
306 };
307 let p = policy(2);
308 let f = Fetcher::new(&http, &p);
309 assert_eq!(f.get("https://x?token=S").unwrap(), br#"{"ok":true}"#);
310 }
311
312 #[test]
313 fn fetcher_gives_up_after_retries() {
314 let http = SeqHttp {
315 calls: RefCell::new(vec![
316 Err(HttpError::Status(503)),
317 Err(HttpError::Status(503)),
318 ]),
319 };
320 let p = policy(1);
321 let f = Fetcher::new(&http, &p);
322 assert!(f.get("https://x").is_err());
323 }
324
325 #[test]
326 fn fetcher_does_not_retry_client_errors() {
327 let http = SeqHttp {
328 calls: RefCell::new(vec![Err(HttpError::Status(404))]),
329 };
330 let p = policy(5);
331 let f = Fetcher::new(&http, &p);
332 assert!(f.get("https://x").unwrap_err().contains("404"));
333 }
334
335 #[test]
336 fn get_json_parses_and_reports_bad_json() {
337 let http = SeqHttp {
338 calls: RefCell::new(vec![Ok(b"{\"a\":1}".to_vec()), Ok(b"not json".to_vec())]),
339 };
340 let p = policy(0);
341 let f = Fetcher::new(&http, &p);
342 assert_eq!(f.get_json("https://x").unwrap()["a"], 1);
343 assert!(f.get_json("https://x").unwrap_err().contains("bad JSON"));
344 }
345
346 #[test]
347 fn throttle_with_rate_limit_runs() {
348 let http = SeqHttp {
349 calls: RefCell::new(vec![Ok(b"{}".to_vec()), Ok(b"{}".to_vec())]),
350 };
351 let p = Policy {
352 rpm: 6000,
353 retries: 0,
354 backoff: Duration::ZERO,
355 };
356 let f = Fetcher::new(&http, &p);
357 assert!(f.get("https://x?token=k").is_ok());
358 assert!(f.get("https://x?token=k").is_ok());
359 }
360
361 #[test]
362 fn get_rows_array_object_error_and_non_array() {
363 let http = SeqHttp {
364 calls: RefCell::new(vec![
365 Ok(br#"[{"a":1}]"#.to_vec()),
366 Ok(br#"{"Error Message":"bad key"}"#.to_vec()),
367 Ok(br#"{"message":"denied"}"#.to_vec()),
368 Ok(b"42".to_vec()),
369 ]),
370 };
371 let p = policy(0);
372 let f = Fetcher::new(&http, &p);
373 assert_eq!(f.get_rows("https://x").unwrap().len(), 1);
374 assert!(f
375 .get_rows("https://x?apikey=S")
376 .unwrap_err()
377 .contains("API error: bad key"));
378 assert!(f
379 .get_rows("https://x")
380 .unwrap_err()
381 .contains("API error: denied"));
382 assert!(f
383 .get_rows("https://x")
384 .unwrap_err()
385 .contains("expected a JSON array"));
386 }
387
388 #[test]
389 fn write_mode_variants() {
390 assert_ne!(WriteMode::Overwrite, WriteMode::Append);
391 assert_eq!(WriteMode::Resume, WriteMode::Resume);
392 }
393}