mas_storage_pg/pagination.rs
1// Copyright 2024 New Vector Ltd.
2// Copyright 2022-2024 The Matrix.org Foundation C.I.C.
3//
4// SPDX-License-Identifier: AGPL-3.0-only
5// Please see LICENSE in the repository root for full details.
6
7//! Utilities to manage paginated queries.
8
9use mas_storage::{Pagination, pagination::PaginationDirection};
10use sea_query::IntoColumnRef;
11use uuid::Uuid;
12
13/// An extension trait to the `sqlx` [`QueryBuilder`], to help adding pagination
14/// to a query
15pub trait QueryBuilderExt {
16 /// Add cursor-based pagination to a query, as used in paginated GraphQL
17 /// connections
18 fn generate_pagination<C: IntoColumnRef>(
19 &mut self,
20 column: C,
21 pagination: Pagination,
22 ) -> &mut Self;
23}
24
25impl QueryBuilderExt for sea_query::SelectStatement {
26 fn generate_pagination<C: IntoColumnRef>(
27 &mut self,
28 column: C,
29 pagination: Pagination,
30 ) -> &mut Self {
31 let id_field = column.into_column_ref();
32
33 // ref: https://github.com/graphql/graphql-relay-js/issues/94#issuecomment-232410564
34 // 1. Start from the greedy query: SELECT * FROM table
35
36 // 2. If the after argument is provided, add `id > parsed_cursor` to the `WHERE`
37 // clause
38 if let Some(after) = pagination.after {
39 self.and_where(sea_query::Expr::col(id_field.clone()).gt(Uuid::from(after)));
40 }
41
42 // 3. If the before argument is provided, add `id < parsed_cursor` to the
43 // `WHERE` clause
44 if let Some(before) = pagination.before {
45 self.and_where(sea_query::Expr::col(id_field.clone()).lt(Uuid::from(before)));
46 }
47
48 match pagination.direction {
49 // 4. If the first argument is provided, add `ORDER BY id ASC LIMIT first+1` to the
50 // query
51 PaginationDirection::Forward => {
52 self.order_by(id_field, sea_query::Order::Asc)
53 .limit((pagination.count + 1) as u64);
54 }
55 // 5. If the first argument is provided, add `ORDER BY id DESC LIMIT last+1` to the
56 // query
57 PaginationDirection::Backward => {
58 self.order_by(id_field, sea_query::Order::Desc)
59 .limit((pagination.count + 1) as u64);
60 }
61 }
62
63 self
64 }
65}