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

origincloud_core/
pages.rs
