sqlx_batch/

lib.rs

use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse_macro_input, Attribute, Data, DeriveInput, Fields, Meta};

/// Get an optionally defined key-value pair from annotation
fn get_attr(attrs: &[Attribute], name: &str) -> Option<String> {
    attrs
        .iter()
        .flat_map(|elem| {
            if let Meta::NameValue(nv) = &elem.meta {
                if nv.path.is_ident(name) {
                    let value = &nv.value;
                    return Some(quote!(#value).to_string());
                }
            }
            if let Meta::Path(p) = &elem.meta {
                if p.get_ident().unwrap() == name {
                    return Some(String::new());
                }
            }
            None
        })
        .last()
}

/// Generates a {stuct_name}Inserter struct. This struct allows you to easily
/// add many of this struct to a given Postgres table.
///
/// Use #\[pgtable = "{table_name}"\] on the main struct to define the postgres
/// table name.
///
/// Use #\[colname = "{renamed_elem}"\] on a field to specify that this field
/// will be renamed when added to the database.
///
/// Use #\[key\] on a field to specify that you want to upsert when inserting
/// this struct in the database, and that this field is a primary key / unique
/// constraint which is used for checking if the row already exists.
#[proc_macro_derive(BatchInserter, attributes(pgtable, colname, key))]
pub fn derive_batch_inserter_from(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    let table_name = get_attr(&input.attrs, "pgtable");

    let struct_name = input.ident;
    let new_struct_name = format_ident!("{}Inserter", struct_name);

    let mut field_names = Vec::new();
    let mut field_types = Vec::new();

    let mut psql_names = Vec::new();
    let mut psql_types = Vec::new();

    let mut keyed_names = Vec::new();

    match &input.data {
        Data::Struct(data) => {
            match &data.fields {
                Fields::Named(fields) => {
                    for field in &fields.named {
                        let ident = &field.ident;
                        let ty = &field.ty;
                        let attrs = &field.attrs;

                        field_names.push(ident);
                        field_types.push(ty);

                        // See if we have specified a custom postgres column name for this field
                        if let Some(name) = get_attr(attrs, "colname") {
                            psql_names.push(name.clone());
                            if get_attr(attrs, "key").is_some() {
                                keyed_names.push(name);
                            }
                        } else {
                            psql_names.push(quote!(#ident).to_string());
                            if get_attr(attrs, "key").is_some() {
                                keyed_names.push(quote!(#ident).to_string());
                            }
                        }

                        match quote!(#ty).to_string().as_ref() {
                            "String" | "&str" | "Option < &str >" | "Option < String >" => {
                                psql_types.push("text[]");
                            }
                            "bool" | "Option < bool >" => psql_types.push("bool[]"),
                            "f8" | "f16" | "f32" | "f64" | "Option < f8 >" | "Option < f16 >"
                            | "Option < f32 >" | "Option < f64 >" => psql_types.push("float[]"),
                            "i16" | "i32" | "i64" | "Option < i16 >" | "Option < i32 >"
                            | "Option < i64 >" => psql_types.push("integer[]"),
                            "NaiveDateTime" | "Option < NaiveDateTime >" => {
                                psql_types.push("timestamp[]")
                            }
                            other => panic!("Type {other} can not be directly converted to a Postgres array type!"),
                        }
                    }
                }
                _ => unimplemented!(),
            }
        }
        _ => unimplemented!(),
    }

    let query_build_fn = if table_name.is_some() {
        let mut content = format!("INSERT INTO {} (", table_name.as_deref().unwrap());
        for name in psql_names.iter() {
            content.push_str(name);
            content.push(',');
        }
        content.pop();

        content.push_str(") SELECT * FROM UNNEST (");

        for (idx, ty) in psql_types.iter().enumerate() {
            content.push_str(&format!("${}::{},", idx + 1, ty));
        }
        content.pop();
        content.push(')');

        if !keyed_names.is_empty() {
            content.push_str(" ON CONFLICT (");
            for name in &keyed_names {
                content.push_str(name);
                content.push(',');
            }
            content.pop();
            content.push_str(") DO ");

            if keyed_names.len() == psql_names.len() {
                content.push_str("NOTHING");
            } else {
                content.push_str("UPDATE SET ");
                for name in &psql_names {
                    if !keyed_names.contains(name) {
                        content.push_str(&format!("{name}=excluded.{name},"));
                    }
                }
                content.pop();
            }
        }

        let cast_tokens = field_types.iter().map(|&ty| {
            if quote!(#ty).to_string().starts_with("Option") {
                quote!(as &[#ty])
            } else {
                quote!()
            }
        });

        quote!(
            fn build(
                self,
            ) -> sqlx::query::Query<'static, sqlx::Postgres, sqlx::postgres::PgArguments>
            {
                sqlx::query!(#content, #(&self.#field_names[..] #cast_tokens),*)
            }
        )
    } else {
        quote!()
    };

    let expanded = quote! {
        #[derive(Default, Debug, PartialEq)]
        struct #new_struct_name {
            #(#field_names: Vec<#field_types>),*
        }

        impl #new_struct_name {
            fn new() -> Self {
                Self::default()
            }

            fn from(items: Vec<#struct_name>) -> Self {
                items.into_iter().fold(Self::default(),|mut inserter, item| {
                    inserter.add(item);
                    inserter
                })
            }

            fn add(&mut self, item: #struct_name) {
                #(self.#field_names.push(item.#field_names));*
            }

            #query_build_fn
        }
    };

    expanded.into()
}