Merge pull request #2185 from SeaQL/pk-arity-trait

`PrimaryKeyArity` trait

src/entity/active_model.rs

@@ -1,5 +1,6 @@
 use crate::{
-    error::*, ConnectionTrait, DeleteResult, EntityTrait, Iterable, PrimaryKeyToColumn, Value,
+    error::*, ConnectionTrait, DeleteResult, EntityTrait, Iterable, PrimaryKeyArity,
+    PrimaryKeyToColumn, PrimaryKeyTrait, Value,
 };
 use async_trait::async_trait;
 use sea_query::{Nullable, ValueTuple};
@@ -139,7 +140,7 @@ pub trait ActiveModelTrait: Clone + Debug {
                 }
             };
         }
-        match <Self::Entity as EntityTrait>::PrimaryKey::iter().count() {
+        match <<<Self::Entity as EntityTrait>::PrimaryKey as PrimaryKeyTrait>::ValueType as PrimaryKeyArity>::ARITY {
             1 => {
                 let s1 = next!();
                 Some(ValueTuple::One(s1))

src/entity/prelude.rs

@@ -4,8 +4,8 @@ pub use crate::{
     ActiveEnum, ActiveModelBehavior, ActiveModelTrait, ColumnDef, ColumnTrait, ColumnType,
     ColumnTypeTrait, ConnectionTrait, CursorTrait, DatabaseConnection, DbConn, EntityName,
     EntityTrait, EnumIter, ForeignKeyAction, Iden, IdenStatic, Linked, LoaderTrait, ModelTrait,
-    PaginatorTrait, PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, QueryResult, Related,
-    RelationDef, RelationTrait, Select, Value,
+    PaginatorTrait, PrimaryKeyArity, PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, QueryResult,
+    Related, RelationDef, RelationTrait, Select, Value,
 };
 
 #[cfg(feature = "macros")]

src/entity/primary_key.rs

@@ -46,9 +46,10 @@ pub trait PrimaryKeyTrait: IdenStatic + Iterable {
         + IntoValueTuple
         + FromValueTuple
         + TryGetableMany
-        + TryFromU64;
+        + TryFromU64
+        + PrimaryKeyArity;
 
-    /// Method to call to perform `AUTOINCREMENT` operation on a Primary Kay
+    /// Method to call to perform `AUTOINCREMENT` operation on a Primary Key
     fn auto_increment() -> bool;
 }
 
@@ -66,6 +67,40 @@ pub trait PrimaryKeyToColumn {
         Self: Sized;
 }
 
+/// How many columns this Primary Key comprises
+pub trait PrimaryKeyArity {
+    /// Arity of the Primary Key
+    const ARITY: usize;
+}
+
+impl<V> PrimaryKeyArity for V
+where
+    V: crate::TryGetable,
+{
+    const ARITY: usize = 1;
+}
+
+macro_rules! impl_pk_arity {
+    ($len:expr, $($tuple_arg:ident),*) => {
+        impl<$($tuple_arg: crate::TryGetableMany,)*> PrimaryKeyArity for ($($tuple_arg,)*) {
+            const ARITY: usize = $len;
+        }
+    }
+}
+
+impl_pk_arity!(1, T1);
+impl_pk_arity!(2, T1, T2);
+impl_pk_arity!(3, T1, T2, T3);
+impl_pk_arity!(4, T1, T2, T3, T4);
+impl_pk_arity!(5, T1, T2, T3, T4, T5);
+impl_pk_arity!(6, T1, T2, T3, T4, T5, T6);
+impl_pk_arity!(7, T1, T2, T3, T4, T5, T6, T7);
+impl_pk_arity!(8, T1, T2, T3, T4, T5, T6, T7, T8);
+impl_pk_arity!(9, T1, T2, T3, T4, T5, T6, T7, T8, T9);
+impl_pk_arity!(10, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10);
+impl_pk_arity!(11, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11);
+impl_pk_arity!(12, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12);
+
 #[cfg(test)]
 mod tests {
     #[test]

src/executor/insert.rs

@@ -116,7 +116,7 @@ where
     {
         // so that self is dropped before entering await
         let mut query = self.query;
-        if db.support_returning() && <A::Entity as EntityTrait>::PrimaryKey::iter().count() > 0 {
+        if db.support_returning() {
             let db_backend = db.get_database_backend();
             let returning =
                 Query::returning().exprs(<A::Entity as EntityTrait>::PrimaryKey::iter().map(|c| {

src/executor/select.rs

@@ -1,7 +1,8 @@
 use crate::{
     error::*, ConnectionTrait, DbBackend, EntityTrait, FromQueryResult, IdenStatic, Iterable,
-    ModelTrait, PartialModelTrait, PrimaryKeyToColumn, QueryResult, QuerySelect, Select, SelectA,
-    SelectB, SelectTwo, SelectTwoMany, Statement, StreamTrait, TryGetableMany,
+    ModelTrait, PartialModelTrait, PrimaryKeyArity, PrimaryKeyToColumn, PrimaryKeyTrait,
+    QueryResult, QuerySelect, Select, SelectA, SelectB, SelectTwo, SelectTwoMany, Statement,
+    StreamTrait, TryGetableMany,
 };
 use futures::{Stream, TryStreamExt};
 use sea_query::{SelectStatement, Value};
@@ -990,6 +991,7 @@ where
     }
 }
 
+#[allow(clippy::unwrap_used)]
 fn consolidate_query_result<L, R>(
     rows: Vec<(L::Model, Option<R::Model>)>,
 ) -> Vec<(L::Model, Vec<R::Model>)>
@@ -997,15 +999,26 @@ where
     L: EntityTrait,
     R: EntityTrait,
 {
-    // This is a strong point to consider adding a trait associated constant
-    // to PrimaryKeyTrait to indicate the arity
-    let pkcol: Vec<_> = <L::PrimaryKey as Iterable>::iter()
-        .map(|pk| pk.into_column())
-        .collect();
-    if pkcol.len() == 1 {
-        consolidate_query_result_of::<L, R, UnitPk<L>>(rows, UnitPk(pkcol[0]))
-    } else {
-        consolidate_query_result_of::<L, R, TuplePk<L>>(rows, TuplePk(pkcol))
+    match <<L::PrimaryKey as PrimaryKeyTrait>::ValueType as PrimaryKeyArity>::ARITY {
+        1 => {
+            let col = <L::PrimaryKey as Iterable>::iter()
+                .next()
+                .unwrap()
+                .into_column();
+            consolidate_query_result_of::<L, R, UnitPk<L>>(rows, UnitPk(col))
+        }
+        2 => {
+            let mut iter = <L::PrimaryKey as Iterable>::iter();
+            let col1 = iter.next().unwrap().into_column();
+            let col2 = iter.next().unwrap().into_column();
+            consolidate_query_result_of::<L, R, PairPk<L>>(rows, PairPk(col1, col2))
+        }
+        _ => {
+            let cols: Vec<_> = <L::PrimaryKey as Iterable>::iter()
+                .map(|pk| pk.into_column())
+                .collect();
+            consolidate_query_result_of::<L, R, TuplePk<L>>(rows, TuplePk(cols))
+        }
     }
 }
 
@@ -1014,8 +1027,9 @@ trait ModelKey<E: EntityTrait> {
     fn get(&self, model: &E::Model) -> Self::Type;
 }
 
-// This could have been an array of [E::Column; <E::PrimaryKey as PrimaryKeyTrait>::ARITY]
+// This could have been an array of [E::Column; <E::PrimaryKey as PrimaryKeyTrait>::ARITY], but it still doesn't compile
 struct UnitPk<E: EntityTrait>(E::Column);
+struct PairPk<E: EntityTrait>(E::Column, E::Column);
 struct TuplePk<E: EntityTrait>(Vec<E::Column>);
 
 impl<E: EntityTrait> ModelKey<E> for UnitPk<E> {
@@ -1025,6 +1039,13 @@ impl<E: EntityTrait> ModelKey<E> for UnitPk<E> {
     }
 }
 
+impl<E: EntityTrait> ModelKey<E> for PairPk<E> {
+    type Type = (Value, Value);
+    fn get(&self, model: &E::Model) -> Self::Type {
+        (model.get(self.0), model.get(self.1))
+    }
+}
+
 impl<E: EntityTrait> ModelKey<E> for TuplePk<E> {
     type Type = Vec<Value>;
     fn get(&self, model: &E::Model) -> Self::Type {

src/schema/entity.rs

@@ -1,6 +1,6 @@
 use crate::{
-    ActiveEnum, ColumnTrait, ColumnType, DbBackend, EntityTrait, Iterable, PrimaryKeyToColumn,
-    PrimaryKeyTrait, RelationTrait, Schema,
+    ActiveEnum, ColumnTrait, ColumnType, DbBackend, EntityTrait, Iterable, PrimaryKeyArity,
+    PrimaryKeyToColumn, PrimaryKeyTrait, RelationTrait, Schema,
 };
 use sea_query::{
     extension::postgres::{Type, TypeCreateStatement},
@@ -172,7 +172,7 @@ where
         stmt.col(&mut column_def);
     }
 
-    if E::PrimaryKey::iter().count() > 1 {
+    if <<E::PrimaryKey as PrimaryKeyTrait>::ValueType as PrimaryKeyArity>::ARITY > 1 {
         let mut idx_pk = Index::create();
         for primary_key in E::PrimaryKey::iter() {
             idx_pk.col(primary_key);
@@ -228,7 +228,7 @@ where
             if E::PrimaryKey::auto_increment() {
                 column_def.auto_increment();
             }
-            if E::PrimaryKey::iter().count() == 1 {
+            if <<E::PrimaryKey as PrimaryKeyTrait>::ValueType as PrimaryKeyArity>::ARITY == 1 {
                 column_def.primary_key();
             }
         }