use super::*;
use crate::{Entity, Relation, RelationType};
use heck::ToSnakeCase;
use sea_query::ForeignKeyAction;
fn relation_column_name(a: &syn::Ident) -> String {
let a = a.to_string();
let b = a.to_snake_case();
if a != b.to_upper_camel_case() {
a
} else {
b
}
}
fn relation_column_list(punctuated: Vec<String>) -> String {
let len = punctuated.len();
let punctuated = punctuated.join(", ");
match len {
0..=1 => punctuated,
_ => format!("({punctuated})"),
}
}
fn foreign_key_action_attr(
attr: TokenStream,
action: Option<&ForeignKeyAction>,
) -> Option<TokenStream> {
action.map(|action| {
let action = Relation::get_foreign_key_action(action);
quote!(, #attr = #action)
})
}
struct DenseRelationField<'a> {
entity: &'a Entity,
rel: &'a Relation,
via_entities: &'a [syn::Ident],
}
impl DenseRelationField<'_> {
fn field_tokens(&self) -> Option<TokenStream> {
let (field, target_entity) = if self.rel.self_referencing {
let table_name = self.entity.get_table_name_snake_case_ident();
let suffix = if self.rel.num_suffix > 0 {
format!("_{}", self.rel.num_suffix)
} else {
String::new()
};
let field = format_ident!("{table_name}{suffix}");
(field, quote!(Entity))
} else {
if !self.rel.impl_related {
return None;
}
let to_entity = self.rel.get_module_name()?;
if self.via_entities.contains(&to_entity) {
return None;
}
let field = match self.rel.rel_type {
RelationType::HasMany => {
let to_entity = to_entity.to_string();
let pluralized = pluralizer::pluralize(&to_entity, 2, false);
format_ident!("{pluralized}")
}
RelationType::HasOne | RelationType::BelongsTo => to_entity.clone(),
};
let field = if self.rel.num_suffix == 0 {
field
} else {
format_ident!("{field}_{}", self.rel.num_suffix)
};
(field, quote!(super::#to_entity::Entity))
};
let rel_field_type = match self.rel.rel_type {
RelationType::BelongsTo => {
let is_optional = !self.rel.columns.is_empty()
&& self.rel.columns.iter().all(|name| {
self.entity
.columns
.iter()
.find(|column| column.name == *name)
.is_some_and(|column| column.not_null)
});
if is_optional {
quote!(BelongsTo<#target_entity>)
} else {
quote!(BelongsTo<Option<#target_entity>>)
}
}
RelationType::HasOne => quote!(HasOne<#target_entity>),
RelationType::HasMany => quote!(HasMany<#target_entity>),
};
let sea_orm_attr = if self.rel.self_referencing {
let (from, to) = self.rel.get_src_ref_columns(
relation_column_name,
relation_column_name,
relation_column_list,
);
let on_update = foreign_key_action_attr(quote!(on_update), self.rel.on_update.as_ref());
let on_delete = foreign_key_action_attr(quote!(on_delete), self.rel.on_delete.as_ref());
let relation_enum = self.rel.get_enum_name().to_string();
quote!(#[sea_orm(self_ref, relation_enum = #relation_enum, from = #from, to = #to #on_update #on_delete)])
} else {
match self.rel.rel_type {
RelationType::HasOne => quote!(#[sea_orm(has_one)]),
RelationType::HasMany => quote!(#[sea_orm(has_many)]),
RelationType::BelongsTo => {
let (from, to) = self.rel.get_src_ref_columns(
relation_column_name,
relation_column_name,
relation_column_list,
);
let on_update =
foreign_key_action_attr(quote!(on_update), self.rel.on_update.as_ref());
let on_delete =
foreign_key_action_attr(quote!(on_delete), self.rel.on_delete.as_ref());
let relation_enum = if self.rel.num_suffix > 0 {
let relation_enum = self.rel.get_enum_name().to_string();
Some(quote!(relation_enum = #relation_enum,))
} else {
None
};
quote!(#[sea_orm(belongs_to, #relation_enum from = #from, to = #to #on_update #on_delete)])
}
}
};
Some(quote! {
#sea_orm_attr
pub #field: #rel_field_type
})
}
}
impl EntityWriter {
#[allow(clippy::too_many_arguments)]
pub fn gen_dense_code_blocks(
entity: &Entity,
with_serde: &WithSerde,
column_option: &ColumnOption,
schema_name: &Option<String>,
serde_skip_deserializing_primary_key: bool,
serde_skip_hidden_column: bool,
model_extra_derives: &TokenStream,
model_extra_attributes: &TokenStream,
_column_extra_derives: &TokenStream,
_seaography: bool,
impl_active_model_behavior: bool,
) -> Vec<TokenStream> {
let mut imports = Self::gen_import(with_serde);
let active_enums = Self::gen_import_active_enum(entity);
imports.extend(active_enums.imports);
let mut code_blocks = vec![
imports,
Self::gen_dense_model_struct(
entity,
with_serde,
column_option,
schema_name,
serde_skip_deserializing_primary_key,
serde_skip_hidden_column,
model_extra_derives,
model_extra_attributes,
&active_enums.type_idents,
),
];
if impl_active_model_behavior {
code_blocks.push(Self::impl_active_model_behavior());
}
code_blocks
}
#[allow(clippy::too_many_arguments)]
pub fn gen_dense_model_struct(
entity: &Entity,
with_serde: &WithSerde,
column_option: &ColumnOption,
schema_name: &Option<String>,
serde_skip_deserializing_primary_key: bool,
serde_skip_hidden_column: bool,
model_extra_derives: &TokenStream,
model_extra_attributes: &TokenStream,
active_enum_type_idents: &ActiveEnumTypeIdents,
) -> TokenStream {
let table_name = entity.table_name.as_str();
let column_names_snake_case = entity.get_column_names_snake_case();
let column_rs_types = Self::get_column_rs_types_with_enum_idents(
entity,
column_option,
active_enum_type_idents,
);
let if_eq_needed = entity.get_eq_needed();
let primary_keys: Vec<String> = entity
.primary_keys
.iter()
.map(|pk| pk.name.clone())
.collect();
let attrs: Vec<TokenStream> = entity
.columns
.iter()
.map(|col| {
let mut attrs: Punctuated<_, Comma> = Punctuated::new();
let is_primary_key = primary_keys.contains(&col.name);
if !col.is_snake_case_name() {
let column_name = &col.name;
attrs.push(quote! { column_name = #column_name });
}
if is_primary_key {
attrs.push(quote! { primary_key });
if !col.auto_increment {
attrs.push(quote! { auto_increment = false });
}
}
if let Some(ts) = col.get_col_type_attrs() {
attrs.extend([ts]);
if !col.not_null {
attrs.push(quote! { nullable });
}
};
if col.unique {
attrs.push(quote! { unique });
} else if let Some(unique_key) = &col.unique_key {
attrs.push(quote! { unique_key = #unique_key });
}
let mut ts = quote! {};
if !attrs.is_empty() {
for (i, attr) in attrs.into_iter().enumerate() {
if i > 0 {
ts = quote! { #ts, };
}
ts = quote! { #ts #attr };
}
ts = quote! { #[sea_orm(#ts)] };
}
let serde_attribute = col.get_serde_attribute(
is_primary_key,
serde_skip_deserializing_primary_key,
serde_skip_hidden_column,
);
ts = quote! {
#ts
#serde_attribute
};
ts
})
.collect();
let schema_name = match Self::gen_schema_name(schema_name) {
Some(schema_name) => quote! {
schema_name = #schema_name,
},
None => quote! {},
};
let extra_derive = with_serde.extra_derive();
let mut compound_objects: Punctuated<_, Comma> = Punctuated::new();
let via_entities = entity.get_conjunct_relations_via_snake_case();
for rel in entity.relations.iter() {
let relation_field = DenseRelationField {
entity,
rel,
via_entities: &via_entities,
};
if let Some(field) = relation_field.field_tokens() {
compound_objects.push(field);
}
}
for (to_entity, via_entity) in entity
.get_conjunct_relations_to_snake_case()
.into_iter()
.zip(via_entities)
{
let field = format_ident!(
"{}",
pluralizer::pluralize(&to_entity.to_string(), 2, false)
);
let via_entity = via_entity.to_string();
compound_objects.push(quote! {
#[sea_orm(has_many, via = #via_entity)]
pub #field: HasMany<super::#to_entity::Entity>
});
}
if !compound_objects.is_empty() {
compound_objects.push_punct(<syn::Token![,]>::default());
}
quote! {
#[sea_orm::model]
#[derive(Clone, Debug, PartialEq #if_eq_needed, DeriveEntityModel #extra_derive #model_extra_derives)]
#[sea_orm(
#schema_name
table_name = #table_name
)]
#model_extra_attributes
pub struct Model {
#(
#attrs
pub #column_names_snake_case: #column_rs_types,
)*
#compound_objects
}
}
}
#[allow(dead_code)]
fn gen_dense_related_entity(entity: &Entity) -> TokenStream {
let via_entities = entity.get_conjunct_relations_via_snake_case();
let related_modules = entity.get_related_entity_modules();
let related_attrs = entity.get_related_entity_attrs();
let related_enum_names = entity.get_related_entity_enum_name();
let items: Vec<_> = related_modules
.into_iter()
.zip(related_attrs)
.zip(related_enum_names)
.filter_map(|((related_module, related_attr), related_enum_name)| {
if !via_entities.contains(&related_module) {
Some(quote!(#related_attr #related_enum_name))
} else {
None
}
})
.collect();
quote! {
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelatedEntity)]
pub enum RelatedEntity {
#(#items),*
}
}
}
}
#[cfg(test)]
mod test {
#[test]
#[ignore]
fn test_name() {
panic!("{}", pluralizer::pluralize("filling", 2, false));
}
}