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}