s2_api/v1/stream/
sse.rs

1use std::str::FromStr;
2
3use s2_common::http::ParseableHeader;
4use serde::Serialize;
5
6use super::{ReadBatch, StreamPosition};
7
8static LAST_EVENT_ID_HEADER: http::HeaderName = http::HeaderName::from_static("last-event-id");
9
10#[derive(Debug, Clone, Copy)]
11pub struct LastEventId {
12    /// Sequence number of the last delivered record.
13    pub seq_num: u64,
14    /// Total records delivered.
15    pub count: usize,
16    /// Total metered bytes delivered.
17    pub bytes: usize,
18}
19
20impl ParseableHeader for LastEventId {
21    fn name() -> &'static http::HeaderName {
22        &LAST_EVENT_ID_HEADER
23    }
24}
25
26impl Serialize for LastEventId {
27    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
28    where
29        S: serde::Serializer,
30    {
31        self.to_string().serialize(serializer)
32    }
33}
34
35impl std::fmt::Display for LastEventId {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        let Self {
38            seq_num,
39            count,
40            bytes,
41        } = self;
42        write!(f, "{seq_num},{count},{bytes}")
43    }
44}
45
46impl FromStr for LastEventId {
47    type Err = s2_common::ValidationError;
48
49    fn from_str(s: &str) -> Result<Self, Self::Err> {
50        let mut iter = s.splitn(3, ",");
51
52        fn get_next<T>(
53            iter: &mut std::str::SplitN<&str>,
54            field: &str,
55        ) -> Result<T, s2_common::ValidationError>
56        where
57            T: FromStr,
58            <T as FromStr>::Err: std::fmt::Display,
59        {
60            let item = iter
61                .next()
62                .ok_or_else(|| format!("missing {field} in Last-Event-Id"))?;
63            item.parse()
64                .map_err(|e| format!("invalid {field} in Last-Event-ID: {e}").into())
65        }
66
67        let seq_num = get_next(&mut iter, "seq_num")?;
68        let count = get_next(&mut iter, "count")?;
69        let bytes = get_next(&mut iter, "bytes")?;
70
71        Ok(Self {
72            seq_num,
73            count,
74            bytes,
75        })
76    }
77}
78
79macro_rules! event {
80    ($name:ident, $val:expr) => {
81        #[derive(Serialize)]
82        #[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
83        #[serde(rename_all = "snake_case")]
84        pub enum $name {
85            $name,
86        }
87
88        impl AsRef<str> for $name {
89            fn as_ref(&self) -> &str {
90                $val
91            }
92        }
93    };
94}
95
96event!(Batch, "batch");
97event!(Error, "error");
98event!(Ping, "ping");
99
100#[derive(Serialize)]
101#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
102#[serde(untagged)]
103pub enum ReadEvent {
104    #[cfg_attr(feature = "utoipa", schema(title = "batch"))]
105    Batch {
106        #[cfg_attr(feature = "utoipa", schema(inline))]
107        event: Batch,
108        data: ReadBatch,
109        #[cfg_attr(feature = "utoipa", schema(value_type = String, pattern = "^[0-9]+,[0-9]+,[0-9]+$"))]
110        id: LastEventId,
111    },
112    #[cfg_attr(feature = "utoipa", schema(title = "error"))]
113    Error {
114        #[cfg_attr(feature = "utoipa", schema(inline))]
115        event: Error,
116        data: String,
117    },
118    #[cfg_attr(feature = "utoipa", schema(title = "ping"))]
119    Ping {
120        #[cfg_attr(feature = "utoipa", schema(inline))]
121        event: Ping,
122        data: PingEventData,
123    },
124    #[cfg_attr(feature = "utoipa", schema(title = "done"))]
125    #[serde(skip)]
126    Done {
127        #[cfg_attr(feature = "utoipa", schema(value_type = String, pattern = r"^\[DONE\]$"))]
128        data: DoneEventData,
129    },
130}
131
132#[cfg(feature = "axum")]
133fn elapsed_since_epoch() -> std::time::Duration {
134    std::time::SystemTime::now()
135        .duration_since(std::time::SystemTime::UNIX_EPOCH)
136        .expect("healthy clock")
137}
138
139#[cfg(feature = "axum")]
140pub fn read_batch_event(
141    format: crate::data::Format,
142    batch: &s2_common::stream::ReadBatch,
143    id: LastEventId,
144) -> Result<axum::response::sse::Event, axum::Error> {
145    axum::response::sse::Event::default()
146        .event(Batch::Batch)
147        .id(id.to_string())
148        .json_data(super::json::serialize_read_batch(format, batch))
149}
150
151#[cfg(feature = "axum")]
152pub fn error_event(data: String) -> Result<axum::response::sse::Event, axum::Error> {
153    Ok(axum::response::sse::Event::default()
154        .event(Error::Error)
155        .data(data))
156}
157
158#[cfg(feature = "axum")]
159pub fn ping_event(
160    tail: s2_common::record::StreamPosition,
161) -> Result<axum::response::sse::Event, axum::Error> {
162    axum::response::sse::Event::default()
163        .event(Ping::Ping)
164        .json_data(PingEventData {
165            timestamp: elapsed_since_epoch().as_millis() as u64,
166            tail: tail.into(),
167        })
168}
169
170#[cfg(feature = "axum")]
171pub fn done_event() -> Result<axum::response::sse::Event, axum::Error> {
172    Ok(axum::response::sse::Event::default().data(DoneEventData))
173}
174
175#[derive(Debug, Clone, Serialize)]
176#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
177#[serde(rename = "[DONE]")]
178pub struct DoneEventData;
179
180impl AsRef<str> for DoneEventData {
181    fn as_ref(&self) -> &str {
182        "[DONE]"
183    }
184}
185
186#[rustfmt::skip]
187#[derive(Debug, Clone, Serialize)]
188#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
189pub struct PingEventData {
190    /// Time the ping was emitted, as Unix epoch milliseconds.
191    pub timestamp: u64,
192    /// Sequence number that will be assigned to the next record on the stream, and timestamp of the last record.
193    pub tail: StreamPosition,
194}
s2_api/v1/stream/sse.rs

s2_api/v1/stream/
sse.rs
