From 21216f3a97b4ccf24bc3ff41de3469596556215a Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Sat, 5 Feb 2022 21:14:10 +0800
Subject: [PATCH] Codegen `column_name` proc_macro attribute (#433)

* feat: codegen `column_name` proc_macro attribute

* test: codegen `column_name`
---
 sea-orm-codegen/src/entity/column.rs     |  4 ++++
 sea-orm-codegen/src/entity/writer.rs     | 25 +++++++++++++++++++-----
 sea-orm-codegen/tests/compact/vendor.rs  |  2 ++
 sea-orm-codegen/tests/expanded/vendor.rs |  2 ++
 4 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/sea-orm-codegen/src/entity/column.rs b/sea-orm-codegen/src/entity/column.rs
index 5890ab8c..3b374b02 100644
--- a/sea-orm-codegen/src/entity/column.rs
+++ b/sea-orm-codegen/src/entity/column.rs
@@ -22,6 +22,10 @@ impl Column {
         format_ident!("{}", escape_rust_keyword(self.name.to_camel_case()))
     }
 
+    pub fn is_snake_case_name(&self) -> bool {
+        self.name.to_snake_case() == self.name
+    }
+
     pub fn get_rs_type(&self) -> TokenStream {
         #[allow(unreachable_patterns)]
         let ident: TokenStream = match &self.col_type {
diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
index 60a01eeb..beb0a042 100644
--- a/sea-orm-codegen/src/entity/writer.rs
+++ b/sea-orm-codegen/src/entity/writer.rs
@@ -306,11 +306,22 @@ impl EntityWriter {
     }
 
     pub fn gen_column_enum(entity: &Entity) -> TokenStream {
-        let column_names_camel_case = entity.get_column_names_camel_case();
+        let column_variants = entity.columns.iter().map(|col| {
+            let variant = col.get_name_camel_case();
+            let mut variant = quote! { #variant };
+            if !col.is_snake_case_name() {
+                let column_name = &col.name;
+                variant = quote! {
+                    #[sea_orm(column_name = #column_name)]
+                    #variant
+                };
+            }
+            variant
+        });
         quote! {
             #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
             pub enum Column {
-                #(#column_names_camel_case,)*
+                #(#column_variants,)*
             }
         }
     }
@@ -474,6 +485,10 @@ impl EntityWriter {
             .iter()
             .map(|col| {
                 let mut attrs: Punctuated<_, Comma> = Punctuated::new();
+                if !col.is_snake_case_name() {
+                    let column_name = &col.name;
+                    attrs.push(quote! { column_name = #column_name });
+                }
                 if primary_keys.contains(&col.name) {
                     attrs.push(quote! { primary_key });
                     if !col.auto_increment {
@@ -724,14 +739,14 @@ mod tests {
                         unique: false,
                     },
                     Column {
-                        name: "name".to_owned(),
+                        name: "_name_".to_owned(),
                         col_type: ColumnType::String(Some(255)),
                         auto_increment: false,
                         not_null: true,
                         unique: false,
                     },
                     Column {
-                        name: "fruit_id".to_owned(),
+                        name: "fruitId".to_owned(),
                         col_type: ColumnType::Integer(Some(11)),
                         auto_increment: false,
                         not_null: false,
@@ -740,7 +755,7 @@ mod tests {
                 ],
                 relations: vec![Relation {
                     ref_table: "fruit".to_owned(),
-                    columns: vec!["fruit_id".to_owned()],
+                    columns: vec!["fruitId".to_owned()],
                     ref_columns: vec!["id".to_owned()],
                     rel_type: RelationType::BelongsTo,
                     on_delete: None,
diff --git a/sea-orm-codegen/tests/compact/vendor.rs b/sea-orm-codegen/tests/compact/vendor.rs
index 314eb396..1351c227 100644
--- a/sea-orm-codegen/tests/compact/vendor.rs
+++ b/sea-orm-codegen/tests/compact/vendor.rs
@@ -7,7 +7,9 @@ use sea_orm::entity::prelude::*;
 pub struct Model {
     #[sea_orm(primary_key)]
     pub id: i32,
+    #[sea_orm(column_name = "_name_")]
     pub name: String,
+    #[sea_orm(column_name = "fruitId")]
     pub fruit_id: Option<i32> ,
 }
 
diff --git a/sea-orm-codegen/tests/expanded/vendor.rs b/sea-orm-codegen/tests/expanded/vendor.rs
index ffc211a8..ab4ae39a 100644
--- a/sea-orm-codegen/tests/expanded/vendor.rs
+++ b/sea-orm-codegen/tests/expanded/vendor.rs
@@ -21,7 +21,9 @@ pub struct Model {
 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
 pub enum Column {
     Id,
+    #[sea_orm(column_name = "_name_")]
     Name,
+    #[sea_orm(column_name = "fruitId")]
     FruitId,
 }