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 amendment #30

Merged
merged 2 commits into from
Jun 26, 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
10 changes: 2 additions & 8 deletions examples/codegen/src/out/cake.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions examples/codegen/src/out/cake_filling.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
5 changes: 1 addition & 4 deletions examples/codegen/src/out/filling.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions examples/codegen/src/out/fruit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,15 +47,15 @@ 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(),
}
}
}

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(),
Expand Down
24 changes: 23 additions & 1 deletion sea-orm-codegen/src/entity/column.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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() },
Expand Down Expand Up @@ -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
}
}

Expand Down Expand Up @@ -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 changes: 4 additions & 0 deletions sea-orm-codegen/src/entity/entity.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
37 changes: 33 additions & 4 deletions sea-orm-codegen/src/entity/relation.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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"),
}
}

Expand All @@ -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,
Expand Down
27 changes: 22 additions & 5 deletions sea-orm-codegen/src/entity/transformer.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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 {
Expand Down
10 changes: 2 additions & 8 deletions sea-orm-codegen/src/entity/writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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! {
Expand Down