Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 3 commits into from
Dec 4, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 14 additions & 74 deletions sea-orm-codegen/src/entity/base_entity.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}

Expand All @@ -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)
Expand Down Expand Up @@ -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(),
Expand All @@ -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![],
Expand Down Expand Up @@ -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());
}
}

Expand All @@ -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();
Expand Down
74 changes: 54 additions & 20 deletions sea-orm-codegen/src/entity/relation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
}
Expand All @@ -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! {
Expand All @@ -65,25 +94,22 @@ 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);
quote! {
on_delete = #action,
}
} else {
TokenStream::new()
quote! {}
};
quote! {
#[sea_orm(
Expand Down Expand Up @@ -144,6 +170,8 @@ impl From<&TableForeignKey> for Relation {
rel_type,
on_delete,
on_update,
self_referencing: false,
num_suffix: 0,
}
}
}
Expand All @@ -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(),
Expand All @@ -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(),
Expand All @@ -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);
}
}

Expand Down
42 changes: 39 additions & 3 deletions sea-orm-codegen/src/entity/transformer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
Expand Down
Loading