Casting columns as a different data type on select, insert and update

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

@@ -63,6 +63,8 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
     // generate Column enum and it's ColumnTrait impl
     let mut columns_enum: Punctuated<_, Comma> = Punctuated::new();
     let mut columns_trait: Punctuated<_, Comma> = Punctuated::new();
+    let mut cast_selects: Punctuated<_, Comma> = Punctuated::new();
+    let mut cast_values: Punctuated<_, Comma> = Punctuated::new();
     let mut primary_keys: Punctuated<_, Comma> = Punctuated::new();
     let mut primary_key_types: Punctuated<_, Comma> = Punctuated::new();
     let mut auto_increment = true;
@@ -90,6 +92,8 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                     let mut nullable = false;
                     let mut default_value = None;
                     let mut default_expr = None;
+                    let mut cast_select = None;
+                    let mut cast_value = None;
                     let mut indexed = false;
                     let mut ignore = false;
                     let mut unique = false;
@@ -169,6 +173,24 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                                                         format!("Invalid enum_name {:?}", nv.lit),
                                                     ));
                                                 }
+                                            } else if name == "cast_select" {
+                                                if let Lit::Str(litstr) = &nv.lit {
+                                                    cast_select = Some(litstr.value());
+                                                } else {
+                                                    return Err(Error::new(
+                                                        field.span(),
+                                                        format!("Invalid cast_select {:?}", nv.lit),
+                                                    ));
+                                                }
+                                            } else if name == "cast_value" {
+                                                if let Lit::Str(litstr) = &nv.lit {
+                                                    cast_value = Some(litstr.value());
+                                                } else {
+                                                    return Err(Error::new(
+                                                        field.span(),
+                                                        format!("Invalid cast_value {:?}", nv.lit),
+                                                    ));
+                                                }
                                             }
                                         }
                                     }
@@ -224,6 +246,23 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                         });
                     }
 
+                    if let Some(cast_select) = cast_select {
+                        cast_selects.push(quote! {
+                            Self::#field_name => sea_orm::sea_query::SimpleExpr::cast_as(
+                                Into::<sea_orm::sea_query::SimpleExpr>::into(expr),
+                                sea_orm::sea_query::Alias::new(&#cast_select),
+                            ),
+                        });
+                    }
+                    if let Some(cast_value) = cast_value {
+                        cast_values.push(quote! {
+                            Self::#field_name => sea_orm::sea_query::SimpleExpr::cast_as(
+                                Into::<sea_orm::sea_query::SimpleExpr>::into(val),
+                                sea_orm::sea_query::Alias::new(&#cast_value),
+                            ),
+                        });
+                    }
+
                     let field_type = &field.ty;
                     let field_type = quote! { #field_type }
                         .to_string() //E.g.: "Option < String >"
@@ -347,6 +386,20 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                     #columns_trait
                 }
             }
+
+            fn cast_select(&self, expr: sea_orm::sea_query::Expr) -> sea_orm::sea_query::SimpleExpr {
+                match self {
+                    #cast_selects
+                    _ => sea_orm::prelude::ColumnTrait::cast_select_enum(self, expr),
+                }
+            }
+
+            fn cast_value(&self, val: sea_orm::sea_query::Expr) -> sea_orm::sea_query::SimpleExpr {
+                match self {
+                    #cast_values
+                    _ => sea_orm::prelude::ColumnTrait::cast_value_enum(self, val),
+                }
+            }
         }
 
         #entity_def

src/entity/column.rs

@@ -1,5 +1,7 @@
-use crate::{cast_text_as_enum, EntityName, IdenStatic, IntoSimpleExpr, Iterable};
-use sea_query::{BinOper, DynIden, Expr, SeaRc, SelectStatement, SimpleExpr, Value};
+use crate::{EntityName, IdenStatic, IntoSimpleExpr, Iterable};
+use sea_query::{
+    Alias, BinOper, DynIden, Expr, Iden, IntoIden, SeaRc, SelectStatement, SimpleExpr, Value,
+};
 use std::str::FromStr;
 
 /// Defines a Column for an Entity
@@ -134,7 +136,7 @@ macro_rules! bind_oper_with_enum_casting {
         where
             V: Into<Value>,
         {
-            let expr = cast_text_as_enum(Expr::val(v), self);
+            let expr = self.cast_value(Expr::val(v));
             Expr::tbl(self.entity_name(), *self).binary(BinOper::$bin_op, expr)
         }
     };
@@ -338,6 +340,42 @@ pub trait ColumnTrait: IdenStatic + Iterable + FromStr {
     fn into_expr(self) -> Expr {
         Expr::expr(self.into_simple_expr())
     }
+
+    /// Cast column expression used in select statement.
+    /// By default it only cast database enum as text.
+    fn cast_select(&self, expr: Expr) -> SimpleExpr {
+        self.cast_select_enum(expr)
+    }
+
+    /// Cast enum column as text
+    fn cast_select_enum(&self, expr: Expr) -> SimpleExpr {
+        cast_enum_text_inner(expr, self, |col, _, col_type| {
+            let type_name = match col_type {
+                ColumnType::Array(_) => TextArray.into_iden(),
+                _ => Text.into_iden(),
+            };
+            col.as_enum(type_name)
+        })
+    }
+
+    /// Cast value of a column into the correct type for database storage.
+    /// By default it only cast text as enum type if it's an enum column.
+    fn cast_value(&self, val: Expr) -> SimpleExpr {
+        self.cast_value_enum(val)
+    }
+
+    /// Cast value of a enum column as enum type
+    fn cast_value_enum(&self, val: Expr) -> SimpleExpr {
+        cast_enum_text_inner(val, self, |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)
+        })
+    }
 }
 
 impl ColumnType {
@@ -512,6 +550,26 @@ impl From<sea_query::ColumnType> for ColumnType {
     }
 }
 
+#[derive(Iden)]
+struct Text;
+
+#[derive(Iden)]
+#[iden = "text[]"]
+struct TextArray;
+
+fn cast_enum_text_inner<C, F>(expr: Expr, col: &C, f: F) -> SimpleExpr
+where
+    C: ColumnTrait,
+    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), col_type),
+        None => expr.into(),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{

src/executor/insert.rs

@@ -1,6 +1,6 @@
 use crate::{
-    cast_enum_as_text, error::*, ActiveModelTrait, ConnectionTrait, EntityTrait, Insert,
-    IntoActiveModel, Iterable, PrimaryKeyTrait, SelectModel, SelectorRaw, Statement, TryFromU64,
+    error::*, ActiveModelTrait, ColumnTrait, ConnectionTrait, EntityTrait, Insert, IntoActiveModel,
+    Iterable, PrimaryKeyTrait, SelectModel, SelectorRaw, Statement, TryFromU64,
 };
 use sea_query::{Expr, FromValueTuple, Iden, InsertStatement, IntoColumnRef, Query, ValueTuple};
 use std::{future::Future, marker::PhantomData};
@@ -186,8 +186,7 @@ where
     let found = match db.support_returning() {
         true => {
             let returning = Query::returning().exprs(
-                <A::Entity as EntityTrait>::Column::iter()
-                    .map(|c| cast_enum_as_text(Expr::col(c), &c)),
+                <A::Entity as EntityTrait>::Column::iter().map(|c| c.cast_select(Expr::col(c))),
             );
             insert_statement.returning(returning);
             SelectorRaw::<SelectModel<<A::Entity as EntityTrait>::Model>>::from_statement(

src/executor/update.rs

@@ -1,5 +1,5 @@
 use crate::{
-    cast_enum_as_text, error::*, ActiveModelTrait, ConnectionTrait, EntityTrait, IntoActiveModel,
+    error::*, ActiveModelTrait, ColumnTrait, ConnectionTrait, EntityTrait, IntoActiveModel,
     Iterable, SelectModel, SelectorRaw, Statement, UpdateMany, UpdateOne,
 };
 use sea_query::{Expr, FromValueTuple, Query, UpdateStatement};
@@ -92,8 +92,7 @@ where
     match db.support_returning() {
         true => {
             let returning = Query::returning().exprs(
-                <A::Entity as EntityTrait>::Column::iter()
-                    .map(|c| cast_enum_as_text(Expr::col(c), &c)),
+                <A::Entity as EntityTrait>::Column::iter().map(|c| c.cast_select(Expr::col(c))),
             );
             query.returning(returning);
             let db_backend = db.get_database_backend();

src/query/combine.rs

@@ -1,6 +1,5 @@
 use crate::{
-    cast_enum_as_text, ColumnTrait, EntityTrait, IdenStatic, Iterable, QueryTrait, Select,
-    SelectTwo, SelectTwoMany,
+    ColumnTrait, EntityTrait, IdenStatic, Iterable, QueryTrait, Select, SelectTwo, SelectTwoMany,
 };
 use core::marker::PhantomData;
 pub use sea_query::JoinType;
@@ -149,7 +148,7 @@ where
     for col in <F::Column as Iterable>::iter() {
         let alias = format!("{}{}", SelectB.as_str(), col.as_str());
         selector.query().expr(SelectExpr {
-            expr: cast_enum_as_text(col.into_expr(), &col),
+            expr: col.cast_select(col.into_expr()),
             alias: Some(SeaRc::new(Alias::new(&alias))),
             window: None,
         });

src/query/helper.rs

@@ -1,6 +1,6 @@
 use crate::{
-    ColumnTrait, ColumnType, EntityTrait, Identity, IntoIdentity, IntoSimpleExpr, Iterable,
-    ModelTrait, PrimaryKeyToColumn, RelationDef,
+    ColumnTrait, EntityTrait, Identity, IntoIdentity, IntoSimpleExpr, Iterable, ModelTrait,
+    PrimaryKeyToColumn, RelationDef,
 };
 use sea_query::{
     Alias, Expr, Iden, IntoCondition, IntoIden, LockType, SeaRc, SelectExpr, SelectStatement,
@@ -67,7 +67,7 @@ pub trait QuerySelect: Sized {
     where
         C: ColumnTrait,
     {
-        self.query().expr(cast_enum_as_text(col.into_expr(), &col));
+        self.query().expr(col.cast_select(col.into_expr()));
         self
     }
 
@@ -678,49 +678,3 @@ pub(crate) fn unpack_table_alias(table_ref: &TableRef) -> Option<DynIden> {
         | TableRef::DatabaseSchemaTableAlias(_, _, _, alias) => Some(SeaRc::clone(alias)),
     }
 }
-
-#[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_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_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, &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), col_type),
-        None => expr.into(),
-    }
-}

src/query/insert.rs

@@ -1,5 +1,5 @@
 use crate::{
-    cast_text_as_enum, ActiveModelTrait, EntityName, EntityTrait, IntoActiveModel, Iterable,
+    ActiveModelTrait, ColumnTrait, EntityName, EntityTrait, IntoActiveModel, Iterable,
     PrimaryKeyTrait, QueryTrait,
 };
 use core::marker::PhantomData;
@@ -134,7 +134,7 @@ where
             }
             if av_has_val {
                 columns.push(col);
-                values.push(cast_text_as_enum(Expr::val(av.into_value().unwrap()), &col));
+                values.push(col.cast_value(Expr::val(av.into_value().unwrap())));
             }
         }
         self.query.columns(columns);

src/query/join.rs

@@ -1,5 +1,5 @@
 use crate::{
-    cast_enum_as_text, join_tbl_on_condition, unpack_table_ref, EntityTrait, IdenStatic, Iterable,
+    join_tbl_on_condition, unpack_table_ref, ColumnTrait, EntityTrait, IdenStatic, Iterable,
     Linked, QuerySelect, Related, Select, SelectA, SelectB, SelectTwo, SelectTwoMany,
 };
 pub use sea_query::JoinType;
@@ -100,7 +100,7 @@ where
                 col.into_iden(),
             );
             select_two.query().expr(SelectExpr {
-                expr: cast_enum_as_text(expr, &col),
+                expr: col.cast_select(expr),
                 alias: Some(SeaRc::new(Alias::new(&alias))),
                 window: None,
             });

src/query/select.rs

@@ -1,7 +1,4 @@
-use crate::{
-    cast_enum_as_text, ColumnTrait, EntityTrait, Iterable, QueryFilter, QueryOrder, QuerySelect,
-    QueryTrait,
-};
+use crate::{ColumnTrait, EntityTrait, Iterable, QueryFilter, QueryOrder, QuerySelect, QueryTrait};
 use core::fmt::Debug;
 use core::marker::PhantomData;
 pub use sea_query::JoinType;
@@ -123,7 +120,7 @@ where
 
     fn column_list(&self) -> Vec<SimpleExpr> {
         E::Column::iter()
-            .map(|col| cast_enum_as_text(col.into_expr(), &col))
+            .map(|col| col.cast_select(col.into_expr()))
             .collect()
     }
 

src/query/update.rs

@@ -1,6 +1,6 @@
 use crate::{
-    cast_text_as_enum, ActiveModelTrait, ColumnTrait, EntityTrait, Iterable, PrimaryKeyToColumn,
-    QueryFilter, QueryTrait,
+    ActiveModelTrait, ColumnTrait, EntityTrait, Iterable, PrimaryKeyToColumn, QueryFilter,
+    QueryTrait,
 };
 use core::marker::PhantomData;
 use sea_query::{Expr, IntoIden, SimpleExpr, UpdateStatement};
@@ -109,7 +109,7 @@ where
             }
             let av = self.model.get(col);
             if av.is_set() {
-                let expr = cast_text_as_enum(Expr::val(av.into_value().unwrap()), &col);
+                let expr = col.cast_value(Expr::val(av.into_value().unwrap()));
                 self.query.value(col, expr);
             }
         }