Merge pull request #1898 from SeaQL/rework-active-enum

Rework ActiveEnum to support TryGetable for Vec<Json>

sea-orm-macros/src/derives/active_enum.rs

@@ -337,6 +337,16 @@ impl ActiveEnum {
                 }
             }
 
+            #[automatically_derived]
+            impl sea_orm::TryGetableArray for #ident {
+                fn try_get_by<I: sea_orm::ColIdx>(res: &sea_orm::QueryResult, index: I) -> std::result::Result<Vec<Self>, sea_orm::TryGetError> {
+                    <<Self as sea_orm::ActiveEnum>::Value as sea_orm::ActiveEnumValue>::try_get_vec_by(res, index)?
+                        .into_iter()
+                        .map(|value| <Self as sea_orm::ActiveEnum>::try_from_value(&value).map_err(Into::into))
+                        .collect()
+                }
+            }
+
             #[automatically_derived]
             #[allow(clippy::from_over_into)]
             impl Into<sea_orm::sea_query::Value> for #ident {

sea-orm-macros/src/derives/try_getable_from_json.rs

@@ -2,6 +2,15 @@ use proc_macro2::{Ident, TokenStream};
 use quote::quote;
 
 pub fn expand_derive_from_json_query_result(ident: Ident) -> syn::Result<TokenStream> {
+    let impl_not_u8 = if cfg!(feature = "postgres-array") {
+        quote!(
+            #[automatically_derived]
+            impl sea_orm::sea_query::value::with_array::NotU8 for #ident {}
+        )
+    } else {
+        quote!()
+    };
+
     Ok(quote!(
         #[automatically_derived]
         impl sea_orm::TryGetableFromJson for #ident {}
@@ -43,5 +52,7 @@ pub fn expand_derive_from_json_query_result(ident: Ident) -> syn::Result<TokenSt
                 sea_orm::Value::Json(None)
             }
         }
+
+        #impl_not_u8
     ))
 }

src/entity/active_enum.rs

@@ -1,4 +1,4 @@
-use crate::{ColumnDef, DbErr, Iterable, QueryResult, TryFromU64, TryGetError, TryGetable};
+use crate::{ColIdx, ColumnDef, DbErr, Iterable, QueryResult, TryFromU64, TryGetError, TryGetable};
 use sea_query::{DynIden, Expr, Nullable, SimpleExpr, Value, ValueType};
 
 /// A Rust representation of enum defined in database.
@@ -110,11 +110,11 @@ use sea_query::{DynIden, Expr, Nullable, SimpleExpr, Value, ValueType};
 /// impl ActiveModelBehavior for ActiveModel {}
 /// ```
 pub trait ActiveEnum: Sized + Iterable {
-    /// Define the Rust type that each enum variant represents.
-    type Value: Into<Value> + ValueType + Nullable + TryGetable;
+    /// Define the Rust type that each enum variant corresponds.
+    type Value: ActiveEnumValue;
 
-    /// Define the enum value in Vector type.
-    type ValueVec: IntoIterator<Item = Self::Value>;
+    /// This has no purpose. It will be removed in the next major version.
+    type ValueVec;
 
     /// Get the name of enum
     fn name() -> DynIden;
@@ -144,19 +144,53 @@ pub trait ActiveEnum: Sized + Iterable {
     }
 }
 
-impl<T> TryGetable for Vec<T>
-where
-    T: ActiveEnum,
-    T::ValueVec: TryGetable,
-{
-    fn try_get_by<I: crate::ColIdx>(res: &QueryResult, index: I) -> Result<Self, TryGetError> {
-        <T::ValueVec as TryGetable>::try_get_by(res, index)?
-            .into_iter()
-            .map(|value| T::try_from_value(&value).map_err(Into::into))
-            .collect()
-    }
+/// The Rust Value backing ActiveEnums
+pub trait ActiveEnumValue: Into<Value> + ValueType + Nullable + TryGetable {
+    /// For getting an array of enum. Postgres only
+    fn try_get_vec_by<I: ColIdx>(res: &QueryResult, index: I) -> Result<Vec<Self>, TryGetError>;
+}
+
+macro_rules! impl_active_enum_value {
+    ($type:ident) => {
+        impl ActiveEnumValue for $type {
+            fn try_get_vec_by<I: ColIdx>(
+                _res: &QueryResult,
+                _index: I,
+            ) -> Result<Vec<Self>, TryGetError> {
+                panic!("Not supported by `postgres-array`")
+            }
+        }
+    };
 }
 
+macro_rules! impl_active_enum_value_with_pg_array {
+    ($type:ident) => {
+        impl ActiveEnumValue for $type {
+            fn try_get_vec_by<I: ColIdx>(
+                _res: &QueryResult,
+                _index: I,
+            ) -> Result<Vec<Self>, TryGetError> {
+                #[cfg(feature = "postgres-array")]
+                {
+                    <Vec<Self>>::try_get_by(_res, _index)
+                }
+                #[cfg(not(feature = "postgres-array"))]
+                panic!("`postgres-array` is not enabled")
+            }
+        }
+    };
+}
+
+impl_active_enum_value!(u8);
+impl_active_enum_value!(u16);
+impl_active_enum_value!(u32);
+impl_active_enum_value!(u64);
+impl_active_enum_value_with_pg_array!(String);
+impl_active_enum_value_with_pg_array!(i8);
+impl_active_enum_value_with_pg_array!(i16);
+impl_active_enum_value_with_pg_array!(i32);
+impl_active_enum_value_with_pg_array!(i64);
+
 impl<T> TryFromU64 for T
 where
     T: ActiveEnum,

src/entity/column.rs

@@ -281,6 +281,7 @@ pub trait ColumnTrait: IdenStatic + Iterable + FromStr {
     }
 
     /// Cast value of an enum column as enum type; do nothing if `self` is not an enum.
+    /// Will also transform `Array(Vec<Json>)` into `Json(Vec<Json>)` if the column type is `Json`.
     fn save_enum_as(&self, val: Expr) -> SimpleExpr {
         cast_enum_as(val, self, |col, enum_name, col_type| {
             let type_name = match col_type {
@@ -412,9 +413,41 @@ where
 {
     let col_def = col.def();
     let col_type = col_def.get_column_type();
-    match col_type.get_enum_name() {
-        Some(enum_name) => f(expr, SeaRc::clone(enum_name), col_type),
-        None => expr.into(),
+
+    match col_type {
+        #[cfg(all(feature = "with-json", feature = "postgres-array"))]
+        ColumnType::Json | ColumnType::JsonBinary => {
+            use sea_query::ArrayType;
+            use serde_json::Value as Json;
+
+            #[allow(clippy::boxed_local)]
+            fn unbox<T>(boxed: Box<T>) -> T {
+                *boxed
+            }
+
+            let expr = expr.into();
+            match expr {
+                SimpleExpr::Value(Value::Array(ArrayType::Json, Some(json_vec))) => {
+                    // flatten Array(Vec<Json>) into Json
+                    let json_vec: Vec<Json> = json_vec
+                        .into_iter()
+                        .filter_map(|val| match val {
+                            Value::Json(Some(json)) => Some(unbox(json)),
+                            _ => None,
+                        })
+                        .collect();
+                    SimpleExpr::Value(Value::Json(Some(Box::new(json_vec.into()))))
+                }
+                SimpleExpr::Value(Value::Array(ArrayType::Json, None)) => {
+                    SimpleExpr::Value(Value::Json(None))
+                }
+                _ => expr,
+            }
+        }
+        _ => match col_type.get_enum_name() {
+            Some(enum_name) => f(expr, SeaRc::clone(enum_name), col_type),
+            None => expr.into(),
+        },
     }
 }
 

src/executor/query.rs

@@ -960,6 +960,25 @@ fn try_get_many_with_slice_len_of(len: usize, cols: &[String]) -> Result<(), Try
     }
 }
 
+/// An interface to get an array of values from the query result.
+/// A type can only implement `ActiveEnum` or `TryGetableFromJson`, but not both.
+/// A blanket impl is provided for `TryGetableFromJson`, while the impl for `ActiveEnum`
+/// is provided by the `DeriveActiveEnum` macro. So as an end user you won't normally
+/// touch this trait.
+pub trait TryGetableArray: Sized {
+    /// Just a delegate
+    fn try_get_by<I: ColIdx>(res: &QueryResult, index: I) -> Result<Vec<Self>, TryGetError>;
+}
+
+impl<T> TryGetable for Vec<T>
+where
+    T: TryGetableArray,
+{
+    fn try_get_by<I: ColIdx>(res: &QueryResult, index: I) -> Result<Self, TryGetError> {
+        T::try_get_by(res, index)
+    }
+}
+
 // TryGetableFromJson //
 
 /// An interface to get a JSON from the query result
@@ -999,6 +1018,22 @@ where
             _ => unreachable!(),
         }
     }
+
+    /// Get a Vec<Self> from an Array of Json
+    fn from_json_vec(value: serde_json::Value) -> Result<Vec<Self>, TryGetError> {
+        match value {
+            serde_json::Value::Array(values) => {
+                let mut res = Vec::new();
+                for item in values {
+                    res.push(serde_json::from_value(item).map_err(json_err)?);
+                }
+                Ok(res)
+            }
+            _ => Err(TryGetError::DbErr(DbErr::Json(
+                "Value is not an Array".to_owned(),
+            ))),
+        }
+    }
 }
 
 #[cfg(feature = "with-json")]
@@ -1011,6 +1046,16 @@ where
     }
 }
 
+#[cfg(feature = "with-json")]
+impl<T> TryGetableArray for T
+where
+    T: TryGetableFromJson,
+{
+    fn try_get_by<I: ColIdx>(res: &QueryResult, index: I) -> Result<Vec<T>, TryGetError> {
+        T::from_json_vec(serde_json::Value::try_get_by(res, index)?)
+    }
+}
+
 // TryFromU64 //
 /// Try to convert a type to a u64
 pub trait TryFromU64: Sized {

tests/active_enum_tests.rs

@@ -22,8 +22,13 @@ async fn main() -> Result<(), DbErr> {
     create_tables(&ctx.db).await?;
     insert_active_enum(&ctx.db).await?;
     insert_active_enum_child(&ctx.db).await?;
+
+    #[cfg(feature = "sqlx-postgres")]
+    insert_active_enum_vec(&ctx.db).await?;
+
     find_related_active_enum(&ctx.db).await?;
     find_linked_active_enum(&ctx.db).await?;
+
     ctx.delete().await;
 
     Ok(())
@@ -205,6 +210,72 @@ pub async fn insert_active_enum_child(db: &DatabaseConnection) -> Result<(), DbE
     Ok(())
 }
 
+pub async fn insert_active_enum_vec(db: &DatabaseConnection) -> Result<(), DbErr> {
+    use categories::*;
+
+    let model = Model {
+        id: 1,
+        categories: None,
+    };
+
+    assert_eq!(
+        model,
+        ActiveModel {
+            id: Set(1),
+            categories: Set(None),
+            ..Default::default()
+        }
+        .insert(db)
+        .await?
+    );
+    assert_eq!(model, Entity::find().one(db).await?.unwrap());
+    assert_eq!(
+        model,
+        Entity::find()
+            .filter(Column::Id.is_not_null())
+            .filter(Column::Categories.is_null())
+            .one(db)
+            .await?
+            .unwrap()
+    );
+
+    let _ = ActiveModel {
+        id: Set(1),
+        categories: Set(Some(vec![Category::Big, Category::Small])),
+        ..model.into_active_model()
+    }
+    .save(db)
+    .await?;
+
+    let model = Entity::find().one(db).await?.unwrap();
+    assert_eq!(
+        model,
+        Model {
+            id: 1,
+            categories: Some(vec![Category::Big, Category::Small]),
+        }
+    );
+    assert_eq!(
+        model,
+        Entity::find()
+            .filter(Column::Id.eq(1))
+            .filter(Expr::cust_with_values(
+                r#"$1 = ANY("categories")"#,
+                vec![Category::Big]
+            ))
+            .one(db)
+            .await?
+            .unwrap()
+    );
+
+    let res = model.delete(db).await?;
+
+    assert_eq!(res.rows_affected, 1);
+    assert_eq!(Entity::find().one(db).await?, None);
+
+    Ok(())
+}
+
 pub async fn find_related_active_enum(db: &DatabaseConnection) -> Result<(), DbErr> {
     assert_eq!(
         active_enum::Model {

tests/common/features/active_enum_vec.rs

@@ -0,0 +1,16 @@
+use super::sea_orm_active_enums::*;
+use sea_orm::entity::prelude::*;
+
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+#[cfg_attr(feature = "sqlx-postgres", sea_orm(schema_name = "public"))]
+#[sea_orm(table_name = "active_enum")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: i32,
+    pub categories: Option<Vec<Category>>,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}

tests/common/features/categories.rs

@@ -0,0 +1,15 @@
+use super::sea_orm_active_enums::*;
+use sea_orm::entity::prelude::*;
+
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+#[sea_orm(table_name = "categories")]
+pub struct Model {
+    #[sea_orm(primary_key, auto_increment = false)]
+    pub id: i32,
+    pub categories: Option<Vec<Category>>,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}