Support Vector of enum for Postgres (#1210)

* Support Vector of enum for Postgres

* Fix clippy

* Add tests for integer array

* Fix dependency

* Bump sea-query to 0.27.2

Co-authored-by: Chris Tsang <chris.2y3@outlook.com>

Cargo.toml

@@ -34,8 +34,8 @@ log = { version = "^0.4" }
 tracing = { version = "^0.1", features = ["log"] }
 rust_decimal = { version = "^1", optional = true }
 sea-orm-macros = { version = "^0.10.3", path = "sea-orm-macros", optional = true }
-sea-query = { version = "^0.27.1", features = ["thread-safe"] }
-sea-query-binder = { version = "^0.2.1", optional = true }
+sea-query = { version = "^0.27.2", features = ["thread-safe"] }
+sea-query-binder = { version = "^0.2.2", optional = true }
 sea-strum = { version = "^0.23", features = ["derive", "sea-orm"] }
 serde = { version = "^1.0", features = ["derive"] }
 serde_json = { version = "^1.0", optional = true }
@@ -76,7 +76,7 @@ with-chrono = ["chrono", "sea-query/with-chrono", "sea-query-binder?/with-chrono
 with-rust_decimal = ["rust_decimal", "sea-query/with-rust_decimal", "sea-query-binder?/with-rust_decimal", "sqlx?/decimal"]
 with-uuid = ["uuid", "sea-query/with-uuid", "sea-query-binder?/with-uuid", "sqlx?/uuid"]
 with-time = ["time", "sea-query/with-time", "sea-query-binder?/with-time", "sqlx?/time"]
-postgres-array = ["sea-query/postgres-array", "sea-query-binder?/postgres-array"]
+postgres-array = ["sea-query/postgres-array", "sea-query-binder?/postgres-array", "sea-orm-macros?/postgres-array"]
 sqlx-dep = []
 sqlx-all = ["sqlx-mysql", "sqlx-postgres", "sqlx-sqlite"]
 sqlx-mysql = ["sqlx-dep", "sea-query-binder/sqlx-mysql", "sqlx/mysql"]

sea-orm-macros/Cargo.toml

@@ -27,3 +27,6 @@ proc-macro2 = "^1"
 [dev-dependencies]
 sea-orm = { path = "../", features = ["macros"] }
 serde = { version = "^1.0", features = ["derive"] }
+
+[features]
+postgres-array = []

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

@@ -287,6 +287,15 @@ impl ActiveEnum {
             quote!()
         };
 
+        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!()
+        };
+
         quote!(
             #[derive(Debug, Clone, PartialEq, Eq)]
             pub struct #enum_name_iden;
@@ -304,6 +313,8 @@ impl ActiveEnum {
             impl sea_orm::ActiveEnum for #ident {
                 type Value = #rs_type;
 
+                type ValueVec = Vec<#rs_type>;
+
                 fn name() -> sea_orm::sea_query::DynIden {
                     sea_orm::sea_query::SeaRc::new(#enum_name_iden) as sea_orm::sea_query::DynIden
                 }
@@ -359,7 +370,7 @@ impl ActiveEnum {
                 }
 
                 fn array_type() -> sea_orm::sea_query::ArrayType {
-                    unimplemented!("Array of Enum is not supported")
+                    <<Self as sea_orm::ActiveEnum>::Value as sea_orm::sea_query::ValueType>::array_type()
                 }
 
                 fn column_type() -> sea_orm::sea_query::ColumnType {
@@ -384,6 +395,8 @@ impl ActiveEnum {
                     write!(f, "{}", v)
                 }
             }
+
+            #impl_not_u8
         )
     }
 }

src/entity/active_enum.rs

@@ -1,4 +1,4 @@
-use crate::{ColumnDef, DbErr, Iterable, TryGetable};
+use crate::{ColumnDef, DbErr, Iterable, QueryResult, TryGetError, TryGetable};
 use sea_query::{DynIden, Expr, Nullable, SimpleExpr, Value, ValueType};
 
 /// A Rust representation of enum defined in database.
@@ -47,6 +47,8 @@ use sea_query::{DynIden, Expr, Nullable, SimpleExpr, Value, ValueType};
 ///     // The macro attribute `rs_type` is being pasted here
 ///     type Value = String;
 ///
+///     type ValueVec = Vec<String>;
+///
 ///     // Will be atomically generated by `DeriveActiveEnum`
 ///     fn name() -> DynIden {
 ///         SeaRc::new(CategoryEnum)
@@ -114,6 +116,9 @@ pub trait ActiveEnum: Sized + Iterable {
     /// Define the Rust type that each enum variant represents.
     type Value: Into<Value> + ValueType + Nullable + TryGetable;
 
+    /// Define the enum value in Vector type.
+    type ValueVec: IntoIterator<Item = Self::Value>;
+
     /// Get the name of enum
     fn name() -> DynIden;
 
@@ -142,6 +147,19 @@ pub trait ActiveEnum: Sized + Iterable {
     }
 }
 
+impl<T> TryGetable for Vec<T>
+where
+    T: ActiveEnum,
+    T::ValueVec: TryGetable,
+{
+    fn try_get(res: &QueryResult, pre: &str, col: &str) -> Result<Self, TryGetError> {
+        <T::ValueVec as TryGetable>::try_get(res, pre, col)?
+            .into_iter()
+            .map(|value| T::try_from_value(&value).map_err(TryGetError::DbErr))
+            .collect()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate as sea_orm;
@@ -163,6 +181,8 @@ mod tests {
         impl ActiveEnum for Category {
             type Value = String;
 
+            type ValueVec = Vec<String>;
+
             fn name() -> DynIden {
                 SeaRc::new(CategoryEnum)
             }

src/entity/column.rs

@@ -353,10 +353,14 @@ impl ColumnType {
     }
 
     pub(crate) fn get_enum_name(&self) -> Option<&DynIden> {
-        match self {
-            ColumnType::Enum { name, .. } => Some(name),
-            _ => None,
+        fn enum_name(col_type: &ColumnType) -> Option<&DynIden> {
+            match col_type {
+                ColumnType::Enum { name, .. } => Some(name),
+                ColumnType::Array(col_type) => enum_name(col_type),
+                _ => None,
+            }
         }
+        enum_name(self)
     }
 }
 

src/query/helper.rs

@@ -1,6 +1,6 @@
 use crate::{
-    ColumnTrait, EntityTrait, Identity, IntoIdentity, IntoSimpleExpr, Iterable, ModelTrait,
-    PrimaryKeyToColumn, RelationDef,
+    ColumnTrait, ColumnType, EntityTrait, Identity, IntoIdentity, IntoSimpleExpr, Iterable,
+    ModelTrait, PrimaryKeyToColumn, RelationDef,
 };
 use sea_query::{
     Alias, Expr, Iden, IntoCondition, IntoIden, LockType, SeaRc, SelectExpr, SelectStatement,
@@ -616,29 +616,45 @@ pub(crate) fn unpack_table_alias(table_ref: &TableRef) -> Option<DynIden> {
 #[derive(Iden)]
 struct Text;
 
+#[derive(Iden)]
+#[iden = "text[]"]
+struct TextArray;
+
 pub(crate) fn cast_enum_as_text<C>(expr: Expr, col: &C) -> SimpleExpr
 where
     C: ColumnTrait,
 {
-    cast_enum_text_inner(expr, col, |col, _| col.as_enum(Text))
+    cast_enum_text_inner(expr, col, |col, _, col_type| {
+        let type_name = match col_type {
+            ColumnType::Array(_) => TextArray.into_iden(),
+            _ => Text.into_iden(),
+        };
+        col.as_enum(type_name)
+    })
 }
 
 pub(crate) fn cast_text_as_enum<C>(expr: Expr, col: &C) -> SimpleExpr
 where
     C: ColumnTrait,
 {
-    cast_enum_text_inner(expr, col, |col, enum_name| col.as_enum(enum_name))
+    cast_enum_text_inner(expr, col, |col, enum_name, col_type| {
+        let type_name = match col_type {
+            ColumnType::Array(_) => Alias::new(&format!("{}[]", enum_name.to_string())).into_iden(),
+            _ => enum_name,
+        };
+        col.as_enum(type_name)
+    })
 }
 
 fn cast_enum_text_inner<C, F>(expr: Expr, col: &C, f: F) -> SimpleExpr
 where
     C: ColumnTrait,
-    F: Fn(Expr, DynIden) -> SimpleExpr,
+    F: Fn(Expr, DynIden, &ColumnType) -> SimpleExpr,
 {
     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)),
+        Some(enum_name) => f(expr, SeaRc::clone(enum_name), col_type),
         None => expr.into(),
     }
 }

tests/collection_tests.rs

@@ -24,6 +24,10 @@ pub async fn insert_collection(db: &DatabaseConnection) -> Result<(), DbErr> {
             id: 1,
             integers: vec![1, 2, 3],
             integers_opt: Some(vec![1, 2, 3]),
+            teas: vec![Tea::BreakfastTea],
+            teas_opt: Some(vec![Tea::BreakfastTea]),
+            colors: vec![Color::Black],
+            colors_opt: Some(vec![Color::Black]),
         }
         .into_active_model()
         .insert(db)
@@ -32,6 +36,10 @@ pub async fn insert_collection(db: &DatabaseConnection) -> Result<(), DbErr> {
             id: 1,
             integers: vec![1, 2, 3],
             integers_opt: Some(vec![1, 2, 3]),
+            teas: vec![Tea::BreakfastTea],
+            teas_opt: Some(vec![Tea::BreakfastTea]),
+            colors: vec![Color::Black],
+            colors_opt: Some(vec![Color::Black]),
         }
     );
 
@@ -40,6 +48,10 @@ pub async fn insert_collection(db: &DatabaseConnection) -> Result<(), DbErr> {
             id: 2,
             integers: vec![10, 9],
             integers_opt: None,
+            teas: vec![Tea::BreakfastTea],
+            teas_opt: None,
+            colors: vec![Color::Black],
+            colors_opt: None,
         }
         .into_active_model()
         .insert(db)
@@ -48,6 +60,10 @@ pub async fn insert_collection(db: &DatabaseConnection) -> Result<(), DbErr> {
             id: 2,
             integers: vec![10, 9],
             integers_opt: None,
+            teas: vec![Tea::BreakfastTea],
+            teas_opt: None,
+            colors: vec![Color::Black],
+            colors_opt: None,
         }
     );
 
@@ -56,6 +72,10 @@ pub async fn insert_collection(db: &DatabaseConnection) -> Result<(), DbErr> {
             id: 3,
             integers: vec![],
             integers_opt: Some(vec![]),
+            teas: vec![],
+            teas_opt: Some(vec![]),
+            colors: vec![],
+            colors_opt: Some(vec![]),
         }
         .into_active_model()
         .insert(db)
@@ -64,6 +84,10 @@ pub async fn insert_collection(db: &DatabaseConnection) -> Result<(), DbErr> {
             id: 3,
             integers: vec![],
             integers_opt: Some(vec![]),
+            teas: vec![],
+            teas_opt: Some(vec![]),
+            colors: vec![],
+            colors_opt: Some(vec![]),
         }
     );
 
@@ -78,6 +102,10 @@ pub async fn update_collection(db: &DatabaseConnection) -> Result<(), DbErr> {
     ActiveModel {
         integers: Set(vec![4, 5, 6]),
         integers_opt: Set(Some(vec![4, 5, 6])),
+        teas: Set(vec![Tea::EverydayTea]),
+        teas_opt: Set(Some(vec![Tea::EverydayTea])),
+        colors: Set(vec![Color::White]),
+        colors_opt: Set(Some(vec![Color::White])),
         ..model.into_active_model()
     }
     .update(db)
@@ -87,6 +115,10 @@ pub async fn update_collection(db: &DatabaseConnection) -> Result<(), DbErr> {
         id: Unchanged(3),
         integers: Set(vec![3, 1, 4]),
         integers_opt: Set(None),
+        teas: Set(vec![Tea::EverydayTea]),
+        teas_opt: Set(None),
+        colors: Set(vec![Color::White]),
+        colors_opt: Set(None),
     }
     .update(db)
     .await?;

tests/common/features/collection.rs

@@ -1,3 +1,4 @@
+use super::sea_orm_active_enums::*;
 use sea_orm::entity::prelude::*;
 
 #[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
@@ -7,6 +8,10 @@ pub struct Model {
     pub id: i32,
     pub integers: Vec<i32>,
     pub integers_opt: Option<Vec<i32>>,
+    pub teas: Vec<Tea>,
+    pub teas_opt: Option<Vec<Tea>>,
+    pub colors: Vec<Color>,
+    pub colors_opt: Option<Vec<Color>>,
 }
 
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]

tests/common/features/schema.rs

@@ -6,7 +6,7 @@ use sea_orm::{
     error::*, sea_query, ConnectionTrait, DatabaseConnection, DbBackend, DbConn, EntityName,
     ExecResult, Schema,
 };
-use sea_query::{extension::postgres::Type, Alias, ColumnDef, ForeignKeyCreateStatement};
+use sea_query::{extension::postgres::Type, Alias, ColumnDef, ForeignKeyCreateStatement, IntoIden};
 
 pub async fn create_tables(db: &DatabaseConnection) -> Result<(), DbErr> {
     let db_backend = db.get_database_backend();
@@ -350,6 +350,35 @@ pub async fn create_collection_table(db: &DbConn) -> Result<ExecResult, DbErr> {
             ColumnDef::new(collection::Column::IntegersOpt)
                 .array(sea_query::ColumnType::Integer(None)),
         )
+        .col(
+            ColumnDef::new(collection::Column::Teas)
+                .array(sea_query::ColumnType::Enum {
+                    name: TeaEnum.into_iden(),
+                    variants: vec![
+                        TeaVariant::EverydayTea.into_iden(),
+                        TeaVariant::BreakfastTea.into_iden(),
+                    ],
+                })
+                .not_null(),
+        )
+        .col(
+            ColumnDef::new(collection::Column::TeasOpt).array(sea_query::ColumnType::Enum {
+                name: TeaEnum.into_iden(),
+                variants: vec![
+                    TeaVariant::EverydayTea.into_iden(),
+                    TeaVariant::BreakfastTea.into_iden(),
+                ],
+            }),
+        )
+        .col(
+            ColumnDef::new(collection::Column::Colors)
+                .array(sea_query::ColumnType::Integer(None))
+                .not_null(),
+        )
+        .col(
+            ColumnDef::new(collection::Column::ColorsOpt)
+                .array(sea_query::ColumnType::Integer(None)),
+        )
         .to_owned();
 
     create_table(db, &stmt, Collection).await