Merge pull request #4245 from Aethelflaed/combination-limit-offset

Implement Offset and Limit Dsl on CombinationClause

CHANGELOG.md

@@ -12,6 +12,7 @@ Increasing the minimal supported Rust version will always be coupled at least wi
 
 ### Added
 
+* Added `limit()` and `offset()` DSL to combination clauses such as `UNION`
 * Fixed `#[derive(Identifiable)]` ignoring attribute `#[diesel(serialize_as)]` on primary keys
 * Added embedded struct support for `AsChangeset` via `#[diesel(embed)]`
 * Support for libsqlite3-sys 0.30.0

diesel/src/query_builder/combination_clause.rs

@@ -1,22 +1,43 @@
 //! Combine queries using a combinator like `UNION`, `INTERSECT` or `EXPECT`
 //! with or without `ALL` rule for duplicates
+//!
+//! Within this module, types commonly use the following abbreviations:
+//!
+//! L: Limit Clause
+//! Of: Offset Clause
+//! LOf: Limit Offset Clause
 
 use crate::backend::{Backend, DieselReserveSpecialization};
+use crate::dsl::AsExprOf;
 use crate::expression::subselect::ValidSubselect;
+use crate::expression::IntoSql;
 use crate::expression::NonAggregate;
 use crate::query_builder::insert_statement::InsertFromSelect;
+use crate::query_builder::limit_clause::{LimitClause, NoLimitClause};
+use crate::query_builder::limit_offset_clause::LimitOffsetClause;
+use crate::query_builder::offset_clause::{NoOffsetClause, OffsetClause};
 use crate::query_builder::{AsQuery, AstPass, Query, QueryFragment, QueryId, SelectQuery};
-use crate::{CombineDsl, Insertable, QueryResult, RunQueryDsl, Table};
+use crate::query_dsl::methods::*;
+use crate::sql_types::BigInt;
+use crate::{CombineDsl, Insertable, QueryDsl, QueryResult, RunQueryDsl, Table};
 
 #[derive(Debug, Copy, Clone, QueryId)]
 #[must_use = "Queries are only executed when calling `load`, `get_result` or similar."]
 /// Combine queries using a combinator like `UNION`, `INTERSECT` or `EXPECT`
 /// with or without `ALL` rule for duplicates
-pub struct CombinationClause<Combinator, Rule, Source, Rhs> {
+pub struct CombinationClause<
+    Combinator,
+    Rule,
+    Source,
+    Rhs,
+    LimitOffset = LimitOffsetClause<NoLimitClause, NoOffsetClause>,
+> {
     combinator: Combinator,
     duplicate_rule: Rule,
     source: ParenthesisWrapper<Source>,
     rhs: ParenthesisWrapper<Rhs>,
+    /// The combined limit/offset clause of the query
+    limit_offset: LimitOffset,
 }
 
 impl<Combinator, Rule, Source, Rhs> CombinationClause<Combinator, Rule, Source, Rhs> {
@@ -32,41 +53,52 @@ impl<Combinator, Rule, Source, Rhs> CombinationClause<Combinator, Rule, Source,
             duplicate_rule,
             source: ParenthesisWrapper(source),
             rhs: ParenthesisWrapper(rhs),
+            limit_offset: LimitOffsetClause {
+                limit_clause: NoLimitClause,
+                offset_clause: NoOffsetClause,
+            },
         }
     }
 }
 
-impl<Combinator, Rule, Source, Rhs> Query for CombinationClause<Combinator, Rule, Source, Rhs>
+impl<Combinator, Rule, Source, Rhs, LOf> QueryDsl
+    for CombinationClause<Combinator, Rule, Source, Rhs, LOf>
+{
+}
+
+impl<Combinator, Rule, Source, Rhs, LOf> Query
+    for CombinationClause<Combinator, Rule, Source, Rhs, LOf>
 where
     Source: Query,
     Rhs: Query<SqlType = Source::SqlType>,
 {
     type SqlType = Source::SqlType;
 }
 
-impl<Combinator, Rule, Source, Rhs> SelectQuery for CombinationClause<Combinator, Rule, Source, Rhs>
+impl<Combinator, Rule, Source, Rhs, LOf> SelectQuery
+    for CombinationClause<Combinator, Rule, Source, Rhs, LOf>
 where
     Source: SelectQuery,
     Rhs: SelectQuery<SqlType = Source::SqlType>,
 {
     type SqlType = Source::SqlType;
 }
 
-impl<Combinator, Rule, Source, Rhs, QS> ValidSubselect<QS>
-    for CombinationClause<Combinator, Rule, Source, Rhs>
+impl<Combinator, Rule, Source, Rhs, LOf, QS> ValidSubselect<QS>
+    for CombinationClause<Combinator, Rule, Source, Rhs, LOf>
 where
     Source: ValidSubselect<QS>,
     Rhs: ValidSubselect<QS>,
 {
 }
 
-impl<Combinator, Rule, Source, Rhs, Conn> RunQueryDsl<Conn>
-    for CombinationClause<Combinator, Rule, Source, Rhs>
+impl<Combinator, Rule, Source, Rhs, LOf, Conn> RunQueryDsl<Conn>
+    for CombinationClause<Combinator, Rule, Source, Rhs, LOf>
 {
 }
 
-impl<Combinator, Rule, Source, Rhs, T> Insertable<T>
-    for CombinationClause<Combinator, Rule, Source, Rhs>
+impl<Combinator, Rule, Source, Rhs, LOf, T> Insertable<T>
+    for CombinationClause<Combinator, Rule, Source, Rhs, LOf>
 where
     T: Table,
     T::AllColumns: NonAggregate,
@@ -79,8 +111,8 @@ where
     }
 }
 
-impl<Combinator, Rule, Source, OriginRhs> CombineDsl
-    for CombinationClause<Combinator, Rule, Source, OriginRhs>
+impl<Combinator, Rule, Source, OriginRhs, LOf> CombineDsl
+    for CombinationClause<Combinator, Rule, Source, OriginRhs, LOf>
 where
     Self: Query,
 {
@@ -129,20 +161,83 @@ where
     }
 }
 
-impl<Combinator, Rule, Source, Rhs, DB: Backend> QueryFragment<DB>
-    for CombinationClause<Combinator, Rule, Source, Rhs>
+impl<Combinator, Rule, Source, Rhs, LOf, DB: Backend> QueryFragment<DB>
+    for CombinationClause<Combinator, Rule, Source, Rhs, LOf>
 where
     Combinator: QueryFragment<DB>,
     Rule: QueryFragment<DB>,
     ParenthesisWrapper<Source>: QueryFragment<DB>,
     ParenthesisWrapper<Rhs>: QueryFragment<DB>,
+    LOf: QueryFragment<DB>,
     DB: Backend + SupportsCombinationClause<Combinator, Rule> + DieselReserveSpecialization,
 {
     fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, DB>) -> QueryResult<()> {
         self.source.walk_ast(out.reborrow())?;
         self.combinator.walk_ast(out.reborrow())?;
         self.duplicate_rule.walk_ast(out.reborrow())?;
-        self.rhs.walk_ast(out)
+        self.rhs.walk_ast(out.reborrow())?;
+        self.limit_offset.walk_ast(out)
+    }
+}
+
+#[doc(hidden)]
+type Limit = AsExprOf<i64, BigInt>;
+
+impl<ST, Combinator, Rule, Source, Rhs, L, Of> LimitDsl
+    for CombinationClause<Combinator, Rule, Source, Rhs, LimitOffsetClause<L, Of>>
+where
+    Self: SelectQuery<SqlType = ST>,
+    CombinationClause<Combinator, Rule, Source, Rhs, LimitOffsetClause<LimitClause<Limit>, Of>>:
+        SelectQuery<SqlType = ST>,
+{
+    type Output =
+        CombinationClause<Combinator, Rule, Source, Rhs, LimitOffsetClause<LimitClause<Limit>, Of>>;
+
+    fn limit(self, limit: i64) -> Self::Output {
+        let limit_clause = LimitClause(limit.into_sql::<BigInt>());
+        CombinationClause {
+            combinator: self.combinator,
+            duplicate_rule: self.duplicate_rule,
+            source: self.source,
+            rhs: self.rhs,
+            limit_offset: LimitOffsetClause {
+                limit_clause,
+                offset_clause: self.limit_offset.offset_clause,
+            },
+        }
+    }
+}
+
+#[doc(hidden)]
+type Offset = Limit;
+
+impl<ST, Combinator, Rule, Source, Rhs, L, Of> OffsetDsl
+    for CombinationClause<Combinator, Rule, Source, Rhs, LimitOffsetClause<L, Of>>
+where
+    Self: SelectQuery<SqlType = ST>,
+    CombinationClause<Combinator, Rule, Source, Rhs, LimitOffsetClause<L, OffsetClause<Offset>>>:
+        SelectQuery<SqlType = ST>,
+{
+    type Output = CombinationClause<
+        Combinator,
+        Rule,
+        Source,
+        Rhs,
+        LimitOffsetClause<L, OffsetClause<Offset>>,
+    >;
+
+    fn offset(self, offset: i64) -> Self::Output {
+        let offset_clause = OffsetClause(offset.into_sql::<BigInt>());
+        CombinationClause {
+            combinator: self.combinator,
+            duplicate_rule: self.duplicate_rule,
+            source: self.source,
+            rhs: self.rhs,
+            limit_offset: LimitOffsetClause {
+                limit_clause: self.limit_offset.limit_clause,
+                offset_clause,
+            },
+        }
     }
 }
 

diesel_compile_tests/tests/fail/exists_can_only_take_subselects.stderr

@@ -7,7 +7,7 @@ error[E0277]: the trait bound `bool: SelectQuery` is not satisfied
    = help: the following other types implement trait `SelectQuery`:
              BoxedSelectStatement<'a, ST, QS, DB, GB>
              SelectStatement<F, S, D, W, O, LOf, G, H, LC>
-             diesel::query_builder::combination_clause::CombinationClause<Combinator, Rule, Source, Rhs>
+             diesel::query_builder::combination_clause::CombinationClause<Combinator, Rule, Source, Rhs, LOf>
    = note: required for `diesel::expression::subselect::Subselect<bool, Bool>` to implement `diesel::Expression`
    = note: 1 redundant requirement hidden
    = note: required for `Exists<bool>` to implement `diesel::Expression`
@@ -22,7 +22,7 @@ error[E0277]: the trait bound `users::columns::id: SelectQuery` is not satisfied
    = help: the following other types implement trait `SelectQuery`:
              BoxedSelectStatement<'a, ST, QS, DB, GB>
              SelectStatement<F, S, D, W, O, LOf, G, H, LC>
-             diesel::query_builder::combination_clause::CombinationClause<Combinator, Rule, Source, Rhs>
+             diesel::query_builder::combination_clause::CombinationClause<Combinator, Rule, Source, Rhs, LOf>
    = note: required for `diesel::expression::subselect::Subselect<users::columns::id, Bool>` to implement `diesel::Expression`
    = note: 1 redundant requirement hidden
    = note: required for `Exists<users::columns::id>` to implement `diesel::Expression`

diesel_tests/tests/combination.rs

@@ -117,6 +117,81 @@ fn except() {
     assert_eq!(expected_data, data);
 }
 
+#[test]
+fn union_with_limit() {
+    use crate::schema::users::dsl::*;
+
+    let conn = &mut connection();
+    let data = vec![
+        NewUser::new("Sean", None),
+        NewUser::new("Tess", None),
+        NewUser::new("Jim", None),
+    ];
+    insert_into(users).values(&data).execute(conn).unwrap();
+    let data = users.order(id).load::<User>(conn).unwrap();
+    let _sean = &data[0];
+    let tess = &data[1];
+    let _jim = &data[2];
+
+    let data: Vec<User> = users
+        .filter(id.le(tess.id))
+        .union(users.filter(id.ge(tess.id)))
+        .limit(2)
+        .load(conn)
+        .unwrap();
+
+    assert_eq!(2, data.len());
+
+    // Cannot determine any `expected_data` to match against because the database might return
+    // different rows
+    //
+    // The following code works with sqlite but fails with postgres:
+    // ```rust
+    // let expected_data = vec![User::new(sean.id, "Sean"), User::new(tess.id, "Tess")];
+    // assert_eq!(expected_data, data);
+    // ```
+    //
+    // thread 'combination::union_with_limit' panicked at diesel_tests/tests/combination.rs:144:5:
+    // assertion `left == right` failed
+    //  left: [User { id: 236, name: "Sean", hair_color: None }, User { id: 237, name: "Tess", hair_color: None }]
+    //  right: [User { id: 237, name: "Tess", hair_color: None }, User { id: 236, name: "Sean", hair_color: None }]
+    //
+    // Using positional_order_by to get predictable result is not possible because the ORDER
+    // fragment it produce would need to be placed in the middle of the fragment generated by the
+    // combination clause to be valid (and ORDER should be before LIMIT in an SQL query).
+}
+
+#[test]
+fn union_with_offset() {
+    use crate::schema::users::dsl::*;
+
+    let conn = &mut connection();
+    let data = vec![
+        NewUser::new("Sean", None),
+        NewUser::new("Tess", None),
+        NewUser::new("Jim", None),
+    ];
+    insert_into(users).values(&data).execute(conn).unwrap();
+    let data = users.order(id).load::<User>(conn).unwrap();
+    let _sean = &data[0];
+    let tess = &data[1];
+    let _jim = &data[2];
+
+    let data: Vec<User> = users
+        .filter(id.le(tess.id))
+        .union(users.filter(id.ge(tess.id)))
+        .limit(3)
+        .offset(1)
+        .load(conn)
+        .unwrap();
+
+    // Cannot determine any `expected_data` to match against because the database might return
+    // different rows
+    // The positional_order_by fix is not applicable, because ORDER should appear before LIMIT
+    // cf. union_with_limit for a lengthier explanation
+    assert_eq!(2, data.len());
+}
+
 #[test]
 fn union_with_order() {
     let conn = &mut connection();