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

Fix: generate relation for related entity of composite foreign key #1693

Merged
merged 2 commits into from
Sep 21, 2023
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.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
93 changes: 75 additions & 18 deletions sea-orm-codegen/src/entity/relation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use heck::{ToSnakeCase, ToUpperCamelCase};
use proc_macro2::{Ident, TokenStream};
use quote::{format_ident, quote};
use sea_query::{ForeignKeyAction, TableForeignKey};
use syn::{punctuated::Punctuated, token::Comma};

use crate::util::escape_rust_keyword;

Expand Down Expand Up @@ -66,17 +67,27 @@ impl Relation {
}
}
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 }
let map_src_column = |src_column: &Ident| {
quote! { Column::#src_column }
};
let map_ref_column = |ref_column: &Ident| {
if module_name.is_some() {
quote! { super::#module_name::Column::#ref_column }
} else {
quote! { Column::#ref_column }
}
};
let map_punctuated =
|punctuated: Punctuated<TokenStream, Comma>| match punctuated.len() {
0..=1 => quote! { #punctuated },
_ => quote! { (#punctuated) },
};
let (from, to) =
self.get_src_ref_columns(map_src_column, map_ref_column, map_punctuated);
quote! {
Entity::#rel_type(#ref_entity)
.from(Column::#column_camel_case)
.to(#to_col)
.from(#from)
.to(#to)
.into()
}
}
Expand All @@ -98,10 +109,20 @@ impl Relation {
}
}
RelationType::BelongsTo => {
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!("{module_name}Column::{ref_column_camel_case}");
let map_src_column = |src_column: &Ident| format!("Column::{src_column}");
let map_ref_column =
|ref_column: &Ident| format!("{module_name}Column::{ref_column}");
let map_punctuated = |punctuated: Vec<String>| {
let len = punctuated.len();
let punctuated = punctuated.join(", ");
match len {
0..=1 => punctuated,
_ => format!("({})", punctuated),
}
};
let (from, to) =
self.get_src_ref_columns(map_src_column, map_ref_column, map_punctuated);

let on_update = if let Some(action) = &self.on_update {
let action = Self::get_foreign_key_action(action);
quote! {
Expand Down Expand Up @@ -139,12 +160,18 @@ impl Relation {
}
}

pub fn get_column_camel_case(&self) -> Ident {
format_ident!("{}", self.columns[0].to_upper_camel_case())
pub fn get_column_camel_case(&self) -> Vec<Ident> {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ident -> Vec<Ident>

self.columns
.iter()
.map(|col| format_ident!("{}", col.to_upper_camel_case()))
.collect()
}

pub fn get_ref_column_camel_case(&self) -> Ident {
format_ident!("{}", self.ref_columns[0].to_upper_camel_case())
pub fn get_ref_column_camel_case(&self) -> Vec<Ident> {
self.ref_columns
.iter()
.map(|col| format_ident!("{}", col.to_upper_camel_case()))
.collect()
}

pub fn get_foreign_key_action(action: &ForeignKeyAction) -> String {
Expand All @@ -157,6 +184,36 @@ impl Relation {
}
.to_owned()
}

pub fn get_src_ref_columns<F1, F2, F3, T, I>(
&self,
map_src_column: F1,
map_ref_column: F2,
map_punctuated: F3,
) -> (T, T)
where
F1: Fn(&Ident) -> T,
F2: Fn(&Ident) -> T,
F3: Fn(I) -> T,
I: Extend<T> + Default,
{
let from: I =
self.get_column_camel_case()
.iter()
.fold(I::default(), |mut acc, src_column| {
acc.extend([map_src_column(src_column)]);
acc
});
let to: I =
self.get_ref_column_camel_case()
.iter()
.fold(I::default(), |mut acc, ref_column| {
acc.extend([map_ref_column(ref_column)]);
acc
});

(map_punctuated(from), map_punctuated(to))
}
}

impl From<&TableForeignKey> for Relation {
Expand Down Expand Up @@ -278,7 +335,7 @@ mod tests {
let relations = setup();
let cols = vec!["Id", "FillingId", "FillingId"];
for (rel, col) in relations.into_iter().zip(cols) {
assert_eq!(rel.get_column_camel_case(), col);
assert_eq!(rel.get_column_camel_case(), [col]);
}
}

Expand All @@ -287,7 +344,7 @@ mod tests {
let relations = setup();
let ref_cols = vec!["CakeId", "Id", "Id"];
for (rel, ref_col) in relations.into_iter().zip(ref_cols) {
assert_eq!(rel.get_ref_column_camel_case(), ref_col);
assert_eq!(rel.get_ref_column_camel_case(), [ref_col]);
}
}
}
58 changes: 54 additions & 4 deletions sea-orm-codegen/src/entity/writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -908,6 +908,52 @@ mod tests {
},
],
},
Entity {
table_name: "cake_filling_price".to_owned(),
columns: vec![
Column {
name: "cake_id".to_owned(),
col_type: ColumnType::Integer,
auto_increment: false,
not_null: true,
unique: false,
},
Column {
name: "filling_id".to_owned(),
col_type: ColumnType::Integer,
auto_increment: false,
not_null: true,
unique: false,
},
Column {
name: "price".to_owned(),
col_type: ColumnType::Decimal(None),
auto_increment: false,
not_null: true,
unique: false,
},
],
relations: vec![Relation {
ref_table: "cake_filling".to_owned(),
columns: vec!["cake_id".to_owned(), "filling_id".to_owned()],
ref_columns: vec!["cake_id".to_owned(), "filling_id".to_owned()],
rel_type: RelationType::BelongsTo,
on_delete: None,
on_update: None,
self_referencing: false,
num_suffix: 0,
impl_related: true,
}],
conjunct_relations: vec![],
primary_keys: vec![
PrimaryKey {
name: "cake_id".to_owned(),
},
PrimaryKey {
name: "filling_id".to_owned(),
},
],
},
Entity {
table_name: "filling".to_owned(),
columns: vec![
Expand Down Expand Up @@ -1361,9 +1407,10 @@ mod tests {
#[test]
fn test_gen_expanded_code_blocks() -> io::Result<()> {
let entities = setup();
const ENTITY_FILES: [&str; 10] = [
const ENTITY_FILES: [&str; 11] = [
include_str!("../../tests/expanded/cake.rs"),
include_str!("../../tests/expanded/cake_filling.rs"),
include_str!("../../tests/expanded/cake_filling_price.rs"),
include_str!("../../tests/expanded/filling.rs"),
include_str!("../../tests/expanded/fruit.rs"),
include_str!("../../tests/expanded/vendor.rs"),
Expand All @@ -1373,9 +1420,10 @@ mod tests {
include_str!("../../tests/expanded/collection.rs"),
include_str!("../../tests/expanded/collection_float.rs"),
];
const ENTITY_FILES_WITH_SCHEMA_NAME: [&str; 10] = [
const ENTITY_FILES_WITH_SCHEMA_NAME: [&str; 11] = [
include_str!("../../tests/expanded_with_schema_name/cake.rs"),
include_str!("../../tests/expanded_with_schema_name/cake_filling.rs"),
include_str!("../../tests/expanded_with_schema_name/cake_filling_price.rs"),
include_str!("../../tests/expanded_with_schema_name/filling.rs"),
include_str!("../../tests/expanded_with_schema_name/fruit.rs"),
include_str!("../../tests/expanded_with_schema_name/vendor.rs"),
Expand Down Expand Up @@ -1460,9 +1508,10 @@ mod tests {
#[test]
fn test_gen_compact_code_blocks() -> io::Result<()> {
let entities = setup();
const ENTITY_FILES: [&str; 10] = [
const ENTITY_FILES: [&str; 11] = [
include_str!("../../tests/compact/cake.rs"),
include_str!("../../tests/compact/cake_filling.rs"),
include_str!("../../tests/compact/cake_filling_price.rs"),
include_str!("../../tests/compact/filling.rs"),
include_str!("../../tests/compact/fruit.rs"),
include_str!("../../tests/compact/vendor.rs"),
Expand All @@ -1472,9 +1521,10 @@ mod tests {
include_str!("../../tests/compact/collection.rs"),
include_str!("../../tests/compact/collection_float.rs"),
];
const ENTITY_FILES_WITH_SCHEMA_NAME: [&str; 10] = [
const ENTITY_FILES_WITH_SCHEMA_NAME: [&str; 11] = [
include_str!("../../tests/compact_with_schema_name/cake.rs"),
include_str!("../../tests/compact_with_schema_name/cake_filling.rs"),
include_str!("../../tests/compact_with_schema_name/cake_filling_price.rs"),
include_str!("../../tests/compact_with_schema_name/filling.rs"),
include_str!("../../tests/compact_with_schema_name/fruit.rs"),
include_str!("../../tests/compact_with_schema_name/vendor.rs"),
Expand Down
31 changes: 31 additions & 0 deletions sea-orm-codegen/tests/compact/cake_filling_price.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0

use sea_orm::entity::prelude::*;

#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
#[sea_orm(table_name = "cake_filling_price")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub cake_id: i32,
#[sea_orm(primary_key, auto_increment = false)]
pub filling_id: i32,
pub price: Decimal,
}

#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::cake_filling::Entity",
from = "(Column::CakeId, Column::FillingId)",
to = "(super::cake_filling::Column::CakeId, super::cake_filling::Column::FillingId)",
)]
CakeFilling,
}

impl Related<super::cake_filling::Entity> for Entity {
fn to() -> RelationDef {
Relation::CakeFilling.def()
}
}

impl ActiveModelBehavior for ActiveModel {}
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0

use sea_orm::entity::prelude::*;

#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq)]
#[sea_orm(schema_name = "schema_name", table_name = "cake_filling_price")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
pub cake_id: i32,
#[sea_orm(primary_key, auto_increment = false)]
pub filling_id: i32,
pub price: Decimal,
}

#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::cake_filling::Entity",
from = "(Column::CakeId, Column::FillingId)",
to = "(super::cake_filling::Column::CakeId, super::cake_filling::Column::FillingId)",
)]
CakeFilling,
}

impl Related<super::cake_filling::Entity> for Entity {
fn to() -> RelationDef {
Relation::CakeFilling.def()
}
}

impl ActiveModelBehavior for ActiveModel {}
79 changes: 79 additions & 0 deletions sea-orm-codegen/tests/expanded/cake_filling_price.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
//! SeaORM Entity. Generated by sea-orm-codegen 0.1.0

use sea_orm::entity::prelude::*;

#[derive(Copy, Clone, Default, Debug, DeriveEntity)]
pub struct Entity;

impl EntityName for Entity {
fn table_name(&self) -> &str {
"cake_filling_price"
}
}

#[derive(Clone, Debug, PartialEq, DeriveModel, DeriveActiveModel, Eq)]
pub struct Model {
pub cake_id: i32,
pub filling_id: i32,
pub price: Decimal,
}

#[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
pub enum Column {
CakeId,
FillingId,
Price,
}

#[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)]
pub enum PrimaryKey {
CakeId,
FillingId,
}

impl PrimaryKeyTrait for PrimaryKey {
type ValueType = (i32, i32);

fn auto_increment() -> bool {
false
}
}

#[derive(Copy, Clone, Debug, EnumIter)]
pub enum Relation {
CakeFilling,
}

impl ColumnTrait for Column {
type EntityName = Entity;

fn def(&self) -> ColumnDef {
match self {
Self::CakeId => ColumnType::Integer.def(),
Self::FillingId => ColumnType::Integer.def(),
Self::Price => ColumnType::Decimal(None).def(),
}
}
}

impl RelationTrait for Relation {
fn def(&self) -> RelationDef {
match self {
Self::CakeFilling => Entity::belongs_to(super::cake_filling::Entity)
.from((Column::CakeId, Column::FillingId))
.to((
super::cake_filling::Column::CakeId,
super::cake_filling::Column::FillingId
))
.into(),
}
}
}

impl Related<super::cake_filling::Entity> for Entity {
fn to() -> RelationDef {
Relation::CakeFilling.def()
}
}

impl ActiveModelBehavior for ActiveModel {}