use std::collections::HashSet;
use std::pin::Pin;
use std::sync::Arc;
use futures::future::BoxFuture;
use futures::{Stream, StreamExt, TryStreamExt, stream};
use crate::Error;
use crate::http::continuation;
pub type ListStream<'a, T> = Pin<Box<dyn Stream<Item = Result<T, Error>> + Send + 'a>>;
#[derive(Debug, Clone, serde::Deserialize)]
pub struct Page<T> {
pub items: Vec<T>,
pub next_page_token: Option<String>,
}
#[derive(Clone, Default)]
pub(crate) struct PageOptions {
pub size: Option<u32>,
pub token: Option<String>,
}
type Fetch<'a, T> =
Arc<dyn Fn(PageOptions) -> BoxFuture<'a, Result<Page<T>, Error>> + Send + Sync + 'a>;
pub struct List<'a, T> {
fetch: Fetch<'a, T>,
options: PageOptions,
}
impl<'a, T: Send + 'a> List<'a, T> {
pub(crate) fn new(
fetch: impl Fn(PageOptions) -> BoxFuture<'a, Result<Page<T>, Error>> + Send + Sync + 'a,
) -> Self {
Self {
fetch: Arc::new(fetch),
options: PageOptions::default(),
}
}
pub(crate) fn with_options(mut self, options: PageOptions) -> Self {
self.options = options;
self
}
pub fn page_size(mut self, size: u32) -> Self {
self.options.size = Some(size);
self
}
pub fn page_token(mut self, token: impl Into<String>) -> Self {
self.options.token = Some(token.into());
self
}
pub async fn send(self) -> Result<Page<T>, Error> {
if self.options.size == Some(0) || self.options.token.as_deref() == Some("") {
return Err(Error::invalid(
"page size and continuation token must be nonempty",
));
}
(self.fetch)(self.options).await
}
pub fn stream(self) -> ListStream<'a, T> {
let mut seen = HashSet::new();
if let Some(token) = &self.options.token {
seen.insert(token.clone());
}
stream::try_unfold((Some(self), seen), |(request, mut seen)| async move {
let Some(request) = request else {
return Ok::<_, Error>(None);
};
let fetch = Arc::clone(&request.fetch);
let mut options = request.options.clone();
let page = request.send().await?;
options.token = continuation(page.next_page_token, &mut seen)?;
let next = options.token.as_ref().map(|_| Self {
fetch,
options: options.clone(),
});
Ok(Some((
stream::iter(page.items.into_iter().map(Ok)),
(next, seen),
)))
})
.try_flatten()
.boxed()
}
}