object_store/client/
pagination.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::Result;
19use futures::Stream;
20use std::future::Future;
21
22/// Takes a paginated operation `op` that when called with:
23///
24/// - A state `S`
25/// - An optional next token `Option<String>`
26///
27/// Returns
28///
29/// - A response value `T`
30/// - The next state `S`
31/// - The next continuation token `Option<String>`
32///
33/// And converts it into a `Stream<Result<T>>` which will first call `op(state, None)`, and yield
34/// the returned response `T`. If the returned continuation token was `None` the stream will then
35/// finish, otherwise it will continue to call `op(state, token)` with the values returned by the
36/// previous call to `op`, until a continuation token of `None` is returned
37///
38pub fn stream_paginated<F, Fut, S, T>(state: S, op: F) -> impl Stream<Item = Result<T>>
39where
40    F: Fn(S, Option<String>) -> Fut + Copy,
41    Fut: Future<Output = Result<(T, S, Option<String>)>>,
42{
43    enum PaginationState<T> {
44        Start(T),
45        HasMore(T, String),
46        Done,
47    }
48
49    futures::stream::unfold(PaginationState::Start(state), move |state| async move {
50        let (s, page_token) = match state {
51            PaginationState::Start(s) => (s, None),
52            PaginationState::HasMore(s, page_token) if !page_token.is_empty() => {
53                (s, Some(page_token))
54            }
55            _ => {
56                return None;
57            }
58        };
59
60        let (resp, s, continuation) = match op(s, page_token).await {
61            Ok(resp) => resp,
62            Err(e) => return Some((Err(e), PaginationState::Done)),
63        };
64
65        let next_state = match continuation {
66            Some(token) => PaginationState::HasMore(s, token),
67            None => PaginationState::Done,
68        };
69
70        Some((Ok(resp), next_state))
71    })
72}
object_store/client/pagination.rs

object_store/client/
pagination.rs
