use super::case_style::{CaseStyle, CaseStyleHelpers};
use super::util::{escape_rust_keyword, trim_starting_raw_identifier};
use heck::{ToSnakeCase, ToUpperCamelCase};
use proc_macro2::{Ident, Span, TokenStream};
use quote::quote;
use syn::{
punctuated::Punctuated, spanned::Spanned, token::Comma, Attribute, Data, Expr, Fields, Lit,
};
pub fn expand_derive_entity_model(data: Data, attrs: Vec<Attribute>) -> syn::Result<TokenStream> {
let mut table_name = None;
let mut comment = quote! {None};
let mut schema_name = quote! { None };
let mut table_iden = false;
let mut rename_all: Option<CaseStyle> = None;
attrs
.iter()
.filter(|attr| attr.path().is_ident("sea_orm"))
.try_for_each(|attr| {
attr.parse_nested_meta(|meta| {
if meta.path.is_ident("comment") {
let name: Lit = meta.value()?.parse()?;
comment = quote! { Some(#name) };
} else if meta.path.is_ident("table_name") {
table_name = Some(meta.value()?.parse::<Lit>()?);
} else if meta.path.is_ident("schema_name") {
let name: Lit = meta.value()?.parse()?;
schema_name = quote! { Some(#name) };
} else if meta.path.is_ident("table_iden") {
table_iden = true;
} else if meta.path.is_ident("rename_all") {
rename_all = Some((&meta).try_into()?);
} else {
let _: Option<Expr> = meta.value().and_then(|v| v.parse()).ok();
}
Ok(())
})
})?;
let entity_def = table_name
.as_ref()
.map(|table_name| {
quote! {
#[doc = " Generated by sea-orm-macros"]
#[derive(Copy, Clone, Default, Debug, sea_orm::prelude::DeriveEntity)]
pub struct Entity;
#[automatically_derived]
impl sea_orm::prelude::EntityName for Entity {
fn schema_name(&self) -> Option<&str> {
#schema_name
}
fn table_name(&self) -> &str {
#table_name
}
fn comment(&self) -> Option<&str> {
#comment
}
}
}
})
.unwrap_or_default();
let mut columns_enum: Punctuated<_, Comma> = Punctuated::new();
let mut columns_trait: Punctuated<_, Comma> = Punctuated::new();
let mut columns_enum_type_name: Punctuated<_, Comma> = Punctuated::new();
let mut columns_select_as: Punctuated<_, Comma> = Punctuated::new();
let mut columns_save_as: Punctuated<_, Comma> = Punctuated::new();
let mut primary_keys: Punctuated<_, Comma> = Punctuated::new();
let mut primary_key_types: Punctuated<_, Comma> = Punctuated::new();
let mut auto_increment = true;
if table_iden {
if let Some(table_name) = table_name {
let table_field_name = Ident::new("Table", Span::call_site());
columns_enum.push(quote! {
#[doc = " Generated by sea-orm-macros"]
#[sea_orm(table_name=#table_name)]
#[strum(disabled)]
#table_field_name
});
columns_trait.push(
quote! { Self::#table_field_name => panic!("Table cannot be used as a column") },
);
}
}
if let Data::Struct(item_struct) = data {
if let Fields::Named(fields) = item_struct.fields {
for field in fields.named {
if let Some(ident) = &field.ident {
let original_field_name = trim_starting_raw_identifier(ident);
let mut field_name = Ident::new(
&original_field_name.to_upper_camel_case(),
Span::call_site(),
);
let mut nullable = false;
let mut default_value = None;
let mut comment = None;
let mut default_expr = None;
let mut select_as = None;
let mut save_as = None;
let mut indexed = false;
let mut ignore = false;
let mut unique = false;
let mut sql_type = None;
let mut column_name = if let Some(case_style) = rename_all {
Some(field_name.convert_case(Some(case_style)))
} else if original_field_name
!= original_field_name.to_upper_camel_case().to_snake_case()
{
Some(original_field_name.to_snake_case())
} else {
None
};
let mut enum_name = None;
let mut is_primary_key = false;
for attr in field.attrs.iter() {
if !attr.path().is_ident("sea_orm") {
continue;
}
attr.parse_nested_meta(|meta| {
if meta.path.is_ident("column_type") {
let lit = meta.value()?.parse()?;
if let Lit::Str(litstr) = lit {
let ty: TokenStream = syn::parse_str(&litstr.value())?;
sql_type = Some(ty);
} else {
return Err(meta.error(format!("Invalid column_type {lit:?}")));
}
} else if meta.path.is_ident("auto_increment") {
let lit = meta.value()?.parse()?;
if let Lit::Bool(litbool) = lit {
auto_increment = litbool.value();
} else {
return Err(
meta.error(format!("Invalid auto_increment = {lit:?}"))
);
}
} else if meta.path.is_ident("comment") {
comment = Some(meta.value()?.parse::<Lit>()?);
} else if meta.path.is_ident("default_value") {
default_value = Some(meta.value()?.parse::<Lit>()?);
} else if meta.path.is_ident("default_expr") {
let lit = meta.value()?.parse()?;
if let Lit::Str(litstr) = lit {
let value_expr: TokenStream = syn::parse_str(&litstr.value())?;
default_expr = Some(value_expr);
} else {
return Err(meta.error(format!("Invalid column_type {lit:?}")));
}
} else if meta.path.is_ident("column_name") {
let lit = meta.value()?.parse()?;
if let Lit::Str(litstr) = lit {
column_name = Some(litstr.value());
} else {
return Err(meta.error(format!("Invalid column_name {lit:?}")));
}
} else if meta.path.is_ident("enum_name") {
let lit = meta.value()?.parse()?;
if let Lit::Str(litstr) = lit {
let ty: Ident = syn::parse_str(&litstr.value())?;
enum_name = Some(ty);
} else {
return Err(meta.error(format!("Invalid enum_name {lit:?}")));
}
} else if meta.path.is_ident("select_as") {
let lit = meta.value()?.parse()?;
if let Lit::Str(litstr) = lit {
select_as = Some(litstr.value());
} else {
return Err(meta.error(format!("Invalid select_as {lit:?}")));
}
} else if meta.path.is_ident("save_as") {
let lit = meta.value()?.parse()?;
if let Lit::Str(litstr) = lit {
save_as = Some(litstr.value());
} else {
return Err(meta.error(format!("Invalid save_as {lit:?}")));
}
} else if meta.path.is_ident("ignore") {
ignore = true;
} else if meta.path.is_ident("primary_key") {
is_primary_key = true;
primary_key_types.push(field.ty.clone());
} else if meta.path.is_ident("nullable") {
nullable = true;
} else if meta.path.is_ident("indexed") {
indexed = true;
} else if meta.path.is_ident("unique") {
unique = true;
} else {
let _: Option<Expr> = meta.value().and_then(|v| v.parse()).ok();
}
Ok(())
})?;
}
if let Some(enum_name) = enum_name {
field_name = enum_name;
}
field_name = Ident::new(&escape_rust_keyword(field_name), Span::call_site());
let variant_attrs = match &column_name {
Some(column_name) => quote! {
#[sea_orm(column_name = #column_name)]
#[doc = " Generated by sea-orm-macros"]
},
None => quote! {
#[doc = " Generated by sea-orm-macros"]
},
};
if ignore {
continue;
} else {
columns_enum.push(quote! {
#variant_attrs
#field_name
});
}
if is_primary_key {
primary_keys.push(quote! {
#variant_attrs
#field_name
});
}
if let Some(select_as) = select_as {
columns_select_as.push(quote! {
Self::#field_name => expr.cast_as(#select_as)
});
}
if let Some(save_as) = save_as {
columns_save_as.push(quote! {
Self::#field_name => val.cast_as(#save_as)
});
}
let field_type = &field.ty;
let field_type = quote! { #field_type }
.to_string() .replace(' ', ""); let field_type = if field_type.starts_with("Option<") {
nullable = true;
&field_type[7..(field_type.len() - 1)] } else {
field_type.as_str()
};
let field_span = field.span();
let sea_query_col_type = crate::derives::sql_type_match::col_type_match(
sql_type, field_type, field_span,
);
let col_def =
quote! { sea_orm::prelude::ColumnTypeTrait::def(#sea_query_col_type) };
let mut match_row = quote! { Self::#field_name => #col_def };
if nullable {
match_row = quote! { #match_row.nullable() };
}
if indexed {
match_row = quote! { #match_row.indexed() };
}
if unique {
match_row = quote! { #match_row.unique() };
}
if let Some(default_value) = default_value {
match_row = quote! { #match_row.default_value(#default_value) };
}
if let Some(comment) = comment {
match_row = quote! { #match_row.comment(#comment) };
}
if let Some(default_expr) = default_expr {
match_row = quote! { #match_row.default(#default_expr) };
}
columns_trait.push(match_row);
let ty: syn::Type = syn::LitStr::new(field_type, field_span)
.parse()
.expect("field type error");
let enum_type_name = quote::quote_spanned! { field_span =>
<#ty as sea_orm::sea_query::ValueType>::enum_type_name()
};
columns_enum_type_name.push(quote! {
Self::#field_name => #enum_type_name
});
}
}
}
}
if !columns_select_as.is_empty() {
columns_select_as.push_punct(Comma::default());
}
if !columns_save_as.is_empty() {
columns_save_as.push_punct(Comma::default());
}
let primary_key = {
let auto_increment = auto_increment && primary_keys.len() == 1;
let primary_key_types = if primary_key_types.len() == 1 {
let first = primary_key_types.first();
quote! { #first }
} else {
quote! { (#primary_key_types) }
};
quote! {
#[doc = " Generated by sea-orm-macros"]
#[derive(Copy, Clone, Debug, EnumIter, DerivePrimaryKey)]
pub enum PrimaryKey {
#primary_keys
}
#[automatically_derived]
impl PrimaryKeyTrait for PrimaryKey {
type ValueType = #primary_key_types;
fn auto_increment() -> bool {
#auto_increment
}
}
}
};
Ok(quote! {
#[doc = " Generated by sea-orm-macros"]
#[derive(Copy, Clone, Debug, sea_orm::prelude::EnumIter, sea_orm::prelude::DeriveColumn)]
pub enum Column {
#columns_enum
}
#[automatically_derived]
impl sea_orm::prelude::ColumnTrait for Column {
type EntityName = Entity;
fn def(&self) -> sea_orm::prelude::ColumnDef {
match self {
#columns_trait
}
}
fn enum_type_name(&self) -> Option<&'static str> {
match self {
#columns_enum_type_name
}
}
fn select_as(&self, expr: sea_orm::sea_query::Expr) -> sea_orm::sea_query::SimpleExpr {
match self {
#columns_select_as
_ => sea_orm::prelude::ColumnTrait::select_enum_as(self, expr),
}
}
fn save_as(&self, val: sea_orm::sea_query::Expr) -> sea_orm::sea_query::SimpleExpr {
match self {
#columns_save_as
_ => sea_orm::prelude::ColumnTrait::save_enum_as(self, val),
}
}
}
#entity_def
#primary_key
})
}