From 6551ebb6c23ce0db27b958216f24ee75f51c87e2 Mon Sep 17 00:00:00 2001 From: Karol Gorecki <2426927+karolgorecki@users.noreply.github.com> Date: Thu, 2 May 2024 01:21:24 +0200 Subject: [PATCH] feat: Support --column-extra-derives --- sea-orm-cli/src/cli.rs | 7 ++ sea-orm-cli/src/commands/generate.rs | 2 + sea-orm-codegen/src/entity/transformer.rs | 1 + sea-orm-codegen/src/entity/writer.rs | 87 ++++++++++++++++++- .../cake_multiple.rs | 77 ++++++++++++++++ .../expanded_with_column_derives/cake_none.rs | 77 ++++++++++++++++ .../expanded_with_column_derives/cake_one.rs | 77 ++++++++++++++++ 7 files changed, 325 insertions(+), 3 deletions(-) create mode 100644 sea-orm-codegen/tests/expanded_with_column_derives/cake_multiple.rs create mode 100644 sea-orm-codegen/tests/expanded_with_column_derives/cake_none.rs create mode 100644 sea-orm-codegen/tests/expanded_with_column_derives/cake_one.rs diff --git a/sea-orm-cli/src/cli.rs b/sea-orm-cli/src/cli.rs index 4c771da00..a589ab9a4 100644 --- a/sea-orm-cli/src/cli.rs +++ b/sea-orm-cli/src/cli.rs @@ -298,6 +298,13 @@ pub enum GenerateSubcommands { )] enum_extra_attributes: Vec, + #[arg( + long, + value_delimiter = ',', + help = "Add extra derive macros to generated column enum (comma separated), e.g. `--column-extra-derives 'async_graphql::Enum','CustomDerive'`" + )] + column_extra_derives: Vec, + #[arg( long, default_value = "false", diff --git a/sea-orm-cli/src/commands/generate.rs b/sea-orm-cli/src/commands/generate.rs index a07ca2e38..0b7c56c7b 100644 --- a/sea-orm-cli/src/commands/generate.rs +++ b/sea-orm-cli/src/commands/generate.rs @@ -33,6 +33,7 @@ pub async fn run_generate_command( model_extra_attributes, enum_extra_derives, enum_extra_attributes, + column_extra_derives, seaography, } => { if verbose { @@ -187,6 +188,7 @@ pub async fn run_generate_command( model_extra_attributes, enum_extra_derives, enum_extra_attributes, + column_extra_derives, seaography, ); let output = EntityTransformer::transform(table_stmts)?.generate(&writer_context); diff --git a/sea-orm-codegen/src/entity/transformer.rs b/sea-orm-codegen/src/entity/transformer.rs index 20f8047cb..0f4a0cb5c 100644 --- a/sea-orm-codegen/src/entity/transformer.rs +++ b/sea-orm-codegen/src/entity/transformer.rs @@ -386,6 +386,7 @@ mod tests { false, &Default::default(), &Default::default(), + &Default::default(), false, ) .into_iter() diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index 1660c2221..7ef971079 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/sea-orm-codegen/src/entity/writer.rs @@ -49,6 +49,7 @@ pub struct EntityWriterContext { pub(crate) model_extra_attributes: TokenStream, pub(crate) enum_extra_derives: TokenStream, pub(crate) enum_extra_attributes: TokenStream, + pub(crate) column_extra_derives: TokenStream, pub(crate) seaography: bool, } @@ -147,6 +148,7 @@ impl EntityWriterContext { model_extra_attributes: Vec, enum_extra_derives: Vec, enum_extra_attributes: Vec, + column_extra_derives: Vec, seaography: bool, ) -> Self { Self { @@ -162,6 +164,7 @@ impl EntityWriterContext { model_extra_attributes: bonus_attributes(model_extra_attributes), enum_extra_derives: bonus_derive(enum_extra_derives), enum_extra_attributes: bonus_attributes(enum_extra_attributes), + column_extra_derives: bonus_derive(column_extra_derives), seaography, } } @@ -221,6 +224,7 @@ impl EntityWriter { serde_skip_hidden_column, &context.model_extra_derives, &context.model_extra_attributes, + &context.column_extra_derives, context.seaography, ) } else { @@ -233,6 +237,7 @@ impl EntityWriter { serde_skip_hidden_column, &context.model_extra_derives, &context.model_extra_attributes, + &context.column_extra_derives, context.seaography, ) }; @@ -346,6 +351,7 @@ impl EntityWriter { serde_skip_hidden_column: bool, model_extra_derives: &TokenStream, model_extra_attributes: &TokenStream, + column_extra_derives: &TokenStream, seaography: bool, ) -> Vec { let mut imports = Self::gen_import(with_serde); @@ -363,7 +369,7 @@ impl EntityWriter { model_extra_derives, model_extra_attributes, ), - Self::gen_column_enum(entity), + Self::gen_column_enum(entity, column_extra_derives), Self::gen_primary_key_enum(entity), Self::gen_impl_primary_key(entity, date_time_crate), Self::gen_relation_enum(entity), @@ -389,6 +395,7 @@ impl EntityWriter { serde_skip_hidden_column: bool, model_extra_derives: &TokenStream, model_extra_attributes: &TokenStream, + _column_extra_derives: &TokenStream, seaography: bool, ) -> Vec { let mut imports = Self::gen_import(with_serde); @@ -527,7 +534,7 @@ impl EntityWriter { } } - pub fn gen_column_enum(entity: &Entity) -> TokenStream { + pub fn gen_column_enum(entity: &Entity, column_extra_derives: &TokenStream) -> TokenStream { let column_variants = entity.columns.iter().map(|col| { let variant = col.get_name_camel_case(); let mut variant = quote! { #variant }; @@ -541,7 +548,7 @@ impl EntityWriter { variant }); quote! { - #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)] + #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn #column_extra_derives)] pub enum Column { #(#column_variants,)* } @@ -1552,6 +1559,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false ) .into_iter() @@ -1573,6 +1581,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, ) .into_iter() @@ -1636,6 +1645,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, ) .into_iter() @@ -1657,6 +1667,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, ) .into_iter() @@ -1690,6 +1701,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1706,6 +1718,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1722,6 +1735,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1736,6 +1750,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1752,6 +1767,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1768,6 +1784,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1784,6 +1801,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1798,6 +1816,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1879,6 +1898,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), true, )) ); @@ -1895,6 +1915,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), true, )) ); @@ -1922,6 +1943,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1936,6 +1958,7 @@ mod tests { false, &bonus_derive(["ts_rs::TS"]), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1952,6 +1975,7 @@ mod tests { false, &bonus_derive(["ts_rs::TS", "utoipa::ToSchema"]), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1970,6 +1994,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -1986,6 +2011,7 @@ mod tests { false, &bonus_derive(["ts_rs::TS"]), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -2002,6 +2028,7 @@ mod tests { false, &bonus_derive(["ts_rs::TS", "utoipa::ToSchema"]), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -2033,6 +2060,50 @@ mod tests { Ok(()) } + #[test] + fn test_gen_with_column_derives() -> io::Result<()> { + let cake_entity = setup().get_mut(0).unwrap().clone(); + + assert_eq!(cake_entity.get_table_name_snake_case(), "cake"); + + assert_eq!( + comparable_file_string(include_str!( + "../../tests/expanded_with_column_derives/cake_one.rs" + ))?, + generated_to_string(EntityWriter::gen_expanded_code_blocks( + &cake_entity, + &WithSerde::None, + &DateTimeCrate::Chrono, + &None, + false, + false, + &TokenStream::new(), + &TokenStream::new(), + &bonus_derive(["async_graphql::Enum"]), + false, + )) + ); + assert_eq!( + comparable_file_string(include_str!( + "../../tests/expanded_with_column_derives/cake_multiple.rs" + ))?, + generated_to_string(EntityWriter::gen_expanded_code_blocks( + &cake_entity, + &WithSerde::None, + &DateTimeCrate::Chrono, + &None, + false, + false, + &TokenStream::new(), + &TokenStream::new(), + &bonus_derive(["async_graphql::Enum", "Eq", "PartialEq"]), + false, + )) + ); + + Ok(()) + } + #[allow(clippy::type_complexity)] fn assert_serde_variant_results( cake_entity: &Entity, @@ -2047,6 +2118,7 @@ mod tests { bool, &TokenStream, &TokenStream, + &TokenStream, bool, ) -> Vec, >, @@ -2078,6 +2150,7 @@ mod tests { serde_skip_hidden_column, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, ) .into_iter() @@ -2110,6 +2183,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -2126,6 +2200,7 @@ mod tests { false, &TokenStream::new(), &bonus_attributes([r#"serde(rename_all = "camelCase")"#]), + &TokenStream::new(), false, )) ); @@ -2142,6 +2217,7 @@ mod tests { false, &TokenStream::new(), &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]), + &TokenStream::new(), false, )) ); @@ -2160,6 +2236,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, )) ); @@ -2176,6 +2253,7 @@ mod tests { false, &TokenStream::new(), &bonus_attributes([r#"serde(rename_all = "camelCase")"#]), + &TokenStream::new(), false, )) ); @@ -2192,6 +2270,7 @@ mod tests { false, &TokenStream::new(), &bonus_attributes([r#"serde(rename_all = "camelCase")"#, "ts(export)"]), + &TokenStream::new(), false, )) ); @@ -2285,6 +2364,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, ) .into_iter() @@ -2306,6 +2386,7 @@ mod tests { false, &TokenStream::new(), &TokenStream::new(), + &TokenStream::new(), false, ) .into_iter() diff --git a/sea-orm-codegen/tests/expanded_with_column_derives/cake_multiple.rs b/sea-orm-codegen/tests/expanded_with_column_derives/cake_multiple.rs new file mode 100644 index 000000000..7f601b405 --- /dev/null +++ b/sea-orm-codegen/tests/expanded_with_column_derives/cake_multiple.rs @@ -0,0 +1,77 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +use sea_orm::entity::prelude:: * ; + +#[derive(Copy, Clone, Default, Debug, DeriveEntity)] +pub struct Entity; + +impl EntityName for Entity { + fn table_name(&self) -> &str { + "cake" + } +} + +#[derive(Clone, Debug, PartialEq, DeriveModel, DeriveActiveModel, Eq)] +pub struct Model { + pub id: i32, + pub name: Option , +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveColumn, async_graphql::Enum, Eq, PartialEq)] +pub enum Column { + Id, + Name, +} + +#[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)] +pub enum PrimaryKey { + Id, +} + +impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + + fn auto_increment() -> bool { + true + } +} + +#[derive(Copy, Clone, Debug, EnumIter)] +pub enum Relation { + Fruit, +} + +impl ColumnTrait for Column { + type EntityName = Entity; + fn def(&self) -> ColumnDef { + match self { + Self::Id => ColumnType::Integer.def(), + Self::Name => ColumnType::Text.def().null(), + } + } +} + +impl RelationTrait for Relation { + fn def(&self) -> RelationDef { + match self { + Self::Fruit => Entity::has_many(super::fruit::Entity).into(), + } + } +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Fruit.def() + } +} + +impl Related for Entity { + fn to() -> RelationDef { + super::cake_filling::Relation::Filling.def() + } + fn via() -> Option { + Some(super::cake_filling::Relation::Cake.def().rev()) + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/sea-orm-codegen/tests/expanded_with_column_derives/cake_none.rs b/sea-orm-codegen/tests/expanded_with_column_derives/cake_none.rs new file mode 100644 index 000000000..a540fad13 --- /dev/null +++ b/sea-orm-codegen/tests/expanded_with_column_derives/cake_none.rs @@ -0,0 +1,77 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +use sea_orm::entity::prelude:: * ; + +#[derive(Copy, Clone, Default, Debug, DeriveEntity)] +pub struct Entity; + +impl EntityName for Entity { + fn table_name(&self) -> &str { + "cake" + } +} + +#[derive(Clone, Debug, PartialEq, DeriveModel, DeriveActiveModel, Eq)] +pub struct Model { + pub id: i32, + pub name: Option , +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)] +pub enum Column { + Id, + Name, +} + +#[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)] +pub enum PrimaryKey { + Id, +} + +impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + + fn auto_increment() -> bool { + true + } +} + +#[derive(Copy, Clone, Debug, EnumIter)] +pub enum Relation { + Fruit, +} + +impl ColumnTrait for Column { + type EntityName = Entity; + fn def(&self) -> ColumnDef { + match self { + Self::Id => ColumnType::Integer.def(), + Self::Name => ColumnType::Text.def().null(), + } + } +} + +impl RelationTrait for Relation { + fn def(&self) -> RelationDef { + match self { + Self::Fruit => Entity::has_many(super::fruit::Entity).into(), + } + } +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Fruit.def() + } +} + +impl Related for Entity { + fn to() -> RelationDef { + super::cake_filling::Relation::Filling.def() + } + fn via() -> Option { + Some(super::cake_filling::Relation::Cake.def().rev()) + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/sea-orm-codegen/tests/expanded_with_column_derives/cake_one.rs b/sea-orm-codegen/tests/expanded_with_column_derives/cake_one.rs new file mode 100644 index 000000000..657daa312 --- /dev/null +++ b/sea-orm-codegen/tests/expanded_with_column_derives/cake_one.rs @@ -0,0 +1,77 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +use sea_orm::entity::prelude:: * ; + +#[derive(Copy, Clone, Default, Debug, DeriveEntity)] +pub struct Entity; + +impl EntityName for Entity { + fn table_name(&self) -> &str { + "cake" + } +} + +#[derive(Clone, Debug, PartialEq, DeriveModel, DeriveActiveModel, Eq)] +pub struct Model { + pub id: i32, + pub name: Option , +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveColumn, async_graphql::Enum)] +pub enum Column { + Id, + Name, +} + +#[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)] +pub enum PrimaryKey { + Id, +} + +impl PrimaryKeyTrait for PrimaryKey { + type ValueType = i32; + + fn auto_increment() -> bool { + true + } +} + +#[derive(Copy, Clone, Debug, EnumIter)] +pub enum Relation { + Fruit, +} + +impl ColumnTrait for Column { + type EntityName = Entity; + fn def(&self) -> ColumnDef { + match self { + Self::Id => ColumnType::Integer.def(), + Self::Name => ColumnType::Text.def().null(), + } + } +} + +impl RelationTrait for Relation { + fn def(&self) -> RelationDef { + match self { + Self::Fruit => Entity::has_many(super::fruit::Entity).into(), + } + } +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Fruit.def() + } +} + +impl Related for Entity { + fn to() -> RelationDef { + super::cake_filling::Relation::Filling.def() + } + fn via() -> Option { + Some(super::cake_filling::Relation::Cake.def().rev()) + } +} + +impl ActiveModelBehavior for ActiveModel {}