1use futures::{Future, Stream, TryStreamExt};
5
6#[derive(Debug, Clone, serde::Deserialize)]
8#[serde(rename_all = "camelCase")]
9pub struct PageInfo {
10 pub has_next_page: bool,
12 pub has_previous_page: bool,
14 pub start_cursor: Option<String>,
16 pub end_cursor: Option<String>,
19}
20
21#[derive(Debug, Clone)]
23pub struct Page<T> {
24 pub nodes: Vec<T>,
26 pub page_info: PageInfo,
28}
29
30pub fn paginate<'a, T, F, Fut>(fetch: F) -> impl Stream<Item = crate::Result<T>> + 'a
41where
42 T: 'a,
43 F: FnMut(Option<String>) -> Fut + 'a,
44 Fut: Future<Output = crate::Result<Page<T>>> + 'a,
45{
46 futures::stream::try_unfold(
47 (fetch, None::<String>, false),
48 |(mut fetch, cursor, done)| async move {
49 if done {
50 return Ok(None);
51 }
52 let page = fetch(cursor.clone()).await?;
53 let next_cursor = page.page_info.end_cursor.clone();
54 let done = !page.page_info.has_next_page || next_cursor == cursor;
57 Ok(Some((page, (fetch, next_cursor, done))))
58 },
59 )
60 .map_ok(|page| futures::stream::iter(page.nodes.into_iter().map(Ok)))
61 .try_flatten()
62}
63
64pub async fn collect_all<'a, T, F, Fut>(fetch: F, limit: Option<usize>) -> crate::Result<Vec<T>>
70where
71 T: 'a,
72 F: FnMut(Option<String>) -> Fut + 'a,
73 Fut: Future<Output = crate::Result<Page<T>>> + 'a,
74{
75 let stream = paginate(fetch);
76 futures::pin_mut!(stream);
77 let mut items = Vec::new();
78 while let Some(item) = stream.try_next().await? {
79 items.push(item);
80 if limit.is_some_and(|limit| items.len() >= limit) {
81 break;
82 }
83 }
84 Ok(items)
85}
86
87#[cfg(test)]
88mod tests {
89 use super::*;
90 use std::sync::atomic::{AtomicUsize, Ordering};
91
92 fn page(nodes: Vec<u32>, end_cursor: Option<&str>, has_next: bool) -> Page<u32> {
93 Page {
94 nodes,
95 page_info: PageInfo {
96 has_next_page: has_next,
97 has_previous_page: false,
98 start_cursor: None,
99 end_cursor: end_cursor.map(str::to_owned),
100 },
101 }
102 }
103
104 #[tokio::test]
105 async fn paginate_walks_all_pages() {
106 let calls = AtomicUsize::new(0);
107 let items = collect_all(
108 |cursor| {
109 calls.fetch_add(1, Ordering::SeqCst);
110 async move {
111 Ok(match cursor.as_deref() {
112 None => page(vec![1, 2], Some("a"), true),
113 Some("a") => page(vec![3], Some("b"), false),
114 other => panic!("unexpected cursor {other:?}"),
115 })
116 }
117 },
118 None,
119 )
120 .await
121 .unwrap();
122 assert_eq!(items, vec![1, 2, 3]);
123 assert_eq!(calls.load(Ordering::SeqCst), 2);
124 }
125
126 #[tokio::test]
127 async fn paginate_stops_on_repeated_cursor() {
128 let items = collect_all(
129 |_cursor| async move { Ok(page(vec![1], Some("stuck"), true)) },
130 None,
131 )
132 .await
133 .unwrap();
134 assert_eq!(items, vec![1, 1]);
137 }
138
139 #[tokio::test]
140 async fn collect_all_honors_limit() {
141 let items = collect_all(
142 |_cursor| async move { Ok(page(vec![1, 2, 3], Some("x"), false)) },
143 Some(2),
144 )
145 .await
146 .unwrap();
147 assert_eq!(items, vec![1, 2]);
148 }
149}