use std::{collections::VecDeque, future::Future};
use futures_util::{Stream, stream};
use crate::{
Client, Typed,
api::{builder, types},
error::Result,
};
#[derive(Debug, Clone)]
pub struct Page<T> {
pub items: Vec<T>,
pub next_cursor: Option<String>,
}
impl From<types::AssetPage> for Page<types::Asset> {
fn from(page: types::AssetPage) -> Self {
Self {
items: page.items,
next_cursor: page.next_cursor,
}
}
}
impl From<types::JobPage> for Page<types::Job> {
fn from(page: types::JobPage) -> Self {
Self {
items: page.items,
next_cursor: page.next_cursor,
}
}
}
pub fn paginate<T, F, Fut>(fetch: F) -> impl Stream<Item = Result<T>>
where
F: FnMut(Option<String>) -> Fut,
Fut: Future<Output = Result<Page<T>>>,
{
struct State<T, F> {
fetch: F,
cursor: Option<String>,
buffer: VecDeque<T>,
done: bool,
}
stream::unfold(
State {
fetch,
cursor: None,
buffer: VecDeque::new(),
done: false,
},
|mut state| async move {
loop {
if let Some(item) = state.buffer.pop_front() {
return Some((Ok(item), state));
}
if state.done {
return None;
}
match (state.fetch)(state.cursor.take()).await {
Ok(page) => {
state.buffer.extend(page.items);
state.cursor = page.next_cursor.filter(|c| !c.is_empty());
state.done = state.cursor.is_none();
}
Err(err) => {
state.done = true;
return Some((Err(err), state));
}
}
}
},
)
}
impl Client {
pub fn paginate_assets<'a>(
&'a self,
configure: impl Fn(builder::ListAssets<'a>) -> builder::ListAssets<'a> + 'a,
) -> impl Stream<Item = Result<types::Asset>> + 'a {
paginate(move |cursor| {
let mut list = configure(self.api.list_assets());
if let Some(cursor) = cursor {
list = list.cursor(cursor);
}
async move { list.send().typed().await.map(Page::from) }
})
}
pub fn paginate_jobs<'a>(
&'a self,
configure: impl Fn(builder::ListJobs<'a>) -> builder::ListJobs<'a> + 'a,
) -> impl Stream<Item = Result<types::Job>> + 'a {
paginate(move |cursor| {
let mut list = configure(self.api.list_jobs());
if let Some(cursor) = cursor {
list = list.cursor(cursor);
}
async move { list.send().typed().await.map(Page::from) }
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures_util::StreamExt;
#[tokio::test]
async fn follows_cursors_to_the_end() {
let pages = paginate(|cursor: Option<String>| async move {
Ok(match cursor.as_deref() {
None => Page {
items: vec![1, 2],
next_cursor: Some("b".into()),
},
Some("b") => Page {
items: vec![3],
next_cursor: Some(String::new()),
},
Some(other) => panic!("unexpected cursor {other}"),
})
});
let items: Vec<i32> = pages.map(|r| r.unwrap()).collect().await;
assert_eq!(items, vec![1, 2, 3]);
}
}