1use crate::{Error, ErrorKind, RequestOptions, Result, Session};
8use futures_util::future::BoxFuture;
9use serde::Serialize;
10use std::{
11 collections::{HashSet, VecDeque},
12 fmt,
13 future::Future,
14 hash::Hash,
15};
16
17pub struct Page<T, C> {
18 pub items: Vec<T>,
19 pub next: Option<C>,
20}
21impl<T, C> fmt::Debug for Page<T, C> {
22 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
23 f.debug_struct("Page")
24 .field("item_count", &self.items.len())
25 .field("has_next", &self.next.is_some())
26 .finish()
27 }
28}
29
30#[derive(Clone, PartialEq, Eq)]
32pub struct Resume<C> {
33 pub cursor: C,
34 pub item_offset: usize,
35}
36impl<C> fmt::Debug for Resume<C> {
37 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38 f.debug_struct("Resume")
39 .field("cursor", &"[REDACTED]")
40 .field("item_offset", &self.item_offset)
41 .finish()
42 }
43}
44
45#[derive(Clone, Copy, Debug)]
46#[non_exhaustive]
47pub struct CollectOptions {
48 pub max_items: usize,
49 pub max_bytes: usize,
51 pub max_pages: usize,
52}
53impl Default for CollectOptions {
54 fn default() -> Self {
55 Self {
56 max_items: 10_000,
57 max_bytes: 64 * 1024 * 1024,
58 max_pages: 1_000,
59 }
60 }
61}
62impl CollectOptions {
63 pub fn with_max_items(mut self, value: usize) -> Result<Self> {
64 self.max_items = value;
65 self.validate()?;
66 Ok(self)
67 }
68 pub fn with_max_bytes(mut self, value: usize) -> Result<Self> {
69 self.max_bytes = value;
70 self.validate()?;
71 Ok(self)
72 }
73 pub fn with_max_pages(mut self, value: usize) -> Result<Self> {
74 self.max_pages = value;
75 self.validate()?;
76 Ok(self)
77 }
78 pub fn validate(&self) -> Result<()> {
79 if self.max_items == 0
80 || self.max_items > 1_000_000
81 || self.max_bytes < 2
82 || self.max_bytes > 512 * 1024 * 1024
83 || self.max_pages == 0
84 || self.max_pages > 100_000
85 {
86 Err(Error::new(ErrorKind::Validation))
87 } else {
88 Ok(())
89 }
90 }
91}
92
93#[derive(Clone, Copy, Debug)]
95#[non_exhaustive]
96pub struct CursorLimits {
97 pub max_bytes: usize,
98 pub max_retained_bytes: usize,
99}
100impl Default for CursorLimits {
101 fn default() -> Self {
102 Self {
103 max_bytes: 4096,
104 max_retained_bytes: 4 * 1024 * 1024,
105 }
106 }
107}
108impl CursorLimits {
109 pub fn validate(&self) -> Result<()> {
110 if self.max_bytes == 0
111 || self.max_bytes > 1024 * 1024
112 || self.max_retained_bytes < self.max_bytes
113 || self.max_retained_bytes > 64 * 1024 * 1024
114 {
115 Err(Error::new(ErrorKind::Validation))
116 } else {
117 Ok(())
118 }
119 }
120}
121#[derive(Clone, Copy, Debug, PartialEq, Eq)]
122#[non_exhaustive]
123pub enum Truncation {
124 Items,
125 Bytes,
126 Pages,
127}
128#[derive(Clone, Copy, Debug, PartialEq, Eq)]
129#[non_exhaustive]
130pub enum Completion {
131 Complete,
132 Truncated(Truncation),
133 Interrupted,
134}
135
136pub struct Collection<T, C> {
137 pub items: Vec<T>,
138 pub completion: Completion,
139 pub resume: Option<Resume<C>>,
140 pub pages_read: usize,
141 pub wire_bytes: usize,
142}
143impl<T, C> Collection<T, C> {
144 pub fn is_complete(&self) -> bool {
145 self.completion == Completion::Complete
146 }
147}
148impl<T, C> fmt::Debug for Collection<T, C> {
149 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
150 f.debug_struct("Collection")
151 .field("item_count", &self.items.len())
152 .field("completion", &self.completion)
153 .field("has_resume", &self.resume.is_some())
154 .field("pages_read", &self.pages_read)
155 .field("wire_bytes", &self.wire_bytes)
156 .finish()
157 }
158}
159pub type CollectionResult<T, C> = std::result::Result<Collection<T, C>, CollectionError<T, C>>;
160
161pub struct CollectionError<T, C> {
163 error: Error,
164 partial: Box<Collection<T, C>>,
165}
166impl<T, C> CollectionError<T, C> {
167 pub fn error(&self) -> &Error {
168 &self.error
169 }
170 pub fn partial(&self) -> &Collection<T, C> {
171 &self.partial
172 }
173 pub fn into_error(self) -> Error {
174 self.error
175 }
176 pub fn into_parts(self) -> (Error, Collection<T, C>) {
177 (self.error, *self.partial)
178 }
179}
180impl<T, C> From<Error> for CollectionError<T, C> {
181 fn from(error: Error) -> Self {
182 Self {
183 error,
184 partial: Box::new(Collection {
185 items: Vec::new(),
186 completion: Completion::Interrupted,
187 resume: None,
188 pages_read: 0,
189 wire_bytes: 2,
190 }),
191 }
192 }
193}
194impl<T, C> fmt::Debug for CollectionError<T, C> {
195 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
196 f.debug_struct("CollectionError")
197 .field("error", &self.error)
198 .field("partial", &self.partial)
199 .finish()
200 }
201}
202impl<T, C> fmt::Display for CollectionError<T, C> {
203 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
204 self.error.fmt(f)
205 }
206}
207impl<T, C> std::error::Error for CollectionError<T, C> {
208 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
209 Some(&self.error)
210 }
211}
212
213type Fetch<'a, T, C> =
214 Box<dyn FnMut(C, RequestOptions) -> BoxFuture<'a, Result<Page<T, C>>> + Send + 'a>;
215pub struct Pages<'a, T, C> {
216 fetch: Fetch<'a, T, C>,
217 cursor: Option<C>,
218 seen: HashSet<C>,
219 options: RequestOptions,
220 session: Session,
221 max_pages: usize,
222 pages_read: usize,
223 failure: Option<Error>,
224 deferred_failure: Option<Error>,
226 rejected_page_items: usize,
227 cursor_limits: CursorLimits,
228 retained_cursor_bytes: usize,
229}
230impl<T, C> fmt::Debug for Pages<'_, T, C> {
231 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
232 f.debug_struct("Pages")
233 .field("pages_read", &self.pages_read)
234 .field("max_pages", &self.max_pages)
235 .field("has_resume", &self.cursor.is_some())
236 .field(
237 "failed",
238 &(self.failure.is_some() || self.deferred_failure.is_some()),
239 )
240 .field("cursor_limits", &self.cursor_limits)
241 .field("retained_cursor_bytes", &self.retained_cursor_bytes)
242 .finish()
243 }
244}
245impl<'a, T: 'a, C: Clone + Eq + Hash + Serialize + Send + 'a> Pages<'a, T, C> {
246 pub fn new<F, Fut>(
247 initial: C,
248 options: RequestOptions,
249 session: Session,
250 max_pages: usize,
251 fetch: F,
252 ) -> Result<Self>
253 where
254 F: FnMut(C, RequestOptions) -> Fut + Send + 'a,
255 Fut: Future<Output = Result<Page<T, C>>> + Send + 'a,
256 {
257 Self::new_with_cursor_limits(
258 initial,
259 options,
260 session,
261 max_pages,
262 CursorLimits::default(),
263 fetch,
264 )
265 }
266 pub fn new_with_cursor_limits<F, Fut>(
267 initial: C,
268 options: RequestOptions,
269 session: Session,
270 max_pages: usize,
271 cursor_limits: CursorLimits,
272 mut fetch: F,
273 ) -> Result<Self>
274 where
275 F: FnMut(C, RequestOptions) -> Fut + Send + 'a,
276 Fut: Future<Output = Result<Page<T, C>>> + Send + 'a,
277 {
278 if max_pages == 0 || max_pages > 100_000 {
279 return Err(Error::new(ErrorKind::Validation));
280 }
281 cursor_limits.validate()?;
282 crate::engine::serialize_bounded(&initial, cursor_limits.max_bytes)?;
283 Ok(Self {
284 fetch: Box::new(move |cursor, options| Box::pin(fetch(cursor, options))),
285 cursor: Some(initial),
286 seen: HashSet::new(),
287 options,
288 session,
289 max_pages,
290 pages_read: 0,
291 failure: None,
292 deferred_failure: None,
293 rejected_page_items: 0,
294 cursor_limits,
295 retained_cursor_bytes: 0,
296 })
297 }
298 pub fn pages_read(&self) -> usize {
299 self.pages_read
300 }
301 pub fn with_cursor_limits(mut self, limits: CursorLimits) -> Result<Self> {
304 if self.pages_read != 0 || !self.seen.is_empty() {
305 return Err(Error::new(ErrorKind::Validation));
306 }
307 limits.validate()?;
308 if let Some(cursor) = &self.cursor {
309 crate::engine::serialize_bounded(cursor, limits.max_bytes)?;
310 }
311 self.cursor_limits = limits;
312 Ok(self)
313 }
314 pub fn resume(&self) -> Option<Resume<C>> {
315 self.cursor.clone().map(|cursor| Resume {
316 cursor,
317 item_offset: self.rejected_page_items,
318 })
319 }
320 async fn check_active(&mut self) -> Result<()> {
321 if let Some(error) = &self.failure {
322 return Err(error.clone());
323 }
324 let result = self.options.run(&self.session, async { Ok(()) }).await;
325 if let Err(error) = &result {
326 self.failure = Some(error.clone());
327 }
328 result
329 }
330 pub async fn next_page(&mut self) -> Result<Option<Page<T, C>>> {
331 if let Some(error) = &self.failure {
332 return Err(error.clone());
333 }
334 if let Some(error) = &self.deferred_failure {
335 self.failure = Some(error.clone());
336 return Err(error.clone());
337 }
338 let Some(_) = &self.cursor else {
339 return Ok(None);
340 };
341 self.check_active().await?;
342 let cursor = self
343 .cursor
344 .as_ref()
345 .ok_or_else(|| Error::new(ErrorKind::Pagination))?;
346 if self.pages_read >= self.max_pages || self.seen.contains(cursor) {
347 let error = Error::new(ErrorKind::Pagination).with_operation("pagination.pages");
348 self.failure = Some(error.clone());
349 return Err(error);
350 }
351 let bytes = match crate::engine::serialize_bounded(cursor, self.cursor_limits.max_bytes) {
352 Ok(bytes) => bytes,
353 Err(error) => {
354 self.failure = Some(error.clone());
355 return Err(error);
356 }
357 };
358 if self.retained_cursor_bytes.saturating_add(bytes.len())
359 > self.cursor_limits.max_retained_bytes
360 {
361 let error = Error::new(ErrorKind::Limit).with_operation("pagination.cursor_history");
362 self.failure = Some(error.clone());
363 return Err(error);
364 }
365 self.retained_cursor_bytes += bytes.len();
366 let cursor = cursor.clone();
367 self.seen.insert(cursor.clone());
368 let result = self
369 .options
370 .run(&self.session, (self.fetch)(cursor, self.options.clone()))
371 .await;
372 let page = match result {
373 Ok(page) => page,
374 Err(error) => {
375 self.failure = Some(error.clone());
376 return Err(error);
377 }
378 };
379 self.pages_read += 1;
380 if let Some(next) = &page.next
381 && let Err(error) = crate::engine::serialize_bounded(next, self.cursor_limits.max_bytes)
382 {
383 let error = error.with_operation("pagination.cursor");
384 self.deferred_failure = Some(error);
385 self.rejected_page_items = page.items.len();
386 return Ok(Some(page));
389 }
390 if page.items.is_empty() && page.next.is_some() {
391 let error = Error::new(ErrorKind::Pagination).with_operation("pagination.empty_page");
392 self.failure = Some(error.clone());
393 return Err(error);
394 }
395 self.cursor = page.next.clone();
396 Ok(Some(page))
397 }
398 pub fn items(self) -> Items<'a, T, C> {
399 Items {
400 pages: self,
401 pending: VecDeque::new(),
402 active_cursor: None,
403 item_offset: 0,
404 initial_skip: 0,
405 }
406 }
407 pub async fn collect(self, options: CollectOptions) -> CollectionResult<T, C>
408 where
409 T: Serialize,
410 {
411 self.items().collect(options).await
412 }
413}
414
415pub struct Items<'a, T, C> {
416 pages: Pages<'a, T, C>,
417 pending: VecDeque<T>,
418 active_cursor: Option<C>,
419 item_offset: usize,
420 initial_skip: usize,
421}
422impl<T, C> fmt::Debug for Items<'_, T, C> {
423 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
424 f.debug_struct("Items")
425 .field("pages", &self.pages)
426 .field("buffered_items", &self.pending.len())
427 .field("item_offset", &self.item_offset)
428 .finish()
429 }
430}
431impl<'a, T: 'a, C: Clone + Eq + Hash + Serialize + Send + 'a> Items<'a, T, C> {
432 pub fn resume_from(mut self, position: Resume<C>) -> Result<Self> {
434 if self.pages.pages_read != 0 || !self.pages.seen.is_empty() || !self.pending.is_empty() {
435 return Err(Error::new(ErrorKind::Validation));
436 }
437 crate::engine::serialize_bounded(&position.cursor, self.pages.cursor_limits.max_bytes)?;
438 self.pages.cursor = Some(position.cursor);
439 self.initial_skip = position.item_offset;
440 Ok(self)
441 }
442 pub fn resume(&self) -> Option<Resume<C>> {
443 if !self.pending.is_empty()
444 || (self.pages.deferred_failure.is_some() && self.active_cursor.is_some())
445 {
446 self.active_cursor.clone().map(|cursor| Resume {
447 cursor,
448 item_offset: self.item_offset,
449 })
450 } else {
451 self.pages.resume().map(|mut position| {
452 position.item_offset = position.item_offset.max(self.initial_skip);
453 position
454 })
455 }
456 }
457 pub fn pages_read(&self) -> usize {
458 self.pages.pages_read
459 }
460 async fn load_page(&mut self) -> Result<bool> {
461 if !self.pending.is_empty() {
462 return Ok(true);
463 }
464 let cursor = self.pages.cursor.clone();
465 let Some(page) = self.pages.next_page().await? else {
466 return Ok(false);
467 };
468 if self.initial_skip > page.items.len() {
469 let error = Error::new(ErrorKind::Pagination).with_operation("pagination.resume");
470 self.pages.failure = Some(error.clone());
471 self.pages.cursor = cursor;
472 return Err(error);
473 }
474 self.item_offset = self.initial_skip;
475 self.pending = page.items.into_iter().skip(self.initial_skip).collect();
476 self.initial_skip = 0;
477 self.active_cursor = cursor;
478 Ok(true)
479 }
480 async fn ensure_item(&mut self) -> Result<bool> {
481 while self.pending.is_empty() {
482 if !self.load_page().await? {
483 return Ok(false);
484 }
485 }
486 Ok(true)
487 }
488 pub async fn next(&mut self) -> Result<Option<T>> {
489 self.pages.check_active().await?;
491 if !self.ensure_item().await? {
492 return Ok(None);
493 }
494 self.item_offset += 1;
495 Ok(self.pending.pop_front())
496 }
497 pub async fn collect(mut self, options: CollectOptions) -> CollectionResult<T, C>
498 where
499 T: Serialize,
500 {
501 options.validate()?;
502 let mut items = Vec::new();
503 let mut wire_bytes = 2usize;
504 let max_pages = options.max_pages.min(self.pages.max_pages);
505 loop {
506 if self.pending.is_empty()
507 && let Some(error) = &self.pages.deferred_failure
508 {
509 self.pages.failure = Some(error.clone());
510 }
511 if let Some(error) = &self.pages.failure {
512 return Err(CollectionError {
513 error: error.clone(),
514 partial: Box::new(Collection {
515 items,
516 completion: Completion::Interrupted,
517 resume: self.resume(),
518 pages_read: self.pages.pages_read,
519 wire_bytes,
520 }),
521 });
522 }
523 let completion = if self.resume().is_none() {
524 Some(Completion::Complete)
525 } else if items.len() >= options.max_items {
526 Some(Completion::Truncated(Truncation::Items))
527 } else if self.pending.is_empty() && self.pages.pages_read >= max_pages {
528 Some(Completion::Truncated(Truncation::Pages))
529 } else {
530 None
531 };
532 if let Some(completion) = completion {
533 if let Some(error) = &self.pages.deferred_failure {
534 return Err(CollectionError {
535 error: error.clone(),
536 partial: Box::new(Collection {
537 items,
538 completion: Completion::Interrupted,
539 resume: self.resume(),
540 pages_read: self.pages.pages_read,
541 wire_bytes,
542 }),
543 });
544 }
545 return Ok(Collection {
546 items,
547 completion,
548 resume: self.resume(),
549 pages_read: self.pages.pages_read,
550 wire_bytes,
551 });
552 }
553 let active = self.pages.check_active().await;
554 let ready = match active {
555 Ok(()) => self.load_page().await,
556 Err(error) => Err(error),
557 };
558 match ready {
559 Ok(false) => {
560 return Ok(Collection {
561 items,
562 completion: Completion::Complete,
563 resume: None,
564 pages_read: self.pages.pages_read,
565 wire_bytes,
566 });
567 }
568 Err(error) => {
569 return Err(CollectionError {
570 error,
571 partial: Box::new(Collection {
572 items,
573 completion: Completion::Interrupted,
574 resume: self.resume(),
575 pages_read: self.pages.pages_read,
576 wire_bytes,
577 }),
578 });
579 }
580 Ok(true) if self.pending.is_empty() => continue,
581 Ok(true) => {}
582 }
583 let comma = usize::from(!items.is_empty());
584 let remaining = options
585 .max_bytes
586 .saturating_sub(wire_bytes)
587 .saturating_sub(comma);
588 let encoded = crate::engine::serialize_bounded(
589 self.pending
590 .front()
591 .ok_or_else(|| Error::new(ErrorKind::Pagination))?,
592 remaining,
593 );
594 match encoded {
595 Ok(bytes) => wire_bytes += bytes.len() + comma,
596 Err(error) if error.kind() == ErrorKind::Limit => {
597 if let Some(error) = &self.pages.deferred_failure {
598 return Err(CollectionError {
599 error: error.clone(),
600 partial: Box::new(Collection {
601 items,
602 completion: Completion::Interrupted,
603 resume: self.resume(),
604 pages_read: self.pages.pages_read,
605 wire_bytes,
606 }),
607 });
608 }
609 return Ok(Collection {
610 items,
611 completion: Completion::Truncated(Truncation::Bytes),
612 resume: self.resume(),
613 pages_read: self.pages.pages_read,
614 wire_bytes,
615 });
616 }
617 Err(error) => {
618 return Err(CollectionError {
619 error,
620 partial: Box::new(Collection {
621 items,
622 completion: Completion::Interrupted,
623 resume: self.resume(),
624 pages_read: self.pages.pages_read,
625 wire_bytes,
626 }),
627 });
628 }
629 }
630 self.item_offset += 1;
631 if let Some(item) = self.pending.pop_front() {
632 items.push(item);
633 }
634 }
635 }
636}
637
638#[cfg(test)]
639mod tests;