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 {