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

Merge branch 'codegen-amendment'

Browse Source
This commit is contained in:
Chris Tsang 2021-06-26 22:08:01 +08:00
commit c36b94592d
6 changed files with 85 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
examples/cli/src/entity/fruit.rs
View File

@ -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(),
}
}
}

24
sea-orm-codegen/src/entity/column.rs
View File

@ -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<bool> = 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,
}
}
}

4
sea-orm-codegen/src/entity/entity.rs
View File

@ -86,6 +86,10 @@ impl Entity {
.collect()
}
pub fn get_relation_defs(&self) -> Vec<TokenStream> {
self.relations.iter().map(|rel| rel.get_def()).collect()
}
pub fn get_relation_rel_types(&self) -> Vec<Ident> {
self.relations
.iter()

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

@ -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,

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

@ -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<Column> = table_create
.get_columns()
.iter()
.map(|col_def| col_def.into())
.collect();
let unique_columns: Vec<String> = 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 {

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

@ -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! {