Merge pull request #360 from caido/feat/column-table-iden

Add feature to generate table Iden

.github/workflows/rust.yml

@@ -316,7 +316,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest]
-        path: [86, 249, 262, 319, 324, 352]
+        path: [86, 249, 262, 319, 324, 352, 356]
     steps:
       - uses: actions/checkout@v2
 

issues/356/Cargo.toml

@@ -0,0 +1,11 @@
+[workspace]
+# A separate workspace
+
+[package]
+name = "sea-orm-issues-356"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[dependencies]
+sea-orm = { path = "../../", features = [ "sqlx-mysql", "runtime-async-std-native-tls" ]}

issues/356/src/main.rs

@@ -0,0 +1,3 @@
+mod model;
+
+pub fn main() {}

issues/356/src/model.rs

@@ -0,0 +1,42 @@
+use sea_orm::entity::prelude::*;
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
+#[sea_orm(table_name = "model", table_iden)]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: i32,
+    pub owner: String,
+    pub name: String,
+    pub description: String,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use sea_orm::*;
+
+    #[test]
+    fn test_columns_1() {
+        assert_eq!(
+            Column::iter()
+                .map(|col| col.to_string())
+                .collect::<Vec<_>>(),
+            vec![
+                "id".to_owned(),
+                "owner".to_owned(),
+                "name".to_owned(),
+                "description".to_owned(),
+            ]
+        );
+        assert_eq!(Column::Table.to_string().as_str(), "model");
+        assert_eq!(Column::Id.to_string().as_str(), "id");
+        assert_eq!(Column::Owner.to_string().as_str(), "owner");
+        assert_eq!(Column::Name.to_string().as_str(), "name");
+        assert_eq!(Column::Description.to_string().as_str(), "description");
+    }
+}

sea-orm-macros/src/attributes.rs

@@ -11,6 +11,7 @@ pub mod derive_attr {
         pub relation: Option<syn::Ident>,
         pub schema_name: Option<syn::Lit>,
         pub table_name: Option<syn::Lit>,
+        pub table_iden: Option<()>,
     }
 }
 

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();
+                                    }
+                                }
                             }
                         }
                     }

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

@@ -12,6 +12,7 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
     // if #[sea_orm(table_name = "foo", schema_name = "bar")] specified, create Entity struct
     let mut table_name = None;
     let mut schema_name = quote! { None };
+    let mut table_iden = false;
     attrs.iter().for_each(|attr| {
         if attr.path.get_ident().map(|i| i == "sea_orm") != Some(true) {
             return;
@@ -28,11 +29,18 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                             schema_name = quote! { Some(#name) };
                         }
                     }
+                } else if let Meta::Path(path) = meta {
+                    if let Some(ident) = path.get_ident() {
+                        if ident == "table_iden" {
+                            table_iden = true;
+                        }
+                    }
                 }
             }
         }
     });
     let entity_def = table_name
+        .as_ref()
         .map(|table_name| {
             quote! {
                 #[derive(Copy, Clone, Default, Debug, sea_orm::prelude::DeriveEntity)]
@@ -58,6 +66,19 @@ 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;
+    if 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 {