1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

mod parse;

use darling::FromDeriveInput;
use proc_macro::{self, TokenStream};
use quote::quote;
use syn::parse_macro_input;

use parse::ModelInput;

/// Automatically implement `Model` for your struct.
/// 
/// ## Attributes
///  * `table` - for structs:
///      - `table_name`: String, optional
///  * `column` - for struct fields:
///      - `dtype`: String, required,
///      - `unique`: bool, optional, default: `false`
#[proc_macro_derive(Model, attributes(table, column))]
pub fn derive(input: TokenStream) -> TokenStream {
    let opts = ModelInput::from_derive_input(&parse_macro_input!(input)).unwrap();

    let ident = opts.ident();

    let table_name = opts.table_name();

    // Retrieve the struct's fields
    let fields = opts.fields();

    if opts.n_fields() == 0 {
        panic!("struct must have at least one field to become a `Model`")
    }

    // Get the fields' idents
    let field_idents = fields
        .map(|f| f.clone().ident()).collect::<Vec<_>>();

    let column_fields = opts.column_fields().collect::<Vec<_>>();
    let column_counter = (1..=column_fields.len()).map(|i| format!("${i}")).collect::<Vec<_>>();
    let column_counter_one = column_counter.join(", ");

    let column_names = column_fields.clone().into_iter().map(|f| f.column_name()).collect::<Vec<_>>();
    let column_names_one = column_names.join(", ");

    let column_field_idents = column_fields.clone().into_iter().map(|f| f.ident()).collect::<Vec<_>>();
    let column_field_types = column_fields.clone().into_iter().map(|f| f.ty()).collect::<Vec<_>>();

    let create_table_sql = opts.get_create_sql();

    // Generate the needed impl code
    let output = quote!(
        impl<'a> TryFrom<&'a pg_worm::Row> for #ident {
            type Error = pg_worm::pg::Error;

            fn try_from(value: &'a pg_worm::Row) -> Result<#ident, Self::Error> {
                // Parse each column into the corresponding field
                Ok(#ident {
                    #(#field_idents: value.try_get(stringify!(#field_idents))?),*
                })
            }
        }

        #[pg_worm::async_trait]
        impl Model<#ident> for #ident {
            fn _create_table_sql() -> &'static str {
                #create_table_sql
            }

            async fn select() -> Vec<#ident> {
                let client = pg_worm::_get_client().expect("not connected to db");
                let rows = client.query(format!("SELECT * FROM {}", #table_name).as_str(), &[]).await.unwrap();
                rows.iter().map(|r| #ident::try_from(r).expect("couldn't parse data")).collect()
            }

            async fn select_one() -> Option<#ident> {
                let client = pg_worm::_get_client().expect("not connected to db");
                let rows = client.query(format!("SELECT * FROM {} LIMIT 1", #table_name).as_str(), &[]).await.unwrap();
                if rows.len() != 1 {
                    return None;
                }
                Some(#ident::try_from(&rows[0]).unwrap())
            }
        }

        impl #ident {
            /// Insert a new entity into the database. 
            /// 
            /// For columns which are autogenerated (like in the example below, `id`),
            /// no value has to be specified.
            /// 
            /// # Example
            /// 
            /// ```ignore
            /// use pg_worm::Model;
            /// 
            /// #[derive(Model)]
            /// struct Book {
            ///     #[column(dtype = "BIGSERIAL")]
            ///     id: i64,
            ///     #[column(dtype = "TEXT")]
            ///     title: String
            /// }
            /// 
            /// async fn some_func() -> Result<(), pg_worm::Error> {
            ///     Book::insert("Foo".to_string()).await?;
            /// }
            /// ```
            pub async fn insert(
                #(#column_field_idents: #column_field_types),*
            ) -> Result<(), pg_worm::Error> {
                let stmt = format!(
                    "INSERT INTO {} ({}) VALUES ({})",
                    #table_name,
                    #column_names_one,
                    #column_counter_one
                );

                let client = pg_worm::_get_client()?;

                client.execute(
                    stmt.as_str(),
                    &[
                        #(&#column_field_idents), *
                    ]
                ).await?;

                Ok(())
            }
        }
    );

    output.into()
}