argos_arpa_item_macro/

lib.rs

//! Derive macro crate for `argos-arpa`.

#![warn(
    missing_docs,
    clippy::all,
    clippy::pedantic,
    clippy::nursery,
    clippy::cargo
)]
#![allow(clippy::must_use_candidate)]

extern crate proc_macro;
extern crate quote;
extern crate syn;

use proc_macro2::TokenStream;
use quote::quote;
use syn::{DataStruct, DeriveInput, parse_macro_input, spanned::Spanned};

mod item;
use item::Item;

type Field = (syn::Ident, Vec<syn::Attribute>, syn::Type);

/// Implements a set of functions to ease the SQL data handling. A struct using
/// `derive(CacheTableItem)` must have a field `id: i32`, and all other fields
/// must be named and implement `sqlx::FromRow`.
///
/// - Use `#[unique]` for values that may not be repeated in any row.
/// - Use `#[derived]` for values that are set automatically by the DB and thus
///     not `insert`ed.
/// - Use `#[secret]` for values that should not be selected when downloading.
#[proc_macro_derive(TableItem, attributes(table, unique, derived, secret))]
pub fn item_macro(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    let item = match Item::make(&input) {
        Ok(i) => i,
        Err(err) => return err,
    };

    let Some(attribute) = input
        .attrs
        .iter()
        .find(|a| a.path().segments.len() == 1 && a.path().segments[0].ident == "table")
    else {
        return syn::Error::new(
            input.span(),
            "'const TABLE: &'static str' must be set! This can be done with \
            '#[table(TABLE)]'.",
        )
        .to_compile_error()
        .into();
    };

    let Ok(table) = attribute.parse_args::<syn::Ident>() else {
        return syn::Error::new(input.span(), "Malformed table argument!")
            .to_compile_error()
            .into();
    };

    let name = input.ident;
    let fields = match input.data {
        syn::Data::Struct(DataStruct {
            fields: syn::Fields::Named(ref fields),
            ..
        }) => fields
            .named
            .iter()
            .filter_map(|field| {
                field
                    .ident
                    .clone()
                    .map(|id| (id, field.attrs.clone(), field.ty.clone()))
            })
            .collect::<Vec<_>>(),

        _ => unimplemented!(),
    };

    let exists = exists(&item);
    let select_all = select_all(&item);
    let select_by_id = select_by_id(&item);
    let delete = delete(&item);
    let check_unique = check_unique(&item);
    let insert = insert(&item);
    let update = update(&item);
    // panic!("{select_all}");

    let insert_columns = insert_columns(&fields);
    let insert_values = insert_values(&fields);
    let unique_values = unique_values(&fields);
    let select = select(&fields);

    let table_name = format!("{table}");
    let test_query = format!("select * from {table_name};");

    let expanded = quote! {
        impl TableItem for #name {
            const TABLE: &'static str = #table_name;

            fn id(&self) -> i32 {
                self.id
            }
            #insert_columns
            #insert_values
            #unique_values
            #select

            #exists
            #select_all
            #select_by_id
            #delete
            #check_unique
            #insert
            #update
        }

        impl #name {
            fn test_table_name_existance() {
                _ = sqlx::query!(#test_query);
            }
        }
    };

    proc_macro::TokenStream::from(expanded)
}

fn exists(item: &Item) -> TokenStream {
    let query = format!("select exists (select 1 from {} where id=$1);", item.table,);

    quote! {
        async fn exists(pool: &sqlx::PgPool, id: i32) -> sqlx::Result<bool> {
            sqlx::query!(#query, id,)
            .fetch_one(pool)
            .await
            .map(|r| r.exists.unwrap_or_default())
        }
    }
}

fn select_all(item: &Item) -> TokenStream {
    let items = item.columns.iter().filter(|c| !c.secret).map(|f| &f.ident);

    let get_string = items
        .clone()
        .map(ToString::to_string)
        .collect::<Vec<_>>()
        .join(",");

    let ids = items;
    let default_ids = item.columns.iter().filter(|c| c.secret).map(|f| &f.ident);
    // .map(|id| quote! {
    //     self.#id: data.#id
    // });

    let query = format!("select {} from {};", get_string, item.table);
    quote! {
        async fn select_all(pool: &sqlx::PgPool) -> sqlx::Result<Vec<Self>> {
            sqlx::query_as!(Self, #query)
            .fetch_all(pool)
            .await
            .map(|ok| ok
                .into_iter()
                .map(|data| Self {
                    #(#ids: data.#ids,)*
                    #(#default_ids: Default::default(),)*
                })
                .collect::<Vec<_>>()
            )
        }
    }
}

fn select_by_id(item: &Item) -> TokenStream {
    let items = item
        .columns
        .iter()
        .filter(|c| !c.secret)
        .map(|f| f.ident.to_string())
        .collect::<Vec<_>>()
        .join(",");

    let query = format!("select {} from {} where id=$1;", items, item.table,);

    quote! {
        async fn select_by_id(
            pool: &sqlx::PgPool,
            id: i32,
        ) -> sqlx::Result<Self> {
            sqlx::query_as!(
                Self,
                #query,
                id,
            )
            .fetch_one(pool)
            .await
        }
    }
}

fn check_unique(item: &Item) -> TokenStream {
    let items = item.columns.iter().filter(|c| c.unique);

    if !items.clone().any(|_| true) {
        return quote! {
            async fn check_unique(
                &self,
                _: &sqlx::PgPool,
            ) -> sqlx::Result<Option<i32>> {
                Ok(None)
            }
        };
    }

    let ids = items.clone().map(|f| &f.ident);

    let checks = items
        .enumerate()
        .map(|(i, f)| format!("{}=${}", f.ident, i + 1,))
        .collect::<Vec<_>>()
        .join(" or ");

    let query = format!("select id from {} where {};", item.table, checks,);

    quote! {
        async fn check_unique(
            &self,
            pool: &sqlx::PgPool,
        ) -> sqlx::Result<Option<i32>> {
            struct Id{id:i32}
            sqlx::query_as!(Id, #query, #(self.#ids,)*)
            .fetch_optional(pool)
            .await
            .map(|i| i.map(|j| j.id))
        }
    }
}

fn delete(item: &Item) -> TokenStream {
    let query = format!("delete from {} where id=$1;", item.table,);

    quote! {
        async fn delete(
            tx: &mut sqlx::PgConnection,
            id: i32,
        ) -> sqlx::Result<()> {
            sqlx::query!(
                #query,
                id,
            )
            .execute(tx)
            .await?;
            Ok(())
        }
    }
}

fn insert(item: &Item) -> TokenStream {
    let items = item.columns.iter().filter(|c| !c.derived);

    let (numbers, columns): (Vec<String>, Vec<String>) = items
        .clone()
        .enumerate()
        .map(|(n, i)| (format!("${}", n + 1), i.ident.to_string()))
        .unzip();

    let ids = items.map(|i| &i.ident).collect::<Vec<_>>();

    let query = format!(
        "insert into {}({}) values({}) returning id;",
        item.table,
        columns.join(","),
        numbers.join(","),
    );
    quote! {
        async fn insert(
            &self,
            tx: &mut sqlx::PgConnection,
        ) -> sqlx::Result<i32> {
            struct Id{id:i32}
            sqlx::query_as!(Id, #query, #(self.#ids,)*)
            .fetch_one(tx)
            .await
            .map(|i| i.id)
        }
    }
}

fn update(item: &Item) -> TokenStream {
    let items = item.columns.iter().filter(|c| !c.derived);

    let columns = items
        .clone()
        .enumerate()
        .map(|(i, f)| format!("{}=${}", f.ident, i + 1,))
        .collect::<Vec<_>>()
        .join(",");

    let ids = items.map(|i| &i.ident).collect::<Vec<_>>();

    let query = format!(
        "update {} set {} where id=${};",
        item.table,
        columns,
        ids.len() + 1
    );

    quote! {
        async fn update(
            &self,
            tx: &mut sqlx::PgConnection,
            id: i32,
        ) -> sqlx::Result<()> {
            sqlx::query!(#query, #(self.#ids,)* id,)
            .execute(tx)
            .await?;
            Ok(())
        }
    }
}

fn insert_columns(fields: &[Field]) -> TokenStream {
    let line = collect_named_data_inverse(fields, Some("derived"), ",", |f| f.0.to_string());

    quote! {
        fn insert_columns() -> &'static str {
            #line
        }
    }
}

fn insert_values(fields: &[Field]) -> TokenStream {
    let (brackets, ids): (Vec<_>, Vec<_>) = fields
        .iter()
        .filter(|(_, atts, _)| atts.iter().all(|a| !is_attribute(a, Some("derived"))))
        .map(|(id, _, t)| match get_typetype(t) {
            TypeType::OptionStringish => (
                false,
                quote! {
                    self.#id
                        .as_ref()
                        .map_or_else(|| "NULL".into(), |v| format!("'{v}'"))
                },
            ),

            TypeType::Option => (
                false,
                quote! {
                    self.#id
                        .as_ref()
                        .map_or_else(|| "NULL".into(), |v| format!("{v}"))
                },
            ),

            TypeType::Stringish => (true, quote! {self.#id}),
            TypeType::Casual => (false, quote! {self.#id}),
        })
        .unzip();

    let line = brackets
        .into_iter()
        .map(|to_quote| if to_quote { "'{}' " } else { "{} " })
        .collect::<Vec<_>>()
        .join(",");

    quote! {
        fn insert_values(&self) -> String {
            format!(#line, #(#ids,)*)
        }
    }
}

enum TypeType {
    Casual,
    Stringish,
    Option,
    OptionStringish,
}

fn get_typetype(t: &syn::Type) -> TypeType {
    if is_stringish(t) {
        return TypeType::Stringish;
    }

    if let syn::Type::Path(path) = t {
        let option = path.path.segments.iter().find(|seg| seg.ident.eq("Option"));

        option.map_or(TypeType::Casual, |op| match &op.arguments {
            syn::PathArguments::None => unimplemented!("Empty options have no support"),

            syn::PathArguments::AngleBracketed(args) => match &args.args[0] {
                // Options only have one argument
                syn::GenericArgument::Type(t) => {
                    if is_stringish(t) {
                        TypeType::OptionStringish
                    } else {
                        TypeType::Option
                    }
                }
                _ => unimplemented!("Only types go in options"),
            },

            syn::PathArguments::Parenthesized(_) => unimplemented!("Tuples have no support"),
        })
    } else {
        TypeType::Casual
    }
}

fn is_stringish(t: &syn::Type) -> bool {
    match t {
        syn::Type::Path(path) => path.path.segments.iter().any(|seg| {
            seg.ident.eq("String") || seg.ident.eq("Uuid") || seg.ident.eq("OffsetDateTime")
        }),

        _ => false,
    }
}

fn unique_values(fields: &[Field]) -> TokenStream {
    let (brackets, ids): (Vec<_>, Vec<_>) = fields
        .iter()
        .filter(|(_, atts, _)| atts.iter().any(|a| is_attribute(a, Some("unique"))))
        .map(|(id, _, t)| ((is_stringish(t), id.to_string()), id))
        .unzip();

    let line = brackets
        .into_iter()
        .map(|(is_str, name)| format!("{}={}", name, if is_str { "'{}'" } else { "{}" }))
        .collect::<Vec<_>>()
        .join(" or ");

    quote! {
        fn unique_values(&self) -> String {
            format!(#line, #(self.#ids,)*)
        }
    }
}

fn select(fields: &[Field]) -> TokenStream {
    let line = collect_named_data_inverse(fields, Some("secret"), ",", |f| f.0.to_string());

    quote! {
        fn select() -> &'static str {
            #line
        }
    }
}

fn collect_named_data_inverse<F>(
    fields: &[Field],
    attribute: Option<&'static str>,
    joiner: &'static str,
    treatment: F,
) -> String
where
    F: Fn(&Field) -> String,
{
    fields
        .iter()
        .filter(|&field| field.1.iter().all(|a| !is_attribute(a, attribute)))
        .map(treatment)
        .collect::<Vec<_>>()
        .join(joiner)
}

fn is_attribute(a: &syn::Attribute, attr: Option<&'static str>) -> bool {
    attr.is_none_or(|attr| match &a.meta {
        syn::Meta::Path(path) => path.is_ident(attr),
        _ => false,
    })
}