clientix_core/client/asynchronous/stream/
sse.rs

1use crate::client::asynchronous::stream::ClientixStream;
2use crate::client::response::{ClientixError, ClientixResult};
3use futures_core::Stream;
4use futures_util::{StreamExt, TryStreamExt};
5use http::{HeaderMap, StatusCode, Version};
6use reqwest::Url;
7use std::net::SocketAddr;
8use std::pin::Pin;
9use std::str::FromStr;
10use std::task::{Context, Poll};
11use encoding_rs::UTF_8;
12use serde::de::DeserializeOwned;
13
14const ID_PROPERTY: &str = "id:";
15const EVENT_PROPERTY: &str = "event:";
16const COMMENT_PROPERTY: &str = ":";
17const RETRY_PROPERTY: &str = "retry:";
18const DATA_PROPERTY: &str = "data:";
19
20#[derive(Debug, Clone)]
21pub struct SSE<T> {
22    id: Option<String>,
23    event: Option<String>,
24    comment: Option<String>,
25    retry: Option<u64>,
26    data: Option<T>
27}
28
29impl<T> SSE<T> {
30
31    fn new() -> Self {
32        Self {
33            id: None,
34            event: None,
35            comment: None,
36            retry: None,
37            data: None,
38        }
39    }
40
41    pub fn id(&self) -> &Option<String> {
42        &self.id
43    }
44
45    pub fn event(&self) -> &Option<String> {
46        &self.event
47    }
48
49    pub fn comment(&self) -> &Option<String> {
50        &self.comment
51    }
52
53    pub fn retry(&self) -> &Option<u64> {
54        &self.retry
55    }
56
57    pub fn data(self) -> Option<T> {
58        self.data
59    }
60
61}
62
63pub struct ClientixSSEStream<T> {
64    version: Version,
65    content_length: Option<u64>,
66    status: StatusCode,
67    url: Url,
68    remote_addr: Option<SocketAddr>,
69    headers: HeaderMap,
70    stream: Pin<Box<dyn Stream<Item = ClientixResult<SSE<T>>>>>,
71}
72
73impl<T> ClientixSSEStream<T> {
74
75    pub fn new(
76        version: Version,
77        content_length: Option<u64>,
78        status: StatusCode,
79        url: Url,
80        remote_addr: Option<SocketAddr>,
81        headers: HeaderMap,
82        stream: impl Stream<Item = ClientixResult<SSE<T>>> + 'static
83    ) -> Self {
84        Self {
85            version,
86            content_length,
87            status,
88            url,
89            remote_addr,
90            headers,
91            stream: Box::pin(stream)
92        }
93    }
94
95    pub fn version(&self) -> Version {
96        self.version
97    }
98
99    pub fn content_length(&self) -> Option<u64> {
100        self.content_length
101    }
102
103    pub fn status(&self) -> StatusCode {
104        self.status
105    }
106
107    pub fn url(&self) -> &Url {
108        &self.url
109    }
110
111    pub fn remote_addr(&self) -> Option<SocketAddr> {
112        self.remote_addr
113    }
114
115    pub fn headers(&self) -> &HeaderMap {
116        &self.headers
117    }
118
119    pub async fn execute<F>(mut self, mut handle: F) where F: FnMut(ClientixResult<SSE<T>>) {
120        while let Some(result) = self.stream.next().await {
121            handle(result);
122        }
123    }
124
125    pub async fn collect(self) -> ClientixResult<Vec<SSE<T>>> {
126        self.stream.try_collect().await
127    }
128
129}
130
131impl ClientixSSEStream<String> {
132
133    pub fn object_stream<T, F>(self, mut convert: F) -> ClientixSSEStream<T> 
134    where T: DeserializeOwned + Clone, F: FnMut(&str) -> ClientixResult<T> + 'static {
135        let version = self.version();
136        let content_length = self.content_length();
137        let status = self.status();
138        let url = self.url().clone();
139        let remote_addr = self.remote_addr();
140        let headers = self.headers().clone();
141        let stream = self
142            .filter(|line| match line {
143                Ok(line) if !line.data.clone().unwrap_or(String::new()).contains("[DONE]") => futures_util::future::ready(true),
144                _ => futures_util::future::ready(false)
145            })
146            .map(move |line| match line {
147                Ok(line) => {
148                    let mut sse = SSE::new();
149                    sse.id = line.id.clone();
150                    sse.event = line.event.clone();
151                    sse.comment = line.comment.clone();
152                    sse.retry = line.retry;
153                    sse.data =  Some(convert(line.data.clone().unwrap_or(String::new()).as_str())?);
154
155                    Ok(sse)
156                },
157                Err(err) => Err(err),
158            });
159        
160        ClientixSSEStream::new(version, content_length, status, url, remote_addr, headers, stream)
161    }
162
163    pub fn json_stream<T>(self) -> ClientixSSEStream<T> where T: DeserializeOwned + Clone {
164        self.object_stream(|string| {
165            serde_json::from_str::<T>(string).map_err(ClientixError::from)
166        })
167    }
168    
169    pub fn xml_stream<T>(self) -> ClientixSSEStream<T> where T: DeserializeOwned + Clone {
170        self.object_stream(|string| serde_xml_rs::from_str::<T>(string).map_err(ClientixError::from))
171    }
172    
173}
174
175impl<T> Stream for ClientixSSEStream<T> {
176    type Item = ClientixResult<SSE<T>>;
177
178    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
179        self.stream.poll_next_unpin(cx)
180    }
181}
182
183impl From<ClientixStream> for ClientixSSEStream<String> {
184    fn from(stream: ClientixStream) -> Self {
185        let version = stream.version();
186        let content_length = stream.content_length();
187        let status = stream.status();
188        let url = stream.url().clone();
189        let remote_addr = stream.remote_addr();
190        let headers = stream.headers().clone();
191
192        let mut buffer = String::new();
193        let stream = stream
194            .map(|chunk| match chunk {
195                Ok(chunk) => {
196                    let (text, _, _) = UTF_8.decode(&chunk);
197                    Ok(text.to_string())
198                },
199                Err(error) => Err(error),
200            })
201            .flat_map(move |text| match text {
202                Ok(text) => {
203                    let mut events = Vec::new();
204                    for line in text.lines() {
205                        let mut sse = SSE::new();
206                        match line {
207                            line if line.starts_with(ID_PROPERTY) => {
208                                sse.id = line.strip_prefix(ID_PROPERTY)
209                                    .map(str::trim)
210                                    .map(str::to_string);
211                            }
212                            line if line.starts_with(EVENT_PROPERTY) => {
213                                sse.event = line.strip_prefix(EVENT_PROPERTY)
214                                    .map(str::trim)
215                                    .map(str::to_string);
216                            }
217                            line if line.starts_with(COMMENT_PROPERTY) => {
218                                sse.comment = line.strip_prefix(COMMENT_PROPERTY)
219                                    .map(str::trim)
220                                    .map(str::to_string);
221                            }
222                            line if line.starts_with(RETRY_PROPERTY) => {
223                                sse.retry = line.strip_prefix(RETRY_PROPERTY)
224                                    .map(str::trim)
225                                    .map(u64::from_str)
226                                    .map(|result| match result {
227                                        Ok(value) => Some(value),
228                                        Err(_) => None
229                                    })
230                                    .unwrap_or(None);
231                            }
232                            line if line.starts_with(DATA_PROPERTY) => {
233                                buffer.push_str(line.trim_start_matches(DATA_PROPERTY).trim());
234                            }
235                            _ => {
236                                if !buffer.is_empty() {
237                                    sse.data = Some(buffer.to_string());
238                                    buffer.clear();
239
240                                    events.push(Ok(sse))
241                                }
242                            }
243                        }
244                    }
245
246                    futures_util::stream::iter(events)
247                }
248                Err(error) => futures_util::stream::iter(vec![Err(error)])
249            });
250
251        ClientixSSEStream::new(
252            version,
253            content_length,
254            status,
255            url,
256            remote_addr,
257            headers,
258            stream
259        )
260    }
261}
clientix_core/client/asynchronous/stream/sse.rs

clientix_core/client/asynchronous/stream/
sse.rs
