sea-orm-macros 2.0.0-rc.38

use super::case_style::{CaseStyle, CaseStyleHelpers};
use super::util::{escape_rust_keyword, trim_starting_raw_identifier};
use heck::{
    ToKebabCase, ToLowerCamelCase, ToShoutySnakeCase, ToSnakeCase, ToTitleCase, ToUpperCamelCase,
};
use proc_macro2::{Ident, Span, TokenStream};
use quote::quote;
use std::str::FromStr;
use syn::meta::ParseNestedMeta;
use syn::{
    Attribute, Data, Fields, Lit, LitStr, Visibility, punctuated::Punctuated, spanned::Spanned,
    token::Comma,
};

const NOT_AUTO_INCRE_TYPE_SUFFIX: [&str; 2] = ["String", "Uuid"];

#[allow(dead_code)]
fn convert_case(s: &str, case_style: CaseStyle) -> String {
    match case_style {
        CaseStyle::PascalCase => s.to_upper_camel_case(),
        CaseStyle::KebabCase => s.to_kebab_case(),
        CaseStyle::MixedCase => s.to_lower_camel_case(),
        CaseStyle::ShoutySnakeCase => s.to_shouty_snake_case(),
        CaseStyle::SnakeCase => s.to_snake_case(),
        CaseStyle::TitleCase => s.to_title_case(),
        CaseStyle::UpperCase => s.to_uppercase(),
        CaseStyle::LowerCase => s.to_lowercase(),
        CaseStyle::ScreamingKebabCase => s.to_kebab_case().to_uppercase(),
        CaseStyle::CamelCase => {
            let camel_case = s.to_upper_camel_case();
            let mut result = String::with_capacity(camel_case.len());
            let mut it = camel_case.chars();
            if let Some(ch) = it.next() {
                result.extend(ch.to_lowercase());
            }
            result.extend(it);
            result
        }
    }
}

#[cfg(feature = "with-json")]
fn serde_deserialize_name(
    orig: &str,
    serde_rename: Option<&str>,
    serde_rename_all: Option<CaseStyle>,
) -> String {
    if let Some(rename) = serde_rename.as_ref() {
        return rename.to_string();
    }

    if let Some(case_style) = serde_rename_all {
        convert_case(orig, case_style)
    } else {
        orig.to_string()
    }
}

fn consume_meta(meta: ParseNestedMeta<'_>) {
    let _ = meta.value().and_then(|v| v.parse::<syn::Expr>());
}

/// Method to derive an Model
pub fn expand_derive_entity_model(
    vis: &Visibility,
    data: &Data,
    attrs: &[Attribute],
) -> syn::Result<TokenStream> {
    // if #[sea_orm(table_name = "foo", schema_name = "bar")] specified, create Entity struct
    let mut table_name = None;
    let mut comment = quote! {None};
    let mut schema_name = quote! { None };
    let mut table_iden = false;
    let mut model_ex = false;
    let mut rename_all: Option<CaseStyle> = None;
    let mut serde_rename_all: Option<CaseStyle> = None;

    // Parse #[serde(rename_all = "...")] at struct level
    attrs
        .iter()
        .filter(|attr| attr.path().is_ident("serde"))
        .try_for_each(|attr| {
            attr.parse_nested_meta(|meta| {
                if meta.path.is_ident("rename_all") {
                    if let Ok(lit) = meta.value().and_then(|v| v.parse::<LitStr>()) {
                        // #[serde(rename_all = "camelCase")]
                        serde_rename_all = CaseStyle::from_str(&lit.value()).ok();
                    } else {
                        // #[serde(rename_all(serialize = "...", deserialize = "..."))]
                        meta.parse_nested_meta(|nested| {
                            if nested.path.is_ident("deserialize") {
                                let lit: LitStr = nested.value()?.parse()?;
                                serde_rename_all = CaseStyle::from_str(&lit.value()).ok();
                            } else {
                                consume_meta(nested);
                            }
                            Ok(())
                        })?;
                    }
                } else {
                    consume_meta(meta);
                }
                Ok(())
            })
        })?;

    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::<LitStr>()?);
                } 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("model_ex") {
                    model_ex = true;
                } else if meta.path.is_ident("rename_all") {
                    rename_all = Some((&meta).try_into()?);
                } else {
                    consume_meta(meta);
                }

                Ok(())
            })
        })?;

    let entity_def = table_name
        .as_ref()
        .map(|table_name| {
            let entity_extra_attr = if model_ex {
                quote!(#[sea_orm(model_ex = ModelEx, active_model_ex = ActiveModelEx)])
            } else {
                quote!()
            };
            quote! {
                #[doc = " Generated by sea-orm-macros"]
                #[derive(Copy, Clone, Default, Debug, sea_orm::prelude::DeriveEntity)]
                #entity_extra_attr
                #vis struct Entity;

                #[automatically_derived]
                impl sea_orm::prelude::EntityName for Entity {
                    fn schema_name(&self) -> Option<&str> {
                        #schema_name
                    }

                    fn table_name(&self) -> &'static str {
                        #table_name
                    }

                    fn comment(&self) -> Option<&str> {
                        #comment
                    }
                }
            }
        })
        .unwrap_or_default();

    // generate Column enum and it's ColumnTrait impl
    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: Option<bool> = None;
    #[cfg(feature = "with-json")]
    let mut columns_json_keys: Punctuated<_, Comma> = Punctuated::new();

    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(), ident.span());

                    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 unique_key = None;
                    let mut renamed_from = None;
                    let mut indexed = false;
                    let mut ignore = false;
                    let mut unique = false;
                    let mut sql_type = None;
                    let mut enum_name = None;
                    let mut is_primary_key = false;
                    let mut is_auto_increment = false;
                    let mut extra = None;
                    let mut seaography_ignore = false;
                    #[cfg(feature = "with-json")]
                    let mut serde_rename: Option<String> = 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()
                    {
                        // `to_snake_case` was used to trim prefix and tailing underscore
                        Some(original_field_name.to_snake_case())
                    } else {
                        None
                    };

                    // search for #[sea_orm(primary_key, auto_increment = false, column_type = "String(StringLen::N(255))", default_value = "new user", default_expr = "gen_random_uuid()", column_name = "name", enum_name = "Name", nullable, indexed, unique)]
                    for attr in field.attrs.iter() {
                        if attr.path().is_ident("sea_orm") {
                            // single param
                            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 {
                                        is_auto_increment = litbool.value();
                                        auto_increment = Some(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 if meta.path.is_ident("unique_key") {
                                    let lit = meta.value()?.parse()?;
                                    if let Lit::Str(litstr) = lit {
                                        unique_key = Some(litstr.value());
                                    } else {
                                        return Err(
                                            meta.error(format!("Invalid unique_key {lit:?}"))
                                        );
                                    }
                                } else if meta.path.is_ident("renamed_from") {
                                    let lit = meta.value()?.parse()?;
                                    if let Lit::Str(litstr) = lit {
                                        renamed_from = Some(litstr.value());
                                    } else {
                                        return Err(
                                            meta.error(format!("Invalid renamed_from {lit:?}"))
                                        );
                                    }
                                } else if meta.path.is_ident("extra") {
                                    let lit = meta.value()?.parse()?;
                                    if let Lit::Str(litstr) = lit {
                                        extra = Some(litstr.value());
                                    } else {
                                        return Err(meta.error(format!("Invalid extra {lit:?}")));
                                    }
                                } else {
                                    consume_meta(meta);
                                }

                                Ok(())
                            })?;
                        } else if attr.path().is_ident("seaography") {
                            attr.parse_nested_meta(|meta| {
                                if meta.path.is_ident("ignore") {
                                    seaography_ignore = true;
                                }

                                Ok(())
                            })?;
                        } else if cfg!(feature = "with-json") && attr.path().is_ident("serde") {
                            #[cfg(feature = "with-json")]
                            attr.parse_nested_meta(|meta| {
                                if meta.path.is_ident("rename") {
                                    if let Ok(lit) = meta.value().and_then(|v| v.parse::<LitStr>())
                                    {
                                        // #[serde(rename = "xxx")]
                                        serde_rename = Some(lit.value());
                                    } else {
                                        // #[serde(rename(serialize = "...", deserialize = "..."))]
                                        meta.parse_nested_meta(|nested| {
                                            if nested.path.is_ident("deserialize") {
                                                let lit: LitStr = nested.value()?.parse()?;
                                                serde_rename = Some(lit.value());
                                            } else {
                                                consume_meta(nested);
                                            }
                                            Ok(())
                                        })?;
                                    }
                                } else {
                                    consume_meta(meta);
                                }
                                Ok(())
                            })?;
                        }
                    }

                    #[cfg(feature = "with-json")]
                    let json_key_name = serde_deserialize_name(
                        &original_field_name,
                        serde_rename.as_deref(),
                        serde_rename_all,
                    );

                    if let Some(enum_name) = enum_name {
                        field_name = enum_name;
                    }

                    field_name = Ident::new(&escape_rust_keyword(field_name), ident.span());

                    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"]
                        },
                    };

                    let field_type = &field.ty;
                    let field_type: String = quote! { #field_type }
                        .to_string() // e.g.: "Option < String >"
                        .split_whitespace()
                        .collect(); // Remove all whitespace

                    if ignore {
                        continue;
                    } else {
                        columns_enum.push(quote! {
                            #variant_attrs
                            #field_name
                        });
                        #[cfg(feature = "with-json")]
                        columns_json_keys.push(quote! {
                            Self::#field_name => #json_key_name
                        });
                    }

                    if is_primary_key {
                        primary_keys.push(quote! {
                            #variant_attrs
                            #field_name
                        });
                    }

                    if !is_primary_key && is_auto_increment {
                        return Err(syn::Error::new_spanned(
                            ident,
                            "auto_increment can only be used on primary_key",
                        ));
                    }

                    if primary_key_types.len() > 1 && is_auto_increment {
                        return Err(syn::Error::new_spanned(
                            ident,
                            "auto_increment cannot be used on composite primary_key",
                        ));
                    }

                    if let Some(select_as) = select_as {
                        columns_select_as.push(quote! {
                            Self::#field_name => sea_orm::sea_query::ExprTrait::cast_as(expr, #select_as)
                        });
                    }
                    if let Some(save_as) = save_as {
                        columns_save_as.push(quote! {
                            Self::#field_name => sea_orm::sea_query::ExprTrait::cast_as(val, #save_as)
                        });
                    }

                    let field_type = if field_type.starts_with("Option<") {
                        nullable = true;
                        &field_type["Option<".len()..(field_type.len() - 1)] // Extract `T` out of `Option<T>`
                    } else {
                        field_type.as_str()
                    };
                    let field_span = field.span();

                    if is_primary_key && auto_increment.is_none() {
                        for suffix in NOT_AUTO_INCRE_TYPE_SUFFIX {
                            if field_type.ends_with(suffix) {
                                auto_increment = Some(false);
                                break;
                            }
                        }
                    }

                    let sea_query_col_type =
                        super::value_type_match::column_type_expr(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 seaography_ignore {
                        match_row = quote! { #match_row.seaography_ignore() };
                    }
                    if unique_key.is_some() {
                        match_row = quote! { #match_row.unique_key(#unique_key) };
                    }
                    if renamed_from.is_some() {
                        match_row = quote! { #match_row.renamed_from(#renamed_from) };
                    }
                    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) };
                    }
                    if let Some(extra) = extra {
                        match_row = quote! { #match_row.extra(#extra) };
                    }
                    // match_row = quote! { #match_row.comment() };
                    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
                    });
                }
            }
        }
    }

    // Add tailing comma
    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 = match auto_increment {
            Some(value) => value && primary_keys.len() == 1,
            None => 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)]
            #vis enum PrimaryKey {
                #primary_keys
            }

            #[automatically_derived]
            impl PrimaryKeyTrait for PrimaryKey {
                type ValueType = #primary_key_types;

                fn auto_increment() -> bool {
                    #auto_increment
                }
            }
        }
    };

    let impl_model_ex = if model_ex {
        quote!()
    } else {
        quote! {
            impl Model {
                #[doc = " Generated by sea-orm-macros"]
                pub fn into_ex(self) -> Self {
                    self
                }
            }
        }
    };

    let with_json_impls = {
        #[cfg(feature = "with-json")]
        quote! {
            fn json_key(&self) -> &'static str {
                match self {
                    #columns_json_keys
                }
            }
        }

        #[cfg(not(feature = "with-json"))]
        quote! {}
    };

    Ok(quote! {
        #impl_model_ex

        #[doc = " Generated by sea-orm-macros"]
        #[derive(Copy, Clone, Debug, sea_orm::prelude::EnumIter, sea_orm::prelude::DeriveColumn)]
        #vis 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),
                }
            }

            #with_json_impls
        }

        #entity_def

        #primary_key
    })
}