sea-orm-macros 2.0.0-rc.38

Derive macros for SeaORM
Documentation
use super::util::{
    escape_rust_keyword, field_not_ignored, format_field_ident, trim_starting_raw_identifier,
};
use heck::ToUpperCamelCase;
use proc_macro2::{Ident, TokenStream};
use quote::{format_ident, quote};
use syn::{Data, DataStruct, Expr, Fields, LitStr, Type, Visibility};

pub(crate) struct DeriveActiveModel {
    model: Ident,
    vis: Visibility,
    fields: Vec<Ident>,
    names: Vec<Ident>,
    types: Vec<Type>,
}

impl DeriveActiveModel {
    pub fn new(vis: &Visibility, ident: &Ident, data: &Data) -> syn::Result<Self> {
        let all_fields = match data {
            Data::Struct(DataStruct {
                fields: Fields::Named(named),
                ..
            }) => &named.named,
            _ => {
                return Err(syn::Error::new_spanned(
                    ident,
                    "You can only derive DeriveActiveModel on structs",
                ));
            }
        };

        let mut fields = Vec::new();
        let mut names = Vec::new();
        let mut types = Vec::new();

        for field in all_fields.iter().filter(|f| field_not_ignored(f)) {
            fields.push(format_field_ident(field));

            let ident = field.ident.as_ref().unwrap().to_string();
            let ident = trim_starting_raw_identifier(ident).to_upper_camel_case();
            let ident = escape_rust_keyword(ident);
            let mut ident = format_ident!("{}", &ident);
            field
                .attrs
                .iter()
                .filter(|attr| attr.path().is_ident("sea_orm"))
                .try_for_each(|attr| {
                    attr.parse_nested_meta(|meta| {
                        if meta.path.is_ident("enum_name") {
                            let litstr: LitStr = meta.value()?.parse()?;
                            ident = syn::parse_str(&litstr.value()).unwrap();
                        } else {
                            // Reads the value expression to advance the parse stream.
                            // Some parameters, such as `primary_key`, do not have any value,
                            // so ignoring an error occurred here.
                            let _: Option<Expr> = meta.value().and_then(|v| v.parse()).ok();
                        }

                        Ok(())
                    })
                })?;

            names.push(ident);
            types.push(field.ty.clone());
        }

        Ok(DeriveActiveModel {
            model: ident.clone(),
            vis: vis.clone(),
            fields,
            names,
            types,
        })
    }
}

impl DeriveActiveModel {
    fn define_active_model(&self) -> TokenStream {
        let vis = &self.vis;
        let fields = &self.fields;
        let types = &self.types;
        quote!(
            #[doc = " Generated by sea-orm-macros"]
            #[derive(Clone, Debug, PartialEq)]
            #vis struct ActiveModel {
                #(
                    #[doc = " Generated by sea-orm-macros"]
                    pub #fields: sea_orm::ActiveValue<#types>
                ),*
            }
        )
    }

    fn impl_active_model(&self) -> TokenStream {
        let mut ts = self.impl_active_model_convert();
        ts.extend(self.impl_active_model_trait());
        ts
    }

    fn impl_active_model_convert(&self) -> TokenStream {
        let model = &self.model;
        let fields = &self.fields;

        quote!(
            #[automatically_derived]
            impl std::default::Default for ActiveModel {
                fn default() -> Self {
                    <Self as sea_orm::ActiveModelBehavior>::new()
                }
            }

            #[automatically_derived]
            impl std::convert::From<#model> for ActiveModel {
                fn from(m: #model) -> Self {
                    Self {
                        #(#fields: sea_orm::ActiveValue::Unchanged(m.#fields)),*
                    }
                }
            }

            #[automatically_derived]
            impl sea_orm::IntoActiveModel<ActiveModel> for #model {
                fn into_active_model(self) -> ActiveModel {
                    self.into()
                }
            }
        )
    }

    fn impl_active_model_trait(&self) -> TokenStream {
        let fields = &self.fields;
        let methods = self.impl_active_model_trait_methods();

        quote! {
            #[automatically_derived]
            impl sea_orm::ActiveModelTrait for ActiveModel {
                type Entity = Entity;

                #methods

                fn default() -> Self {
                    Self {
                        #(#fields: sea_orm::ActiveValue::NotSet),*
                    }
                }
            }
        }
    }

    pub fn impl_active_model_trait_methods(&self) -> TokenStream {
        let fields = &self.fields;
        let names = &self.names;

        quote!(
            fn take(&mut self, c: <Self::Entity as sea_orm::EntityTrait>::Column) -> sea_orm::ActiveValue<sea_orm::Value> {
                match c {
                    #(<Self::Entity as sea_orm::EntityTrait>::Column::#names => {
                        let mut value = sea_orm::ActiveValue::NotSet;
                        std::mem::swap(&mut value, &mut self.#fields);
                        value.into_wrapped_value()
                    },)*
                    _ => sea_orm::ActiveValue::NotSet,
                }
            }

            fn get(&self, c: <Self::Entity as sea_orm::EntityTrait>::Column) -> sea_orm::ActiveValue<sea_orm::Value> {
                match c {
                    #(<Self::Entity as sea_orm::EntityTrait>::Column::#names => self.#fields.clone().into_wrapped_value(),)*
                    _ => sea_orm::ActiveValue::NotSet,
                }
            }

            fn set_if_not_equals(&mut self, c: <Self::Entity as sea_orm::EntityTrait>::Column, v: sea_orm::Value) {
                match c {
                    #(<Self::Entity as sea_orm::EntityTrait>::Column::#names => self.#fields.set_if_not_equals(v.unwrap()),)*
                    _ => (),
                }
            }

            fn try_set(&mut self, c: <Self::Entity as sea_orm::EntityTrait>::Column, v: sea_orm::Value) -> Result<(), sea_orm::DbErr> {
                match c {
                    #(<Self::Entity as sea_orm::EntityTrait>::Column::#names => self.#fields = sea_orm::ActiveValue::Set(sea_orm::sea_query::ValueType::try_from(v).map_err(|e| sea_orm::DbErr::Type(e.to_string()))?),)*
                    _ => return Err(sea_orm::DbErr::Type(format!("ActiveModel does not have this field: {:?}", sea_orm::ColumnTrait::as_column_ref(&c)))),
                }
                Ok(())
            }

            fn not_set(&mut self, c: <Self::Entity as sea_orm::EntityTrait>::Column) {
                match c {
                    #(<Self::Entity as sea_orm::EntityTrait>::Column::#names => self.#fields = sea_orm::ActiveValue::NotSet,)*
                    _ => (),
                }
            }

            fn is_not_set(&self, c: <Self::Entity as sea_orm::EntityTrait>::Column) -> bool {
                match c {
                    #(<Self::Entity as sea_orm::EntityTrait>::Column::#names => self.#fields.is_not_set(),)*
                    _ => panic!("This ActiveModel does not have this field"),
                }
            }

            fn reset(&mut self, c: <Self::Entity as sea_orm::EntityTrait>::Column) {
                match c {
                    #(<Self::Entity as sea_orm::EntityTrait>::Column::#names => self.#fields.reset(),)*
                    _ => panic!("This ActiveModel does not have this field"),
                }
            }

            fn default_values() -> Self {
                use sea_orm::value::{DefaultActiveValue, DefaultActiveValueNone, DefaultActiveValueNotSet};
                let mut default = <Self as sea_orm::ActiveModelTrait>::default();
                #(default.#fields = (&default.#fields).default_value();)*
                default
            }
        )
    }
}

fn derive_into_model(ident: &Ident, data: &Data) -> syn::Result<TokenStream> {
    let model_fields = match data {
        Data::Struct(DataStruct {
            fields: Fields::Named(named),
            ..
        }) => &named.named,
        _ => {
            return Err(syn::Error::new_spanned(
                ident,
                "You can only derive DeriveActiveModel on structs",
            ));
        }
    };

    let active_model_field: Vec<Ident> = model_fields
        .iter()
        .filter(|f| field_not_ignored(f))
        .map(format_field_ident)
        .collect();

    let model_field: Vec<Ident> = model_fields.iter().map(format_field_ident).collect();

    let ignore_attr: Vec<bool> = model_fields.iter().map(|f| !field_not_ignored(f)).collect();

    let model_field_value: Vec<TokenStream> = model_field
        .iter()
        .zip(ignore_attr)
        .map(|(field, ignore)| {
            if ignore {
                quote! {
                    Default::default()
                }
            } else {
                quote! {
                    a.#field.unwrap()
                }
            }
        })
        .collect();

    Ok(quote!(
        #[automatically_derived]
        impl std::convert::TryFrom<ActiveModel> for #ident {
            type Error = sea_orm::DbErr;
            fn try_from(a: ActiveModel) -> Result<Self, sea_orm::DbErr> {
                #(if a.#active_model_field.is_not_set() {
                    return Err(sea_orm::DbErr::AttrNotSet(stringify!(#active_model_field).to_owned()));
                })*
                Ok(
                    Self {
                        #(#model_field: #model_field_value),*
                    }
                )
            }
        }

        #[automatically_derived]
        impl sea_orm::TryIntoModel<#ident> for ActiveModel {
            fn try_into_model(self) -> Result<#ident, sea_orm::DbErr> {
                self.try_into()
            }
        }
    ))
}

pub fn expand_derive_active_model(
    vis: &Visibility,
    ident: &Ident,
    data: &Data,
) -> syn::Result<TokenStream> {
    let derive_active_model = DeriveActiveModel::new(vis, ident, data)?;

    let define_active_model = derive_active_model.define_active_model();
    let impl_active_model = derive_active_model.impl_active_model();
    let derive_into_model = derive_into_model(ident, data)?;

    Ok(quote!(
        #define_active_model

        #impl_active_model

        #derive_into_model
    ))
}