From f066b7a4d282e4cd9a94b4bd9f5a35b4405b9c21 Mon Sep 17 00:00:00 2001 From: Billy Chan Date: Wed, 3 Apr 2024 12:00:46 +0800 Subject: [PATCH 1/2] `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 c0e344840..abd687a3f 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 ::PrimaryKey::iter().count() { + match <<::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 ee73f3d06..265f0c564 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 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 34b803378..2b7168388 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() && ::PrimaryKey::iter().count() > 0 { + if db.support_returning() { let db_backend = db.get_database_backend(); let returning = Query::returning().exprs(::PrimaryKey::iter().map(|c| { diff --git a/src/schema/entity.rs b/src/schema/entity.rs index a99226631..7acabab5c 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 <::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 <::ValueType as PrimaryKeyArity>::ARITY == 1 { column_def.primary_key(); } } From fe3e92bfdfed17732c3949c89684001180dd2e57 Mon Sep 17 00:00:00 2001 From: Chris Tsang Date: Wed, 3 Apr 2024 20:51:39 +0100 Subject: [PATCH 2/2] Make use of ARITY --- src/entity/prelude.rs | 4 ++-- src/executor/select.rs | 45 +++++++++++++++++++++++++++++++----------- 2 files changed, 35 insertions(+), 14 deletions(-) diff --git a/src/entity/prelude.rs b/src/entity/prelude.rs index fd1d91b1d..532b52d39 100644 --- 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 index 493f861fe..84a23ba67 100644 --- 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,6 +991,7 @@ where } } +#[allow(clippy::unwrap_used)] fn consolidate_query_result( rows: Vec<(L::Model, Option)>, ) -> Vec<(L::Model, Vec)> @@ -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<_> = ::iter() - .map(|pk| pk.into_column()) - .collect(); - if pkcol.len() == 1 { - consolidate_query_result_of::>(rows, UnitPk(pkcol[0])) - } else { - consolidate_query_result_of::>(rows, TuplePk(pkcol)) + match <::ValueType as PrimaryKeyArity>::ARITY { + 1 => { + let col = ::iter() + .next() + .unwrap() + .into_column(); + consolidate_query_result_of::>(rows, UnitPk(col)) + } + 2 => { + let mut iter = ::iter(); + let col1 = iter.next().unwrap().into_column(); + let col2 = iter.next().unwrap().into_column(); + consolidate_query_result_of::>(rows, PairPk(col1, col2)) + } + _ => { + let cols: Vec<_> = ::iter() + .map(|pk| pk.into_column()) + .collect(); + consolidate_query_result_of::>(rows, TuplePk(cols)) + } } } @@ -1014,8 +1027,9 @@ trait ModelKey { fn get(&self, model: &E::Model) -> Self::Type; } -// This could have been an array of [E::Column; ::ARITY] +// This could have been an array of [E::Column; ::ARITY], but it still doesn't compile struct UnitPk(E::Column); +struct PairPk(E::Column, E::Column); struct TuplePk(Vec); impl ModelKey for UnitPk { @@ -1025,6 +1039,13 @@ impl ModelKey for UnitPk { } } +impl ModelKey for PairPk { + type Type = (Value, Value); + fn get(&self, model: &E::Model) -> Self::Type { + (model.get(self.0), model.get(self.1)) + } +} + impl ModelKey for TuplePk { type Type = Vec; fn get(&self, model: &E::Model) -> Self::Type {