origincloud_core/
pages.rs1use 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
13pub struct Pages<T> {
28 stream: BoxStream<Result<T, Error>>,
29}
30
31struct State<R, F, T> {
32 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 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 #[allow(clippy::should_implement_trait)] 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}