databento/live/
protocol.rs

1//! Lower-level live interfaces exposed for those who want more customization or
2//! control.
3//!
4//! As these are not part of the primary live API, they are less documented and
5//! subject to change without warning.
6
7use std::{
8    borrow::Cow,
9    collections::HashMap,
10    fmt::{Debug, Display},
11};
12
13use dbn::{Compression, SType, Schema};
14use hex::ToHex;
15use sha2::{Digest, Sha256};
16use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
17use tracing::{debug, error, instrument};
18
19use crate::{ApiKey, Error, USER_AGENT};
20
21use super::{SlowReaderBehavior, Subscription};
22
23/// Returns the host and port for the live gateway for the given dataset.
24///
25/// Performs no validation on `dataset`.
26pub fn determine_gateway(dataset: &str) -> String {
27    const DEFAULT_PORT: u16 = 13_000;
28
29    let dataset_subdomain: String = dataset.replace('.', "-").to_ascii_lowercase();
30    format!("{dataset_subdomain}.lsg.databento.com:{DEFAULT_PORT}")
31}
32
33/// The core live API protocol.
34pub struct Protocol<W> {
35    sender: W,
36}
37
38impl<W> Protocol<W>
39where
40    W: AsyncWriteExt + Unpin,
41{
42    /// Creates a new instance of the live API protocol that will send raw API messages
43    /// to `sender`.
44    pub fn new(sender: W) -> Self {
45        Self { sender }
46    }
47
48    /// Conducts CRAM authentication with the live gateway. Returns the session ID.
49    ///
50    /// # Errors
51    /// This function returns an error if the gateway fails to respond or the authentication
52    /// request is rejected.
53    ///
54    /// # Cancel safety
55    /// This method is not cancellation safe. If this method is used in a
56    /// [`tokio::select!`] statement and another branch completes first, the
57    /// authentication may have been only partially sent, resulting in the gateway
58    /// rejecting the authentication and closing the connection.
59    #[instrument(skip(self, recver, key, options))]
60    pub async fn authenticate<R>(
61        &mut self,
62        recver: &mut R,
63        key: &ApiKey,
64        dataset: &str,
65        options: SessionOptions<'_>,
66    ) -> crate::Result<String>
67    where
68        R: AsyncBufReadExt + Unpin,
69    {
70        let mut greeting = String::new();
71        // Greeting
72        recver.read_line(&mut greeting).await?;
73        greeting.pop(); // remove newline
74
75        debug!(greeting);
76        let mut response = String::new();
77        // Challenge
78        recver.read_line(&mut response).await?;
79        response.pop(); // remove newline
80
81        // Parse challenge
82        let challenge = Challenge::parse(&response).inspect_err(|_| {
83            error!(?response, "No CRAM challenge in response from gateway");
84        })?;
85        debug!(%challenge, "Received CRAM challenge");
86
87        // Send CRAM reply/auth request
88        let auth_req = AuthRequest::new(key, dataset, &challenge, options);
89        debug!(?auth_req, "Sending CRAM reply");
90        self.sender.write_all(auth_req.as_bytes()).await?;
91
92        response.clear();
93        recver.read_line(&mut response).await?;
94        if response.is_empty() {
95            error!("Received empty auth response");
96        } else {
97            debug!(
98                auth_resp = &response[..response.len() - 1],
99                "Received auth response"
100            );
101        }
102        response.pop(); // remove newline
103
104        let auth_resp = AuthResponse::parse(&response)?;
105        Ok(auth_resp
106            .session_id()
107            .map(ToOwned::to_owned)
108            .unwrap_or_default())
109    }
110
111    /// Sends one or more subscription messages for `sub` depending on the number of symbols.
112    ///
113    /// # Errors
114    /// This function returns an error if it's unable to communicate with the gateway.
115    ///
116    /// # Cancel safety
117    /// This method is not cancellation safe. If this method is used in a
118    /// [`tokio::select!`] statement and another branch completes first, the subscription
119    /// may have been partially sent, resulting in the gateway rejecting the
120    /// subscription, sending an error, and closing the connection.
121    pub async fn subscribe(&mut self, sub: &Subscription) -> crate::Result<()> {
122        let Subscription {
123            schema,
124            stype_in,
125            start,
126            use_snapshot,
127            ..
128        } = &sub;
129
130        if *use_snapshot && start.is_some() {
131            return Err(Error::BadArgument {
132                param_name: "use_snapshot",
133                desc: "cannot request snapshot with start time".to_owned(),
134            });
135        }
136        let start_nanos = sub.start.as_ref().map(|start| start.unix_timestamp_nanos());
137
138        let symbol_chunks = sub.symbols.to_chunked_api_string();
139        let last_chunk_idx = symbol_chunks.len() - 1;
140        for (i, sym_str) in symbol_chunks.into_iter().enumerate() {
141            let sub_req = SubRequest::new(
142                *schema,
143                *stype_in,
144                start_nanos,
145                *use_snapshot,
146                sub.id,
147                &sym_str,
148                i == last_chunk_idx,
149            );
150            debug!(?sub_req, "Sending subscription request");
151            self.sender.write_all(sub_req.as_bytes()).await?;
152        }
153        Ok(())
154    }
155
156    /// Sends a start session message to the live gateway.
157    ///
158    /// # Errors
159    /// This function returns an error if it's unable to communicate with
160    /// the gateway.
161    ///
162    /// # Cancel safety
163    /// This method is not cancellation safe. If this method is used in a
164    /// [`tokio::select!`] statement and another branch completes first, the live
165    /// gateway may only receive a partial message, resulting in it sending an error and
166    /// closing the connection.
167    pub async fn start_session(&mut self) -> crate::Result<()> {
168        Ok(self.sender.write_all(StartRequest.as_bytes()).await?)
169    }
170
171    /// Shuts down the inner writer.
172    ///
173    /// # Errors
174    /// This function returns an error if the shut down did not complete successfully.
175    pub async fn shutdown(&mut self) -> crate::Result<()> {
176        Ok(self.sender.shutdown().await?)
177    }
178
179    /// Consumes the protocol instance and returns the inner sender.
180    pub fn into_inner(self) -> W {
181        self.sender
182    }
183}
184
185/// A challenge request from the live gateway.
186///
187/// See the [raw API documentation](https://databento.com/docs/api-reference-live/gateway-control-messages/challenge-request?live=raw)
188/// for more information.
189#[derive(Debug, Clone)]
190pub struct Challenge<'a>(&'a str);
191
192impl<'a> Challenge<'a> {
193    /// Parses a challenge request from the given raw response.
194    ///
195    /// # Errors
196    /// Returns an error if the response does not begin with "cram=".
197    // Can't use `FromStr` with lifetime
198    pub fn parse(response: &'a str) -> crate::Result<Self> {
199        if let Some(challenge) = response.strip_prefix("cram=") {
200            Ok(Self(challenge))
201        } else {
202            Err(Error::internal(
203                "no CRAM challenge in response from gateway",
204            ))
205        }
206    }
207}
208
209impl Display for Challenge<'_> {
210    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
211        write!(f, "{}", self.0)
212    }
213}
214
215/// Optional configuration options for live sessions sent on authentication requests.
216#[derive(Clone, Debug)]
217pub struct SessionOptions<'a> {
218    /// The compression mode for the session.
219    pub compression: Compression,
220    /// Whether to `ts_out` should be appended to each record.
221    pub send_ts_out: bool,
222    /// The heartbeat interval in seconds.
223    pub heartbeat_interval_s: Option<i64>,
224    /// Extension string to append to the user agent.
225    pub user_agent_ext: Option<&'a str>,
226    /// The behavior of the gateway when the client falls behind real time.
227    pub slow_reader_behavior: Option<SlowReaderBehavior>,
228}
229
230impl Default for SessionOptions<'_> {
231    fn default() -> Self {
232        Self {
233            compression: Compression::None,
234            send_ts_out: false,
235            heartbeat_interval_s: None,
236            user_agent_ext: None,
237            slow_reader_behavior: None,
238        }
239    }
240}
241
242/// Parses a pipe-delimited string of key=value pairs into an iterator.
243fn parse_kv_pairs(s: &str) -> impl Iterator<Item = (&str, &str)> {
244    s.split('|').filter_map(|kvp| kvp.split_once('='))
245}
246
247/// A raw API message to be sent to the live gateway.
248pub trait RawApiMsg {
249    /// Returns the request as a string slice.
250    fn as_str(&self) -> &str;
251
252    /// Returns the request as a byte slice.
253    fn as_bytes(&self) -> &[u8] {
254        self.as_str().as_bytes()
255    }
256}
257
258/// An authentication request to be sent to the live gateway.
259///
260/// See the [raw API documentation](https://databento.com/docs/api-reference-live/client-control-messages/authentication-request?live=raw)
261/// for more information.
262#[derive(Clone)]
263pub struct AuthRequest(String);
264
265impl AuthRequest {
266    /// Creates the raw API authentication request message from the given parameters.
267    pub fn new(
268        key: &ApiKey,
269        dataset: &str,
270        challenge: &Challenge,
271        options: SessionOptions,
272    ) -> Self {
273        let challenge_key = format!("{challenge}|{}", key.0);
274        let mut hasher = Sha256::new();
275        hasher.update(challenge_key.as_bytes());
276        let hashed = hasher.finalize();
277        let bucket_id = key.bucket_id();
278        let encoded_response = hashed.encode_hex::<String>();
279        let send_ts_out = options.send_ts_out as u8;
280        let user_agent: Cow<'_, str> = match options.user_agent_ext {
281            Some(ext) => Cow::Owned(format!("{} {ext}", *USER_AGENT)),
282            None => Cow::Borrowed(&USER_AGENT),
283        };
284        let mut req = format!(
285            "auth={encoded_response}-{bucket_id}|dataset={dataset}|encoding=dbn|compression={compression}|ts_out={send_ts_out}|client={user_agent}",
286            compression = options.compression,
287        );
288        if let Some(heartbeat_interval_s) = options.heartbeat_interval_s {
289            req = format!("{req}|heartbeat_interval_s={heartbeat_interval_s}");
290        }
291        if let Some(slow_reader_behavior) = options.slow_reader_behavior {
292            req = format!("{req}|slow_reader_behavior={slow_reader_behavior}");
293        }
294        req.push('\n');
295        Self(req)
296    }
297}
298
299impl RawApiMsg for AuthRequest {
300    fn as_str(&self) -> &str {
301        self.0.as_str()
302    }
303}
304
305impl Debug for AuthRequest {
306    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
307        // Should never be empty
308        write!(f, "{}", &self.0[..self.0.len() - 1])
309    }
310}
311
312/// An authentication response from the live gateway.
313///
314/// See the [raw API documentation](https://databento.com/docs/api-reference-live/gateway-control-messages/authentication-response?live=raw)
315/// for more information.
316pub struct AuthResponse<'a>(HashMap<&'a str, &'a str>);
317
318impl<'a> AuthResponse<'a> {
319    /// Parses a challenge request from the given raw response.
320    ///
321    /// # Errors
322    /// Returns an error if the response does not begin with "cram=".
323    // Can't use `FromStr` with lifetime
324    pub fn parse(response: &'a str) -> crate::Result<Self> {
325        let auth_keys: HashMap<&'a str, &'a str> = parse_kv_pairs(response).collect();
326        // Lack of success key also indicates something went wrong
327        if auth_keys.get("success").map(|v| *v != "1").unwrap_or(true) {
328            return Err(Error::Auth(
329                auth_keys
330                    .get("error")
331                    .map(|msg| (*msg).to_owned())
332                    .unwrap_or_else(|| response.to_owned()),
333            ));
334        }
335        Ok(Self(auth_keys))
336    }
337
338    /// Returns the session ID if present or `None`.
339    pub fn session_id(&self) -> Option<&str> {
340        self.0.get("session_id").copied()
341    }
342
343    /// Returns a reference to the key-value pairs.
344    pub fn get_ref(&self) -> &HashMap<&'a str, &'a str> {
345        &self.0
346    }
347}
348
349/// A subscription request to be sent to the live gateway.
350///
351/// See the [raw API documentation](https://databento.com/docs/api-reference-live/client-control-messages/subscription-request?live=raw)
352/// for more information.
353#[derive(Clone)]
354pub struct SubRequest(String);
355
356impl SubRequest {
357    /// Creates the raw API authentication request message from the given parameters.
358    /// `symbols` is expected to already be a valid length, such as from
359    /// [`Symbols::to_chunked_api_string()`](crate::Symbols::to_chunked_api_string).
360    pub fn new(
361        schema: Schema,
362        stype_in: SType,
363        start_nanos: Option<i128>,
364        use_snapshot: bool,
365        id: Option<u32>,
366        symbols: &str,
367        is_last: bool,
368    ) -> Self {
369        let use_snapshot = use_snapshot as u8;
370        let is_last = is_last as u8;
371        let mut args = format!(
372            "schema={schema}|stype_in={stype_in}|symbols={symbols}|snapshot={use_snapshot}|is_last={is_last}"
373        );
374
375        if let Some(start) = start_nanos {
376            args = format!("{args}|start={start}");
377        }
378        if let Some(id) = id {
379            args = format!("{args}|id={id}");
380        }
381        args.push('\n');
382        Self(args)
383    }
384}
385
386impl RawApiMsg for SubRequest {
387    fn as_str(&self) -> &str {
388        self.0.as_str()
389    }
390}
391
392impl Debug for SubRequest {
393    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
394        // Should never be empty
395        write!(f, "{}", &self.0[..self.0.len() - 1])
396    }
397}
398
399/// A request to begin the session to be sent to the live gateway.
400///
401/// See the [raw API documentation](https://databento.com/docs/api-reference-live/client-control-messages/session-start?live=raw)
402/// for more information.
403#[derive(Debug, Clone, Copy)]
404pub struct StartRequest;
405
406impl RawApiMsg for StartRequest {
407    fn as_str(&self) -> &str {
408        "start_session\n"
409    }
410}
databento/live/protocol.rs

databento/live/
protocol.rs
