From ec0deaebef967ffc53460c620e1a1b493d65b2b8 Mon Sep 17 00:00:00 2001 From: Huliiiiii Date: Thu, 25 Jun 2026 12:36:43 +0800 Subject: [PATCH 1/2] Move `count()` to `CountTrait` --- sea-orm-sync/src/entity/prelude.rs | 2 +- sea-orm-sync/src/executor/paginator.rs | 14 +---- sea-orm-sync/src/executor/select_ext.rs | 69 +++++++++++++++++++- sea-orm-sync/src/query/mod.rs | 2 +- src/entity/prelude.rs | 8 +-- src/executor/paginator.rs | 14 +---- src/executor/select_ext.rs | 84 ++++++++++++++++++++++++- src/query/mod.rs | 2 +- 8 files changed, 157 insertions(+), 38 deletions(-) diff --git a/sea-orm-sync/src/entity/prelude.rs b/sea-orm-sync/src/entity/prelude.rs index 6e3e2cbc13..34fd5d7703 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 5215ba1815..ab54b300b9 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 8213c52e28..02a63342ac 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 4fa9b0a82b..7ce6ec4bd7 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 6e3e2cbc13..65ffcd843c 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 a3e88a6572..b267e64072 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 4ff81d310a..d82673a565 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 4fa9b0a82b..7ce6ec4bd7 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; From 7721692071cf8fef9a7885f3a594ef4f99940c62 Mon Sep 17 00:00:00 2001 From: Chris Tsang Date: Sun, 19 Jul 2026 15:22:52 +0100 Subject: [PATCH 2/2] Keep count() on PaginatorTrait Instead of a dedicated CountTrait, keep count() as a default method on PaginatorTrait. It wraps the query as-is in a SELECT COUNT(*) FROM (..) subquery, so it honors any limit/offset set on the query, unlike num_items which resets them. --- sea-orm-sync/src/entity/prelude.rs | 2 +- sea-orm-sync/src/executor/paginator.rs | 26 ++++++++ sea-orm-sync/src/executor/select_ext.rs | 68 +------------------- sea-orm-sync/src/query/mod.rs | 2 +- sea-orm-sync/tests/paginator_tests.rs | 7 ++- src/entity/prelude.rs | 8 +-- src/executor/paginator.rs | 26 ++++++++ src/executor/select_ext.rs | 83 +------------------------ src/query/mod.rs | 2 +- tests/paginator_tests.rs | 7 ++- 10 files changed, 74 insertions(+), 157 deletions(-) diff --git a/sea-orm-sync/src/entity/prelude.rs b/sea-orm-sync/src/entity/prelude.rs index 34fd5d7703..6e3e2cbc13 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, CountTrait, CursorTrait, DatabaseConnection, DbConn, EntityName, + 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, diff --git a/sea-orm-sync/src/executor/paginator.rs b/sea-orm-sync/src/executor/paginator.rs index ab54b300b9..7a0bc04f04 100644 --- a/sea-orm-sync/src/executor/paginator.rs +++ b/sea-orm-sync/src/executor/paginator.rs @@ -267,6 +267,32 @@ where /// Paginate the result of a select operation. fn paginate(self, db: &'db C, page_size: u64) -> Paginator<'db, C, Self::Selector>; + + /// Count the number of rows this query selects, honoring any `limit` and + /// `offset` set on it. + /// + /// This differs from [`Paginator::num_items`], which resets `limit`/`offset` + /// to report the grand total for pagination. `count` wraps the query as-is in + /// a `SELECT COUNT(*) FROM (..)` subquery, so a query with a `limit` counts at + /// most that many rows. + fn count(self, db: &'db C) -> Result + where + Self: Sized, + { + let paginator = self.paginate(db, 1); + let query = SelectStatement::new() + .expr(Expr::cust("COUNT(*) AS num_items")) + .from_subquery( + paginator.query.clone().clear_order_by().to_owned(), + "sub_query", + ) + .to_owned(); + let result = match db.query_one(&query)? { + Some(res) => res, + None => return Ok(0), + }; + Ok(result.try_get::("", "num_items")? as u64) + } } 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 02a63342ac..fa43ef547c 100644 --- a/sea-orm-sync/src/executor/select_ext.rs +++ b/sea-orm-sync/src/executor/select_ext.rs @@ -1,6 +1,6 @@ use crate::{ ConnectionTrait, DbErr, EntityTrait, Select, SelectFive, SelectFour, SelectSix, SelectThree, - QueryTrait, SelectTwo, Selector, SelectorRaw, SelectorTrait, Statement, Topology, + SelectTwo, Selector, SelectorRaw, SelectorTrait, Topology, }; use sea_query::{Expr, SelectStatement}; @@ -21,14 +21,6 @@ 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. @@ -39,37 +31,6 @@ 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, @@ -79,15 +40,6 @@ 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, @@ -107,15 +59,6 @@ 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, @@ -189,15 +132,6 @@ 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 7ce6ec4bd7..4fa9b0a82b 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, CountTrait, CursorTrait, InsertResult, PaginatorTrait, SelectExt, Statement, + ConnectionTrait, CursorTrait, InsertResult, PaginatorTrait, SelectExt, Statement, TransactionTrait, UpdateResult, Value, Values, }; pub use sea_query::ExprTrait; diff --git a/sea-orm-sync/tests/paginator_tests.rs b/sea-orm-sync/tests/paginator_tests.rs index 1e91016adc..950268e34c 100644 --- a/sea-orm-sync/tests/paginator_tests.rs +++ b/sea-orm-sync/tests/paginator_tests.rs @@ -4,7 +4,7 @@ pub mod common; pub use common::{TestContext, features::*, setup::*}; use pretty_assertions::assert_eq; -use sea_orm::{PaginatorTrait, QueryOrder, Set, entity::prelude::*}; +use sea_orm::{PaginatorTrait, QueryOrder, QuerySelect, Set, entity::prelude::*}; #[sea_orm_macros::test] fn paginator_tests() -> Result<(), DbErr> { @@ -145,5 +145,10 @@ pub fn paginator_count(db: &DatabaseConnection) -> Result<(), DbErr> { assert_eq!(Entity::find().filter(Column::Id.gt(100)).count(db)?, 0); + // `count` honors `limit`/`offset` set on the query, unlike `num_items`. + assert_eq!(Entity::find().limit(3).count(db)?, 3); + + assert_eq!(Entity::find().limit(4).offset(8).count(db)?, 2); + Ok(()) } diff --git a/src/entity/prelude.rs b/src/entity/prelude.rs index 65ffcd843c..6e3e2cbc13 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, 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, + 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, error::*, sea_query::{DynIden, Expr, RcOrArc, SeaRc, StringLen}, }; diff --git a/src/executor/paginator.rs b/src/executor/paginator.rs index b267e64072..83d7166823 100644 --- a/src/executor/paginator.rs +++ b/src/executor/paginator.rs @@ -274,6 +274,32 @@ where /// Paginate the result of a select operation. fn paginate(self, db: &'db C, page_size: u64) -> Paginator<'db, C, Self::Selector>; + + /// Count the number of rows this query selects, honoring any `limit` and + /// `offset` set on it. + /// + /// This differs from [`Paginator::num_items`], which resets `limit`/`offset` + /// to report the grand total for pagination. `count` wraps the query as-is in + /// a `SELECT COUNT(*) FROM (..)` subquery, so a query with a `limit` counts at + /// most that many rows. + async fn count(self, db: &'db C) -> Result + where + Self: Send + Sized, + { + let paginator = self.paginate(db, 1); + let query = SelectStatement::new() + .expr(Expr::cust("COUNT(*) AS num_items")) + .from_subquery( + paginator.query.clone().clear_order_by().to_owned(), + "sub_query", + ) + .to_owned(); + let result = match db.query_one(&query).await? { + Some(res) => res, + None => return Ok(0), + }; + Ok(result.try_get::("", "num_items")? as u64) + } } impl<'db, C, S> PaginatorTrait<'db, C> for Selector diff --git a/src/executor/select_ext.rs b/src/executor/select_ext.rs index d82673a565..6cbbeba55b 100644 --- a/src/executor/select_ext.rs +++ b/src/executor/select_ext.rs @@ -1,6 +1,6 @@ use crate::{ - ConnectionTrait, DbErr, EntityTrait, QueryTrait, Select, SelectFive, SelectFour, SelectSix, - SelectThree, SelectTwo, Selector, SelectorRaw, SelectorTrait, Statement, Topology, + ConnectionTrait, DbErr, EntityTrait, Select, SelectFive, SelectFour, SelectSix, SelectThree, + SelectTwo, Selector, SelectorRaw, SelectorTrait, Topology, }; use sea_query::{Expr, SelectStatement}; @@ -22,15 +22,6 @@ 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. @@ -41,37 +32,6 @@ 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, @@ -81,19 +41,6 @@ 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 @@ -114,19 +61,6 @@ 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, @@ -200,19 +134,6 @@ 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 7ce6ec4bd7..4fa9b0a82b 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, CountTrait, CursorTrait, InsertResult, PaginatorTrait, SelectExt, Statement, + ConnectionTrait, CursorTrait, InsertResult, PaginatorTrait, SelectExt, Statement, TransactionTrait, UpdateResult, Value, Values, }; pub use sea_query::ExprTrait; diff --git a/tests/paginator_tests.rs b/tests/paginator_tests.rs index 14deca3cf0..c52d007138 100644 --- a/tests/paginator_tests.rs +++ b/tests/paginator_tests.rs @@ -4,7 +4,7 @@ pub mod common; pub use common::{TestContext, features::*, setup::*}; use pretty_assertions::assert_eq; -use sea_orm::{PaginatorTrait, QueryOrder, Set, entity::prelude::*}; +use sea_orm::{PaginatorTrait, QueryOrder, QuerySelect, Set, entity::prelude::*}; #[sea_orm_macros::test] async fn paginator_tests() -> Result<(), DbErr> { @@ -149,5 +149,10 @@ pub async fn paginator_count(db: &DatabaseConnection) -> Result<(), DbErr> { 0 ); + // `count` honors `limit`/`offset` set on the query, unlike `num_items`. + assert_eq!(Entity::find().limit(3).count(db).await?, 3); + + assert_eq!(Entity::find().limit(4).offset(8).count(db).await?, 2); + Ok(()) }