diff --git a/sea-orm-sync/src/entity/prelude.rs b/sea-orm-sync/src/entity/prelude.rs index 6e3e2cbc1..34fd5d770 100644 --- a/sea-orm-sync/src/entity/prelude.rs +++ b/sea-orm-sync/src/entity/prelude.rs @@ -1,6 +1,6 @@ pub use crate::{ ActiveEnum, ActiveModelBehavior, ActiveModelTrait, ColumnDef, ColumnTrait, ColumnType, - ColumnTypeTrait, ConnectionTrait, CursorTrait, DatabaseConnection, DbConn, EntityName, + ColumnTypeTrait, ConnectionTrait, CountTrait, CursorTrait, DatabaseConnection, DbConn, EntityName, EntityTrait, EnumIter, ForeignKeyAction, Iden, IdenStatic, Linked, LoaderTrait, ModelTrait, PaginatorTrait, PrimaryKeyArity, PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, QueryResult, Related, RelatedSelfVia, RelationDef, RelationTrait, Select, SelectExt, Value, diff --git a/sea-orm-sync/src/executor/paginator.rs b/sea-orm-sync/src/executor/paginator.rs index 5215ba181..ab54b300b 100644 --- a/sea-orm-sync/src/executor/paginator.rs +++ b/sea-orm-sync/src/executor/paginator.rs @@ -87,11 +87,7 @@ where Some(res) => res, None => return Ok(0), }; - #[allow(clippy::match_single_binding)] - let num_items = match self.db.get_database_backend() { - _ => result.try_get::("", "num_items")? as u64, - }; - Ok(num_items) + Ok(result.try_get::("", "num_items")? as u64) } /// Get the total number of pages @@ -271,14 +267,6 @@ where /// Paginate the result of a select operation. fn paginate(self, db: &'db C, page_size: u64) -> Paginator<'db, C, Self::Selector>; - - /// Perform a count on the paginated results - fn count(self, db: &'db C) -> Result - where - Self: Sized, - { - self.paginate(db, 1).num_items() - } } impl<'db, C, S> PaginatorTrait<'db, C> for Selector diff --git a/sea-orm-sync/src/executor/select_ext.rs b/sea-orm-sync/src/executor/select_ext.rs index 8213c52e2..02a63342a 100644 --- a/sea-orm-sync/src/executor/select_ext.rs +++ b/sea-orm-sync/src/executor/select_ext.rs @@ -1,10 +1,9 @@ use crate::{ ConnectionTrait, DbErr, EntityTrait, Select, SelectFive, SelectFour, SelectSix, SelectThree, - SelectTwo, Selector, SelectorRaw, SelectorTrait, Topology, + QueryTrait, SelectTwo, Selector, SelectorRaw, SelectorTrait, Statement, Topology, }; use sea_query::{Expr, SelectStatement}; -// TODO: Move count here /// Helper trait for selectors with convenient methods pub trait SelectExt { /// This method is unstable and is only used for internal testing. @@ -22,6 +21,14 @@ pub trait SelectExt { } } +/// Helper trait for counting rows selected by a query. +pub trait CountTrait { + /// Count the number of rows selected by this query. + fn count(self, db: &impl ConnectionTrait) -> Result + where + Self: Sized; +} + fn into_exists_query(mut stmt: SelectStatement) -> SelectStatement { stmt.clear_selects(); // Expr::Custom has fewer branches, but this may not have any significant impact on performance. @@ -32,6 +39,37 @@ fn into_exists_query(mut stmt: SelectStatement) -> SelectStatement { stmt } +fn build_count_query(stmt: SelectStatement) -> SelectStatement { + SelectStatement::new() + .expr(Expr::cust("COUNT(*) AS count")) + .from_subquery(stmt, "sub_query") + .to_owned() +} + +fn build_count_query_raw(stmt: Statement) -> SelectStatement { + let sub_query_sql = stmt.sql.trim().trim_end_matches(';').trim(); + let count_sql = format!("COUNT(*) AS count FROM ({sub_query_sql}) AS sub_query"); + + let mut query = SelectStatement::new(); + query.expr(if let Some(values) = stmt.values { + Expr::cust_with_values(count_sql, values.0) + } else { + Expr::cust(count_sql) + }); + query +} + +fn exec_count(db: &C, stmt: SelectStatement) -> Result +where + C: ConnectionTrait, +{ + let result = match db.query_one(&stmt)? { + Some(res) => res, + None => return Ok(0), + }; + Ok(result.try_get::("", "count")? as u64) +} + impl SelectExt for Selector where S: SelectorTrait, @@ -41,6 +79,15 @@ where } } +impl CountTrait for Selector +where + S: SelectorTrait, +{ + fn count(self, db: &impl ConnectionTrait) -> Result { + exec_count(db, build_count_query(self.query)) + } +} + impl SelectExt for SelectorRaw where S: SelectorTrait, @@ -60,6 +107,15 @@ where } } +impl CountTrait for SelectorRaw +where + S: SelectorTrait, +{ + fn count(self, db: &impl ConnectionTrait) -> Result { + exec_count(db, build_count_query_raw(self.stmt)) + } +} + impl SelectExt for Select where E: EntityTrait, @@ -133,6 +189,15 @@ where } } +impl CountTrait for T +where + T: QueryTrait, +{ + fn count(self, db: &impl ConnectionTrait) -> Result { + exec_count(db, build_count_query(self.into_query())) + } +} + #[cfg(test)] mod tests { use super::SelectExt; diff --git a/sea-orm-sync/src/query/mod.rs b/sea-orm-sync/src/query/mod.rs index 4fa9b0a82..7ce6ec4bd 100644 --- a/sea-orm-sync/src/query/mod.rs +++ b/sea-orm-sync/src/query/mod.rs @@ -47,7 +47,7 @@ pub use update::*; pub(crate) use util::*; pub use crate::{ - ConnectionTrait, CursorTrait, InsertResult, PaginatorTrait, SelectExt, Statement, + ConnectionTrait, CountTrait, CursorTrait, InsertResult, PaginatorTrait, SelectExt, Statement, TransactionTrait, UpdateResult, Value, Values, }; pub use sea_query::ExprTrait; diff --git a/src/entity/prelude.rs b/src/entity/prelude.rs index 6e3e2cbc1..65ffcd843 100644 --- a/src/entity/prelude.rs +++ b/src/entity/prelude.rs @@ -1,9 +1,9 @@ pub use crate::{ ActiveEnum, ActiveModelBehavior, ActiveModelTrait, ColumnDef, ColumnTrait, ColumnType, - ColumnTypeTrait, ConnectionTrait, CursorTrait, DatabaseConnection, DbConn, EntityName, - EntityTrait, EnumIter, ForeignKeyAction, Iden, IdenStatic, Linked, LoaderTrait, ModelTrait, - PaginatorTrait, PrimaryKeyArity, PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, QueryResult, - Related, RelatedSelfVia, RelationDef, RelationTrait, Select, SelectExt, Value, + ColumnTypeTrait, ConnectionTrait, CountTrait, CursorTrait, DatabaseConnection, DbConn, + EntityName, EntityTrait, EnumIter, ForeignKeyAction, Iden, IdenStatic, Linked, LoaderTrait, + ModelTrait, PaginatorTrait, PrimaryKeyArity, PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, + QueryResult, Related, RelatedSelfVia, RelationDef, RelationTrait, Select, SelectExt, Value, error::*, sea_query::{DynIden, Expr, RcOrArc, SeaRc, StringLen}, }; diff --git a/src/executor/paginator.rs b/src/executor/paginator.rs index a3e88a657..b267e6407 100644 --- a/src/executor/paginator.rs +++ b/src/executor/paginator.rs @@ -89,11 +89,7 @@ where Some(res) => res, None => return Ok(0), }; - #[allow(clippy::match_single_binding)] - let num_items = match self.db.get_database_backend() { - _ => result.try_get::("", "num_items")? as u64, - }; - Ok(num_items) + Ok(result.try_get::("", "num_items")? as u64) } /// Get the total number of pages @@ -278,14 +274,6 @@ where /// Paginate the result of a select operation. fn paginate(self, db: &'db C, page_size: u64) -> Paginator<'db, C, Self::Selector>; - - /// Perform a count on the paginated results - async fn count(self, db: &'db C) -> Result - where - Self: Send + Sized, - { - self.paginate(db, 1).num_items().await - } } impl<'db, C, S> PaginatorTrait<'db, C> for Selector diff --git a/src/executor/select_ext.rs b/src/executor/select_ext.rs index 4ff81d310..d82673a56 100644 --- a/src/executor/select_ext.rs +++ b/src/executor/select_ext.rs @@ -1,10 +1,9 @@ use crate::{ - ConnectionTrait, DbErr, EntityTrait, Select, SelectFive, SelectFour, SelectSix, SelectThree, - SelectTwo, Selector, SelectorRaw, SelectorTrait, Topology, + ConnectionTrait, DbErr, EntityTrait, QueryTrait, Select, SelectFive, SelectFour, SelectSix, + SelectThree, SelectTwo, Selector, SelectorRaw, SelectorTrait, Statement, Topology, }; use sea_query::{Expr, SelectStatement}; -// TODO: Move count here #[async_trait::async_trait] /// Helper trait for selectors with convenient methods pub trait SelectExt { @@ -23,6 +22,15 @@ pub trait SelectExt { } } +#[async_trait::async_trait] +/// Helper trait for counting rows selected by a query. +pub trait CountTrait { + /// Count the number of rows selected by this query. + async fn count(self, db: &impl ConnectionTrait) -> Result + where + Self: Send + Sized; +} + fn into_exists_query(mut stmt: SelectStatement) -> SelectStatement { stmt.clear_selects(); // Expr::Custom has fewer branches, but this may not have any significant impact on performance. @@ -33,6 +41,37 @@ fn into_exists_query(mut stmt: SelectStatement) -> SelectStatement { stmt } +fn build_count_query(stmt: SelectStatement) -> SelectStatement { + SelectStatement::new() + .expr(Expr::cust("COUNT(*) AS count")) + .from_subquery(stmt, "sub_query") + .to_owned() +} + +fn build_count_query_raw(stmt: Statement) -> SelectStatement { + let sub_query_sql = stmt.sql.trim().trim_end_matches(';').trim(); + let count_sql = format!("COUNT(*) AS count FROM ({sub_query_sql}) AS sub_query"); + + let mut query = SelectStatement::new(); + query.expr(if let Some(values) = stmt.values { + Expr::cust_with_values(count_sql, values.0) + } else { + Expr::cust(count_sql) + }); + query +} + +async fn exec_count(db: &C, stmt: SelectStatement) -> Result +where + C: ConnectionTrait, +{ + let result = match db.query_one(&stmt).await? { + Some(res) => res, + None => return Ok(0), + }; + Ok(result.try_get::("", "count")? as u64) +} + impl SelectExt for Selector where S: SelectorTrait, @@ -42,6 +81,19 @@ where } } +#[async_trait::async_trait] +impl CountTrait for Selector +where + S: SelectorTrait, +{ + async fn count(self, db: &impl ConnectionTrait) -> Result + where + Self: Send + Sized, + { + exec_count(db, build_count_query(self.query)).await + } +} + #[async_trait::async_trait] impl SelectExt for SelectorRaw where @@ -62,6 +114,19 @@ where } } +#[async_trait::async_trait] +impl CountTrait for SelectorRaw +where + S: SelectorTrait, +{ + async fn count(self, db: &impl ConnectionTrait) -> Result + where + Self: Send + Sized, + { + exec_count(db, build_count_query_raw(self.stmt)).await + } +} + impl SelectExt for Select where E: EntityTrait, @@ -135,6 +200,19 @@ where } } +#[async_trait::async_trait] +impl CountTrait for T +where + T: QueryTrait + Send, +{ + async fn count(self, db: &impl ConnectionTrait) -> Result + where + Self: Send + Sized, + { + exec_count(db, build_count_query(self.into_query())).await + } +} + #[cfg(test)] mod tests { use super::SelectExt; diff --git a/src/query/mod.rs b/src/query/mod.rs index 4fa9b0a82..7ce6ec4bd 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -47,7 +47,7 @@ pub use update::*; pub(crate) use util::*; pub use crate::{ - ConnectionTrait, CursorTrait, InsertResult, PaginatorTrait, SelectExt, Statement, + ConnectionTrait, CountTrait, CursorTrait, InsertResult, PaginatorTrait, SelectExt, Statement, TransactionTrait, UpdateResult, Value, Values, }; pub use sea_query::ExprTrait;