Make use of ARITY

SeaQL · Mar 29, 2024 · cd7305c · cd7305c
1 parent 7c90c44
commit cd7305c
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 14 deletions.
diff --git a/src/entity/active_model.rs b/src/entity/active_model.rs
@@ -1,5 +1,6 @@
 use crate::{
-    error::*, ConnectionTrait, DeleteResult, EntityTrait, Iterable, PrimaryKeyToColumn, Value,
+    error::*, ConnectionTrait, DeleteResult, EntityTrait, Iterable, 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>::ARITY {
             1 => {
                 let s1 = next!();
                 Some(ValueTuple::One(s1))

diff --git a/src/entity/primary_key.rs b/src/entity/primary_key.rs
@@ -55,6 +55,9 @@ pub trait PrimaryKeyTrait: IdenStatic + Iterable {
     fn auto_increment() -> bool;
 }
 
+/// The maximum number of parts primary keys can have
+pub const PRIMARY_KEY_MAX_ARITY: usize = 12;
+
 /// How to map a Primary Key to a column
 pub trait PrimaryKeyToColumn {
     #[allow(missing_docs)]

diff --git a/src/executor/select.rs b/src/executor/select.rs
@@ -1,7 +1,7 @@
 use crate::{
     error::*, ConnectionTrait, DbBackend, EntityTrait, FromQueryResult, IdenStatic, Iterable,
-    ModelTrait, PartialModelTrait, PrimaryKeyToColumn, QueryResult, QuerySelect, Select, SelectA,
-    SelectB, SelectTwo, SelectTwoMany, Statement, StreamTrait, TryGetableMany,
+    ModelTrait, PartialModelTrait, PrimaryKeyToColumn, PrimaryKeyTrait, QueryResult, QuerySelect,
+    Select, SelectA, SelectB, SelectTwo, SelectTwoMany, Statement, StreamTrait, TryGetableMany,
 };
 use futures::{Stream, TryStreamExt};
 use sea_query::{SelectStatement, Value};
@@ -990,22 +990,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>::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 +1026,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 +1038,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 {