Make use of ARITY

SeaQL · Apr 3, 2024 · fe3e92b · fe3e92b
1 parent f066b7a
commit fe3e92b
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 14 deletions.
diff --git a/src/entity/prelude.rs b/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")]

diff --git a/src/executor/select.rs b/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,22 +991,34 @@ where
     }
 }
 
+#[allow(clippy::unwrap_used)]
 fn consolidate_query_result<L, R>(
     rows: Vec<(L::Model, Option<R::Model>)>,
 ) -> Vec<(L::Model, Vec<R::Model>)>
 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 {