symbols/

lib.rs

#![deny(warnings)]
#![deny(missing_docs)]

//! # Symbols
//!
//! This is an utility to build a proc-macro that connects to a database, retrieves data from given table and populates an enum variants with primary keys values  
//! It also generates a method for every non-primary-key field, and, when there are multiple primary keys, a costructor for every possible subset of primary keys

use std::{collections::HashMap, env, fs, future::Future, io};

use heck::{ToSnakeCase, ToUpperCamelCase};

use itertools::Itertools;

use proc_macro2::{Ident, Literal, Span, TokenStream};

use quote::quote;

use sea_orm::{
    DatabaseConnection, EntityName, EntityTrait, Iterable, ModelTrait, PrimaryKeyToColumn, QueryFilter, Value,
};

use serde::{de::DeserializeOwned, Serialize};

pub use symbols_models::EntityFilter;

use syn::{punctuated::Punctuated, token::Comma, Fields, ItemEnum, Lit, LitBool, Meta, NestedMeta, Variant};

use tracing::{error, info};

/// Main function  
/// Given a database model (via generics), an enum item, a list of arguments and an async function to retrieve a database connection
/// it populates the enum using primary key(s) values.  
/// Only string-typed primary keys are supported.
///
/// When a single primary key is present, it simply generate an as_str method and a TryFrom<&str> implementation.  
/// When multiple primary keys are present, it generates a costructor for every possible subset of primary keys.
///
/// For every non-primary key field of a supported type, it generates a const method to retrieve it.
///
/// Replacements can be done on every string-typed field, even primary keys, and are done using annotated parameters.  
/// Two type of replacements are supported:
/// * basic: written in the form #[macro(field = "enum")] or #[macro(field(type = "enum"))], where we are telling to replace string values from `field` with variants from enum `enum`, variant names will be the CamelCase version of field value.
/// * advanced: written in the form #[macro(field(type = "bar", fn = "foo"))], where we are telling to replace string values from `field` with a call to method `foo` from struct/enum `bar`, method output is expected to be of type `bar`.
pub async fn symbols<M, F, Fut>(item: &mut ItemEnum, args: &[NestedMeta], get_conn: F) -> syn::Result<TokenStream>
where
    M: EntityTrait + EntityFilter + Default,
    <M as EntityTrait>::Model: Serialize + DeserializeOwned,
    <M as EntityTrait>::Column: PartialEq,
    F: Fn() -> Fut,
    Fut: Future<Output = syn::Result<DatabaseConnection>>,
{
    let name = &item.ident;
    let primary_keys = <M as EntityTrait>::PrimaryKey::iter().map(|k| k.into_column()).collect::<Vec<_>>();

    let mut constructors = HashMap::new();
    let mut methods = HashMap::new();

    let data = get_data::<M, _, _>(get_conn).await?;

    data.iter().try_for_each(|v| {
        let mut key_s = vec![];

        // scan primary keys
        for k in &primary_keys {
            let val = v.get(*k);
            // only string values are accepted
            if let Value::String(Some(s)) = val {
                key_s.push(s.to_upper_camel_case());

                // if we have a single primary key, create a method as_str and a counter-trait-impl TryFrom<&str>
                if primary_keys.len() == 1 {
                    let key = Ident::new(&s.to_upper_camel_case(), Span::call_site());
                    let v = Literal::string(s.as_str());

                    let (_, method, _) = methods
                        .entry(String::from("as_str"))
                        .or_insert_with(|| (quote! { &'static str }, Punctuated::<_, Comma>::new(), false));
                    method.push(quote! {
                        #name::#key => #v
                    });

                    let (_, method, _) = methods
                        .entry(String::from("try_from"))
                        .or_insert_with(|| (quote! { () }, Punctuated::<_, Comma>::new(), false));
                    method.push(quote! {
                        #v => Ok(#name::#key)
                    });
                }
            } else {
                return Err(syn::Error::new(Span::call_site(), format!("Unrecognized value type {val:?}")));
            }
        }
        // push primary keys into enum variants
        let key_ident = Ident::new(&key_s.join("_"), Span::call_site());
        item.variants.push(Variant {
            attrs: vec![],
            ident: key_ident.clone(),
            fields: Fields::Unit,
            discriminant: None,
        });
        // generate constructors for every combination of primary keys
        if primary_keys.len() > 1 {
            for n in 1..=primary_keys.len() {
                for combo in primary_keys.iter().enumerate().combinations(n) {
                    let cols = combo.iter().map(|(_, col)| **col).collect::<Vec<_>>();
                    let method = combo
                        .iter()
                        .map(|(_, col)| format!("{col:?}").to_snake_case())
                        .collect::<Vec<_>>()
                        .join("_and_");
                    let key = combo.iter().map(|(index, _)| key_s[*index].clone()).collect::<Vec<_>>();
                    let (_, method) = constructors.entry(method).or_insert_with(|| (cols, HashMap::new()));
                    let (_, idents) =
                        method.entry(key.join("_")).or_insert_with(|| (key, Punctuated::<_, Comma>::new()));
                    idents.push(quote! { #name::#key_ident });
                }
            }
        }

        // create a method for every non-primary_key column
        for col in <M as EntityTrait>::Column::iter() {
            let replace = get_replacement::<M>(col, args);

            // skip self-describing methods (would be an as_str clone)
            if primary_keys.len() == 1 && primary_keys.contains(&col) && replace.is_none() {
                continue;
            }

            // keep only managed data types
            let (t, value) = match v.get(col) {
                Value::Bool(b) => (
                    quote! { bool },
                    b.map(|b| {
                        let v = LitBool::new(b, Span::call_site());
                        quote! { #v }
                    }),
                ),
                Value::TinyInt(n) => (
                    quote! { i8 },
                    n.map(|n| {
                        let v = Literal::i8_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::SmallInt(n) => (
                    quote! { i16 },
                    n.map(|n| {
                        let v = Literal::i16_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::Int(n) => (
                    quote! { i32 },
                    n.map(|n| {
                        let v = Literal::i32_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::BigInt(n) => (
                    quote! { i64 },
                    n.map(|n| {
                        let v = Literal::i64_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::TinyUnsigned(n) => (
                    quote! { u8 },
                    n.map(|n| {
                        let v = Literal::u8_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::SmallUnsigned(n) => (
                    quote! { u16 },
                    n.map(|n| {
                        let v = Literal::u16_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::Unsigned(n) => (
                    quote! { u32 },
                    n.map(|n| {
                        let v = Literal::u32_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::BigUnsigned(n) => (
                    quote! { u64 },
                    n.map(|n| {
                        let v = Literal::u64_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::Float(n) => (
                    quote! { f32 },
                    n.map(|n| {
                        let v = Literal::f32_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::Double(n) => (
                    quote! { f64 },
                    n.map(|n| {
                        let v = Literal::f64_unsuffixed(n);
                        quote! { #v }
                    }),
                ),
                Value::String(s) => match replace {
                    Some(Replacement::Type(r)) => (
                        r.clone(),
                        s.map(|s| {
                            let ident = Ident::new(&s.to_upper_camel_case(), Span::call_site());
                            quote! { #r::#ident }
                        }),
                    ),
                    Some(Replacement::Fn(f, Some(r))) => (
                        r.clone(),
                        s.map(|s| {
                            let v = Literal::string(s.as_str());
                            quote! { #r::#f(#v) }
                        }),
                    ),
                    Some(Replacement::Fn(_, None)) => {
                        // teoretically we could accept only a function, but we won't know the return type
                        return Err(syn::Error::new(
                            Span::call_site(),
                            format!("Missing parameter type for field {col:?}"),
                        ));
                    }
                    _ => (
                        quote! { &'static str },
                        s.map(|s| {
                            let v = Literal::string(s.as_str());
                            quote! { #v }
                        }),
                    ),
                },
                // disable ChronoDateTime for now, it would only produce methods for created_at and updated_at fields
                // Value::ChronoDateTime(dt) => (quote! { chrono::NaiveDateTime }, Lit::Verbatim(Literal)),
                _ => continue,
            };
            let (_, method, option) =
                methods.entry(format!("{col:?}")).or_insert_with(|| (t, Punctuated::<_, Comma>::new(), false));
            if let Some(v) = value {
                method.push(quote! {
                    #name::#key_ident => #v
                });
            } else {
                *option = true;
            }
        }

        Ok(())
    })?;

    // decorate constructors
    let constructors = constructors.into_iter().map(|(name, (cols, body))| {
        let is_full = cols.len() == primary_keys.len();
        let fn_name = Ident::new(&format!("get_by_{name}"), Span::call_site());
        let signature = cols
            .iter()
            .map(|col| {
                let field_name = Ident::new(&format!("{col:?}").to_snake_case(), Span::call_site());
                match get_replacement::<M>(*col, args) {
                    Some(Replacement::Type(r)) => quote! { #field_name: #r },
                    _ => quote! { #field_name: &str },
                }
            })
            .collect::<Punctuated<_, Comma>>();
        let m = cols
            .iter()
            .map(|col| {
                let field_name = Ident::new(&format!("{col:?}").to_snake_case(), Span::call_site());
                quote! { #field_name }
            })
            .collect::<Punctuated<_, Comma>>();
        let body = body
            .iter()
            .map(|(_, (values, array_body))| {
                let args = cols
                    .iter()
                    .enumerate()
                    .map(|(index, col)| match get_replacement::<M>(*col, args) {
                        Some(Replacement::Type(r)) => {
                            let ident = Ident::new(&values[index].to_upper_camel_case(), Span::call_site());
                            quote! { #r::#ident }
                        }
                        _ => {
                            let v = Literal::string(values[index].as_str());
                            quote! { #v }
                        }
                    })
                    .collect::<Punctuated<_, Comma>>();
                if is_full {
                    quote! {
                        (#args,) => Some(#array_body)
                    }
                } else {
                    quote! {
                        (#args,) => &[#array_body]
                    }
                }
            })
            .collect::<Punctuated<_, Comma>>();
        if is_full {
            quote! {
                pub const fn #fn_name(#signature) -> Option<Self> {
                    match (#m,) {
                        #body,
                        _ => None,
                    }
                }
            }
        } else {
            quote! {
                pub const fn #fn_name(#signature) -> &'static [Self] {
                    match (#m,) {
                        #body,
                        _ => &[],
                    }
                }
            }
        }
    });

    // separate try_from from other methods
    let try_from = methods
        .remove("try_from")
        .map(|(_, matches, _)| {
            quote! {
                impl<'a> TryFrom<&'a str> for #name {
                    type Error = String;
                    fn try_from(s: &'a str) -> Result<Self, Self::Error> {
                        match s {
                            #matches,
                            _ => Err(format!("Unknown {} {}", stringify!(#name), s)),
                        }
                    }
                }
            }
        })
        .unwrap_or_default();

    // decorate methods
    let methods: TokenStream = methods
        .into_iter()
        .map(|(name, (t, matches, option))| {
            let n = Ident::new(&name.to_snake_case(), Span::call_site());
            if option {
                if matches.is_empty() {
                    quote! {
                        pub const fn #n(&self) -> Option<#t> {
                            None
                        }
                    }
                } else {
                    quote! {
                        pub const fn #n(&self) -> Option<#t> {
                            Some(match self {
                                #matches,
                                _ => return None,
                            })
                        }
                    }
                }
            } else {
                quote! {
                    pub const fn #n(&self) -> #t {
                        match self {
                            #matches,
                        }
                    }
                }
            }
        })
        .chain(constructors)
        .collect();

    // output result
    Ok(quote! {
        #item

        impl #name {
            #methods
        }

        #try_from
    })
}

/// Replacement types
enum Replacement {
    Type(TokenStream),
    Fn(TokenStream, Option<TokenStream>),
}

/// Field replacement facility
/// Searches between macro arguments
fn get_replacement<M>(col: M::Column, args: &[NestedMeta]) -> Option<Replacement>
where
    M: EntityTrait,
    M::Column: PartialEq,
{
    let col_name = format!("{col:?}");
    let field_name = col_name.to_snake_case();
    // search for replacements
    args.iter().find_map(|arg| {
        // simple #[macro(field = "enum")]
        if let NestedMeta::Meta(Meta::NameValue(mv)) = arg {
            if mv.path.is_ident(&col_name) || mv.path.is_ident(&field_name) {
                if let Lit::Str(s) = &mv.lit {
                    let ident = Ident::new(&s.value(), Span::call_site());
                    return Some(Replacement::Type(quote! { #ident }));
                }
            }
        }
        // quite complex #[macro(field(type = "enum", fn = "foo"))]
        if let NestedMeta::Meta(Meta::List(ml)) = arg {
            if ml.path.is_ident(&col_name) || ml.path.is_ident(&field_name) {
                return ml.nested.iter().fold(None, |mut acc, nested| {
                    if let NestedMeta::Meta(Meta::NameValue(mv)) = nested {
                        if let Lit::Str(s) = &mv.lit {
                            let ident = Ident::new(&s.value(), Span::call_site());
                            if mv.path.is_ident("type") {
                                if let Some(Replacement::Fn(f, None)) = acc {
                                    acc = Some(Replacement::Fn(f, Some(quote! { #ident })));
                                } else {
                                    acc = Some(Replacement::Type(quote! { #ident }));
                                }
                            } else if mv.path.is_ident("fn") {
                                if let Some(Replacement::Type(t)) = acc {
                                    acc = Some(Replacement::Fn(quote! { #ident }, Some(t)));
                                } else {
                                    acc = Some(Replacement::Fn(quote! { #ident }, None));
                                }
                            }
                        }
                    }
                    acc
                });
            }
        }
        None
    })
}

/// Data retrieve function with cache capabilities
/// File access is sync to not have to depend on an async runtime
async fn get_data<M, F, Fut>(get_conn: F) -> syn::Result<Vec<<M as EntityTrait>::Model>>
where
    M: EntityTrait + EntityFilter + Default,
    <M as EntityTrait>::Model: Serialize + DeserializeOwned,
    F: Fn() -> Fut,
    Fut: Future<Output = syn::Result<DatabaseConnection>>,
{
    let instance = M::default();
    let mut cache = env::temp_dir();
    cache.push(EntityName::table_name(&instance));
    cache.set_extension("cache");
    if cache.exists() {
        info!("Cache file {} exists, loading data from there", cache.display());

        let file = fs::File::open(&cache)
            .map_err(|e| syn::Error::new(Span::call_site(), format!("Error reading {}: {}", cache.display(), e)))?;

        match bincode::deserialize_from(io::BufReader::new(file)) {
            Ok(data) => return Ok(data),
            Err(e) => error!("Error deserializing {}: {}", cache.display(), e),
        }
    } else {
        info!("Cache file {} doesn't exists, creating", cache.display());
    }

    let conn = get_conn().await?;
    let data = <M as EntityTrait>::find()
        .filter(M::filter())
        .all(&conn)
        .await
        .map_err(|e| syn::Error::new(Span::call_site(), e))?;
    let buf = bincode::serialize(&data)
        .map_err(|e| syn::Error::new(Span::call_site(), format!("Error serializing {}: {}", cache.display(), e)))?;
    fs::write(&cache, buf)
        .map_err(|e| syn::Error::new(Span::call_site(), format!("Error writing {}: {}", cache.display(), e)))?;
    Ok(data)
}