finance_query/streaming/polygon/
mod.rs

1//! Polygon.io-backed streaming sources.
2//!
3//! Polygon runs one WebSocket cluster per asset class, so each class gets a
4//! purpose-built source instead of riding a single generic feed. Everything
5//! here reuses the existing [`PolygonStream`] adapter — the only Polygon
6//! WebSocket client in the crate.
7
8mod book;
9mod options;
10mod price;
11mod trades;
12
13use std::collections::{HashMap, HashSet};
14use std::sync::Arc;
15
16use futures::StreamExt;
17use tokio::sync::{RwLock, broadcast, mpsc};
18use tracing::{debug, info, warn};
19
20use crate::adapters::polygon::websocket::{ClusterDTO, PolygonMessage, PolygonStream};
21
22use super::client::{StreamError, StreamResult};
23use super::source::{StreamCommand, apply_command};
24
25pub(crate) use book::PolygonBookSource;
26pub(crate) use options::PolygonOptionsSource;
27pub(crate) use price::PolygonPriceSource;
28pub(crate) use trades::PolygonTradeSource;
29
30/// Asset class of a Polygon real-time cluster.
31///
32/// Each variant is a separate upstream connection with its own channel
33/// vocabulary and symbol format.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
35#[non_exhaustive]
36pub enum AssetClass {
37    /// US equities and ETFs (`AAPL`).
38    #[default]
39    Stocks,
40    /// Options contracts in OCC format (`O:AAPL250117C00150000`).
41    Options,
42    /// Currency pairs (`EUR/USD`).
43    Forex,
44    /// Crypto pairs (`BTC-USD`).
45    Crypto,
46    /// Futures contracts (`ESZ4`).
47    Futures,
48    /// Index values (`I:SPX`).
49    Indices,
50}
51
52impl AssetClass {
53    pub(crate) fn cluster(self) -> ClusterDTO {
54        match self {
55            Self::Stocks => ClusterDTO::Stocks,
56            Self::Options => ClusterDTO::Options,
57            Self::Forex => ClusterDTO::Forex,
58            Self::Crypto => ClusterDTO::Crypto,
59            Self::Futures => ClusterDTO::Futures,
60            Self::Indices => ClusterDTO::Indices,
61        }
62    }
63
64    /// Channels carrying price-forming events for this class.
65    pub(crate) fn price_channels(self) -> &'static [&'static str] {
66        match self {
67            Self::Stocks | Self::Options | Self::Futures => &["T", "Q"],
68            Self::Forex => &["C", "CA"],
69            Self::Crypto => &["XT", "XQ"],
70            Self::Indices => &["V"],
71        }
72    }
73
74    /// Channel carrying individual trade prints, where the class has one.
75    pub(crate) fn trade_channel(self) -> Option<&'static str> {
76        match self {
77            Self::Stocks | Self::Options | Self::Futures => Some("T"),
78            Self::Crypto => Some("XT"),
79            // Forex and indices publish no per-trade prints.
80            Self::Forex | Self::Indices => None,
81        }
82    }
83
84    pub(crate) fn label(self) -> &'static str {
85        match self {
86            Self::Stocks => "polygon-stocks",
87            Self::Options => "polygon-options",
88            Self::Forex => "polygon-forex",
89            Self::Crypto => "polygon-crypto",
90            Self::Futures => "polygon-futures",
91            Self::Indices => "polygon-indices",
92        }
93    }
94
95    /// Put a user-supplied symbol into the cluster's wire format.
96    ///
97    /// A bare options underlying (no digits) becomes an `O:AAPL*` wildcard so
98    /// the whole chain is followed; a full OCC symbol stays exact.
99    pub(crate) fn wire_symbol(self, symbol: &str) -> String {
100        let symbol = symbol.trim().to_uppercase();
101        match self {
102            Self::Indices if !symbol.starts_with("I:") => format!("I:{symbol}"),
103            Self::Options => {
104                let body = symbol.strip_prefix("O:").unwrap_or(&symbol);
105                if body.chars().any(|c| c.is_ascii_digit()) {
106                    format!("O:{body}")
107                } else {
108                    format!("O:{}*", body.trim_end_matches('*'))
109                }
110            }
111            _ => symbol,
112        }
113    }
114}
115
116/// Expand `symbols` into `prefix.symbol` channel names.
117pub(crate) fn channels_for(
118    class: AssetClass,
119    prefixes: &[&str],
120    symbols: impl IntoIterator<Item = String>,
121) -> Vec<String> {
122    let symbols: Vec<String> = symbols
123        .into_iter()
124        .map(|s| class.wire_symbol(&s))
125        .filter(|s| !s.is_empty())
126        .collect();
127    prefixes
128        .iter()
129        .flat_map(|p| symbols.iter().map(move |s| format!("{p}.{s}")))
130        .collect()
131}
132
133/// Decodes wire events into `T`, and drops any per-symbol state on unsubscribe.
134///
135/// Stateful because sources merge several event types into one snapshot per
136/// symbol; that state must not outlive the subscription that created it.
137pub(crate) trait SessionHandler<T>: Send {
138    /// Decode one wire event.
139    fn on_event(&mut self, msg: PolygonMessage) -> Vec<T>;
140
141    /// Forget state for symbols that just left the subscription set.
142    fn on_unsubscribe(&mut self, _removed: &[String]) {}
143}
144
145/// Adapts a stateless decode function to [`SessionHandler`].
146pub(crate) struct Decode<F>(pub(crate) F);
147
148impl<T, F> SessionHandler<T> for Decode<F>
149where
150    F: FnMut(PolygonMessage) -> Vec<T> + Send,
151{
152    fn on_event(&mut self, msg: PolygonMessage) -> Vec<T> {
153        (self.0)(msg)
154    }
155}
156
157/// `true` when a wire symbol — possibly an `O:AAPL*` wildcard — covers `key`.
158fn covers(pattern: &str, key: &str) -> bool {
159    match pattern.strip_suffix('*') {
160        Some(prefix) => key.starts_with(prefix),
161        None => key == pattern,
162    }
163}
164
165/// Drop per-symbol state for symbols that just left the subscription set.
166///
167/// Keys are wire symbols, so user input is normalised the same way the
168/// subscription was; an options wildcard prunes the whole chain it created.
169pub(crate) fn prune_symbols<V>(
170    state: &mut HashMap<String, V>,
171    class: AssetClass,
172    removed: &[String],
173) {
174    let patterns: Vec<String> = removed.iter().map(|s| class.wire_symbol(s)).collect();
175    state.retain(|key, _| !patterns.iter().any(|p| covers(p, key)));
176}
177
178/// Run one connected Polygon session, decoding wire events via `handler`.
179pub(crate) async fn run_polygon_session<T, H>(
180    class: AssetClass,
181    prefixes: &[&str],
182    subscriptions: &Arc<RwLock<HashSet<String>>>,
183    broadcast_tx: &broadcast::Sender<T>,
184    command_rx: &mut mpsc::Receiver<StreamCommand>,
185    mut handler: H,
186) -> StreamResult<()>
187where
188    T: Clone + Send + 'static,
189    H: SessionHandler<T>,
190{
191    let initial: Vec<String> = subscriptions.read().await.iter().cloned().collect();
192    let channels = channels_for(class, prefixes, initial);
193    let channel_refs: Vec<&str> = channels.iter().map(String::as_str).collect();
194
195    let mut stream = PolygonStream::from_singleton()
196        .map_err(|e| StreamError::ConnectionFailed(e.to_string()))?
197        .cluster(class.cluster())
198        .subscribe(&channel_refs)
199        .build()
200        .await
201        .map_err(|e| StreamError::ConnectionFailed(e.to_string()))?;
202
203    let sender = stream.sender();
204    info!("Connected to Polygon {} cluster", class.label());
205
206    loop {
207        tokio::select! {
208            Some(msg) = stream.next() => {
209                match msg {
210                    PolygonMessage::Status(status) => debug!("polygon status: {status}"),
211                    PolygonMessage::Unknown(raw) => warn!("unparsed polygon frame: {raw}"),
212                    msg => {
213                        for item in handler.on_event(msg) {
214                            let _ = broadcast_tx.send(item);
215                        }
216                    }
217                }
218            }
219
220            Some(cmd) = command_rx.recv() => {
221                let Some(changed) = apply_command(&cmd, subscriptions).await else {
222                    return Ok(());
223                };
224                if changed.is_empty() {
225                    continue;
226                }
227                let subscribing = matches!(cmd, StreamCommand::Subscribe(_));
228                if !subscribing {
229                    handler.on_unsubscribe(&changed);
230                }
231                let channels = channels_for(class, prefixes, changed);
232                let result = if subscribing {
233                    sender.subscribe_channels(&channels).await
234                } else {
235                    sender.unsubscribe_channels(&channels).await
236                };
237                if let Err(e) = result {
238                    return Err(StreamError::WebSocketError(e.to_string()));
239                }
240            }
241
242            else => break,
243        }
244    }
245
246    // Upstream ended without a Close command — reconnect.
247    Err(StreamError::WebSocketError(format!(
248        "{} connection closed",
249        class.label()
250    )))
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    #[test]
258    fn wire_symbol_applies_cluster_prefixes() {
259        assert_eq!(AssetClass::Indices.wire_symbol("spx"), "I:SPX");
260        assert_eq!(AssetClass::Indices.wire_symbol("I:SPX"), "I:SPX");
261        assert_eq!(AssetClass::Stocks.wire_symbol("aapl"), "AAPL");
262        assert_eq!(AssetClass::Crypto.wire_symbol("btc-usd"), "BTC-USD");
263        assert_eq!(
264            AssetClass::Options.wire_symbol("AAPL250117C00150000"),
265            "O:AAPL250117C00150000"
266        );
267        assert_eq!(AssetClass::Options.wire_symbol("aapl"), "O:AAPL*");
268        assert_eq!(AssetClass::Options.wire_symbol("O:AAPL*"), "O:AAPL*");
269    }
270
271    #[test]
272    fn pruning_drops_exact_and_wildcard_matches() {
273        let mut state: HashMap<String, u8> = HashMap::from([
274            ("O:AAPL250117C00150000".to_string(), 1),
275            ("O:AAPL250117P00150000".to_string(), 2),
276            ("O:SPY250117C00500000".to_string(), 3),
277        ]);
278
279        // A bare underlying was subscribed as the `O:AAPL*` wildcard, so it
280        // must take the whole chain with it.
281        prune_symbols(&mut state, AssetClass::Options, &["AAPL".to_string()]);
282        assert_eq!(state.len(), 1);
283        assert!(state.contains_key("O:SPY250117C00500000"));
284
285        let mut equities: HashMap<String, u8> =
286            HashMap::from([("AAPL".to_string(), 1), ("NVDA".to_string(), 2)]);
287        prune_symbols(&mut equities, AssetClass::Stocks, &["aapl".to_string()]);
288        assert_eq!(equities.keys().collect::<Vec<_>>(), vec!["NVDA"]);
289    }
290
291    #[test]
292    fn channels_expand_across_prefixes_and_symbols() {
293        let channels = channels_for(
294            AssetClass::Crypto,
295            &["XT", "XQ"],
296            ["btc-usd".to_string(), "eth-usd".to_string()],
297        );
298        assert_eq!(
299            channels,
300            vec!["XT.BTC-USD", "XT.ETH-USD", "XQ.BTC-USD", "XQ.ETH-USD"]
301        );
302    }
303}
finance_query/streaming/polygon/mod.rs

finance_query/streaming/polygon/
mod.rs
