Refactor Query as traits

src/entity/base.rs

@@ -1,7 +1,7 @@
 use crate::{
     ActiveModelOf, ActiveModelTrait, ColumnTrait, Insert, ModelTrait, OneOrManyActiveModel,
-    PrimaryKeyOfModel, PrimaryKeyTrait, RelationBuilder, RelationTrait, RelationType, Select,
+    PrimaryKeyOfModel, PrimaryKeyTrait, QueryFilter, RelationBuilder, RelationTrait, RelationType,
-    SelectHelper,
+    Select,
 };
 use sea_query::{Iden, IntoValueTuple};
 use std::fmt::Debug;

src/entity/column.rs

@@ -57,7 +57,7 @@ pub trait ColumnTrait: IdenStatic + Iterable {
     bind_oper!(lte);
 
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -75,7 +75,7 @@ pub trait ColumnTrait: IdenStatic + Iterable {
     }
 
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -93,7 +93,7 @@ pub trait ColumnTrait: IdenStatic + Iterable {
     }
 
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -108,7 +108,7 @@ pub trait ColumnTrait: IdenStatic + Iterable {
     }
 
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -123,7 +123,7 @@ pub trait ColumnTrait: IdenStatic + Iterable {
     }
 
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -139,7 +139,7 @@ pub trait ColumnTrait: IdenStatic + Iterable {
     }
 
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -155,7 +155,7 @@ pub trait ColumnTrait: IdenStatic + Iterable {
     }
 
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()

src/entity/prelude.rs

@@ -1,6 +1,6 @@
 pub use crate::{
     ActiveModelOf, ActiveModelTrait, ColumnTrait, ColumnType, DeriveActiveModel, DeriveColumn,
     DeriveEntity, DeriveModel, DerivePrimaryKey, EntityName, EntityTrait, EnumIter, Iden,
-    IdenStatic, ModelTrait, PrimaryKeyOfModel, PrimaryKeyTrait, QueryResult, Related, RelationDef,
+    IdenStatic, ModelTrait, PrimaryKeyOfModel, PrimaryKeyTrait, QueryFilter, QueryResult, Related,
-    RelationTrait, Select, TypeErr, Value,
+    RelationDef, RelationTrait, Select, TypeErr, Value,
 };

src/lib.rs

@@ -2,7 +2,7 @@ mod connector;
 mod database;
 mod driver;
 pub mod entity;
-mod query;
+pub mod query;
 pub mod tests_cfg;
 mod util;
 

src/query/combine.rs

@@ -71,7 +71,7 @@ where
 #[cfg(test)]
 mod tests {
     use crate::tests_cfg::{cake, fruit};
-    use crate::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper};
+    use crate::{ColumnTrait, EntityTrait, QueryFilter, QueryTrait, SelectHelper};
     use sea_query::MysqlQueryBuilder;
 
     #[test]

src/query/helper.rs

@@ -1,6 +1,9 @@
-use crate::{ColumnTrait, Identity, IntoSimpleExpr, RelationDef};
+use crate::{
+    ColumnTrait, EntityTrait, Identity, IntoSimpleExpr, Iterable, ModelTrait, PrimaryKeyOfModel,
+    RelationDef,
+};
 use sea_query::{Alias, Expr, IntoCondition, SelectExpr, SelectStatement, SimpleExpr};
-pub use sea_query::{Condition, JoinType, Order};
+pub use sea_query::{Condition, ConditionalStatement, JoinType, Order};
 use std::rc::Rc;
 
 pub trait SelectHelper: Sized {
@@ -14,7 +17,7 @@ pub trait SelectHelper: Sized {
 
     /// Add a select column
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::PostgresQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::PostgresQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -35,7 +38,7 @@ pub trait SelectHelper: Sized {
 
     /// Add a select column with alias
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::PostgresQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::PostgresQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -57,47 +60,9 @@ pub trait SelectHelper: Sized {
         self
     }
 
-    /// Add an AND WHERE expression
-    /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
-    ///
-    /// assert_eq!(
-    ///     cake::Entity::find()
-    ///         .filter(cake::Column::Id.eq(4))
-    ///         .filter(cake::Column::Id.eq(5))
-    ///         .build(MysqlQueryBuilder)
-    ///         .to_string(),
-    ///     "SELECT `cake`.`id`, `cake`.`name` FROM `cake` WHERE `cake`.`id` = 4 AND `cake`.`id` = 5"
-    /// );
-    /// ```
-    ///
-    /// Add a condition tree.
-    /// ```
-    /// use sea_orm::{Condition, ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
-    ///
-    /// assert_eq!(
-    ///     cake::Entity::find()
-    ///         .filter(
-    ///             Condition::any()
-    ///                 .add(cake::Column::Id.eq(4))
-    ///                 .add(cake::Column::Id.eq(5))
-    ///         )
-    ///         .build(MysqlQueryBuilder)
-    ///         .to_string(),
-    ///     "SELECT `cake`.`id`, `cake`.`name` FROM `cake` WHERE `cake`.`id` = 4 OR `cake`.`id` = 5"
-    /// );
-    /// ```
-    fn filter<F>(mut self, filter: F) -> Self
-    where
-        F: IntoCondition,
-    {
-        self.query().cond_where(filter.into_condition());
-        self
-    }
-
     /// Add a group by column
     /// ```
-    /// use sea_orm::{ColumnTrait, EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::PostgresQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::PostgresQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -119,7 +84,7 @@ pub trait SelectHelper: Sized {
 
     /// Add an order_by expression
     /// ```
-    /// use sea_orm::{EntityTrait, Order, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -140,7 +105,7 @@ pub trait SelectHelper: Sized {
 
     /// Add an order_by expression (ascending)
     /// ```
-    /// use sea_orm::{EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -161,7 +126,7 @@ pub trait SelectHelper: Sized {
 
     /// Add an order_by expression (descending)
     /// ```
-    /// use sea_orm::{EntityTrait, QueryTrait, SelectHelper, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
     ///
     /// assert_eq!(
     ///     cake::Entity::find()
@@ -214,6 +179,63 @@ pub trait SelectHelper: Sized {
     }
 }
 
+pub trait QueryFilter: Sized {
+    type QueryStatement: ConditionalStatement;
+
+    fn query(&mut self) -> &mut Self::QueryStatement;
+
+    /// Add an AND WHERE expression
+    /// ```
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    ///
+    /// assert_eq!(
+    ///     cake::Entity::find()
+    ///         .filter(cake::Column::Id.eq(4))
+    ///         .filter(cake::Column::Id.eq(5))
+    ///         .build(MysqlQueryBuilder)
+    ///         .to_string(),
+    ///     "SELECT `cake`.`id`, `cake`.`name` FROM `cake` WHERE `cake`.`id` = 4 AND `cake`.`id` = 5"
+    /// );
+    /// ```
+    ///
+    /// Add a condition tree.
+    /// ```
+    /// use sea_orm::{entity::*, query::*, tests_cfg::cake, sea_query::MysqlQueryBuilder};
+    ///
+    /// assert_eq!(
+    ///     cake::Entity::find()
+    ///         .filter(
+    ///             Condition::any()
+    ///                 .add(cake::Column::Id.eq(4))
+    ///                 .add(cake::Column::Id.eq(5))
+    ///         )
+    ///         .build(MysqlQueryBuilder)
+    ///         .to_string(),
+    ///     "SELECT `cake`.`id`, `cake`.`name` FROM `cake` WHERE `cake`.`id` = 4 OR `cake`.`id` = 5"
+    /// );
+    /// ```
+    fn filter<F>(mut self, filter: F) -> Self
+    where
+        F: IntoCondition,
+    {
+        self.query().cond_where(filter.into_condition());
+        self
+    }
+
+    /// Apply a where condition using the model's primary key
+    fn belongs_to<E>(mut self, model: &E::Model) -> Self
+    where
+        E: EntityTrait,
+        E::PrimaryKey: PrimaryKeyOfModel<E::Model>,
+    {
+        for key in E::PrimaryKey::iter() {
+            let col = key.into_column();
+            self = self.filter(col.eq(model.get(col)));
+        }
+        self
+    }
+}
+
 fn join_condition(rel: RelationDef) -> SimpleExpr {
     let from_tbl = rel.from_tbl.clone();
     let to_tbl = rel.to_tbl.clone();

src/query/insert.rs

@@ -71,7 +71,7 @@ impl<A> QueryTrait for Insert<A>
 where
     A: ActiveModelTrait,
 {
-    type QueryStatementBuilder = InsertStatement;
+    type QueryStatement = InsertStatement;
 
     fn query(&mut self) -> &mut InsertStatement {
         &mut self.query

src/query/join.rs

@@ -1,28 +1,10 @@
-use crate::{
+use crate::{EntityTrait, Related, Select, SelectHelper, SelectTwo};
-    ColumnTrait, EntityTrait, Iterable, ModelTrait, PrimaryKeyOfModel, Related, Select,
-    SelectHelper, SelectTwo,
-};
 pub use sea_query::JoinType;
 
 impl<E> Select<E>
 where
     E: EntityTrait,
 {
-    /// Apply a where condition using the model's primary key
-    pub fn belongs_to<R>(self, model: &R::Model) -> Self
-    where
-        R: EntityTrait + Related<E>,
-        R::PrimaryKey: PrimaryKeyOfModel<R::Model>,
-    {
-        if let Some(key) = R::PrimaryKey::iter().next() {
-            // TODO: supporting composite primary key
-            let col = key.into_column();
-            self.filter(col.eq(model.get(col)))
-        } else {
-            panic!("undefined primary key");
-        }
-    }
-
     /// Left Join with a Related Entity.
     pub fn left_join<R>(self, _: R) -> Self
     where
@@ -71,7 +53,7 @@ where
 #[cfg(test)]
 mod tests {
     use crate::tests_cfg::{cake, filling, fruit};
-    use crate::{QueryTrait, ColumnTrait, EntityTrait, SelectHelper};
+    use crate::{ColumnTrait, EntityTrait, QueryFilter, QueryTrait};
     use sea_query::MysqlQueryBuilder;
 
     #[test]

src/query/select.rs

@@ -1,4 +1,4 @@
-use crate::{ColumnTrait, EntityTrait, Iterable, QueryTrait, SelectHelper};
+use crate::{ColumnTrait, EntityTrait, Iterable, QueryFilter, QueryTrait, SelectHelper};
 use core::fmt::Debug;
 use core::marker::PhantomData;
 pub use sea_query::JoinType;
@@ -37,6 +37,17 @@ where
     }
 }
 
+impl<E> QueryFilter for Select<E>
+where
+    E: EntityTrait,
+{
+    type QueryStatement = SelectStatement;
+
+    fn query(&mut self) -> &mut SelectStatement {
+        &mut self.query
+    }
+}
+
 impl<E, F> SelectHelper for SelectTwo<E, F>
 where
     E: EntityTrait,
@@ -47,6 +58,18 @@ where
     }
 }
 
+impl<E, F> QueryFilter for SelectTwo<E, F>
+where
+    E: EntityTrait,
+    F: EntityTrait,
+{
+    type QueryStatement = SelectStatement;
+
+    fn query(&mut self) -> &mut SelectStatement {
+        &mut self.query
+    }
+}
+
 impl<C> IntoSimpleExpr for C
 where
     C: ColumnTrait,
@@ -95,7 +118,7 @@ impl<E> QueryTrait for Select<E>
 where
     E: EntityTrait,
 {
-    type QueryStatementBuilder = SelectStatement;
+    type QueryStatement = SelectStatement;
     fn query(&mut self) -> &mut SelectStatement {
         &mut self.query
     }
@@ -112,7 +135,7 @@ where
     E: EntityTrait,
     F: EntityTrait,
 {
-    type QueryStatementBuilder = SelectStatement;
+    type QueryStatement = SelectStatement;
     fn query(&mut self) -> &mut SelectStatement {
         &mut self.query
     }

src/query/traits.rs

@@ -2,16 +2,16 @@ use crate::Statement;
 use sea_query::{QueryBuilder, QueryStatementBuilder};
 
 pub trait QueryTrait {
-    type QueryStatementBuilder: QueryStatementBuilder;
+    type QueryStatement: QueryStatementBuilder;
 
     /// Get a mutable ref to the query builder
-    fn query(&mut self) -> &mut Self::QueryStatementBuilder;
+    fn query(&mut self) -> &mut Self::QueryStatement;
 
     /// Get an immutable ref to the query builder
-    fn as_query(&self) -> &Self::QueryStatementBuilder;
+    fn as_query(&self) -> &Self::QueryStatement;
 
     /// Take ownership of the query builder
-    fn into_query(self) -> Self::QueryStatementBuilder;
+    fn into_query(self) -> Self::QueryStatement;
 
     /// Build the query as [`Statement`]
     fn build<B>(&self, builder: B) -> Statement