Skip to main content

origincloud_core/
pages.rs

1use std::collections::VecDeque;
2use std::fmt;
3use std::future::Future;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6
7use futures_core::Stream;
8use futures_util::StreamExt as _;
9
10use crate::Error;
11use crate::runtime::{BoxStream, MaybeSend};
12
13/// Everything a list method lists, page by page: a [`Stream`] that fetches
14/// the next page as it is reached. A page that fails ends it with its
15/// error.
16///
17/// Without a stream library at hand, [`Pages::next`] reads it:
18///
19/// ```no_run
20/// # async fn list(mut pages: origincloud_core::Pages<String>) -> Result<(), origincloud_core::Error> {
21/// while let Some(item) = pages.next().await {
22///     println!("{}", item?);
23/// }
24/// # Ok(())
25/// # }
26/// ```
27pub struct Pages<T> {
28    stream: BoxStream<Result<T, Error>>,
29}
30
31struct State<R, F, T> {
32    /// The request for the next page; `None` after the last one.
33    request: Option<R>,
34    buffered: VecDeque<T>,
35    fetch: F,
36    set_page_token: fn(&mut R, String),
37}
38
39impl<T: MaybeSend + 'static> Pages<T> {
40    /// The items `fetch` lists from `request`'s page on: it answers with a
41    /// page's items and the next page's token, which `set_page_token` puts
42    /// in the request for it, until the token is empty. The generated
43    /// clients call it.
44    pub fn new<R, F, Fut>(request: R, fetch: F, set_page_token: fn(&mut R, String)) -> Self
45    where
46        R: Clone + MaybeSend + 'static,
47        F: Fn(R) -> Fut + MaybeSend + 'static,
48        Fut: Future<Output = Result<(Vec<T>, String), Error>> + MaybeSend + 'static,
49    {
50        let state = State {
51            request: Some(request),
52            buffered: VecDeque::new(),
53            fetch,
54            set_page_token,
55        };
56        let stream = futures_util::stream::unfold(state, |mut state| async move {
57            loop {
58                if let Some(item) = state.buffered.pop_front() {
59                    return Some((Ok(item), state));
60                }
61                let mut request = state.request.take()?;
62                match (state.fetch)(request.clone()).await {
63                    Ok((items, next_page_token)) => {
64                        state.buffered.extend(items);
65                        if !next_page_token.is_empty() {
66                            (state.set_page_token)(&mut request, next_page_token);
67                            state.request = Some(request);
68                        }
69                    }
70                    Err(err) => return Some((Err(err), state)),
71                }
72            }
73        });
74        Self {
75            stream: Box::pin(stream),
76        }
77    }
78
79    /// The next item, fetching the next page when it is needed; `None`
80    /// after the last one.
81    #[allow(clippy::should_implement_trait)] // Stream::next is StreamExt's.
82    pub async fn next(&mut self) -> Option<Result<T, Error>> {
83        self.stream.next().await
84    }
85}
86
87impl<T> Stream for Pages<T> {
88    type Item = Result<T, Error>;
89
90    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
91        self.stream.as_mut().poll_next(cx)
92    }
93}
94
95impl<T> fmt::Debug for Pages<T> {
96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97        f.debug_struct("Pages").finish_non_exhaustive()
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use crate::error::ApiError;
105
106    #[tokio::test]
107    async fn pages_are_followed_until_the_token_is_empty() {
108        let mut pages = Pages::new(
109            String::new(),
110            |token: String| async move {
111                match token.as_str() {
112                    "" => Ok((vec![1, 2], "p2".to_owned())),
113                    "p2" => Ok((vec![], "p3".to_owned())),
114                    "p3" => Ok((vec![3], String::new())),
115                    other => panic!("token {other}"),
116                }
117            },
118            |request, token| *request = token,
119        );
120        let mut got = Vec::new();
121        while let Some(item) = pages.next().await {
122            got.push(item.unwrap());
123        }
124        assert_eq!(got, [1, 2, 3]);
125    }
126
127    #[tokio::test]
128    async fn a_failed_page_ends_the_stream_with_its_error() {
129        let mut pages = Pages::new(
130            String::new(),
131            |token: String| async move {
132                match token.as_str() {
133                    "" => Ok((vec![1], "p2".to_owned())),
134                    _ => Err(Error::from(ApiError::from_response(400, b""))),
135                }
136            },
137            |request, token| *request = token,
138        );
139        assert_eq!(pages.next().await.unwrap().unwrap(), 1);
140        assert!(pages.next().await.unwrap().is_err());
141        assert!(pages.next().await.is_none());
142    }
143}