use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use crate::common::PaginationList;
use crate::error::ApifyClientResult;
type PageFuture<T> = Pin<Box<dyn Future<Output = ApifyClientResult<PaginationList<T>>> + Send>>;
type PageFetcher<T> = Box<dyn Fn(i64, Option<i64>) -> PageFuture<T> + Send + Sync>;
macro_rules! list_iterator {
($self:expr, $options:expr, $list:ident) => {{
let client = $self.clone();
let options = $options;
let start = options.offset.unwrap_or(0);
let total_limit = options.limit;
$crate::clients::pagination::ListIterator::new(
start,
total_limit,
Box::new(move |offset, page_limit| {
let client = client.clone();
let mut options = options.clone();
options.offset = Some(offset);
options.limit = page_limit;
Box::pin(async move { client.$list(options).await })
}),
)
}};
}
pub(crate) use list_iterator;
fn min_positive_limit(a: Option<i64>, b: Option<i64>) -> Option<i64> {
let a = a.filter(|&x| x > 0);
let b = b.filter(|&x| x > 0);
match (a, b) {
(Some(x), Some(y)) => Some(x.min(y)),
(Some(x), None) | (None, Some(x)) => Some(x),
(None, None) => None,
}
}
pub struct ListIterator<T> {
fetch: PageFetcher<T>,
buffer: VecDeque<T>,
next_offset: i64,
remaining: Option<i64>,
chunk_size: Option<i64>,
single_page: bool,
exhausted: bool,
}
impl<T> ListIterator<T> {
pub(crate) fn new(start_offset: i64, total_limit: Option<i64>, fetch: PageFetcher<T>) -> Self {
let cap = total_limit.filter(|&l| l > 0);
Self {
fetch,
buffer: VecDeque::new(),
next_offset: start_offset,
remaining: cap,
chunk_size: None,
single_page: false,
exhausted: false,
}
}
pub(crate) fn new_single_page(fetch: PageFetcher<T>) -> Self {
Self {
single_page: true,
..Self::new(0, None, fetch)
}
}
pub fn with_chunk_size(mut self, chunk_size: i64) -> Self {
self.chunk_size = (chunk_size > 0).then_some(chunk_size);
self
}
pub async fn next(&mut self) -> ApifyClientResult<Option<T>> {
if let Some(item) = self.buffer.pop_front() {
return Ok(Some(item));
}
if self.exhausted {
return Ok(None);
}
let page_limit = min_positive_limit(self.remaining, self.chunk_size);
let page = (self.fetch)(self.next_offset, page_limit).await?;
let received = page.items.len() as i64;
let mut items = page.items;
if let Some(rem) = self.remaining {
if received > rem {
items.truncate(rem as usize);
}
}
self.next_offset += received;
if let Some(rem) = self.remaining.as_mut() {
*rem -= received;
}
if self.single_page {
self.exhausted = true;
} else if received == 0 {
self.exhausted = true;
} else if matches!(self.remaining, Some(r) if r <= 0) {
self.exhausted = true;
} else if page.total > 0 {
if self.next_offset >= page.total {
self.exhausted = true;
}
} else {
if page.limit <= 0 || received < page.limit {
self.exhausted = true;
}
}
self.buffer.extend(items);
Ok(self.buffer.pop_front())
}
pub async fn collect_all(mut self) -> ApifyClientResult<Vec<T>> {
let mut out = Vec::new();
while let Some(item) = self.next().await? {
out.push(item);
}
Ok(out)
}
}
#[cfg(test)]
mod tests {
use super::{ListIterator, PageFetcher};
use crate::common::PaginationList;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
fn slicing_fetcher(
all: Vec<i64>,
default_page: i64,
report_total: bool,
calls: Arc<AtomicUsize>,
) -> PageFetcher<i64> {
let all = Arc::new(all);
Box::new(move |offset, limit| {
calls.fetch_add(1, Ordering::SeqCst);
let all = all.clone();
let page_size = limit.filter(|&l| l > 0).unwrap_or(default_page);
Box::pin(async move {
let start = (offset.max(0) as usize).min(all.len());
let end = (start + page_size.max(0) as usize).min(all.len());
let items = all[start..end].to_vec();
Ok(PaginationList {
total: if report_total { all.len() as i64 } else { 0 },
offset,
limit: page_size,
count: items.len() as i64,
desc: false,
items,
})
})
})
}
#[tokio::test]
async fn walks_all_pages_using_reported_total() {
let calls = Arc::new(AtomicUsize::new(0));
let iter = ListIterator::new(
0,
None,
slicing_fetcher((0..5).collect(), 2, true, calls.clone()),
)
.with_chunk_size(2);
assert_eq!(iter.collect_all().await.unwrap(), vec![0, 1, 2, 3, 4]);
assert_eq!(calls.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn walks_all_pages_when_total_is_zero() {
let calls = Arc::new(AtomicUsize::new(0));
let iter = ListIterator::new(0, None, slicing_fetcher((0..5).collect(), 2, false, calls))
.with_chunk_size(2);
assert_eq!(iter.collect_all().await.unwrap(), vec![0, 1, 2, 3, 4]);
}
#[tokio::test]
async fn total_zero_exact_multiple_terminates_on_empty_page() {
let calls = Arc::new(AtomicUsize::new(0));
let iter = ListIterator::new(
0,
None,
slicing_fetcher((0..4).collect(), 2, false, calls.clone()),
)
.with_chunk_size(2);
assert_eq!(iter.collect_all().await.unwrap(), vec![0, 1, 2, 3]);
assert_eq!(calls.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn reported_total_avoids_extra_request_on_exact_multiple() {
let calls = Arc::new(AtomicUsize::new(0));
let iter = ListIterator::new(
0,
None,
slicing_fetcher((0..4).collect(), 2, true, calls.clone()),
)
.with_chunk_size(2);
assert_eq!(iter.collect_all().await.unwrap(), vec![0, 1, 2, 3]);
assert_eq!(calls.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn non_final_short_page_with_reported_total_does_not_truncate() {
let all: Vec<i64> = (0..6).collect();
let calls = Arc::new(AtomicUsize::new(0));
let counter = calls.clone();
let all = Arc::new(all);
let fetch: PageFetcher<i64> = Box::new(move |offset, _limit| {
counter.fetch_add(1, Ordering::SeqCst);
let all = all.clone();
Box::pin(async move {
let start = (offset.max(0) as usize).min(all.len());
let end = (start + 2).min(all.len());
let items = all[start..end].to_vec();
Ok(PaginationList {
total: 100,
offset,
limit: 4,
count: items.len() as i64,
desc: false,
items,
})
})
});
let iter = ListIterator::new(0, None, fetch).with_chunk_size(4);
let got = iter.collect_all().await.unwrap();
assert_eq!(got, vec![0, 1, 2, 3, 4, 5]);
assert_eq!(calls.load(Ordering::SeqCst), 4);
}
#[tokio::test]
async fn total_limit_caps_items_yielded() {
let calls = Arc::new(AtomicUsize::new(0));
let iter = ListIterator::new(
0,
Some(3),
slicing_fetcher((0..100).collect(), 1000, true, calls.clone()),
);
assert_eq!(iter.collect_all().await.unwrap(), vec![0, 1, 2]);
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn total_limit_with_smaller_chunk_size_pages_and_caps() {
let calls = Arc::new(AtomicUsize::new(0));
let iter = ListIterator::new(
0,
Some(5),
slicing_fetcher((0..100).collect(), 1000, true, calls.clone()),
)
.with_chunk_size(2);
assert_eq!(iter.collect_all().await.unwrap(), vec![0, 1, 2, 3, 4]);
assert_eq!(calls.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn cap_truncates_page_that_exceeds_remaining_budget() {
let calls = Arc::new(AtomicUsize::new(0));
let counter = calls.clone();
let fetch: PageFetcher<i64> = Box::new(move |offset, _limit| {
counter.fetch_add(1, Ordering::SeqCst);
Box::pin(async move {
Ok(PaginationList {
total: 100,
offset,
limit: 5,
count: 5,
desc: false,
items: vec![0, 1, 2, 3, 4],
})
})
});
let iter = ListIterator::new(0, Some(3), fetch);
assert_eq!(iter.collect_all().await.unwrap(), vec![0, 1, 2]);
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"cap reached on the first (over-long) page, so no second fetch"
);
}
#[tokio::test]
async fn honours_caller_start_offset() {
let calls = Arc::new(AtomicUsize::new(0));
let iter = ListIterator::new(2, None, slicing_fetcher((0..5).collect(), 10, true, calls));
assert_eq!(iter.collect_all().await.unwrap(), vec![2, 3, 4]);
}
#[tokio::test]
async fn single_page_fetches_once_and_stops() {
let calls = Arc::new(AtomicUsize::new(0));
let counter = calls.clone();
let fetch: PageFetcher<i64> = Box::new(move |offset, _limit| {
counter.fetch_add(1, Ordering::SeqCst);
Box::pin(async move {
Ok(PaginationList {
total: 0,
offset,
limit: 3,
count: 3,
desc: false,
items: vec![10, 20, 30],
})
})
});
let iter = ListIterator::new_single_page(fetch);
assert_eq!(iter.collect_all().await.unwrap(), vec![10, 20, 30]);
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"single-page must fetch exactly once"
);
}
#[tokio::test]
async fn empty_first_page_yields_nothing() {
let calls = Arc::new(AtomicUsize::new(0));
let mut iter = ListIterator::new(0, None, slicing_fetcher(vec![], 5, true, calls));
assert!(iter.next().await.unwrap().is_none());
}
}