Skip to main content

tango/
pagination.rs

1//! Pagination: [`Page<T>`] envelope + async [`PageStream<T>`].
2//!
3//! Every list endpoint returns a [`Page<T>`] for a single page of results.
4//! For walking every page, resource methods expose an `iterate_*` constructor
5//! that returns a [`PageStream<T>`], which implements [`futures::Stream`].
6//!
7//! ```no_run
8//! use futures::TryStreamExt;
9//! use tango::Client;
10//! # async fn run() -> tango::Result<()> {
11//! let client = Client::builder().api_key("…").build()?;
12//! let opts = tango::ListContractsOptions::builder()
13//!     .awarding_agency("9700")
14//!     .build();
15//! let mut stream = client.iterate_contracts(opts);
16//! while let Some(record) = stream.try_next().await? {
17//!     println!("{:?}", record.get("piid"));
18//! }
19//! # Ok(()) }
20//! ```
21
22use 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/// A single page of results from a list endpoint.
31///
32/// `next`/`previous` are the server-supplied URLs to the adjacent pages;
33/// `cursor` is the keyset cursor extracted from `next` for convenience (empty
34/// on offset-paginated endpoints).
35///
36/// The struct is `#[non_exhaustive]`: the SDK builds pages, callers read them.
37#[derive(Debug, Clone)]
38#[non_exhaustive]
39pub struct Page<T> {
40    /// Total number of items across all pages, when the server reports it.
41    /// For some endpoints this matches `results.len()`.
42    pub count: u64,
43    /// URL of the next page, or `None` when on the last page.
44    pub next: Option<String>,
45    /// URL of the previous page, or `None` when on the first page.
46    pub previous: Option<String>,
47    /// `"exact"` or `"approximate"`, from the response's `X-Results-CountType` header; `None` when the endpoint does not send it.
48    ///
49    /// Past 1,000 matches the contract, IDV, opportunity and notice lists may report a query-planner estimate as [`count`](Self::count), which can be well off; this says which one you got.
50    /// It is read from this page's own response, so concurrent calls on one client never see each other's value the way [`Client::last_response_headers`](crate::Client::last_response_headers) can.
51    pub count_type: Option<String>,
52    /// The response's `meta` object, passed through as the API sends it: agency-filter diagnostics (`resolved_filters`, `agency_warnings`, `unresolved_agency_tokens`, `resolved_agencies`), `upgrade_hints` and the like.
53    ///
54    /// Each entry of `meta["resolved_filters"][<filter name>]` that resolved says how its token matched in `matched_by` (Tango API 5.9.0+): `key` (an organization UUID), `code` (a 3-digit CGAC or 4-digit FPDS code), `name` (the organization's name, including a department's everyday name, a spelling variant or a rename), `alias` (an abbreviation or the organization's own alias) or `fuzzy` (a looser text match, worth checking against the resolved name).
55    /// An entry that did not resolve has no `matched_by`.
56    pub meta: Option<serde_json::Value>,
57    /// Opaque page metadata for keyset endpoints.
58    pub page_metadata: Option<serde_json::Value>,
59    /// Cursor extracted from `next` on cursor-paginated endpoints. Pass it
60    /// back via the `cursor` option to fetch the next page. Empty when the
61    /// endpoint is offset-paginated or no cursor is set.
62    pub cursor: Option<String>,
63    /// The actual data for this page.
64    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    /// Decode a JSON byte slice into a [`Page<T>`].
85    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    /// Decode a list response body and read [`count_type`](Self::count_type) from its headers.
113    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
131/// Response header that says whether a list's `count` is exact or a query-planner estimate.
132const 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
151/// Future produced by a [`FetchFn`].
152type FetchFut<T> = Pin<Box<dyn Future<Output = Result<Page<T>>> + Send>>;
153
154/// Fetcher signature for [`PageStream`]: given the next `(page, cursor)`,
155/// return a future for the next [`Page<T>`].
156pub(crate) type FetchFn<T> =
157    Box<dyn FnMut(Client, Option<u32>, Option<String>) -> FetchFut<T> + Send>;
158
159/// An async stream over every result of a paginated list endpoint.
160///
161/// Yields `T` items one at a time, fetching successive pages on demand. The
162/// stream follows `?cursor=` when the server returns one in `next`, falling
163/// back to `?page=` for offset-paginated endpoints.
164///
165/// Construct with the `iterate_*` methods on [`Client`]; you typically use
166/// `try_next().await` (via [`futures::TryStreamExt`]) or `collect_all().await`.
167///
168/// # Cancellation
169///
170/// Dropping the stream stops fetching mid-page. In-flight requests are
171/// cancelled when their future is dropped (`reqwest`'s standard behavior).
172pub 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
182// All fields are themselves Unpin (Pin<Box<dyn Future>> is Unpin, IntoIter is
183// Unpin, the rest are inert). The Stream impl uses `Pin<&mut Self>` purely as
184// a type-system requirement, not because the struct holds self-referential
185// state. Marking Unpin lets us call `get_mut` in `poll_next` and freely mutate
186// the buffered iterator/page-tracking fields.
187impl<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    /// Drain the stream into a single `Vec<T>`.
216    ///
217    /// Convenience for callers who want every result. For huge result sets,
218    /// prefer iterating with `try_next` to bound memory.
219    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        // PageStream is Unpin (see impl above) — safe to project to &mut Self.
233        let this = self.get_mut();
234        loop {
235            // 1. Drain the buffer FIRST, even after `done` was set. The last
236            //    page legitimately holds results after the server says next=null;
237            //    yielding them before honoring done is what makes single-page
238            //    iteration correct. (Mirrors the F1 Go-iterator bug fix.)
239            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            // 2. Drive any in-flight fetch, or start a new one.
247            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                    // Decide how to advance to the next page.
282                    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                    // Loop back: yield the first buffered item.
303                }
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}