Codegen ActiveEnum & Create Enum From ActiveEnum

Cargo.toml

@@ -30,7 +30,7 @@ futures-util = { version = "^0.3" }
 log = { version = "^0.4", optional = true }
 rust_decimal = { version = "^1", optional = true }
 sea-orm-macros = { version = "^0.4.2", path = "sea-orm-macros", optional = true }
-sea-query = { version = "^0.19.4", features = ["thread-safe"] }
+sea-query = { version = "^0.20.0", features = ["thread-safe"] }
 sea-strum = { version = "^0.21", features = ["derive", "sea-orm"] }
 serde = { version = "^1.0", features = ["derive"] }
 serde_json = { version = "^1", optional = true }

sea-orm-cli/Cargo.toml

@@ -22,7 +22,7 @@ clap = { version = "^2.33.3" }
 dotenv = { version = "^0.15" }
 async-std = { version = "^1.9", features = [ "attributes" ] }
 sea-orm-codegen = { version = "^0.4.2", path = "../sea-orm-codegen" }
-sea-schema = { version = "^0.2.9", default-features = false, features = [
+sea-schema = { version = "0.3.0", default-features = false, features = [
     "debug-print",
     "sqlx-mysql",
     "sqlx-postgres",

sea-orm-codegen/Cargo.toml

@@ -15,7 +15,7 @@ name = "sea_orm_codegen"
 path = "src/lib.rs"
 
 [dependencies]
-sea-query = { version = "^0.16.4" }
+sea-query = { version = "0.20.0" }
 syn = { version = "^1", default-features = false, features = [
     "derive",
     "parsing",

sea-orm-codegen/src/entity/active_enum.rs

@@ -0,0 +1,31 @@
+use heck::CamelCase;
+use proc_macro2::TokenStream;
+use quote::{format_ident, quote};
+
+#[derive(Clone, Debug)]
+pub struct ActiveEnum {
+    pub(crate) enum_name: String,
+    pub(crate) values: Vec<String>,
+}
+
+impl ActiveEnum {
+    pub fn impl_active_enum(&self) -> TokenStream {
+        let enum_name = &self.enum_name;
+        let enum_iden = format_ident!("{}", enum_name.to_camel_case());
+        let values = &self.values;
+        let variants = self
+            .values
+            .iter()
+            .map(|v| format_ident!("{}", v.to_camel_case()));
+        quote! {
+            #[derive(Debug, Clone, PartialEq, EnumIter, DeriveActiveEnum)]
+            #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = #enum_name)]
+            pub enum #enum_iden {
+                #(
+                    #[sea_orm(string_value = #values)]
+                    #variants,
+                )*
+            }
+        }
+    }
+}

sea-orm-codegen/src/entity/column.rs

@@ -24,26 +24,27 @@ impl Column {
 
     pub fn get_rs_type(&self) -> TokenStream {
         #[allow(unreachable_patterns)]
-        let ident: TokenStream = match self.col_type {
+        let ident: TokenStream = match &self.col_type {
             ColumnType::Char(_)
             | ColumnType::String(_)
             | ColumnType::Text
-            | ColumnType::Custom(_) => "String",
-            ColumnType::TinyInteger(_) => "i8",
-            ColumnType::SmallInteger(_) => "i16",
-            ColumnType::Integer(_) => "i32",
-            ColumnType::BigInteger(_) => "i64",
-            ColumnType::Float(_) => "f32",
-            ColumnType::Double(_) => "f64",
-            ColumnType::Json | ColumnType::JsonBinary => "Json",
-            ColumnType::Date => "Date",
-            ColumnType::Time(_) => "Time",
-            ColumnType::DateTime(_) | ColumnType::Timestamp(_) => "DateTime",
-            ColumnType::TimestampWithTimeZone(_) => "DateTimeWithTimeZone",
-            ColumnType::Decimal(_) | ColumnType::Money(_) => "Decimal",
-            ColumnType::Uuid => "Uuid",
-            ColumnType::Binary(_) => "Vec<u8>",
-            ColumnType::Boolean => "bool",
+            | ColumnType::Custom(_) => "String".to_owned(),
+            ColumnType::TinyInteger(_) => "i8".to_owned(),
+            ColumnType::SmallInteger(_) => "i16".to_owned(),
+            ColumnType::Integer(_) => "i32".to_owned(),
+            ColumnType::BigInteger(_) => "i64".to_owned(),
+            ColumnType::Float(_) => "f32".to_owned(),
+            ColumnType::Double(_) => "f64".to_owned(),
+            ColumnType::Json | ColumnType::JsonBinary => "Json".to_owned(),
+            ColumnType::Date => "Date".to_owned(),
+            ColumnType::Time(_) => "Time".to_owned(),
+            ColumnType::DateTime(_) | ColumnType::Timestamp(_) => "DateTime".to_owned(),
+            ColumnType::TimestampWithTimeZone(_) => "DateTimeWithTimeZone".to_owned(),
+            ColumnType::Decimal(_) | ColumnType::Money(_) => "Decimal".to_owned(),
+            ColumnType::Uuid => "Uuid".to_owned(),
+            ColumnType::Binary(_) => "Vec<u8>".to_owned(),
+            ColumnType::Boolean => "bool".to_owned(),
+            ColumnType::Enum(name, _) => name.to_camel_case(),
             _ => unimplemented!(),
         }
         .parse()

sea-orm-codegen/src/entity/mod.rs

@@ -1,3 +1,4 @@
+mod active_enum;
 mod base_entity;
 mod column;
 mod conjunct_relation;
@@ -6,6 +7,7 @@ mod relation;
 mod transformer;
 mod writer;
 
+pub use active_enum::*;
 pub use base_entity::*;
 pub use column::*;
 pub use conjunct_relation::*;

sea-orm-codegen/src/entity/transformer.rs

@@ -1,5 +1,6 @@
 use crate::{
-    Column, ConjunctRelation, Entity, EntityWriter, Error, PrimaryKey, Relation, RelationType,
+    ActiveEnum, Column, ConjunctRelation, Entity, EntityWriter, Error, PrimaryKey, Relation,
+    RelationType,
 };
 use sea_query::TableStatement;
 use std::collections::HashMap;
@@ -9,6 +10,7 @@ pub struct EntityTransformer;
 
 impl EntityTransformer {
     pub fn transform(table_stmts: Vec<TableStatement>) -> Result<EntityWriter, Error> {
+        let mut enums: HashMap<String, ActiveEnum> = HashMap::new();
         let mut inverse_relations: HashMap<String, Vec<Relation>> = HashMap::new();
         let mut conjunct_relations: HashMap<String, Vec<ConjunctRelation>> = HashMap::new();
         let mut entities = HashMap::new();
@@ -22,7 +24,15 @@ impl EntityTransformer {
                 }
             };
             let table_name = match table_create.get_table_name() {
-                Some(s) => s,
+                Some(table_ref) => match table_ref {
+                    sea_query::TableRef::Table(t)
+                    | sea_query::TableRef::SchemaTable(_, t)
+                    | sea_query::TableRef::DatabaseSchemaTable(_, _, t)
+                    | sea_query::TableRef::TableAlias(t, _)
+                    | sea_query::TableRef::SchemaTableAlias(_, t, _)
+                    | sea_query::TableRef::DatabaseSchemaTableAlias(_, _, t, _) => t.to_string(),
+                    _ => unimplemented!(),
+                },
                 None => {
                     return Err(Error::TransformError(
                         "Table name should not be empty".into(),
@@ -44,6 +54,18 @@ impl EntityTransformer {
                         > 0;
                     col
                 })
+                .map(|col| {
+                    if let sea_query::ColumnType::Enum(enum_name, values) = &col.col_type {
+                        enums.insert(
+                            enum_name.clone(),
+                            ActiveEnum {
+                                enum_name: enum_name.clone(),
+                                values: values.clone(),
+                            },
+                        );
+                    }
+                    col
+                })
                 .collect();
             let mut ref_table_counts: HashMap<String, usize> = HashMap::new();
             let relations: Vec<Relation> = table_create
@@ -170,6 +192,7 @@ impl EntityTransformer {
         }
         Ok(EntityWriter {
             entities: entities.into_iter().map(|(_, v)| v).collect(),
+            enums,
         })
     }
 }

sea-orm-codegen/src/entity/writer.rs

@@ -1,13 +1,14 @@
-use std::str::FromStr;
-
-use crate::Entity;
+use crate::{ActiveEnum, Entity};
+use heck::CamelCase;
 use proc_macro2::TokenStream;
-use quote::quote;
+use quote::{format_ident, quote};
+use std::{collections::HashMap, str::FromStr};
 use syn::{punctuated::Punctuated, token::Comma};
 
 #[derive(Clone, Debug)]
 pub struct EntityWriter {
     pub(crate) entities: Vec<Entity>,
+    pub(crate) enums: HashMap<String, ActiveEnum>,
 }
 
 pub struct WriterOutput {
@@ -83,6 +84,9 @@ impl EntityWriter {
         files.extend(self.write_entities(expanded_format, with_serde));
         files.push(self.write_mod());
         files.push(self.write_prelude());
+        if !self.enums.is_empty() {
+            files.push(self.write_sea_orm_active_enums());
+        }
         WriterOutput { files }
     }
 
@@ -112,7 +116,7 @@ impl EntityWriter {
         let code_blocks: Vec<TokenStream> = self
             .entities
             .iter()
-            .map(|entity| Self::gen_mod(entity))
+            .map(Self::gen_mod)
             .collect();
         Self::write(
             &mut lines,
@@ -122,6 +126,14 @@ impl EntityWriter {
         );
         lines.push("".to_owned());
         Self::write(&mut lines, code_blocks);
+        if !self.enums.is_empty() {
+            Self::write(
+                &mut lines,
+                vec![quote! {
+                    pub mod sea_orm_active_enums;
+                }],
+            );
+        }
         OutputFile {
             name: "mod.rs".to_owned(),
             content: lines.join("\n"),
@@ -134,7 +146,7 @@ impl EntityWriter {
         let code_blocks = self
             .entities
             .iter()
-            .map(|entity| Self::gen_prelude_use(entity))
+            .map(Self::gen_prelude_use)
             .collect();
         Self::write(&mut lines, code_blocks);
         OutputFile {
@@ -143,6 +155,28 @@ impl EntityWriter {
         }
     }
 
+    pub fn write_sea_orm_active_enums(&self) -> OutputFile {
+        let mut lines = Vec::new();
+        Self::write_doc_comment(&mut lines);
+        Self::write(
+            &mut lines,
+            vec![quote! {
+                use sea_orm::entity::prelude::*;
+            }],
+        );
+        lines.push("".to_owned());
+        let code_blocks = self
+            .enums
+            .iter()
+            .map(|(_, active_enum)| active_enum.impl_active_enum())
+            .collect();
+        Self::write(&mut lines, code_blocks);
+        OutputFile {
+            name: "sea_orm_active_enums.rs".to_owned(),
+            content: lines.join("\n"),
+        }
+    }
+
     pub fn write(lines: &mut Vec<String>, code_blocks: Vec<TokenStream>) {
         lines.extend(
             code_blocks
@@ -163,8 +197,10 @@ impl EntityWriter {
     }
 
     pub fn gen_expanded_code_blocks(entity: &Entity, with_serde: &WithSerde) -> Vec<TokenStream> {
+        let mut imports = Self::gen_import(with_serde);
+        imports.extend(Self::gen_import_active_enum(entity));
         let mut code_blocks = vec![
-            Self::gen_import(with_serde),
+            imports,
             Self::gen_entity_struct(),
             Self::gen_impl_entity_name(entity),
             Self::gen_model_struct(entity, with_serde),
@@ -182,10 +218,9 @@ impl EntityWriter {
     }
 
     pub fn gen_compact_code_blocks(entity: &Entity, with_serde: &WithSerde) -> Vec<TokenStream> {
-        let mut code_blocks = vec![
-            Self::gen_import(with_serde),
-            Self::gen_compact_model_struct(entity, with_serde),
-        ];
+        let mut imports = Self::gen_import(with_serde);
+        imports.extend(Self::gen_import_active_enum(entity));
+        let mut code_blocks = vec![imports, Self::gen_compact_model_struct(entity, with_serde)];
         let relation_defs = if entity.get_relation_enum_name().is_empty() {
             vec![
                 Self::gen_relation_enum(entity),
@@ -249,6 +284,21 @@ impl EntityWriter {
         }
     }
 
+    pub fn gen_import_active_enum(entity: &Entity) -> TokenStream {
+        entity
+            .columns
+            .iter()
+            .fold(TokenStream::new(), |mut ts, col| {
+                if let sea_query::ColumnType::Enum(enum_name, _) = &col.col_type {
+                    let enum_name = format_ident!("{}", enum_name.to_camel_case());
+                    ts.extend(vec![quote! {
+                        use super::sea_orm_active_enums::#enum_name;
+                    }]);
+                }
+                ts
+            })
+    }
+
     pub fn gen_model_struct(entity: &Entity, with_serde: &WithSerde) -> TokenStream {
         let column_names_snake_case = entity.get_column_names_snake_case();
         let column_rs_types = entity.get_column_rs_types();

src/query/helper.rs

@@ -465,10 +465,12 @@ pub(crate) fn join_tbl_on_condition(
 
 pub(crate) fn unpack_table_ref(table_ref: &TableRef) -> DynIden {
     match table_ref {
-        TableRef::Table(tbl) => SeaRc::clone(tbl),
-        TableRef::SchemaTable(_, tbl) => SeaRc::clone(tbl),
-        TableRef::TableAlias(tbl, _) => SeaRc::clone(tbl),
-        TableRef::SchemaTableAlias(_, tbl, _) => SeaRc::clone(tbl),
-        TableRef::SubQuery(_, tbl) => SeaRc::clone(tbl),
+        TableRef::Table(tbl)
+        | TableRef::SchemaTable(_, tbl)
+        | TableRef::DatabaseSchemaTable(_, _, tbl)
+        | TableRef::TableAlias(tbl, _)
+        | TableRef::SchemaTableAlias(_, tbl, _)
+        | TableRef::DatabaseSchemaTableAlias(_, _, tbl, _)
+        | TableRef::SubQuery(_, tbl) => SeaRc::clone(tbl),
     }
 }