diff --git a/examples/codegen/src/out/cake.rs b/examples/codegen/src/out/cake.rs index 87ee5f1a..9b786a5f 100644 --- a/examples/codegen/src/out/cake.rs +++ b/examples/codegen/src/out/cake.rs @@ -53,14 +53,8 @@ impl ColumnTrait for Column { impl RelationTrait for Relation { fn def(&self) -> RelationDef { match self { - Self::CakeFilling => Entity::has_many(super::cake_filling::Entity) - .from(Column::Id) - .to(super::cake_filling::Column::CakeId) - .into(), - Self::Fruit => Entity::has_many(super::fruit::Entity) - .from(Column::Id) - .to(super::fruit::Column::CakeId) - .into(), + Self::CakeFilling => Entity::has_many(super::cake_filling::Entity).into(), + Self::Fruit => Entity::has_many(super::fruit::Entity).into(), } } } diff --git a/examples/codegen/src/out/cake_filling.rs b/examples/codegen/src/out/cake_filling.rs index b35279d4..d5b4b8b6 100644 --- a/examples/codegen/src/out/cake_filling.rs +++ b/examples/codegen/src/out/cake_filling.rs @@ -54,11 +54,11 @@ impl ColumnTrait for Column { impl RelationTrait for Relation { fn def(&self) -> RelationDef { match self { - Self::Cake => Entity::has_one(super::cake::Entity) + Self::Cake => Entity::belongs_to(super::cake::Entity) .from(Column::CakeId) .to(super::cake::Column::Id) .into(), - Self::Filling => Entity::has_one(super::filling::Entity) + Self::Filling => Entity::belongs_to(super::filling::Entity) .from(Column::FillingId) .to(super::filling::Column::Id) .into(), diff --git a/examples/codegen/src/out/filling.rs b/examples/codegen/src/out/filling.rs index 4134652f..e4563e55 100644 --- a/examples/codegen/src/out/filling.rs +++ b/examples/codegen/src/out/filling.rs @@ -52,10 +52,7 @@ impl ColumnTrait for Column { impl RelationTrait for Relation { fn def(&self) -> RelationDef { match self { - Self::CakeFilling => Entity::has_many(super::cake_filling::Entity) - .from(Column::Id) - .to(super::cake_filling::Column::FillingId) - .into(), + Self::CakeFilling => Entity::has_many(super::cake_filling::Entity).into(), } } } diff --git a/examples/codegen/src/out/fruit.rs b/examples/codegen/src/out/fruit.rs index 17f653ed..f65d5b38 100644 --- a/examples/codegen/src/out/fruit.rs +++ b/examples/codegen/src/out/fruit.rs @@ -55,7 +55,7 @@ impl ColumnTrait for Column { impl RelationTrait for Relation { fn def(&self) -> RelationDef { match self { - Self::Cake => Entity::has_one(super::cake::Entity) + Self::Cake => Entity::belongs_to(super::cake::Entity) .from(Column::CakeId) .to(super::cake::Column::Id) .into(), diff --git a/sea-orm-codegen/src/entity/entity.rs b/sea-orm-codegen/src/entity/entity.rs index 788002ce..2abb4a7a 100644 --- a/sea-orm-codegen/src/entity/entity.rs +++ b/sea-orm-codegen/src/entity/entity.rs @@ -86,6 +86,10 @@ impl Entity { .collect() } + pub fn get_relation_defs(&self) -> Vec { + self.relations.iter().map(|rel| rel.get_def()).collect() + } + pub fn get_relation_rel_types(&self) -> Vec { self.relations .iter() diff --git a/sea-orm-codegen/src/entity/relation.rs b/sea-orm-codegen/src/entity/relation.rs index bbcfd155..8eab372d 100644 --- a/sea-orm-codegen/src/entity/relation.rs +++ b/sea-orm-codegen/src/entity/relation.rs @@ -1,9 +1,15 @@ use heck::{CamelCase, SnakeCase}; -use proc_macro2::Ident; -use quote::format_ident; -use sea_orm::RelationType; +use proc_macro2::{Ident, TokenStream}; +use quote::{format_ident, quote}; use sea_query::TableForeignKey; +#[derive(Clone, Debug)] +pub enum RelationType { + HasOne, + HasMany, + BelongsTo, +} + #[derive(Clone, Debug)] pub struct Relation { pub(crate) ref_table: String, @@ -21,10 +27,33 @@ impl Relation { format_ident!("{}", self.ref_table.to_camel_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(); + match self.rel_type { + RelationType::HasOne | RelationType::HasMany => { + quote! { + Entity::#rel_type(super::#ref_table_snake_case::Entity).into() + } + } + RelationType::BelongsTo => { + let column_camel_case = self.get_column_camel_case(); + let ref_column_camel_case = self.get_ref_column_camel_case(); + quote! { + Entity::#rel_type(super::#ref_table_snake_case::Entity) + .from(Column::#column_camel_case) + .to(super::#ref_table_snake_case::Column::#ref_column_camel_case) + .into() + } + } + } + } + pub fn get_rel_type(&self) -> Ident { match self.rel_type { RelationType::HasOne => format_ident!("has_one"), RelationType::HasMany => format_ident!("has_many"), + RelationType::BelongsTo => format_ident!("belongs_to"), } } @@ -49,7 +78,7 @@ impl From<&TableForeignKey> for Relation { }; let columns = tbl_fk.get_columns(); let ref_columns = tbl_fk.get_ref_columns(); - let rel_type = RelationType::HasOne; + let rel_type = RelationType::BelongsTo; Self { ref_table, columns, diff --git a/sea-orm-codegen/src/entity/transformer.rs b/sea-orm-codegen/src/entity/transformer.rs index d437d0f3..02c7fb82 100644 --- a/sea-orm-codegen/src/entity/transformer.rs +++ b/sea-orm-codegen/src/entity/transformer.rs @@ -1,5 +1,4 @@ -use crate::{Entity, EntityWriter, Error, PrimaryKey, Relation}; -use sea_orm::RelationType; +use crate::{Column, Entity, EntityWriter, Error, PrimaryKey, Relation, RelationType}; use sea_query::TableStatement; use sea_schema::mysql::def::Schema; use std::{collections::HashMap, mem::swap}; @@ -31,11 +30,16 @@ impl EntityTransformer { )) } }; - let columns = table_create + let columns: Vec = table_create .get_columns() .iter() .map(|col_def| col_def.into()) .collect(); + let unique_columns: Vec = columns + .iter() + .filter(|col| col.unique) + .map(|col| col.name.clone()) + .collect(); let relations = table_create .get_foreign_key_create_stmts() .iter() @@ -64,9 +68,22 @@ impl EntityTransformer { entities.push(entity); for mut rel in relations.into_iter() { let ref_table = rel.ref_table; - swap(&mut rel.columns, &mut rel.ref_columns); - rel.rel_type = RelationType::HasMany; + let mut unique = true; + for col in rel.columns.iter() { + if !unique_columns.contains(col) { + unique = false; + break; + } + } + let rel_type = if unique { + RelationType::HasOne + } else { + RelationType::HasMany + }; + rel.rel_type = rel_type; rel.ref_table = table_name.clone(); + rel.columns = Vec::new(); + rel.ref_columns = Vec::new(); if let Some(vec) = inverse_relations.get_mut(&ref_table) { vec.push(rel); } else { diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index 73e12d07..01dd4a4c 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/sea-orm-codegen/src/entity/writer.rs @@ -217,20 +217,14 @@ 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_rel_types = entity.get_relation_rel_types(); - let relation_ref_tables_snake_case = entity.get_relation_ref_tables_snake_case(); - let relation_columns_camel_case = entity.get_relation_columns_camel_case(); - let relation_ref_columns_camel_case = entity.get_relation_ref_columns_camel_case(); + let relation_defs = entity.get_relation_defs(); let quoted = if relation_ref_tables_camel_case.is_empty() { quote! { _ => panic!("No RelationDef"), } } else { quote! { - #(Self::#relation_ref_tables_camel_case => Entity::#relation_rel_types(super::#relation_ref_tables_snake_case::Entity) - .from(Column::#relation_columns_camel_case) - .to(super::#relation_ref_tables_snake_case::Column::#relation_ref_columns_camel_case) - .into()),* + #(Self::#relation_ref_tables_camel_case => #relation_defs),* } }; quote! {