forum/sqly/src/lib.rs

use quote::quote;
use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::{bracketed, parse_macro_input, Expr, Ident, LitStr, Result, Token};

extern crate proc_macro;

#[proc_macro]
pub fn query_parsed(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let macro_params = parse_macro_input!(input as QueryChecked);
    let connection = macro_params.connection;
    let query = macro_params.query;
    let params = macro_params.params;
    let return_type = macro_params.return_type;

    let query = quote! {
        let rows = sqlx::query!(#query, #params)
            .fetch_all(#connection)
            .await?.iter()
    };

    let parse = generate_parse_value(&return_type);

    let result = quote! {
        {
            #query;
            #parse
        }
    };

    proc_macro::TokenStream::from(result)
}

fn generate_parse_value(value: &TypeMapping) -> proc_macro2::TokenStream {
    match value.clone() {
        TypeMapping::Column(name) => generate_parse_column(name),
        TypeMapping::Vec(mapping) => generate_parse_vec(mapping),
        TypeMapping::Object {
            struct_name,
            properties,
        } => generate_parse_object(struct_name, properties),
    }
}

fn generate_parse_column(column: Ident) -> proc_macro2::TokenStream {
    quote! {
        rows[0].#column
    }
}

fn generate_parse_object(
    struct_name: Ident,
    properties: Vec<TypeProperty>,
) -> proc_macro2::TokenStream {
    let property_lines =
        properties
            .iter()
            .map(|property: &TypeProperty| -> proc_macro2::TokenStream {
                let name = &property.name;
                let value = generate_parse_value(&property.mapping);
                quote! {
                    #name: #value
                }
            });

    quote! {
        #struct_name {
            #( #property_lines ),*
        }
    }
}

fn generate_parse_vec(inner_type: Box<TypeMapping>) -> proc_macro2::TokenStream {
    let value_expression = generate_parse_value(&*inner_type);

    quote! {
        {
            let vec = Vec::new();
            let optional_value = #value_expression;

            if optional_value.is_some() {
                vec.push(#value_expression);
            }
            vec
        }
    }
}

#[derive(Clone, Debug)]
struct QueryChecked {
    connection: Expr,
    query: LitStr,
    params: Punctuated<Expr, Token![,]>,
    return_type: TypeMapping,
}

#[derive(Clone, Debug)]
enum TypeMapping {
    Column(Ident),
    Vec(Box<TypeMapping>),
    Object {
        struct_name: Ident,
        properties: Vec<TypeProperty>,
    },
}

#[derive(Clone, Debug)]
struct TypeProperty {
    name: Ident,
    mapping: TypeMapping,
}

impl Parse for QueryChecked {
    fn parse(input: ParseStream) -> Result<Self> {
        let mut connection: Option<Expr> = None;
        let mut query: Option<LitStr> = None;
        let mut params: Option<Punctuated<Expr, Token![,]>> = None;
        let mut return_type: Option<TypeMapping> = None;

        while !input.is_empty() {
            let identifier = input.parse::<syn::Ident>()?;
            input.parse::<Token![=]>()?;

            match identifier.to_string().as_str() {
                "connection" => connection = Some(input.parse()?),
                "query" => query = Some(input.parse()?),
                "params" => {
                    let content;
                    bracketed!(content in input);

                    params = Some(Punctuated::parse_terminated(&content)?)
                }
                "return_type" => return_type = Some(input.parse()?),
                unknown => {
                    return Err(input.error(format!("Unknown property {}", unknown)));
                }
            }

            if !input.is_empty() {
                input.parse::<Token![,]>()?;
            }
        }

        Ok(QueryChecked {
            connection: connection.ok_or(input.error("Expected connection property"))?,
            query: query.ok_or(input.error("Expected query property"))?,
            params: params.unwrap_or(Punctuated::new()),
            return_type: return_type.ok_or(input.error("Expected return type property"))?,
        })
    }
}

impl Parse for TypeMapping {
    fn parse(input: ParseStream) -> Result<Self> {
        let identifier = input.parse::<syn::Ident>()?;

        if !input.is_empty() && input.peek(Token![<]) && identifier.to_string() == "Vec" {
            input.parse::<Token![<]>()?;
            let internal: TypeMapping = input.parse()?;
            input.parse::<Token![>]>()?;

            Ok(TypeMapping::Vec(Box::new(internal)))
        } else if !input.is_empty() && input.peek(syn::token::Brace) {
            let internal;
            syn::braced!(internal in input);

            let properties: Punctuated<TypeProperty, Token![,]> =
                Punctuated::parse_terminated(&internal)?;

            let properties = properties.iter().map(|p| p.clone()).collect();

            Ok(TypeMapping::Object {
                struct_name: identifier,
                properties,
            })
        } else {
            Ok(TypeMapping::Column(identifier))
        }
    }
}
impl Parse for TypeProperty {
    fn parse(input: ParseStream) -> Result<Self> {
        let identifier = input.parse::<syn::Ident>()?;
        input.parse::<Token![:]>()?;
        let property_type: TypeMapping = input.parse()?;

        Ok(TypeProperty {
            name: identifier,
            mapping: property_type,
        })
    }
}