diff --git a/examples/cli/src/entity/fruit.rs b/examples/cli/src/entity/fruit.rs index d3f21bf9..c9b6d27a 100644 --- a/examples/cli/src/entity/fruit.rs +++ b/examples/cli/src/entity/fruit.rs @@ -48,7 +48,7 @@ impl ColumnTrait for Column { match self { Self::Id => ColumnType::Integer.def(), Self::Name => ColumnType::String(Some(255u32)).def(), - Self::CakeId => ColumnType::Integer.def(), + Self::CakeId => ColumnType::Integer.def().null(), } } } diff --git a/sea-orm-codegen/src/entity/column.rs b/sea-orm-codegen/src/entity/column.rs index b575057f..e8a0b543 100644 --- a/sea-orm-codegen/src/entity/column.rs +++ b/sea-orm-codegen/src/entity/column.rs @@ -9,6 +9,7 @@ pub struct Column { pub(crate) col_type: ColumnType, pub(crate) auto_increment: bool, pub(crate) not_null: bool, + pub(crate) unique: bool, } impl Column { @@ -50,7 +51,7 @@ impl Column { } pub fn get_def(&self) -> TokenStream { - match &self.col_type { + let mut col_def = match &self.col_type { ColumnType::Char(s) => match s { Some(s) => quote! { ColumnType::Char(Some(#s)).def() }, None => quote! { ColumnType::Char(None).def() }, @@ -86,7 +87,18 @@ impl Column { let s = s.to_string(); quote! { ColumnType::Custom(#s.to_owned()).def() } } + }; + if !self.not_null { + col_def.extend(quote! { + .null() + }); } + if self.unique { + col_def.extend(quote! { + .unique() + }); + } + col_def } } @@ -115,11 +127,21 @@ impl From<&ColumnDef> for Column { }) .collect(); let not_null = !not_nulls.is_empty(); + let uniques: Vec = col_def + .get_column_spec() + .iter() + .filter_map(|spec| match spec { + ColumnSpec::UniqueKey => Some(true), + _ => None, + }) + .collect(); + let unique = !uniques.is_empty(); Self { name, col_type, auto_increment, not_null, + unique, } } } 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! {