Skip to main content

i_love_jesus/
lib.rs

1use diesel::dsl;
2use diesel::expression::is_aggregate;
3use diesel::expression::ValidGrouping;
4use diesel::pg::Pg;
5use diesel::query_builder::AsQuery;
6use diesel::query_builder::AstPass;
7use diesel::query_builder::Query;
8use diesel::query_builder::QueryFragment;
9use diesel::query_builder::QueryId;
10use diesel::query_dsl::methods;
11use diesel::query_dsl::QueryDsl;
12use diesel::sql_types::SingleValue;
13use diesel::sql_types::{self};
14use diesel::AppearsOnTable;
15use diesel::Expression;
16use diesel::ExpressionMethods;
17use diesel::QueryResult;
18use diesel::SelectableExpression;
19pub use i_love_jesus_macros::CursorKeysModule;
20
21pub trait CursorKey<C> {
22    type SqlType: SingleValue;
23    type CursorValue: QueryFragment<Pg> + 'static + Send + Expression<SqlType = Self::SqlType>;
24    type SqlValue: QueryFragment<Pg> + 'static + Send + Expression<SqlType = Self::SqlType>;
25
26    fn get_cursor_value(cursor: &C) -> Self::CursorValue;
27
28    fn get_sql_value() -> Self::SqlValue;
29}
30
31struct Bound<C> {
32    cursor: C,
33    or_equal: bool,
34    /// [cursor value, sql value for each row]
35    values: Vec<[Box<dyn QueryFragment<Pg> + Send>; 2]>,
36}
37
38impl<C> Bound<C> {
39    fn new(cursor: C, or_equal: bool) -> Self {
40        Bound {
41            cursor,
42            or_equal,
43            values: vec![],
44        }
45    }
46}
47
48#[derive(PartialEq, Eq, Clone, Copy)]
49pub enum SortDirection {
50    /// Ascending - least first
51    Asc,
52    /// Descending - greatest first
53    Desc,
54}
55
56#[inline]
57pub fn asc_if(condition: bool) -> SortDirection {
58    if condition {
59        SortDirection::Asc
60    } else {
61        SortDirection::Desc
62    }
63}
64
65#[inline]
66pub fn desc_if(condition: bool) -> SortDirection {
67    asc_if(!condition)
68}
69
70pub struct PaginatedQueryBuilder<C, Q> {
71    query: Q,
72    after: Option<Bound<C>>,
73    before: Option<Bound<C>>,
74    limit_and_offset_from_end: bool,
75    /// functions to apply sorting to the query (before reversal for limit_and_offset_from_end): [normal, reversed]
76    #[allow(clippy::type_complexity)]
77    sorts: Vec<[Box<dyn FnOnce(Q) -> Q + Send>; 2]>,
78    direction: SortDirection,
79}
80
81impl<C, Q> PaginatedQueryBuilder<C, Q> {
82    /// Query must be boxed. The given sort direction applies to all calls to `then_order_by`.
83    pub fn new(query: Q, direction: SortDirection) -> Self {
84        PaginatedQueryBuilder {
85            query,
86            after: None,
87            before: None,
88            limit_and_offset_from_end: false,
89            sorts: vec![],
90            direction,
91        }
92    }
93
94    pub fn after(mut self, cursor: Option<C>) -> Self {
95        self.after = cursor.map(|cursor| Bound::new(cursor, false));
96        self
97    }
98
99    pub fn after_or_equal(mut self, cursor: Option<C>) -> Self {
100        self.after = cursor.map(|cursor| Bound::new(cursor, true));
101        self
102    }
103
104    pub fn before(mut self, cursor: Option<C>) -> Self {
105        self.before = cursor.map(|cursor| Bound::new(cursor, false));
106        self
107    }
108
109    pub fn before_or_equal(mut self, cursor: Option<C>) -> Self {
110        self.before = cursor.map(|cursor| Bound::new(cursor, true));
111        self
112    }
113
114    pub fn limit_and_offset_from_end(mut self) -> Self {
115        self.limit_and_offset_from_end = true;
116        self
117    }
118
119    pub fn then_order_by<K: CursorKey<C>>(mut self, _key: K) -> Self
120    where
121        Q: methods::ThenOrderDsl<dsl::Desc<K::SqlValue>, Output = Q>
122            + methods::ThenOrderDsl<K::SqlValue, Output = Q>,
123    {
124        self.sorts.push([
125            // descending (normal)
126            Box::new(|q| q.then_order_by(K::get_sql_value().desc())),
127            // ascending (reversed)
128            Box::new(|q| q.then_order_by(K::get_sql_value())),
129        ]);
130        for bound in [&mut self.after, &mut self.before]
131            .into_iter()
132            .flat_map(|i| i.as_mut())
133        {
134            bound.values.push([
135                Box::new(K::get_cursor_value(&bound.cursor)),
136                Box::new(K::get_sql_value()),
137            ]);
138        }
139        self
140    }
141}
142
143impl<T, C, Q: methods::FilterDsl<T, Output = Q>> methods::FilterDsl<T>
144    for PaginatedQueryBuilder<C, Q>
145{
146    type Output = Self;
147
148    fn filter(mut self, predicate: T) -> Self::Output {
149        self.query = self.query.filter(predicate);
150        self
151    }
152}
153
154impl<C, Q: QueryDsl> QueryDsl for PaginatedQueryBuilder<C, Q> {}
155
156struct BoundCompare {
157    /// [cursor value, sql value for each row]
158    ///
159    /// the left tuple has element 0 of each array, and the right tuple has element 1 of each array
160    values: Vec<[Box<dyn QueryFragment<Pg> + Send>; 2]>,
161    op: &'static str,
162}
163
164impl QueryFragment<Pg> for BoundCompare {
165    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> {
166        // left tuple
167        out.push_sql("(");
168        for (i, [value, _]) in self.values.iter().enumerate() {
169            if i != 0 {
170                out.push_sql(", ");
171            }
172            value.walk_ast(out.reborrow())?;
173        }
174        out.push_sql(")");
175
176        // comparison operator
177        out.push_sql(self.op);
178
179        // right tuple
180        out.push_sql("(");
181        for (i, [_, value]) in self.values.iter().enumerate() {
182            if i != 0 {
183                out.push_sql(", ");
184            }
185            value.walk_ast(out.reborrow())?;
186        }
187        out.push_sql(")");
188
189        Ok(())
190    }
191}
192
193impl Expression for BoundCompare {
194    type SqlType = sql_types::Bool;
195}
196
197impl<QS> SelectableExpression<QS> for BoundCompare {}
198impl<QS> AppearsOnTable<QS> for BoundCompare {}
199impl<GB> ValidGrouping<GB> for BoundCompare {
200    type IsAggregate = is_aggregate::No;
201}
202
203impl<C, Q> AsQuery for PaginatedQueryBuilder<C, Q>
204where
205    Q: Query + methods::FilterDsl<BoundCompare, Output = Q>,
206{
207    type Query = PaginatedQuery<Q>;
208    type SqlType = Q::SqlType;
209    fn as_query(mut self) -> Self::Query {
210        // sort the inner query
211        for [sort, sort_reversed] in self.sorts {
212            self.query = if self.limit_and_offset_from_end ^ (self.direction == SortDirection::Asc)
213            {
214                sort_reversed(self.query)
215            } else {
216                sort(self.query)
217            };
218        }
219
220        // filter the inner query with tuple comparison
221        if self.direction == SortDirection::Asc {
222            std::mem::swap(&mut self.after, &mut self.before);
223        }
224        if let Some(bound) = self.after {
225            self.query = self.query.filter(BoundCompare {
226                values: bound.values,
227                op: if bound.or_equal { " >= " } else { " > " },
228            });
229        }
230        if let Some(bound) = self.before {
231            self.query = self.query.filter(BoundCompare {
232                values: bound.values,
233                op: if bound.or_equal { " <= " } else { " < " },
234            });
235        }
236
237        PaginatedQuery {
238            query: self.query,
239            reverse: self.limit_and_offset_from_end,
240        }
241    }
242}
243
244pub struct PaginatedQuery<Q> {
245    query: Q,
246    reverse: bool,
247}
248
249impl<Q: Query> Query for PaginatedQuery<Q> {
250    type SqlType = Q::SqlType;
251}
252
253impl<Q: QueryFragment<Pg>> QueryFragment<Pg> for PaginatedQuery<Q> {
254    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> {
255        out.push_sql("SELECT * FROM (");
256        self.query.walk_ast(out.reborrow())?;
257        out.push_sql(") dullbananas_i_love_jesus_subquery");
258        if self.reverse {
259            // do the 2nd reversal
260            //
261            // `row_number() OVER ()` is each row's index in the output of the inner query
262            out.push_sql(" ORDER BY row_number() OVER () DESC");
263        }
264        Ok(())
265    }
266}
267
268impl<Q> QueryId for PaginatedQuery<Q> {
269    type QueryId = ();
270    const HAS_STATIC_QUERY_ID: bool = false;
271}