create_table_from_entity support comment (#2009)

* support table and column comment * create_table_from_entity support comment * fmt code * fmt code * add comment test * fix test
SeaQL · Dec 14, 2023 · 94eac96 · 94eac96
1 parent 544675f
commit 94eac96
Show file tree

Hide file tree

Showing 11 changed files with 81 additions and 7 deletions.
diff --git a/sea-orm-codegen/src/entity/base_entity.rs b/sea-orm-codegen/src/entity/base_entity.rs
@@ -15,9 +15,14 @@ pub struct Entity {
     pub(crate) relations: Vec<Relation>,
     pub(crate) conjunct_relations: Vec<ConjunctRelation>,
     pub(crate) primary_keys: Vec<PrimaryKey>,
+    pub(crate) comment: Option<String>,
 }
 
 impl Entity {
+    pub fn get_table_comment(&self) -> Option<&String> {
+        self.comment.as_ref()
+    }
+
     pub fn get_table_name_snake_case(&self) -> String {
         self.table_name.to_snake_case()
     }
@@ -318,6 +323,7 @@ mod tests {
             primary_keys: vec![PrimaryKey {
                 name: "id".to_owned(),
             }],
+            comment: None,
         }
     }
 

diff --git a/sea-orm-codegen/src/entity/transformer.rs b/sea-orm-codegen/src/entity/transformer.rs
@@ -117,12 +117,14 @@ impl EntityTransformer {
                             .collect::<Vec<_>>()
                     }),
             );
+            let comment = table_create.get_comment().cloned();
             let entity = Entity {
                 table_name: table_name.clone(),
                 columns,
                 relations: relations.clone(),
                 conjunct_relations: vec![],
                 primary_keys,
+                comment,
             };
             entities.insert(table_name.clone(), entity.clone());
             for mut rel in relations.into_iter() {

diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
@@ -785,6 +785,7 @@ impl EntityWriter {
             },
             None => quote! {},
         };
+
         let extra_derive = with_serde.extra_derive();
 
         quote! {
@@ -882,9 +883,11 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
             Entity {
                 table_name: "_cake_filling_".to_owned(),
+                comment: None,
                 columns: vec![
                     Column {
                         name: "cake_id".to_owned(),
@@ -980,6 +983,7 @@ mod tests {
                         name: "filling_id".to_owned(),
                     },
                 ],
+                comment: None,
             },
             Entity {
                 table_name: "filling".to_owned(),
@@ -1007,6 +1011,7 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
             Entity {
                 table_name: "fruit".to_owned(),
@@ -1061,6 +1066,7 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
             Entity {
                 table_name: "vendor".to_owned(),
@@ -1102,6 +1108,7 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
             Entity {
                 table_name: "rust_keyword".to_owned(),
@@ -1259,6 +1266,7 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
             Entity {
                 table_name: "cake_with_float".to_owned(),
@@ -1303,6 +1311,7 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
             Entity {
                 table_name: "cake_with_double".to_owned(),
@@ -1347,6 +1356,7 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
             Entity {
                 table_name: "collection".to_owned(),
@@ -1378,6 +1388,7 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
             Entity {
                 table_name: "collection_float".to_owned(),
@@ -1409,6 +1420,7 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
         ]
     }
@@ -1770,6 +1782,7 @@ mod tests {
     fn test_gen_with_seaography() -> io::Result<()> {
         let cake_entity = Entity {
             table_name: "cake".to_owned(),
+            comment: None,
             columns: vec![
                 Column {
                     name: "id".to_owned(),
@@ -2225,6 +2238,7 @@ mod tests {
                 primary_keys: vec![PrimaryKey {
                     name: "id".to_owned(),
                 }],
+                comment: None,
             },
         ];
         const ENTITY_FILES: [&str; 1] = [include_str!("../../tests/postgres/binary_json.rs")];
@@ -2308,6 +2322,7 @@ mod tests {
         let entities = vec![
             Entity {
                 table_name: "tea_pairing".to_owned(),
+                comment: None,
                 columns: vec![
                     Column {
                         name: "id".to_owned(),
@@ -2351,6 +2366,7 @@ mod tests {
             },
             Entity {
                 table_name: "tea_pairing_with_size".to_owned(),
+                comment: None,
                 columns: vec![
                     Column {
                         name: "id".to_owned(),

diff --git a/sea-orm-macros/src/derives/attributes.rs b/sea-orm-macros/src/derives/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 comment: Option<syn::Lit>,
         pub table_iden: Option<()>,
     }
 }

diff --git a/sea-orm-macros/src/derives/entity_model.rs b/sea-orm-macros/src/derives/entity_model.rs
@@ -10,14 +10,18 @@ use syn::{
 pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Result<TokenStream> {
     // if #[sea_orm(table_name = "foo", schema_name = "bar")] specified, create Entity struct
     let mut table_name = None;
+    let mut comment = quote! {None};
     let mut schema_name = quote! { None };
     let mut table_iden = false;
     attrs
         .iter()
         .filter(|attr| attr.path().is_ident("sea_orm"))
         .try_for_each(|attr| {
             attr.parse_nested_meta(|meta| {
-                if meta.path.is_ident("table_name") {
+                if meta.path.is_ident("comment") {
+                    let name: Lit = meta.value()?.parse()?;
+                    comment = quote! { Some(#name) };
+                } else if meta.path.is_ident("table_name") {
                     table_name = Some(meta.value()?.parse::<Lit>()?);
                 } else if meta.path.is_ident("schema_name") {
                     let name: Lit = meta.value()?.parse()?;
@@ -51,6 +55,10 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                     fn table_name(&self) -> &str {
                         #table_name
                     }
+
+                    fn comment(&self) -> Option<&str> {
+                        #comment
+                    }
                 }
             }
         })
@@ -90,6 +98,7 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
 
                     let mut nullable = false;
                     let mut default_value = None;
+                    let mut comment = None;
                     let mut default_expr = None;
                     let mut select_as = None;
                     let mut save_as = None;
@@ -134,6 +143,8 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                                         meta.error(format!("Invalid auto_increment = {:?}", lit))
                                     );
                                 }
+                            } else if meta.path.is_ident("comment") {
+                                comment = Some(meta.value()?.parse::<Lit>()?);
                             } else if meta.path.is_ident("default_value") {
                                 default_value = Some(meta.value()?.parse::<Lit>()?);
                             } else if meta.path.is_ident("default_expr") {
@@ -274,9 +285,13 @@ pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Res
                     if let Some(default_value) = default_value {
                         match_row = quote! { #match_row.default_value(#default_value) };
                     }
+                    if let Some(comment) = comment {
+                        match_row = quote! { #match_row.comment(#comment) };
+                    }
                     if let Some(default_expr) = default_expr {
                         match_row = quote! { #match_row.default(#default_expr) };
                     }
+                    // match_row = quote! { #match_row.comment() };
                     columns_trait.push(match_row);
                 }
             }

diff --git a/sea-orm-macros/src/derives/model.rs b/sea-orm-macros/src/derives/model.rs
@@ -182,7 +182,6 @@ impl DeriveModel {
 /// Method to derive an ActiveModel
 pub fn expand_derive_model(input: syn::DeriveInput) -> syn::Result<TokenStream> {
     let ident_span = input.ident.span();
-
     match DeriveModel::new(input) {
         Ok(model) => model.expand(),
         Err(Error::InputNotStruct) => Ok(quote_spanned! {

diff --git a/src/entity/base_entity.rs b/src/entity/base_entity.rs
@@ -20,6 +20,11 @@ pub trait EntityName: IdenStatic + Default {
         None
     }
 
+    /// Method to get the comment for the schema, defaults to [Option::None] if not set
+    fn comment(&self) -> Option<&str> {
+        None
+    }
+
     /// Get the name of the table
     fn table_name(&self) -> &str;
 

diff --git a/src/entity/column.rs b/src/entity/column.rs
@@ -16,6 +16,7 @@ pub struct ColumnDef {
     pub(crate) unique: bool,
     pub(crate) indexed: bool,
     pub(crate) default: Option<SimpleExpr>,
+    pub(crate) comment: Option<String>,
 }
 
 macro_rules! bind_oper {
@@ -312,6 +313,7 @@ impl ColumnTypeTrait for ColumnType {
             unique: false,
             indexed: false,
             default: None,
+            comment: None,
         }
     }
 
@@ -344,6 +346,11 @@ impl ColumnDef {
         self.unique = true;
         self
     }
+    /// Set column comment
+    pub fn comment(mut self, v: &str) -> Self {
+        self.comment = Some(v.into());
+        self
+    }
 
     /// Mark the column as nullable
     pub fn null(self) -> Self {

diff --git a/src/schema/entity.rs b/src/schema/entity.rs
@@ -163,6 +163,10 @@ where
 {
     let mut stmt = TableCreateStatement::new();
 
+    if let Some(comment) = entity.comment() {
+        stmt.comment(comment);
+    }
+
     for column in E::Column::iter() {
         let mut column_def = column_def_from_entity_column::<E>(column, backend);
         stmt.col(&mut column_def);
@@ -213,6 +217,9 @@ where
     if let Some(default) = orm_column_def.default {
         column_def.default(default);
     }
+    if let Some(comment) = orm_column_def.comment {
+        column_def.comment(comment);
+    }
     for primary_key in E::PrimaryKey::iter() {
         if column.to_string() == primary_key.into_column().to_string() {
             if E::PrimaryKey::auto_increment() {

diff --git a/tests/common/features/applog.rs b/tests/common/features/applog.rs
@@ -1,12 +1,15 @@
 use sea_orm::entity::prelude::*;
 
 #[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
-#[sea_orm(table_name = "applog")]
+#[sea_orm(table_name = "applog", comment = "app logs")]
 pub struct Model {
-    #[sea_orm(primary_key)]
+    #[sea_orm(primary_key, comment = "ID")]
     pub id: i32,
+    #[sea_orm(comment = "action")]
     pub action: String,
+    #[sea_orm(comment = "action data")]
     pub json: Json,
+    #[sea_orm(comment = "create time")]
     pub created_at: DateTimeWithTimeZone,
 }
 

diff --git a/tests/common/features/schema.rs b/tests/common/features/schema.rs
@@ -67,19 +67,32 @@ pub async fn create_tables(db: &DatabaseConnection) -> Result<(), DbErr> {
 pub async fn create_log_table(db: &DbConn) -> Result<ExecResult, DbErr> {
     let stmt = sea_query::Table::create()
         .table(applog::Entity)
+        .comment("app logs")
         .col(
             ColumnDef::new(applog::Column::Id)
                 .integer()
                 .not_null()
+                .comment("ID")
                 .auto_increment()
                 .primary_key(),
         )
-        .col(ColumnDef::new(applog::Column::Action).string().not_null())
-        .col(ColumnDef::new(applog::Column::Json).json().not_null())
+        .col(
+            ColumnDef::new(applog::Column::Action)
+                .string()
+                .not_null()
+                .comment("action"),
+        )
+        .col(
+            ColumnDef::new(applog::Column::Json)
+                .json()
+                .not_null()
+                .comment("action data"),
+        )
         .col(
             ColumnDef::new(applog::Column::CreatedAt)
                 .timestamp_with_time_zone()
-                .not_null(),
+                .not_null()
+                .comment("create time"),
         )
         .to_owned();