From f04ef378c6b9478dd2e5291e8c621988787172e7 Mon Sep 17 00:00:00 2001
From: Emile Fugulin <code@efugulin.com>
Date: Sat, 4 Dec 2021 12:05:24 -0500
Subject: [PATCH] Add feature to generate table Iden

---
 Cargo.toml                                 |  1 +
 sea-orm-macros/Cargo.toml                  |  3 +++
 sea-orm-macros/src/derives/column.rs       |  5 +++++
 sea-orm-macros/src/derives/entity_model.rs | 13 +++++++++++++
 4 files changed, 22 insertions(+)

diff --git a/Cargo.toml b/Cargo.toml
index f73cfa51..27380d04 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -67,6 +67,7 @@ with-json = ["serde_json", "sea-query/with-json", "chrono/serde"]
 with-chrono = ["chrono", "sea-query/with-chrono"]
 with-rust_decimal = ["rust_decimal", "sea-query/with-rust_decimal"]
 with-uuid = ["uuid", "sea-query/with-uuid"]
+with-table-iden = ["sea-orm-macros/with-table-iden"]
 sqlx-all = ["sqlx-mysql", "sqlx-postgres", "sqlx-sqlite"]
 sqlx-dep = ["sqlx-json", "sqlx-chrono", "sqlx-decimal", "sqlx-uuid"]
 sqlx-json = ["sqlx/json", "with-json"]
diff --git a/sea-orm-macros/Cargo.toml b/sea-orm-macros/Cargo.toml
index c4c82756..be22fb2f 100644
--- a/sea-orm-macros/Cargo.toml
+++ b/sea-orm-macros/Cargo.toml
@@ -25,3 +25,6 @@ proc-macro2 = "^1"
 [dev-dependencies]
 sea-orm = { path = "../", features = ["macros"] }
 serde = { version = "^1.0", features = ["derive"] }
+
+[features]
+with-table-iden = []
diff --git a/sea-orm-macros/src/derives/column.rs b/sea-orm-macros/src/derives/column.rs
index 8e60c560..97191250 100644
--- a/sea-orm-macros/src/derives/column.rs
+++ b/sea-orm-macros/src/derives/column.rs
@@ -45,6 +45,11 @@ pub fn impl_default_as_str(ident: &Ident, data: &Data) -> syn::Result<TokenStrea
                                         column_name = litstr.value();
                                     }
                                 }
+                                if name == "table_name" {
+                                    if let Lit::Str(litstr) = &nv.lit {
+                                        column_name = litstr.value();
+                                    }
+                                }
                             }
                         }
                     }
diff --git a/sea-orm-macros/src/derives/entity_model.rs b/sea-orm-macros/src/derives/entity_model.rs
index e39fecf1..3a54e145 100644
--- a/sea-orm-macros/src/derives/entity_model.rs
+++ b/sea-orm-macros/src/derives/entity_model.rs
@@ -33,6 +33,7 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
         }
     });
     let entity_def = table_name
+        .clone()
         .map(|table_name| {
             quote! {
                 #[derive(Copy, Clone, Default, Debug, sea_orm::prelude::DeriveEntity)]
@@ -58,6 +59,18 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
     let mut primary_keys: Punctuated<_, Comma> = Punctuated::new();
     let mut primary_key_types: Punctuated<_, Comma> = Punctuated::new();
     let mut auto_increment = true;
+
+    #[cfg(feature = "with-table-iden")]
+    if let Some(table_name) = table_name {
+        let table_field_name = Ident::new("Table", Span::call_site());
+        columns_enum.push(quote! {
+            #[sea_orm(table_name=#table_name)]
+            #[strum(disabled)]
+            #table_field_name
+        });
+        columns_trait
+            .push(quote! { Self::#table_field_name => panic!("Table cannot be used as a column") });
+    }
     if let Data::Struct(item_struct) = data {
         if let Fields::Named(fields) = item_struct.fields {
             for field in fields.named {