1use crate::client::Client;
23use crate::error::{Error, Result};
24use futures::stream::Stream;
25use serde::Deserialize;
26use std::future::Future;
27use std::pin::Pin;
28use std::task::{Context, Poll};
29
30#[derive(Debug, Clone)]
38#[non_exhaustive]
39pub struct Page<T> {
40 pub count: u64,
43 pub next: Option<String>,
45 pub previous: Option<String>,
47 pub count_type: Option<String>,
52 pub meta: Option<serde_json::Value>,
57 pub page_metadata: Option<serde_json::Value>,
59 pub cursor: Option<String>,
63 pub results: Vec<T>,
65}
66
67#[derive(Deserialize)]
68struct RawPage<T> {
69 #[serde(default)]
70 count: Option<serde_json::Value>,
71 #[serde(default)]
72 next: Option<String>,
73 #[serde(default)]
74 previous: Option<String>,
75 #[serde(default)]
76 meta: Option<serde_json::Value>,
77 #[serde(default)]
78 page_metadata: Option<serde_json::Value>,
79 #[serde(default = "Vec::new")]
80 results: Vec<T>,
81}
82
83impl<T> Page<T> {
84 pub(crate) fn decode(bytes: &[u8]) -> Result<Self>
86 where
87 T: for<'de> Deserialize<'de>,
88 {
89 let raw: RawPage<T> = serde_json::from_slice(bytes).map_err(Error::Decode)?;
90 let count = match raw.count {
91 Some(serde_json::Value::Number(n)) => n.as_u64().unwrap_or(0),
92 Some(serde_json::Value::String(s)) => s.parse::<u64>().unwrap_or(0),
93 _ => 0,
94 };
95 let cursor = raw.next.as_deref().and_then(extract_cursor);
96 let mut out = Self {
97 count,
98 next: raw.next.filter(|s| !s.is_empty()),
99 previous: raw.previous.filter(|s| !s.is_empty()),
100 count_type: None,
101 meta: raw.meta.filter(|m| !m.is_null()),
102 page_metadata: raw.page_metadata,
103 cursor,
104 results: raw.results,
105 };
106 if out.count == 0 {
107 out.count = out.results.len() as u64;
108 }
109 Ok(out)
110 }
111
112 pub(crate) fn decode_response(
114 bytes: &[u8],
115 headers: &reqwest::header::HeaderMap,
116 ) -> Result<Self>
117 where
118 T: for<'de> Deserialize<'de>,
119 {
120 let mut page = Self::decode(bytes)?;
121 page.count_type = headers
122 .get(COUNT_TYPE_HEADER)
123 .and_then(|v| v.to_str().ok())
124 .map(str::trim)
125 .filter(|v| !v.is_empty())
126 .map(str::to_ascii_lowercase);
127 Ok(page)
128 }
129}
130
131const COUNT_TYPE_HEADER: &str = "x-results-counttype";
133
134fn extract_cursor(url: &str) -> Option<String> {
135 let parsed = reqwest::Url::parse(url).ok()?;
136 parsed
137 .query_pairs()
138 .find(|(k, _)| k == "cursor")
139 .map(|(_, v)| v.into_owned())
140 .filter(|v| !v.is_empty())
141}
142
143fn extract_page(url: &str) -> Option<u32> {
144 let parsed = reqwest::Url::parse(url).ok()?;
145 parsed
146 .query_pairs()
147 .find(|(k, _)| k == "page")
148 .and_then(|(_, v)| v.parse::<u32>().ok())
149}
150
151type FetchFut<T> = Pin<Box<dyn Future<Output = Result<Page<T>>> + Send>>;
153
154pub(crate) type FetchFn<T> =
157 Box<dyn FnMut(Client, Option<u32>, Option<String>) -> FetchFut<T> + Send>;
158
159pub struct PageStream<T> {
173 client: Client,
174 fetch: FetchFn<T>,
175 next_page: Option<u32>,
176 next_cursor: Option<String>,
177 buf: std::vec::IntoIter<T>,
178 done: bool,
179 in_flight: Option<FetchFut<T>>,
180}
181
182impl<T> Unpin for PageStream<T> {}
188
189impl<T> std::fmt::Debug for PageStream<T> {
190 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
191 f.debug_struct("PageStream")
192 .field("next_page", &self.next_page)
193 .field("next_cursor", &self.next_cursor)
194 .field("done", &self.done)
195 .finish_non_exhaustive()
196 }
197}
198
199impl<T> PageStream<T>
200where
201 T: Send + 'static,
202{
203 pub(crate) fn new(client: Client, fetch: FetchFn<T>) -> Self {
204 Self {
205 client,
206 fetch,
207 next_page: None,
208 next_cursor: None,
209 buf: Vec::new().into_iter(),
210 done: false,
211 in_flight: None,
212 }
213 }
214
215 pub async fn collect_all(self) -> Result<Vec<T>> {
220 use futures::TryStreamExt;
221 self.try_collect().await
222 }
223}
224
225impl<T> Stream for PageStream<T>
226where
227 T: Send + 'static,
228{
229 type Item = Result<T>;
230
231 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
232 let this = self.get_mut();
234 loop {
235 if let Some(item) = this.buf.next() {
240 return Poll::Ready(Some(Ok(item)));
241 }
242 if this.done {
243 return Poll::Ready(None);
244 }
245
246 if this.in_flight.is_none() {
248 let client = this.client.clone();
249 let page = this.next_page;
250 let cursor = this.next_cursor.clone();
251 let fut = (this.fetch)(client, page, cursor);
252 this.in_flight = Some(fut);
253 }
254
255 let fut = this
256 .in_flight
257 .as_mut()
258 .expect("we just set in_flight to Some");
259 match fut.as_mut().poll(cx) {
260 Poll::Pending => return Poll::Pending,
261 Poll::Ready(Err(e)) => {
262 this.in_flight = None;
263 this.done = true;
264 return Poll::Ready(Some(Err(e)));
265 }
266 Poll::Ready(Ok(page)) => {
267 this.in_flight = None;
268 let Page {
269 next,
270 cursor,
271 results,
272 ..
273 } = page;
274
275 if results.is_empty() {
276 this.done = true;
277 return Poll::Ready(None);
278 }
279 this.buf = results.into_iter();
280
281 match next.as_deref() {
283 None => {
284 this.done = true;
285 this.next_page = None;
286 this.next_cursor = None;
287 }
288 Some(next_url) => {
289 if let Some(c) = cursor {
290 this.next_cursor = Some(c);
291 this.next_page = None;
292 } else if let Some(p) = extract_page(next_url) {
293 this.next_page = Some(p);
294 this.next_cursor = None;
295 } else {
296 this.done = true;
297 this.next_page = None;
298 this.next_cursor = None;
299 }
300 }
301 }
302 }
304 }
305 }
306 }
307}
308
309#[cfg(test)]
310mod tests {
311 use super::*;
312
313 #[test]
314 fn decode_basic_page() {
315 let body = br#"{
316 "count": 2,
317 "next": null,
318 "previous": null,
319 "results": [{"a": 1}, {"a": 2}]
320 }"#;
321 let page: Page<serde_json::Value> = Page::decode(body).expect("decode");
322 assert_eq!(page.count, 2);
323 assert!(page.next.is_none());
324 assert_eq!(page.results.len(), 2);
325 assert!(page.cursor.is_none());
326 }
327
328 #[test]
329 fn decode_string_count() {
330 let body = br#"{"count": "42", "results": []}"#;
331 let page: Page<serde_json::Value> = Page::decode(body).expect("decode");
332 assert_eq!(page.count, 42);
333 }
334
335 #[test]
336 fn decode_extracts_cursor_from_next() {
337 let body = br#"{
338 "count": 1,
339 "next": "https://tango.example/api/contracts/?cursor=abc123&limit=25",
340 "results": [{"piid": "X"}]
341 }"#;
342 let page: Page<serde_json::Value> = Page::decode(body).expect("decode");
343 assert_eq!(page.cursor.as_deref(), Some("abc123"));
344 assert!(page.next.is_some());
345 }
346
347 #[test]
348 fn count_defaults_to_results_len_when_zero() {
349 let body = br#"{"results": [{"a": 1}, {"a": 2}, {"a": 3}]}"#;
350 let page: Page<serde_json::Value> = Page::decode(body).expect("decode");
351 assert_eq!(page.count, 3);
352 }
353
354 #[test]
355 fn decode_passes_meta_through() {
356 let body = br#"{
357 "count": 1,
358 "results": [{"a": 1}],
359 "meta": {"resolved_filters": {"agency": [{"token": "dod", "matched_by": "alias"}]}}
360 }"#;
361 let page: Page<serde_json::Value> = Page::decode(body).expect("decode");
362 let meta = page.meta.expect("meta");
363 assert_eq!(
364 meta["resolved_filters"]["agency"][0]["matched_by"],
365 serde_json::json!("alias")
366 );
367 assert!(page.count_type.is_none());
368 }
369
370 #[test]
371 fn decode_response_reads_count_type_header() {
372 let mut headers = reqwest::header::HeaderMap::new();
373 headers.insert("X-Results-CountType", "Approximate".parse().unwrap());
374 let page: Page<serde_json::Value> =
375 Page::decode_response(br#"{"count": 10000, "results": []}"#, &headers).expect("decode");
376 assert_eq!(page.count_type.as_deref(), Some("approximate"));
377 let page: Page<serde_json::Value> = Page::decode_response(
378 br#"{"count": 1, "results": []}"#,
379 &reqwest::header::HeaderMap::new(),
380 )
381 .expect("decode");
382 assert!(page.count_type.is_none());
383 assert!(page.meta.is_none());
384 }
385
386 #[test]
387 fn extract_page_parses_query_param() {
388 assert_eq!(
389 extract_page("https://tango.example/api/contracts/?page=4&limit=25"),
390 Some(4)
391 );
392 assert_eq!(extract_page("https://tango.example/api/contracts/"), None);
393 }
394}