Make use of ARITY

This commit is contained in:
Chris Tsang 2024-04-03 20:51:39 +01:00
parent f066b7a4d2
commit fe3e92bfdf
2 changed files with 35 additions and 14 deletions

View File

@ -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")]

View File

@ -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<L, R>(
rows: Vec<(L::Model, Option<R::Model>)>,
) -> Vec<(L::Model, Vec<R::Model>)>
@ -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<_> = <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 {