xmtp_proto/traits/combinators/
v3_paged.rs

1use std::marker::PhantomData;
2
3use xmtp_common::{MaybeSend, MaybeSync};
4use xmtp_configuration::MAX_PAGE_SIZE;
5
6use crate::{
7    api::{ApiClientError, Client, Endpoint, Pageable, Query},
8    api_client::Paged,
9};
10
11/// Endpoint that is paged with [`PagingInfo`]
12/// implements the v3 backend paging algorithm for endpoints
13/// which implement the [`Pageable`] trait
14pub struct V3Paged<E, T> {
15    endpoint: E,
16    id_cursor: Option<u64>,
17    _marker: PhantomData<T>,
18}
19
20#[xmtp_common::async_trait]
21impl<E, T, C> Query<C> for V3Paged<E, T>
22where
23    E: Query<C, Output = T> + Pageable,
24    C: Client,
25    T: Default + prost::Message + Paged + 'static,
26{
27    type Output = Vec<<T as Paged>::Message>;
28    async fn query(&mut self, client: &C) -> Result<Vec<<T as Paged>::Message>, ApiClientError> {
29        let mut out: Vec<<T as Paged>::Message> = vec![];
30        self.endpoint.set_cursor(self.id_cursor.unwrap_or(0));
31        loop {
32            let result: T = self.endpoint.query(client).await?;
33            let info = *result.info();
34            let mut messages = result.messages();
35            let num_messages = messages.len();
36            out.append(&mut messages);
37
38            if num_messages < MAX_PAGE_SIZE as usize || info.is_none() {
39                break;
40            }
41
42            let paging_info = info.expect("Empty paging info");
43            if paging_info.id_cursor == 0 {
44                break;
45            }
46
47            self.endpoint.set_cursor(paging_info.id_cursor);
48        }
49        Ok(out)
50    }
51}
52
53pub struct V3PagedSpecialized<S> {
54    _marker: PhantomData<S>,
55}
56
57impl<S, E: Endpoint<S>, T: MaybeSend + MaybeSync> Endpoint<V3PagedSpecialized<S>>
58    for V3Paged<E, T>
59{
60    type Output = <E as Endpoint<S>>::Output;
61
62    fn grpc_endpoint(&self) -> std::borrow::Cow<'static, str> {
63        self.endpoint.grpc_endpoint()
64    }
65
66    fn body(&self) -> Result<bytes::Bytes, crate::api::BodyError> {
67        self.endpoint.body()
68    }
69}
70
71/// Set an endpoint to be paged with v3 paging info
72pub fn v3_paged<E, T>(endpoint: E, id_cursor: Option<u64>) -> V3Paged<E, T> {
73    V3Paged {
74        endpoint,
75        id_cursor,
76        _marker: PhantomData,
77    }
78}
79
80#[cfg(test)]
81mod tests {
82
83    use std::borrow::Cow;
84
85    use prost::Message;
86
87    use crate::{
88        api::{self, Endpoint, EndpointExt, mock::MockNetworkClient},
89        mls_v1::{PagingInfo, SortDirection},
90    };
91
92    use super::*;
93    use rstest::*;
94
95    #[derive(prost::Message)]
96    struct TestV3Pageable {
97        #[prost(message, optional, tag = "1")]
98        info: Option<PagingInfo>,
99        #[prost(int32, repeated, tag = "2")]
100        msgs: Vec<i32>,
101    }
102
103    impl Paged for TestV3Pageable {
104        type Message = i32;
105
106        fn info(&self) -> &Option<PagingInfo> {
107            &self.info
108        }
109
110        fn messages(self) -> Vec<Self::Message> {
111            self.msgs
112        }
113    }
114
115    #[derive(Default)]
116    struct PageableTestEndpoint {
117        inner: TestV3Pageable,
118    }
119
120    impl Endpoint for PageableTestEndpoint {
121        type Output = TestV3Pageable;
122
123        fn grpc_endpoint(&self) -> std::borrow::Cow<'static, str> {
124            Cow::Borrowed("")
125        }
126
127        fn body(&self) -> Result<bytes::Bytes, api::BodyError> {
128            Ok(self.inner.encode_to_vec().into())
129        }
130    }
131
132    impl Pageable for PageableTestEndpoint {
133        fn set_cursor(&mut self, cursor: u64) {
134            if let Some(ref mut info) = self.inner.info {
135                info.id_cursor = cursor;
136            }
137        }
138    }
139
140    #[fixture]
141    fn client() -> MockNetworkClient {
142        let mut client = MockNetworkClient::new();
143        client.expect_request().times(1).returning(|_, _, b| {
144            let body = TestV3Pageable::decode(b.clone()).unwrap();
145            assert_eq!(
146                body.info.unwrap().id_cursor,
147                1,
148                "expected 1 got {}",
149                body.info.unwrap().id_cursor
150            );
151            Ok(http::Response::new(
152                TestV3Pageable {
153                    info: Some(PagingInfo {
154                        direction: SortDirection::Ascending as i32,
155                        limit: 100,
156                        id_cursor: 4,
157                    }),
158                    msgs: vec![0; MAX_PAGE_SIZE as usize],
159                }
160                .encode_to_vec()
161                .into(),
162            ))
163        });
164        client.expect_request().times(1).returning(|_, _, b| {
165            let body = TestV3Pageable::decode(b.clone()).unwrap();
166            assert_eq!(
167                body.info.unwrap().id_cursor,
168                4,
169                "expected 4 got {}",
170                body.info.unwrap().id_cursor
171            );
172            Ok(http::Response::new(
173                TestV3Pageable {
174                    info: Some(PagingInfo {
175                        direction: SortDirection::Ascending as i32,
176                        limit: 100,
177                        id_cursor: 6,
178                    }),
179                    msgs: vec![1; MAX_PAGE_SIZE as usize],
180                }
181                .encode_to_vec()
182                .into(),
183            ))
184        });
185        client.expect_request().times(1).returning(|_, _, b| {
186            let body = TestV3Pageable::decode(b.clone()).unwrap();
187            assert_eq!(
188                body.info.unwrap().id_cursor,
189                6,
190                "expected 6 got {}",
191                body.info.unwrap().id_cursor
192            );
193            Ok(http::Response::new(
194                TestV3Pageable {
195                    info: None,
196                    msgs: vec![7],
197                }
198                .encode_to_vec()
199                .into(),
200            ))
201        });
202        client
203    }
204
205    #[rstest]
206    #[xmtp_common::test]
207    async fn pages_endpoint(client: MockNetworkClient) {
208        let endpoint = PageableTestEndpoint {
209            inner: TestV3Pageable {
210                info: Some(PagingInfo {
211                    direction: SortDirection::Ascending as i32,
212                    limit: 100,
213                    id_cursor: 2,
214                }),
215                msgs: vec![],
216            },
217        };
218        // let result = api::v3_paged(endpoint, Some(1)).query(&client).await;
219        let result = endpoint.v3_paged(Some(1)).query(&client).await;
220        assert!(result.is_ok());
221        let result = result.unwrap();
222        let msgs = std::iter::repeat_n(0, MAX_PAGE_SIZE as usize)
223            .chain(std::iter::repeat_n(1, MAX_PAGE_SIZE as usize))
224            .chain(vec![7])
225            .collect::<Vec<_>>();
226        assert_eq!(result, msgs, "{:?}", result);
227    }
228
229    #[rstest]
230    #[xmtp_common::test]
231    async fn pages_endpoint_can_be_retried(client: MockNetworkClient) {
232        let endpoint = PageableTestEndpoint {
233            inner: TestV3Pageable {
234                info: Some(PagingInfo {
235                    direction: SortDirection::Ascending as i32,
236                    limit: 100,
237                    id_cursor: 2,
238                }),
239                msgs: vec![],
240            },
241        };
242        let result = api::v3_paged(api::retry(endpoint), Some(1))
243            .query(&client)
244            .await;
245        assert!(result.is_ok());
246        let result = result.unwrap();
247        let msgs = std::iter::repeat_n(0, MAX_PAGE_SIZE as usize)
248            .chain(std::iter::repeat_n(1, MAX_PAGE_SIZE as usize))
249            .chain(vec![7])
250            .collect::<Vec<_>>();
251        assert_eq!(result, msgs, "{:?}", result);
252    }
253
254    #[xmtp_common::test]
255    fn test_grpc_endpoint_delegates_to_wrapped_endpoint() {
256        let base_endpoint = PageableTestEndpoint::default();
257        let paged_endpoint: V3Paged<PageableTestEndpoint, TestV3Pageable> =
258            v3_paged(base_endpoint, Some(0));
259        assert_eq!(paged_endpoint.grpc_endpoint(), "");
260    }
261
262    #[xmtp_common::test]
263    fn test_body_delegates_to_wrapped_endpoint() {
264        let base_endpoint = PageableTestEndpoint::default();
265        let paged_endpoint: V3Paged<PageableTestEndpoint, TestV3Pageable> =
266            v3_paged(base_endpoint, Some(0));
267        let result = paged_endpoint.body();
268        assert!(result.is_ok());
269        assert_eq!(
270            result.unwrap(),
271            bytes::Bytes::from(TestV3Pageable::default().encode_to_vec())
272        );
273    }
274
275    #[xmtp_common::test]
276    fn test_pageable_test_endpoint_body_encodes_protobuf_message() {
277        let endpoint = PageableTestEndpoint {
278            inner: TestV3Pageable {
279                info: Some(PagingInfo {
280                    direction: SortDirection::Ascending as i32,
281                    limit: 100,
282                    id_cursor: 42,
283                }),
284                msgs: vec![1, 2, 3],
285            },
286        };
287        let result = endpoint.body();
288        assert!(result.is_ok());
289        let expected_bytes = endpoint.inner.encode_to_vec();
290        assert_eq!(result.unwrap(), bytes::Bytes::from(expected_bytes));
291    }
292
293    // this test here to ensure it compiles
294    #[xmtp_common::test]
295    async fn endpoints_can_be_chained() {
296        let client = MockNetworkClient::new();
297        std::mem::drop(
298            PageableTestEndpoint::default()
299                .v3_paged(Some(0))
300                .retry()
301                .query(&client),
302        );
303    }
304}
xmtp_proto/traits/combinators/v3_paged.rs

xmtp_proto/traits/combinators/
v3_paged.rs
