caisin/sea-orm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Codegen Handle Self Referencing & Multiple Relations to the Same Related Entity (#347)

Browse Source 
* [sea-orm-codegen] handle self referencing relation & multiple relations to the same related entity

* Test [cli]

* Test [cli]
This commit is contained in:
Billy Chan 2021-12-04 21:22:10 +08:00 committed by GitHub
parent 1229287fd8
commit ccb8b95324
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 324 additions and 124 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

88
sea-orm-codegen/src/entity/base_entity.rs
View File

@ -73,17 +73,17 @@ impl Entity {
.collect()
}
pub fn get_relation_ref_tables_snake_case(&self) -> Vec<Ident> {
pub fn get_relation_module_name(&self) -> Vec<Option<Ident>> {
self.relations
.iter()
.map(|rel| rel.get_ref_table_snake_case())
.map(|rel| rel.get_module_name())
.collect()
}
pub fn get_relation_ref_tables_camel_case(&self) -> Vec<Ident> {
pub fn get_relation_enum_name(&self) -> Vec<Ident> {
self.relations
.iter()
.map(|rel| rel.get_ref_table_camel_case())
.map(|rel| rel.get_enum_name())
.collect()
}
@ -95,27 +95,6 @@ impl Entity {
self.relations.iter().map(|rel| rel.get_attrs()).collect()
}
pub fn get_relation_rel_types(&self) -> Vec<Ident> {
self.relations
.iter()
.map(|rel| rel.get_rel_type())
.collect()
}
pub fn get_relation_columns_camel_case(&self) -> Vec<Ident> {
self.relations
.iter()
.map(|rel| rel.get_column_camel_case())
.collect()
}
pub fn get_relation_ref_columns_camel_case(&self) -> Vec<Ident> {
self.relations
.iter()
.map(|rel| rel.get_ref_column_camel_case())
.collect()
}
pub fn get_primary_key_auto_increment(&self) -> Ident {
let auto_increment = self.columns.iter().any(|col| col.auto_increment);
format_ident!("{}", auto_increment)
@ -201,6 +180,8 @@ mod tests {
rel_type: RelationType::HasOne,
on_delete: Some(ForeignKeyAction::Cascade),
on_update: Some(ForeignKeyAction::Cascade),
self_referencing: false,
num_suffix: 0,
},
Relation {
ref_table: "filling".to_owned(),
@ -209,6 +190,8 @@ mod tests {
rel_type: RelationType::HasOne,
on_delete: Some(ForeignKeyAction::Cascade),
on_update: Some(ForeignKeyAction::Cascade),
self_referencing: false,
num_suffix: 0,
},
],
conjunct_relations: vec![],
@ -321,28 +304,20 @@ mod tests {
}
#[test]
fn test_get_relation_ref_tables_snake_case() {
fn test_get_relation_module_name() {
let entity = setup();
for (i, elem) in entity
.get_relation_ref_tables_snake_case()
.into_iter()
.enumerate()
{
assert_eq!(elem, entity.relations[i].get_ref_table_snake_case());
for (i, elem) in entity.get_relation_module_name().into_iter().enumerate() {
assert_eq!(elem, entity.relations[i].get_module_name());
}
}
#[test]
fn test_get_relation_ref_tables_camel_case() {
fn test_get_relation_enum_name() {
let entity = setup();
for (i, elem) in entity
.get_relation_ref_tables_camel_case()
.into_iter()
.enumerate()
{
assert_eq!(elem, entity.relations[i].get_ref_table_camel_case());
for (i, elem) in entity.get_relation_enum_name().into_iter().enumerate() {
assert_eq!(elem, entity.relations[i].get_enum_name());
}
}
@ -367,41 +342,6 @@ mod tests {
}
}
#[test]
fn test_get_relation_rel_types() {
let entity = setup();
for (i, elem) in entity.get_relation_rel_types().into_iter().enumerate() {
assert_eq!(elem, entity.relations[i].get_rel_type());
}
}
#[test]
fn test_get_relation_columns_camel_case() {
let entity = setup();
for (i, elem) in entity
.get_relation_columns_camel_case()
.into_iter()
.enumerate()
{
assert_eq!(elem, entity.relations[i].get_column_camel_case());
}
}
#[test]
fn test_get_relation_ref_columns_camel_case() {
let entity = setup();
for (i, elem) in entity
.get_relation_ref_columns_camel_case()
.into_iter()
.enumerate()
{
assert_eq!(elem, entity.relations[i].get_ref_column_camel_case());
}
}
#[test]
fn test_get_primary_key_auto_increment() {
let mut entity = setup();

74
sea-orm-codegen/src/entity/relation.rs
View File

@ -18,33 +18,58 @@ pub struct Relation {
pub(crate) rel_type: RelationType,
pub(crate) on_update: Option<ForeignKeyAction>,
pub(crate) on_delete: Option<ForeignKeyAction>,
pub(crate) self_referencing: bool,
pub(crate) num_suffix: usize,
}
impl Relation {
pub fn get_ref_table_snake_case(&self) -> Ident {
format_ident!("{}", self.ref_table.to_snake_case())
pub fn get_enum_name(&self) -> Ident {
let name = if self.self_referencing {
format_ident!("SelfRef")
} else {
format_ident!("{}", self.ref_table.to_camel_case())
};
if self.num_suffix > 0 {
format_ident!("{}{}", name, self.num_suffix)
} else {
name
}
}
pub fn get_ref_table_camel_case(&self) -> Ident {
format_ident!("{}", self.ref_table.to_camel_case())
pub fn get_module_name(&self) -> Option<Ident> {
if self.self_referencing {
None
} else {
Some(format_ident!("{}", self.ref_table.to_snake_case()))
}
}
pub fn get_def(&self) -> TokenStream {
let rel_type = self.get_rel_type();
let ref_table_snake_case = self.get_ref_table_snake_case();
let module_name = self.get_module_name();
let ref_entity = if module_name.is_some() {
quote! { super::#module_name::Entity }
} else {
quote! { Entity }
};
match self.rel_type {
RelationType::HasOne | RelationType::HasMany => {
quote! {
Entity::#rel_type(super::#ref_table_snake_case::Entity).into()
Entity::#rel_type(#ref_entity).into()
}
}
RelationType::BelongsTo => {
let column_camel_case = self.get_column_camel_case();
let ref_column_camel_case = self.get_ref_column_camel_case();
let to_col = if module_name.is_some() {
quote! { super::#module_name::Column::#ref_column_camel_case }
} else {
quote! { Column::#ref_column_camel_case }
};
quote! {
Entity::#rel_type(super::#ref_table_snake_case::Entity)
Entity::#rel_type(#ref_entity)
.from(Column::#column_camel_case)
.to(super::#ref_table_snake_case::Column::#ref_column_camel_case)
.to(#to_col)
.into()
}
}
@ -53,8 +78,12 @@ 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);
let module_name = if let Some(module_name) = self.get_module_name() {
format!("super::{}::", module_name)
} else {
format!("")
};
let ref_entity = format!("{}Entity", module_name);
match self.rel_type {
RelationType::HasOne | RelationType::HasMany => {
quote! {
@ -65,17 +94,14 @@ impl Relation {
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 to = format!("{}Column::{}", module_name, 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()
quote! {}
};
let on_delete = if let Some(action) = &self.on_delete {
let action = Self::get_foreign_key_action(action);
@ -83,7 +109,7 @@ impl Relation {
on_delete = #action,
}
} else {
TokenStream::new()
quote! {}
};
quote! {
#[sea_orm(
@ -144,6 +170,8 @@ impl From<&TableForeignKey> for Relation {
rel_type,
on_delete,
on_update,
self_referencing: false,
num_suffix: 0,
}
}
}
@ -163,6 +191,8 @@ mod tests {
rel_type: RelationType::HasOne,
on_delete: None,
on_update: None,
self_referencing: false,
num_suffix: 0,
},
Relation {
ref_table: "filling".to_owned(),
@ -171,6 +201,8 @@ mod tests {
rel_type: RelationType::BelongsTo,
on_delete: Some(ForeignKeyAction::Cascade),
on_update: Some(ForeignKeyAction::Cascade),
self_referencing: false,
num_suffix: 0,
},
Relation {
ref_table: "filling".to_owned(),
@ -179,25 +211,27 @@ mod tests {
rel_type: RelationType::HasMany,
on_delete: Some(ForeignKeyAction::Cascade),
on_update: None,
self_referencing: false,
num_suffix: 0,
},
]
}
#[test]
fn test_get_ref_table_snake_case() {
fn test_get_module_name() {
let relations = setup();
let snake_cases = vec!["fruit", "filling", "filling"];
for (rel, snake_case) in relations.into_iter().zip(snake_cases) {
assert_eq!(rel.get_ref_table_snake_case().to_string(), snake_case);
assert_eq!(rel.get_module_name().unwrap().to_string(), snake_case);
}
}
#[test]
fn test_get_ref_table_camel_case() {
fn test_get_enum_name() {
let relations = setup();
let camel_cases = vec!["Fruit", "Filling", "Filling"];
for (rel, camel_case) in relations.into_iter().zip(camel_cases) {
assert_eq!(rel.get_ref_table_camel_case().to_string(), camel_case);
assert_eq!(rel.get_enum_name().to_string(), camel_case);
}
}

42
sea-orm-codegen/src/entity/transformer.rs
View File

@ -45,11 +45,38 @@ impl EntityTransformer {
col
})
.collect();
let relations = table_create
let mut ref_table_counts: HashMap<String, usize> = HashMap::new();
let relations: Vec<Relation> = table_create
.get_foreign_key_create_stmts()
.iter()
.map(|fk_create_stmt| fk_create_stmt.get_foreign_key())
.map(|tbl_fk| tbl_fk.into());
.map(|tbl_fk| {
let ref_tbl = tbl_fk.get_ref_table().unwrap();
if let Some(count) = ref_table_counts.get_mut(&ref_tbl) {
if *count == 0 {
*count = 1;
}
*count += 1;
} else {
ref_table_counts.insert(ref_tbl, 0);
};
tbl_fk.into()
})
.collect::<Vec<_>>()
.into_iter()
.rev()
.map(|mut rel: Relation| {
rel.self_referencing = rel.ref_table == table_name;
if let Some(count) = ref_table_counts.get_mut(&rel.ref_table) {
rel.num_suffix = *count;
if *count > 0 {
*count -= 1;
}
}
rel
})
.rev()
.collect();
let primary_keys = table_create
.get_indexes()
.iter()
@ -67,12 +94,21 @@ impl EntityTransformer {
let entity = Entity {
table_name: table_name.clone(),
columns,
relations: relations.clone().collect(),
relations: relations.clone(),
conjunct_relations: vec![],
primary_keys,
};
entities.insert(table_name.clone(), entity.clone());
for (i, mut rel) in relations.into_iter().enumerate() {
// This will produce a duplicated relation
if rel.self_referencing {
continue;
}
// This will cause compile error on the many side,
// got relation variant but without Related<T> implemented
if rel.num_suffix > 0 {
continue;
}
let is_conjunct_relation = entity.primary_keys.len() == entity.columns.len()
&& rel.columns.len() == 2
&& rel.ref_columns.len() == 2

146
sea-orm-codegen/src/entity/writer.rs
View File

@ -186,7 +186,7 @@ impl EntityWriter {
Self::gen_import(with_serde),
Self::gen_compact_model_struct(entity, with_serde),
];
let relation_defs = if entity.get_relation_ref_tables_camel_case().is_empty() {
let relation_defs = if entity.get_relation_enum_name().is_empty() {
vec![
Self::gen_relation_enum(entity),
Self::gen_impl_relation_trait(entity),
@ -298,11 +298,11 @@ impl EntityWriter {
}
pub fn gen_relation_enum(entity: &Entity) -> TokenStream {
let relation_ref_tables_camel_case = entity.get_relation_ref_tables_camel_case();
let relation_enum_name = entity.get_relation_enum_name();
quote! {
#[derive(Copy, Clone, Debug, EnumIter)]
pub enum Relation {
#(#relation_ref_tables_camel_case,)*
#(#relation_enum_name,)*
}
}
}
@ -324,16 +324,16 @@ impl EntityWriter {
}
pub fn gen_impl_relation_trait(entity: &Entity) -> TokenStream {
let relation_ref_tables_camel_case = entity.get_relation_ref_tables_camel_case();
let relation_enum_name = entity.get_relation_enum_name();
let relation_defs = entity.get_relation_defs();
let quoted = if relation_ref_tables_camel_case.is_empty() {
let quoted = if relation_enum_name.is_empty() {
quote! {
panic!("No RelationDef")
}
} else {
quote! {
match self {
#(Self::#relation_ref_tables_camel_case => #relation_defs,)*
#(Self::#relation_enum_name => #relation_defs,)*
}
}
};
@ -347,17 +347,25 @@ impl EntityWriter {
}
pub fn gen_impl_related(entity: &Entity) -> Vec<TokenStream> {
let camel = entity.get_relation_ref_tables_camel_case();
let snake = entity.get_relation_ref_tables_snake_case();
camel
.into_iter()
.zip(snake)
.map(|(c, s)| {
quote! {
impl Related<super::#s::Entity> for Entity {
fn to() -> RelationDef {
Relation::#c.def()
}
entity
.relations
.iter()
.filter(|rel| !rel.self_referencing && rel.num_suffix == 0)
.map(|rel| {
let enum_name = rel.get_enum_name();
let module_name = rel.get_module_name();
let inner = quote! {
fn to() -> RelationDef {
Relation::#enum_name.def()
}
};
if module_name.is_some() {
quote! {
impl Related<super::#module_name::Entity> for Entity { #inner }
}
} else {
quote! {
impl Related<Entity> for Entity { #inner }
}
}
})
@ -471,14 +479,14 @@ impl EntityWriter {
}
pub fn gen_compact_relation_enum(entity: &Entity) -> TokenStream {
let relation_ref_tables_camel_case = entity.get_relation_ref_tables_camel_case();
let relation_enum_name = entity.get_relation_enum_name();
let attrs = entity.get_relation_attrs();
quote! {
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#(
#attrs
#relation_ref_tables_camel_case,
#relation_enum_name,
)*
}
}
@ -523,6 +531,8 @@ mod tests {
rel_type: RelationType::HasMany,
on_delete: None,
on_update: None,
self_referencing: false,
num_suffix: 0,
}],
conjunct_relations: vec![ConjunctRelation {
via: "cake_filling".to_owned(),
@ -558,6 +568,8 @@ mod tests {
rel_type: RelationType::BelongsTo,
on_delete: Some(ForeignKeyAction::Cascade),
on_update: Some(ForeignKeyAction::Cascade),
self_referencing: false,
num_suffix: 0,
},
Relation {
ref_table: "filling".to_owned(),
@ -566,6 +578,8 @@ mod tests {
rel_type: RelationType::BelongsTo,
on_delete: Some(ForeignKeyAction::Cascade),
on_update: Some(ForeignKeyAction::Cascade),
self_referencing: false,
num_suffix: 0,
},
],
conjunct_relations: vec![],
@ -638,6 +652,8 @@ mod tests {
rel_type: RelationType::BelongsTo,
on_delete: None,
on_update: None,
self_referencing: false,
num_suffix: 0,
},
Relation {
ref_table: "vendor".to_owned(),
@ -646,6 +662,8 @@ mod tests {
rel_type: RelationType::HasMany,
on_delete: None,
on_update: None,
self_referencing: false,
num_suffix: 0,
},
],
conjunct_relations: vec![],
@ -685,6 +703,8 @@ mod tests {
rel_type: RelationType::BelongsTo,
on_delete: None,
on_update: None,
self_referencing: false,
num_suffix: 0,
}],
conjunct_relations: vec![],
primary_keys: vec![PrimaryKey {
@ -750,8 +770,94 @@ mod tests {
not_null: true,
unique: false,
},
Column {
name: "self_id1".to_owned(),
col_type: ColumnType::Integer(Some(11)),
auto_increment: false,
not_null: true,
unique: false,
},
Column {
name: "self_id2".to_owned(),
col_type: ColumnType::Integer(Some(11)),
auto_increment: false,
not_null: true,
unique: false,
},
Column {
name: "fruit_id1".to_owned(),
col_type: ColumnType::Integer(Some(11)),
auto_increment: false,
not_null: true,
unique: false,
},
Column {
name: "fruit_id2".to_owned(),
col_type: ColumnType::Integer(Some(11)),
auto_increment: false,
not_null: true,
unique: false,
},
Column {
name: "cake_id".to_owned(),
col_type: ColumnType::Integer(Some(11)),
auto_increment: false,
not_null: true,
unique: false,
},
],
relations: vec![
Relation {
ref_table: "rust_keyword".to_owned(),
columns: vec!["self_id1".to_owned()],
ref_columns: vec!["id".to_owned()],
rel_type: RelationType::BelongsTo,
on_delete: None,
on_update: None,
self_referencing: true,
num_suffix: 1,
},
Relation {
ref_table: "rust_keyword".to_owned(),
columns: vec!["self_id2".to_owned()],
ref_columns: vec!["id".to_owned()],
rel_type: RelationType::BelongsTo,
on_delete: None,
on_update: None,
self_referencing: true,
num_suffix: 2,
},
Relation {
ref_table: "fruit".to_owned(),
columns: vec!["fruit_id1".to_owned()],
ref_columns: vec!["id".to_owned()],
rel_type: RelationType::BelongsTo,
on_delete: None,
on_update: None,
self_referencing: false,
num_suffix: 1,
},
Relation {
ref_table: "fruit".to_owned(),
columns: vec!["fruit_id2".to_owned()],
ref_columns: vec!["id".to_owned()],
rel_type: RelationType::BelongsTo,
on_delete: None,
on_update: None,
self_referencing: false,
num_suffix: 2,
},
Relation {
ref_table: "cake".to_owned(),
columns: vec!["cake_id".to_owned()],
ref_columns: vec!["id".to_owned()],
rel_type: RelationType::BelongsTo,
on_delete: None,
on_update: None,
self_referencing: false,
num_suffix: 0,
},
],
relations: vec![],
conjunct_relations: vec![],
primary_keys: vec![PrimaryKey {
name: "id".to_owned(),

46
sea-orm-codegen/tests/compact/rust_keyword.rs
View File

@ -14,14 +14,50 @@ pub struct Model {
pub r#typeof: i32,
pub crate_: i32,
pub self_: i32,
pub self_id1: i32,
pub self_id2: i32,
pub fruit_id1: i32,
pub fruit_id2: i32,
pub cake_id: i32,
}
#[derive(Copy, Clone, Debug, EnumIter)]
pub enum Relation {}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "Entity",
from = "Column::SelfId1",
to = "Column::Id",
)]
SelfRef1,
#[sea_orm(
belongs_to = "Entity",
from = "Column::SelfId2",
to = "Column::Id",
)]
SelfRef2,
#[sea_orm(
belongs_to = "super::fruit::Entity",
from = "Column::FruitId1",
to = "super::fruit::Column::Id",
)]
Fruit1,
#[sea_orm(
belongs_to = "super::fruit::Entity",
from = "Column::FruitId2",
to = "super::fruit::Column::Id",
)]
Fruit2,
#[sea_orm(
belongs_to = "super::cake::Entity",
from = "Column::CakeId",
to = "super::cake::Column::Id",
)]
Cake,
}
impl RelationTrait for Relation {
fn def(&self) -> RelationDef {
panic!("No RelationDef")
impl Related<super::cake::Entity> for Entity {
fn to() -> RelationDef {
Relation::Cake.def()
}
}

52
sea-orm-codegen/tests/expanded/rust_keyword.rs
View File

@ -21,6 +21,11 @@ pub struct Model {
pub r#typeof: i32,
pub crate_: i32,
pub self_: i32,
pub self_id1: i32,
pub self_id2: i32,
pub fruit_id1: i32,
pub fruit_id2: i32,
pub cake_id: i32,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
@ -33,6 +38,11 @@ pub enum Column {
Typeof,
Crate,
Self_,
SelfId1,
SelfId2,
FruitId1,
FruitId2,
CakeId,
}
#[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)]
@ -49,7 +59,13 @@ impl PrimaryKeyTrait for PrimaryKey {
}
#[derive(Copy, Clone, Debug, EnumIter)]
pub enum Relation {}
pub enum Relation {
SelfRef1,
SelfRef2,
Fruit1,
Fruit2,
Cake,
}
impl ColumnTrait for Column {
type EntityName = Entity;
@ -64,13 +80,45 @@ impl ColumnTrait for Column {
Self::Typeof => ColumnType::Integer.def(),
Self::Crate => ColumnType::Integer.def(),
Self::Self_ => ColumnType::Integer.def(),
Self::SelfId1 => ColumnType::Integer.def(),
Self::SelfId2 => ColumnType::Integer.def(),
Self::FruitId1 => ColumnType::Integer.def(),
Self::FruitId2 => ColumnType::Integer.def(),
Self::CakeId => ColumnType::Integer.def(),
}
}
}
impl RelationTrait for Relation {
fn def(&self) -> RelationDef {
panic!("No RelationDef")
match self {
Self::SelfRef1 => Entity::belongs_to(Entity)
.from(Column::SelfId1)
.to(Column::Id)
.into(),
Self::SelfRef2 => Entity::belongs_to(Entity)
.from(Column::SelfId2)
.to(Column::Id)
.into(),
Self::Fruit1 => Entity::belongs_to(super::fruit::Entity)
.from(Column::FruitId1)
.to(super::fruit::Column::Id)
.into(),
Self::Fruit2 => Entity::belongs_to(super::fruit::Entity)
.from(Column::FruitId2)
.to(super::fruit::Column::Id)
.into(),
Self::Cake => Entity::belongs_to(super::cake::Entity)
.from(Column::CakeId)
.to(super::cake::Column::Id)
.into(),
}
}
}
impl Related<super::cake::Entity> for Entity {
fn to() -> RelationDef {
Relation::Cake.def()
}
}