Skip to main content

finance_query_core/models/
logo.rs

1use crate::client::FetchClient;
2use std::collections::HashMap;
3use std::env;
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6use tokio::sync::{Mutex, RwLock};
7use tokio::time::timeout;
8use tracing::debug;
9
10/// Simple circuit breaker to protect the external logo service.
11#[derive(Debug)]
12struct CircuitBreaker {
13    failure_threshold: u32,
14    timeout_duration: Duration,
15    failure_count: u32,
16    last_failure_time: Option<Instant>,
17    state: CircuitState,
18}
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21enum CircuitState {
22    Closed,
23    Open,
24    HalfOpen,
25}
26
27impl CircuitBreaker {
28    fn new(failure_threshold: u32, timeout_duration: Duration) -> Self {
29        Self {
30            failure_threshold,
31            timeout_duration,
32            failure_count: 0,
33            last_failure_time: None,
34            state: CircuitState::Closed,
35        }
36    }
37
38    fn allow(&mut self) -> bool {
39        match self.state {
40            CircuitState::Closed => true,
41            CircuitState::Open => {
42                if let Some(last) = self.last_failure_time {
43                    if last.elapsed() > self.timeout_duration {
44                        self.state = CircuitState::HalfOpen;
45                        return true;
46                    }
47                }
48                false
49            }
50            CircuitState::HalfOpen => true,
51        }
52    }
53
54    fn record_success(&mut self) {
55        self.failure_count = 0;
56        self.state = CircuitState::Closed;
57        self.last_failure_time = None;
58    }
59
60    fn record_failure(&mut self) {
61        self.failure_count = self.failure_count.saturating_add(1);
62        self.last_failure_time = Some(Instant::now());
63        if self.failure_count >= self.failure_threshold {
64            self.state = CircuitState::Open;
65        }
66    }
67}
68
69#[derive(Debug, Clone)]
70struct CacheEntry {
71    value: String,
72    expires_at: Instant,
73}
74
75/// Fetches company logos using logo.dev API with caching and a circuit breaker.
76/// 
77/// This fetcher uses logo.dev's API (https://logo.dev) which provides company logos
78/// by ticker symbol. It constructs URLs like:
79/// `https://img.logo.dev/ticker/AAPL?token=...&format=png&fallback=404&size=50&theme=dark`
80/// 
81/// The fetcher includes:
82/// 1. In-memory caching with TTL
83/// 2. Circuit breaker pattern to protect against service failures
84/// 3. Configurable timeouts
85#[derive(Debug, Clone)]
86pub struct LogoFetcher {
87    fetch_client: Arc<FetchClient>,
88    cache: Arc<RwLock<HashMap<String, CacheEntry>>>,
89    circuit_breaker: Arc<Mutex<CircuitBreaker>>,
90    timeout: Duration,
91    cache_ttl: Duration,
92    enabled: bool,
93}
94
95impl LogoFetcher {
96    /// Create a new logo fetcher using the provided fetch client.
97    ///
98    /// Environment variables:
99    /// - `DISABLE_LOGO_FETCHING`: when "true", skip all logo requests.
100    /// - `LOGO_TIMEOUT_SECONDS`: per-request timeout (default 2s).
101    /// - `LOGO_CIRCUIT_BREAKER_THRESHOLD`: failures before opening (default 5).
102    /// - `LOGO_CIRCUIT_BREAKER_TIMEOUT`: cooldown in seconds (default 300s).
103    pub fn new(fetch_client: Arc<FetchClient>) -> Self {
104        let enabled = env::var("DISABLE_LOGO_FETCHING")
105            .map(|v| v.to_lowercase() != "true")
106            .unwrap_or(true);
107
108        let timeout_secs = env::var("LOGO_TIMEOUT_SECONDS")
109            .ok()
110            .and_then(|v| v.parse::<f64>().ok())
111            .unwrap_or(2.0)
112            .max(0.1);
113
114        let breaker_threshold = env::var("LOGO_CIRCUIT_BREAKER_THRESHOLD")
115            .ok()
116            .and_then(|v| v.parse::<u32>().ok())
117            .unwrap_or(5)
118            .max(1);
119
120        let breaker_timeout = env::var("LOGO_CIRCUIT_BREAKER_TIMEOUT")
121            .ok()
122            .and_then(|v| v.parse::<u64>().ok())
123            .unwrap_or(300);
124
125        Self {
126            fetch_client,
127            cache: Arc::new(RwLock::new(HashMap::new())),
128            circuit_breaker: Arc::new(Mutex::new(CircuitBreaker::new(
129                breaker_threshold,
130                Duration::from_secs(breaker_timeout),
131            ))),
132            timeout: Duration::from_secs_f64(timeout_secs),
133            cache_ttl: Duration::from_secs(60 * 60 * 24),
134            enabled,
135        }
136    }
137
138    /// Returns whether logo fetching is enabled.
139    pub fn is_enabled(&self) -> bool {
140        self.enabled
141    }
142
143    /// Fetch a logo for the given ticker symbol using logo.dev API.
144    /// 
145    /// This uses the ticker symbol directly to fetch the logo from logo.dev.
146    /// The website parameter is ignored but kept for backward compatibility.
147    pub async fn fetch_logo(&self, symbol: &str, _website: Option<&str>) -> Option<String> {
148        if !self.enabled {
149            return None;
150        }
151
152        if symbol.is_empty() {
153            return None;
154        }
155
156        if !self.allow_request().await {
157            debug!("Logo circuit breaker is open; skipping fetch");
158            return None;
159        }
160
161        let cache_key = format!("logo:{}", symbol.to_uppercase());
162        
163        if let Some(cached) = self.get_cached(&cache_key).await {
164            return Some(cached);
165        }
166
167        // Construct logo.dev URL with hardcoded token
168        let logo_url = format!(
169            "https://img.logo.dev/ticker/{}?token=pk_NNp9abu9TMm9II6Z0666YA&format=png&fallback=404&size=50&theme=dark",
170            symbol.to_uppercase()
171        );
172
173        match self.try_fetch_logo(&logo_url).await {
174            Ok(url) => {
175                self.cache_value(&cache_key, url.clone()).await;
176                self.record_success().await;
177                Some(url)
178            }
179            Err(err) => {
180                debug!("Logo fetch failed for {}: {}", symbol, err);
181                self.record_failure().await;
182                None
183            }
184        }
185    }
186
187    async fn get_cached(&self, key: &str) -> Option<String> {
188        let now = Instant::now();
189        let mut cache = self.cache.write().await;
190        if let Some(entry) = cache.get(key) {
191            if entry.expires_at > now {
192                return Some(entry.value.clone());
193            }
194        }
195        cache.remove(key);
196        None
197    }
198
199    async fn cache_value(&self, key: &str, value: String) {
200        let expires_at = Instant::now() + self.cache_ttl;
201        let mut cache = self.cache.write().await;
202        cache.insert(key.to_string(), CacheEntry { value, expires_at });
203    }
204
205    async fn allow_request(&self) -> bool {
206        let mut breaker = self.circuit_breaker.lock().await;
207        breaker.allow()
208    }
209
210    async fn record_success(&self) {
211        let mut breaker = self.circuit_breaker.lock().await;
212        breaker.record_success();
213    }
214
215    async fn record_failure(&self) {
216        let mut breaker = self.circuit_breaker.lock().await;
217        breaker.record_failure();
218    }
219
220    async fn try_fetch_logo(&self, url: &str) -> Result<String, String> {
221        let fetch = self.fetch_client.fetch_response(url);
222        match timeout(self.timeout, fetch).await {
223            Ok(Ok(response)) => {
224                let status = response.status();
225                if status.is_success() {
226                    Ok(response.url().to_string())
227                } else {
228                    Err(format!("HTTP {}", status))
229                }
230            }
231            Ok(Err(err)) => Err(err.to_string()),
232            Err(_) => Err(format!("Timed out after {:?}", self.timeout)),
233        }
234    }
235}