From f066b7a4d282e4cd9a94b4bd9f5a35b4405b9c21 Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Wed, 3 Apr 2024 12:00:46 +0800
Subject: [PATCH] `PrimaryKeyArity` trait

---
 src/entity/active_model.rs |  5 +++--
 src/entity/primary_key.rs  | 39 ++++++++++++++++++++++++++++++++++++--
 src/executor/insert.rs     |  2 +-
 src/schema/entity.rs       |  8 ++++----
 4 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/src/entity/active_model.rs b/src/entity/active_model.rs
index c0e34484..abd687a3 100644
--- 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, 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))
diff --git a/src/entity/primary_key.rs b/src/entity/primary_key.rs
index ee73f3d0..265f0c56 100644
--- a/src/entity/primary_key.rs
+++ b/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]
diff --git a/src/executor/insert.rs b/src/executor/insert.rs
index 34b80337..2b716838 100644
--- a/src/executor/insert.rs
+++ b/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| {
diff --git a/src/schema/entity.rs b/src/schema/entity.rs
index a9922663..7acabab5 100644
--- a/src/schema/entity.rs
+++ b/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();
             }
         }