diff --git a/sea-orm-codegen/Cargo.toml b/sea-orm-codegen/Cargo.toml index 56cd72ae..7e564b0d 100644 --- a/sea-orm-codegen/Cargo.toml +++ b/sea-orm-codegen/Cargo.toml @@ -15,7 +15,7 @@ name = "sea_orm_codegen" path = "src/lib.rs" [dependencies] -sea-query = { version = "^0.15" } +sea-query = { version = "0.16.1", git = "https://github.com/SeaQL/sea-query.git", branch = "foreign-key-getters" } syn = { version = "^1", default-features = false, features = [ "derive", "parsing", @@ -25,3 +25,6 @@ syn = { version = "^1", default-features = false, features = [ quote = "^1" heck = "^0.3" proc-macro2 = "^1" + +[dev-dependencies] +pretty_assertions = { version = "^0.7" } diff --git a/sea-orm-codegen/src/entity/base_entity.rs b/sea-orm-codegen/src/entity/base_entity.rs index 7b28f9e4..02f9ee21 100644 --- a/sea-orm-codegen/src/entity/base_entity.rs +++ b/sea-orm-codegen/src/entity/base_entity.rs @@ -91,6 +91,10 @@ impl Entity { self.relations.iter().map(|rel| rel.get_def()).collect() } + pub fn get_relation_attrs(&self) -> Vec { + self.relations.iter().map(|rel| rel.get_attrs()).collect() + } + pub fn get_relation_rel_types(&self) -> Vec { self.relations .iter() @@ -168,7 +172,7 @@ impl Entity { mod tests { use crate::{Column, Entity, PrimaryKey, Relation, RelationType}; use quote::format_ident; - use sea_query::ColumnType; + use sea_query::{ColumnType, ForeignKeyAction}; fn setup() -> Entity { Entity { @@ -195,12 +199,16 @@ mod tests { columns: vec!["id".to_owned()], ref_columns: vec!["cake_id".to_owned()], rel_type: RelationType::HasOne, + on_delete: Some(ForeignKeyAction::Cascade), + on_update: Some(ForeignKeyAction::Cascade), }, Relation { ref_table: "filling".to_owned(), columns: vec!["id".to_owned()], ref_columns: vec!["cake_id".to_owned()], rel_type: RelationType::HasOne, + on_delete: Some(ForeignKeyAction::Cascade), + on_update: Some(ForeignKeyAction::Cascade), }, ], conjunct_relations: vec![], @@ -347,6 +355,18 @@ mod tests { } } + #[test] + fn test_get_relation_attrs() { + let entity = setup(); + + for (i, elem) in entity.get_relation_attrs().into_iter().enumerate() { + assert_eq!( + elem.to_string(), + entity.relations[i].get_attrs().to_string() + ); + } + } + #[test] fn test_get_relation_rel_types() { let entity = setup(); diff --git a/sea-orm-codegen/src/entity/relation.rs b/sea-orm-codegen/src/entity/relation.rs index d675137a..f1e9b441 100644 --- a/sea-orm-codegen/src/entity/relation.rs +++ b/sea-orm-codegen/src/entity/relation.rs @@ -1,7 +1,7 @@ use heck::{CamelCase, SnakeCase}; use proc_macro2::{Ident, TokenStream}; use quote::{format_ident, quote}; -use sea_query::TableForeignKey; +use sea_query::{ForeignKeyAction, TableForeignKey}; #[derive(Clone, Debug)] pub enum RelationType { @@ -16,6 +16,8 @@ pub struct Relation { pub(crate) columns: Vec, pub(crate) ref_columns: Vec, pub(crate) rel_type: RelationType, + pub(crate) on_update: Option, + pub(crate) on_delete: Option, } impl Relation { @@ -49,6 +51,53 @@ impl Relation { } } + pub fn get_attrs(&self) -> TokenStream { + let rel_type = self.get_rel_type(); + let ref_table_snake_case = self.get_ref_table_snake_case(); + let ref_entity = format!("super::{}::Entity", ref_table_snake_case); + match self.rel_type { + RelationType::HasOne | RelationType::HasMany => { + quote! { + #[sea_orm(#rel_type = #ref_entity)] + } + } + RelationType::BelongsTo => { + let column_camel_case = self.get_column_camel_case(); + let ref_column_camel_case = self.get_ref_column_camel_case(); + let from = format!("Column::{}", column_camel_case); + let to = format!( + "super::{}::Column::{}", + ref_table_snake_case, ref_column_camel_case + ); + let on_update = if let Some(action) = &self.on_update { + let action = Self::get_foreign_key_action(action); + quote! { + on_update = #action, + } + } else { + TokenStream::new() + }; + let on_delete = if let Some(action) = &self.on_delete { + let action = Self::get_foreign_key_action(action); + quote! { + on_delete = #action, + } + } else { + TokenStream::new() + }; + quote! { + #[sea_orm( + #rel_type = #ref_entity, + from = #from, + to = #to, + #on_update + #on_delete + )] + } + } + } + } + pub fn get_rel_type(&self) -> Ident { match self.rel_type { RelationType::HasOne => format_ident!("has_one"), @@ -64,6 +113,17 @@ impl Relation { pub fn get_ref_column_camel_case(&self) -> Ident { format_ident!("{}", self.ref_columns[0].to_camel_case()) } + + pub fn get_foreign_key_action(action: &ForeignKeyAction) -> String { + match action { + ForeignKeyAction::Restrict => "Restrict", + ForeignKeyAction::Cascade => "Cascade", + ForeignKeyAction::SetNull => "SetNull", + ForeignKeyAction::NoAction => "NoAction", + ForeignKeyAction::SetDefault => "SetDefault", + } + .to_owned() + } } impl From<&TableForeignKey> for Relation { @@ -75,11 +135,15 @@ impl From<&TableForeignKey> for Relation { let columns = tbl_fk.get_columns(); let ref_columns = tbl_fk.get_ref_columns(); let rel_type = RelationType::BelongsTo; + let on_delete = tbl_fk.get_on_delete(); + let on_update = tbl_fk.get_on_update(); Self { ref_table, columns, ref_columns, rel_type, + on_delete, + on_update, } } } @@ -88,6 +152,7 @@ impl From<&TableForeignKey> for Relation { mod tests { use crate::{Relation, RelationType}; use proc_macro2::TokenStream; + use sea_query::ForeignKeyAction; fn setup() -> Vec { vec![ @@ -96,18 +161,24 @@ mod tests { columns: vec!["id".to_owned()], ref_columns: vec!["cake_id".to_owned()], rel_type: RelationType::HasOne, + on_delete: None, + on_update: None, }, Relation { ref_table: "filling".to_owned(), columns: vec!["filling_id".to_owned()], ref_columns: vec!["id".to_owned()], rel_type: RelationType::BelongsTo, + on_delete: Some(ForeignKeyAction::Cascade), + on_update: Some(ForeignKeyAction::Cascade), }, Relation { ref_table: "filling".to_owned(), columns: vec!["filling_id".to_owned()], ref_columns: vec!["id".to_owned()], rel_type: RelationType::HasMany, + on_delete: Some(ForeignKeyAction::Cascade), + on_update: None, }, ] } diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index ad356872..75229835 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/sea-orm-codegen/src/entity/writer.rs @@ -31,7 +31,7 @@ impl EntityWriter { .map(|entity| { let mut lines = Vec::new(); Self::write_doc_comment(&mut lines); - let code_blocks = Self::gen_code_blocks(entity); + let code_blocks = Self::gen_expanded_code_blocks(entity); Self::write(&mut lines, code_blocks); OutputFile { name: format!("{}.rs", entity.get_table_name_snake_case()), @@ -97,7 +97,7 @@ impl EntityWriter { lines.push("".to_owned()); } - pub fn gen_code_blocks(entity: &Entity) -> Vec { + pub fn gen_expanded_code_blocks(entity: &Entity) -> Vec { let mut code_blocks = vec![ Self::gen_import(), Self::gen_entity_struct(), @@ -116,6 +116,23 @@ impl EntityWriter { code_blocks } + pub fn gen_compact_code_blocks(entity: &Entity) -> Vec { + let mut code_blocks = vec![Self::gen_import(), Self::gen_compact_model_struct(entity)]; + let relation_defs = if entity.get_relation_ref_tables_camel_case().is_empty() { + vec![ + Self::gen_relation_enum(entity), + Self::gen_impl_relation_trait(entity), + ] + } else { + vec![Self::gen_compact_relation_enum(entity)] + }; + code_blocks.extend(relation_defs); + code_blocks.extend(Self::gen_impl_related(entity)); + code_blocks.extend(Self::gen_impl_conjunct_related(entity)); + code_blocks.extend(vec![Self::gen_impl_active_model_behavior()]); + code_blocks + } + pub fn gen_import() -> TokenStream { quote! { use sea_orm::entity::prelude::*; @@ -297,6 +314,54 @@ impl EntityWriter { pub use super::#table_name_snake_case_ident::Entity as #table_name_camel_case_ident; } } + + pub fn gen_compact_model_struct(entity: &Entity) -> TokenStream { + let table_name = entity.table_name.as_str(); + let column_names_snake_case = entity.get_column_names_snake_case(); + let column_rs_types = entity.get_column_rs_types(); + let primary_keys: Vec = entity + .primary_keys + .iter() + .map(|pk| pk.name.clone()) + .collect(); + let attrs: Vec = entity + .columns + .iter() + .map(|col| { + if !primary_keys.contains(&col.name) { + TokenStream::new() + } else { + quote! { + #[sea_orm(primary_key)] + } + } + }) + .collect(); + quote! { + #[derive(Clone, Debug, PartialEq, DeriveEntityModel)] + #[sea_orm(table_name = #table_name)] + pub struct Model { + #( + #attrs + pub #column_names_snake_case: #column_rs_types, + )* + } + } + } + + pub fn gen_compact_relation_enum(entity: &Entity) -> TokenStream { + let relation_ref_tables_camel_case = entity.get_relation_ref_tables_camel_case(); + let attrs = entity.get_relation_attrs(); + quote! { + #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] + pub enum Relation { + #( + #attrs + #relation_ref_tables_camel_case, + )* + } + } + } } #[cfg(test)] @@ -304,18 +369,11 @@ mod tests { use crate::{ Column, ConjunctRelation, Entity, EntityWriter, PrimaryKey, Relation, RelationType, }; + use pretty_assertions::assert_eq; use proc_macro2::TokenStream; - use sea_query::ColumnType; + use sea_query::{ColumnType, ForeignKeyAction}; use std::io::{self, BufRead, BufReader}; - const EXPANDED_ENTITY_FILES: [&str; 5] = [ - include_str!("../../tests/expanded/cake.rs"), - include_str!("../../tests/expanded/cake_filling.rs"), - include_str!("../../tests/expanded/filling.rs"), - include_str!("../../tests/expanded/fruit.rs"), - include_str!("../../tests/expanded/vendor.rs"), - ]; - fn setup() -> Vec { vec![ Entity { @@ -341,6 +399,8 @@ mod tests { columns: vec![], ref_columns: vec![], rel_type: RelationType::HasMany, + on_delete: None, + on_update: None, }], conjunct_relations: vec![ConjunctRelation { via: "cake_filling".to_owned(), @@ -374,12 +434,16 @@ mod tests { columns: vec!["cake_id".to_owned()], ref_columns: vec!["id".to_owned()], rel_type: RelationType::BelongsTo, + on_delete: Some(ForeignKeyAction::Cascade), + on_update: Some(ForeignKeyAction::Cascade), }, Relation { ref_table: "filling".to_owned(), columns: vec!["filling_id".to_owned()], ref_columns: vec!["id".to_owned()], rel_type: RelationType::BelongsTo, + on_delete: Some(ForeignKeyAction::Cascade), + on_update: Some(ForeignKeyAction::Cascade), }, ], conjunct_relations: vec![], @@ -450,12 +514,16 @@ mod tests { columns: vec!["cake_id".to_owned()], ref_columns: vec!["id".to_owned()], rel_type: RelationType::BelongsTo, + on_delete: None, + on_update: None, }, Relation { ref_table: "vendor".to_owned(), columns: vec![], ref_columns: vec![], rel_type: RelationType::HasMany, + on_delete: None, + on_update: None, }, ], conjunct_relations: vec![], @@ -493,6 +561,8 @@ mod tests { columns: vec!["fruit_id".to_owned()], ref_columns: vec!["id".to_owned()], rel_type: RelationType::BelongsTo, + on_delete: None, + on_update: None, }], conjunct_relations: vec![], primary_keys: vec![PrimaryKey { @@ -503,13 +573,20 @@ mod tests { } #[test] - fn test_gen_code_blocks() -> io::Result<()> { + fn test_gen_expanded_code_blocks() -> io::Result<()> { let entities = setup(); + const ENTITY_FILES: [&str; 5] = [ + include_str!("../../tests/expanded/cake.rs"), + include_str!("../../tests/expanded/cake_filling.rs"), + include_str!("../../tests/expanded/filling.rs"), + include_str!("../../tests/expanded/fruit.rs"), + include_str!("../../tests/expanded/vendor.rs"), + ]; - assert_eq!(entities.len(), EXPANDED_ENTITY_FILES.len()); + assert_eq!(entities.len(), ENTITY_FILES.len()); for (i, entity) in entities.iter().enumerate() { - let mut reader = BufReader::new(EXPANDED_ENTITY_FILES[i].as_bytes()); + let mut reader = BufReader::new(ENTITY_FILES[i].as_bytes()); let mut lines: Vec = Vec::new(); reader.read_until(b';', &mut Vec::new())?; @@ -521,7 +598,46 @@ mod tests { } let content = lines.join(""); let expected: TokenStream = content.parse().unwrap(); - let generated = EntityWriter::gen_code_blocks(entity) + let generated = EntityWriter::gen_expanded_code_blocks(entity) + .into_iter() + .skip(1) + .fold(TokenStream::new(), |mut acc, tok| { + acc.extend(tok); + acc + }); + assert_eq!(expected.to_string(), generated.to_string()); + } + + Ok(()) + } + + #[test] + fn test_gen_compact_code_blocks() -> io::Result<()> { + let entities = setup(); + const ENTITY_FILES: [&str; 5] = [ + include_str!("../../tests/compact/cake.rs"), + include_str!("../../tests/compact/cake_filling.rs"), + include_str!("../../tests/compact/filling.rs"), + include_str!("../../tests/compact/fruit.rs"), + include_str!("../../tests/compact/vendor.rs"), + ]; + + assert_eq!(entities.len(), ENTITY_FILES.len()); + + for (i, entity) in entities.iter().enumerate() { + let mut reader = BufReader::new(ENTITY_FILES[i].as_bytes()); + let mut lines: Vec = Vec::new(); + + reader.read_until(b';', &mut Vec::new())?; + + let mut line = String::new(); + while reader.read_line(&mut line)? > 0 { + lines.push(line.to_owned()); + line.clear(); + } + let content = lines.join(""); + let expected: TokenStream = content.parse().unwrap(); + let generated = EntityWriter::gen_compact_code_blocks(entity) .into_iter() .skip(1) .fold(TokenStream::new(), |mut acc, tok| { diff --git a/sea-orm-codegen/tests/compact/cake.rs b/sea-orm-codegen/tests/compact/cake.rs new file mode 100644 index 00000000..4a64611e --- /dev/null +++ b/sea-orm-codegen/tests/compact/cake.rs @@ -0,0 +1,34 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel)] +#[sea_orm(table_name = "cake")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + pub name: String, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm(has_many = "super::fruit::Entity")] + Fruit, +} + +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/compact/cake_filling.rs b/sea-orm-codegen/tests/compact/cake_filling.rs new file mode 100644 index 00000000..af966c66 --- /dev/null +++ b/sea-orm-codegen/tests/compact/cake_filling.rs @@ -0,0 +1,46 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel)] +#[sea_orm(table_name = "_cake_filling_")] +pub struct Model { + #[sea_orm(primary_key)] + pub cake_id: i32, + #[sea_orm(primary_key)] + pub filling_id: i32, +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm( + belongs_to = "super::cake::Entity", + from = "Column::CakeId", + to = "super::cake::Column::Id", + on_update = "Cascade", + on_delete = "Cascade", + )] + Cake, + #[sea_orm( + belongs_to = "super::filling::Entity", + from = "Column::FillingId", + to = "super::filling::Column::Id", + on_update = "Cascade", + on_delete = "Cascade", + )] + Filling, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Cake.def() + } +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Filling.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/sea-orm-codegen/tests/compact/filling.rs b/sea-orm-codegen/tests/compact/filling.rs new file mode 100644 index 00000000..820f65c8 --- /dev/null +++ b/sea-orm-codegen/tests/compact/filling.rs @@ -0,0 +1,33 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel)] +#[sea_orm(table_name = "filling")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + pub name: String, +} + +#[derive(Copy, Clone, Debug, EnumIter)] +pub enum Relation {} + +impl RelationTrait for Relation { + fn def(&self) -> RelationDef { + match self { + _ => panic!("No RelationDef"), + } + } +} + +impl Related for Entity { + fn to() -> RelationDef { + super::cake_filling::Relation::Cake.def() + } + fn via() -> Option { + Some(super::cake_filling::Relation::Filling.def().rev()) + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/sea-orm-codegen/tests/compact/fruit.rs b/sea-orm-codegen/tests/compact/fruit.rs new file mode 100644 index 00000000..6399a51f --- /dev/null +++ b/sea-orm-codegen/tests/compact/fruit.rs @@ -0,0 +1,38 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel)] +#[sea_orm(table_name = "fruit")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + pub name: String, + pub cake_id: Option , +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm( + belongs_to = "super::cake::Entity", + from = "Column::CakeId", + to = "super::cake::Column::Id", + )] + Cake, + #[sea_orm(has_many = "super::vendor::Entity")] + Vendor, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Cake.def() + } +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Vendor.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/sea-orm-codegen/tests/compact/mod.rs b/sea-orm-codegen/tests/compact/mod.rs new file mode 100644 index 00000000..5a8c6c21 --- /dev/null +++ b/sea-orm-codegen/tests/compact/mod.rs @@ -0,0 +1,9 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +pub mod prelude; + +pub mod cake; +pub mod cake_filling; +pub mod filling; +pub mod fruit; +pub mod vendor; diff --git a/sea-orm-codegen/tests/compact/prelude.rs b/sea-orm-codegen/tests/compact/prelude.rs new file mode 100644 index 00000000..b4e85c78 --- /dev/null +++ b/sea-orm-codegen/tests/compact/prelude.rs @@ -0,0 +1,7 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +pub use super::cake::Entity as Cake; +pub use super::cake_filling::Entity as CakeFilling; +pub use super::filling::Entity as Filling; +pub use super::fruit::Entity as Fruit; +pub use super::vendor::Entity as Vendor; diff --git a/sea-orm-codegen/tests/compact/vendor.rs b/sea-orm-codegen/tests/compact/vendor.rs new file mode 100644 index 00000000..314eb396 --- /dev/null +++ b/sea-orm-codegen/tests/compact/vendor.rs @@ -0,0 +1,30 @@ +//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0 + +use sea_orm::entity::prelude::*; + +#[derive(Clone, Debug, PartialEq, DeriveEntityModel)] +#[sea_orm(table_name = "vendor")] +pub struct Model { + #[sea_orm(primary_key)] + pub id: i32, + pub name: String, + pub fruit_id: Option , +} + +#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)] +pub enum Relation { + #[sea_orm( + belongs_to = "super::fruit::Entity", + from = "Column::FruitId", + to = "super::fruit::Column::Id", + )] + Fruit, +} + +impl Related for Entity { + fn to() -> RelationDef { + Relation::Fruit.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/src/query/helper.rs b/src/query/helper.rs index 7aeb9f0b..d93595f1 100644 --- a/src/query/helper.rs +++ b/src/query/helper.rs @@ -3,7 +3,9 @@ use crate::{ PrimaryKeyToColumn, RelationDef, }; pub use sea_query::{Condition, ConditionalStatement, DynIden, JoinType, Order, OrderedStatement}; -use sea_query::{Expr, IntoCondition, LockType, SeaRc, SelectExpr, SelectStatement, SimpleExpr, TableRef}; +use sea_query::{ + Expr, IntoCondition, LockType, SeaRc, SelectExpr, SelectStatement, SimpleExpr, TableRef, +}; // LINT: when the column does not appear in tables selected from // LINT: when there is a group by clause, but some columns don't have aggregate functions