Make use of ARITY

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,
+    PaginatorTrait, PrimaryKeyArity, PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, QueryResult,
-    RelationDef, RelationTrait, Select, Value,
+    Related, RelationDef, RelationTrait, Select, Value,
 };
 
 #[cfg(feature = "macros")]

src/executor/select.rs

@@ -1,7 +1,8 @@
 use crate::{
     error::*, ConnectionTrait, DbBackend, EntityTrait, FromQueryResult, IdenStatic, Iterable,
-    ModelTrait, PartialModelTrait, PrimaryKeyToColumn, QueryResult, QuerySelect, Select, SelectA,
+    ModelTrait, PartialModelTrait, PrimaryKeyArity, PrimaryKeyToColumn, PrimaryKeyTrait,
-    SelectB, SelectTwo, SelectTwoMany, Statement, StreamTrait, TryGetableMany,
+    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
+    match <<L::PrimaryKey as PrimaryKeyTrait>::ValueType as PrimaryKeyArity>::ARITY {
-    // to PrimaryKeyTrait to indicate the arity
+        1 => {
-    let pkcol: Vec<_> = <L::PrimaryKey as Iterable>::iter()
+            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();
-    if pkcol.len() == 1 {
+            consolidate_query_result_of::<L, R, TuplePk<L>>(rows, TuplePk(cols))
-        consolidate_query_result_of::<L, R, UnitPk<L>>(rows, UnitPk(pkcol[0]))
+        }
-    } else {
-        consolidate_query_result_of::<L, R, TuplePk<L>>(rows, TuplePk(pkcol))
     }
 }
 
@@ -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 {