From 8b30ef74a22153ebdb2e0cf65afd411fe68ba8a6 Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Thu, 12 Aug 2021 19:07:03 +0800
Subject: [PATCH] Customize column names

---
 Cargo.toml                           |  2 +-
 sea-orm-macros/src/derives/column.rs | 20 ++++++++++++++------
 sea-orm-macros/src/lib.rs            | 12 +++++++++++-
 src/entity/prelude.rs                |  4 ++--
 src/lib.rs                           |  4 ++--
 src/tests_cfg/filling.rs             | 13 ++++++++++++-
 6 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index ec3d5282..aa40ea2b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -33,7 +33,7 @@ chrono = { version = "^0", optional = true }
 futures = { version = "^0.3" }
 futures-util = { version = "^0.3" }
 rust_decimal = { version = "^1", optional = true }
-sea-orm-macros = { version = "^0.1", optional = true }
+sea-orm-macros = { version = "^0.1", path = "sea-orm-macros", optional = true }
 sea-query = { version = "^0.12.8" }
 sea-strum = { version = "^0.21", features = ["derive", "sea-orm"] }
 serde = { version = "^1.0", features = ["derive"] }
diff --git a/sea-orm-macros/src/derives/column.rs b/sea-orm-macros/src/derives/column.rs
index 78a72259..bf40f58c 100644
--- a/sea-orm-macros/src/derives/column.rs
+++ b/sea-orm-macros/src/derives/column.rs
@@ -3,7 +3,7 @@ use proc_macro2::{Ident, TokenStream};
 use quote::{quote, quote_spanned};
 use syn::{Data, DataEnum, Fields, Variant};
 
-pub fn expand_derive_column(ident: Ident, data: Data) -> syn::Result<TokenStream> {
+pub fn expand_derive_column(ident: &Ident, data: &Data) -> syn::Result<TokenStream> {
     let variants = match data {
         syn::Data::Enum(DataEnum { variants, .. }) => variants,
         _ => {
@@ -30,12 +30,10 @@ pub fn expand_derive_column(ident: Ident, data: Data) -> syn::Result<TokenStream
         })
         .collect();
 
+    let impl_iden = expand_derive_custom_column(ident)?;
+
     Ok(quote!(
-        impl sea_orm::Iden for #ident {
-            fn unquoted(&self, s: &mut dyn std::fmt::Write) {
-                write!(s, "{}", self.as_str()).unwrap();
-            }
-        }
+        #impl_iden
 
         impl sea_orm::IdenStatic for #ident {
             fn as_str(&self) -> &str {
@@ -46,3 +44,13 @@ pub fn expand_derive_column(ident: Ident, data: Data) -> syn::Result<TokenStream
         }
     ))
 }
+
+pub fn expand_derive_custom_column(ident: &Ident) -> syn::Result<TokenStream> {
+    Ok(quote!(
+        impl sea_orm::Iden for #ident {
+            fn unquoted(&self, s: &mut dyn std::fmt::Write) {
+                write!(s, "{}", self.as_str()).unwrap();
+            }
+        }
+    ))
+}
diff --git a/sea-orm-macros/src/lib.rs b/sea-orm-macros/src/lib.rs
index 65b8fd15..7e3f4761 100644
--- a/sea-orm-macros/src/lib.rs
+++ b/sea-orm-macros/src/lib.rs
@@ -29,7 +29,17 @@ pub fn derive_primary_key(input: TokenStream) -> TokenStream {
 pub fn derive_column(input: TokenStream) -> TokenStream {
     let DeriveInput { ident, data, .. } = parse_macro_input!(input);
 
-    match derives::expand_derive_column(ident, data) {
+    match derives::expand_derive_column(&ident, &data) {
+        Ok(ts) => ts.into(),
+        Err(e) => e.to_compile_error().into(),
+    }
+}
+
+#[proc_macro_derive(DeriveCustomColumn)]
+pub fn derive_custom_column(input: TokenStream) -> TokenStream {
+    let DeriveInput { ident, .. } = parse_macro_input!(input);
+
+    match derives::expand_derive_custom_column(&ident) {
         Ok(ts) => ts.into(),
         Err(e) => e.to_compile_error().into(),
     }
diff --git a/src/entity/prelude.rs b/src/entity/prelude.rs
index c7182287..447117b7 100644
--- a/src/entity/prelude.rs
+++ b/src/entity/prelude.rs
@@ -1,7 +1,7 @@
 pub use crate::{
     error::*, ActiveModelBehavior, ActiveModelTrait, ColumnDef, ColumnTrait, ColumnType,
-    DeriveActiveModel, DeriveActiveModelBehavior, DeriveColumn, DeriveEntity, DeriveModel,
-    DerivePrimaryKey, EntityName, EntityTrait, EnumIter, Iden, IdenStatic, ModelTrait,
+    DeriveActiveModel, DeriveActiveModelBehavior, DeriveColumn, DeriveCustomColumn, DeriveEntity,
+    DeriveModel, DerivePrimaryKey, EntityName, EntityTrait, EnumIter, Iden, IdenStatic, ModelTrait,
     PrimaryKeyToColumn, PrimaryKeyTrait, QueryFilter, QueryResult, Related, RelationDef,
     RelationTrait, Select, Value,
 };
diff --git a/src/lib.rs b/src/lib.rs
index a81c620b..becf73eb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -219,8 +219,8 @@ pub use executor::*;
 pub use query::*;
 
 pub use sea_orm_macros::{
-    DeriveActiveModel, DeriveActiveModelBehavior, DeriveColumn, DeriveEntity, DeriveModel,
-    DerivePrimaryKey, FromQueryResult,
+    DeriveActiveModel, DeriveActiveModelBehavior, DeriveColumn, DeriveCustomColumn, DeriveEntity,
+    DeriveModel, DerivePrimaryKey, FromQueryResult,
 };
 
 pub use sea_query;
diff --git a/src/tests_cfg/filling.rs b/src/tests_cfg/filling.rs
index a246d146..838f96b9 100644
--- a/src/tests_cfg/filling.rs
+++ b/src/tests_cfg/filling.rs
@@ -16,12 +16,23 @@ pub struct Model {
     pub name: String,
 }
 
-#[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
+// If your column names are not in snake-case, derive `DeriveCustomColumn` here.
+#[derive(Copy, Clone, Debug, EnumIter, DeriveCustomColumn)]
 pub enum Column {
     Id,
     Name,
 }
 
+// Then, customize each column names here.
+impl IdenStatic for Column {
+    fn as_str(&self) -> &str {
+        match self {
+            Self::Id => "id",
+            Self::Name => "name",
+        }
+    }
+}
+
 #[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)]
 pub enum PrimaryKey {
     Id,