Skip to main content

postgresql_derive/
postgresql_derive.rs

1//! `#[derive(FromRow)]` for the `postgresql` crate.
2//!
3//! Decodes a [`postgresql::Row`] into a struct by looking up each field's
4//! column by name. Column names default to the field name, and can be changed
5//! with `#[postgresql(rename = "...")]` or transformed with a container-level
6//! `#[postgresql(rename_all = "...")]` (`snake_case`, `camelCase`,
7//! `PascalCase`, `SCREAMING_SNAKE_CASE`, `lowercase`, `UPPERCASE`).
8//!
9//! `#[postgresql(default)]` makes a missing or `NULL` column fall back to the
10//! field type's `Default` instead of erroring.
11
12use proc_macro::TokenStream;
13use quote::quote;
14use syn::{Data, DeriveInput, Fields, WherePredicate, parse_macro_input};
15
16/// Derives [`postgresql::FromRow`] for a struct with named fields.
17#[proc_macro_derive(FromRow, attributes(postgresql, pg))]
18pub fn derive_from_row(input: TokenStream) -> TokenStream {
19    let input = parse_macro_input!(input as DeriveInput);
20    match expand(&input) {
21        Ok(tokens) => tokens.into(),
22        Err(err) => err.to_compile_error().into(),
23    }
24}
25
26fn is_pg_attr(attr: &syn::Attribute) -> bool {
27    attr.path().is_ident("postgresql") || attr.path().is_ident("pg")
28}
29
30/// Builds a `where` clause from the type's own predicates plus `extra`.
31fn merge_where(generics: &syn::Generics, extra: impl IntoIterator<Item = WherePredicate>) -> proc_macro2::TokenStream {
32    let mut predicates: syn::punctuated::Punctuated<WherePredicate, syn::Token![,]> = generics
33        .where_clause
34        .as_ref()
35        .map(|w| w.predicates.clone())
36        .unwrap_or_default();
37    for predicate in extra {
38        predicates.push(predicate);
39    }
40    if predicates.is_empty() {
41        quote! {}
42    } else {
43        quote! { where #predicates }
44    }
45}
46
47fn get_rename_all(attrs: &[syn::Attribute]) -> syn::Result<Option<String>> {
48    let mut rename_all = None;
49    for attr in attrs.iter().filter(|a| is_pg_attr(a)) {
50        attr.parse_nested_meta(|meta| {
51            if meta.path.is_ident("rename_all") {
52                let value = meta.value()?;
53                let lit: syn::LitStr = value.parse()?;
54                rename_all = Some(lit.value());
55                Ok(())
56            } else {
57                // Ignore unknown container attributes so field-only keys don't error.
58                let _ = meta.input;
59                Ok(())
60            }
61        })?;
62    }
63    Ok(rename_all)
64}
65
66fn expand(input: &DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
67    let name = &input.ident;
68    let rename_all = get_rename_all(&input.attrs)?;
69
70    let fields = match &input.data {
71        Data::Struct(data) => match &data.fields {
72            Fields::Named(fields) => &fields.named,
73            _ => {
74                return Err(syn::Error::new_spanned(
75                    input,
76                    "FromRow can only be derived for structs with named fields",
77                ));
78            }
79        },
80        _ => {
81            return Err(syn::Error::new_spanned(input, "FromRow can only be derived for structs"));
82        }
83    };
84
85    let mut assignments = Vec::with_capacity(fields.len());
86    let mut shape_entries = Vec::with_capacity(fields.len());
87    let mut shape_bounds: Vec<syn::WherePredicate> = Vec::with_capacity(fields.len());
88
89    for field in fields {
90        let ident = field.ident.as_ref().expect("named field");
91        let mut column: Option<String> = None;
92        let mut default = false;
93
94        for attr in field.attrs.iter().filter(|a| is_pg_attr(a)) {
95            attr.parse_nested_meta(|meta| {
96                if meta.path.is_ident("rename") || meta.path.is_ident("column") {
97                    let value = meta.value()?;
98                    let lit: syn::LitStr = value.parse()?;
99                    column = Some(lit.value());
100                    Ok(())
101                } else if meta.path.is_ident("default") {
102                    default = true;
103                    Ok(())
104                } else {
105                    Ok(())
106                }
107            })?;
108        }
109
110        let column = column.unwrap_or_else(|| apply_rename(&ident.to_string(), rename_all.as_deref()));
111        let ty = &field.ty;
112
113        if default {
114            assignments.push(quote! {
115                #ident: ::postgresql::Row::try_get::<#ty>(row, #column)
116                    .or_else(|_| ::core::result::Result::Ok(::core::default::Default::default()))?,
117            });
118        } else {
119            assignments.push(quote! {
120                #ident: ::postgresql::Row::try_get::<#ty>(row, #column)?,
121            });
122            shape_bounds.push(syn::parse_quote!(#ty: ::postgresql::StaticType));
123            shape_entries.push(quote! {
124                ::postgresql::ColumnSpec {
125                    name: #column,
126                    type_oid: <#ty as ::postgresql::StaticType>::OID,
127                    nullable: <#ty as ::postgresql::StaticType>::NULLABLE,
128                    known: true,
129                },
130            });
131        }
132    }
133
134    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
135
136    let mut row_shape_predicates: syn::punctuated::Punctuated<syn::WherePredicate, syn::Token![,]> = input
137        .generics
138        .where_clause
139        .as_ref()
140        .map(|w| w.predicates.clone())
141        .unwrap_or_default();
142    for bound in shape_bounds {
143        row_shape_predicates.push(bound);
144    }
145    let row_shape_where = if row_shape_predicates.is_empty() {
146        quote! {}
147    } else {
148        quote! { where #row_shape_predicates }
149    };
150
151    let expanded = quote! {
152        impl #impl_generics ::postgresql::FromRow for #name #ty_generics #where_clause {
153            fn from_row(
154                row: &::postgresql::Row,
155            ) -> ::core::result::Result<Self, ::postgresql::Error> {
156                ::core::result::Result::Ok(Self {
157                    #(#assignments)*
158                })
159            }
160        }
161
162        impl #impl_generics ::postgresql::RowShape for #name #ty_generics #row_shape_where {
163            const SHAPE: &'static [::postgresql::ColumnSpec] = &[ #(#shape_entries)* ];
164        }
165    };
166    Ok(expanded)
167}
168
169fn apply_rename(field: &str, rename_all: Option<&str>) -> String {
170    match rename_all {
171        Some("camelCase") => to_camel_case(field),
172        Some("PascalCase") => to_pascal_case(field),
173        Some("SCREAMING_SNAKE_CASE") | Some("UPPERCASE") => field.to_ascii_uppercase(),
174        Some("lowercase") => field.to_ascii_lowercase(),
175        Some("snake_case") | None => to_snake_case(field),
176        Some(_) => field.to_string(),
177    }
178}
179
180fn to_camel_case(s: &str) -> String {
181    let mut out = String::with_capacity(s.len());
182    let mut upper = false;
183    for c in s.chars() {
184        if c == '_' {
185            upper = true;
186        } else if upper {
187            out.extend(c.to_uppercase());
188            upper = false;
189        } else {
190            out.push(c);
191        }
192    }
193    out
194}
195
196fn to_pascal_case(s: &str) -> String {
197    let camel = to_camel_case(s);
198    let mut chars = camel.chars();
199    match chars.next() {
200        Some(c) => c.to_uppercase().collect::<String>() + chars.as_str(),
201        None => String::new(),
202    }
203}
204
205fn to_snake_case(s: &str) -> String {
206    let mut out = String::with_capacity(s.len() + 4);
207    for (i, c) in s.chars().enumerate() {
208        if c.is_uppercase() && i > 0 {
209            out.push('_');
210        }
211        out.extend(c.to_lowercase());
212    }
213    out
214}
215
216/// Derives [`postgresql::ToRow`] for a struct with named fields.
217///
218/// Fields are encoded in declaration order, matching a binary `COPY` column
219/// list.
220#[proc_macro_derive(ToRow, attributes(postgresql, pg))]
221pub fn derive_to_row(input: TokenStream) -> TokenStream {
222    let input = parse_macro_input!(input as DeriveInput);
223    match expand_to_row(&input) {
224        Ok(tokens) => tokens.into(),
225        Err(err) => err.to_compile_error().into(),
226    }
227}
228
229fn expand_to_row(input: &DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
230    let name = &input.ident;
231    let fields = match &input.data {
232        Data::Struct(data) => match &data.fields {
233            Fields::Named(fields) => &fields.named,
234            _ => {
235                return Err(syn::Error::new_spanned(
236                    input,
237                    "ToRow can only be derived for structs with named fields",
238                ));
239            }
240        },
241        _ => {
242            return Err(syn::Error::new_spanned(input, "ToRow can only be derived for structs"));
243        }
244    };
245
246    let encoders = fields.iter().map(|field| {
247        let ident = field.ident.as_ref().expect("named field");
248        quote! { &self.#ident as &dyn ::postgresql::ToSql }
249    });
250
251    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
252    Ok(quote! {
253        impl #impl_generics ::postgresql::ToRow for #name #ty_generics #where_clause {
254            fn encode_row(
255                &self,
256                buf: &mut ::postgresql::__private::BytesMut,
257            ) -> ::core::result::Result<(), ::postgresql::Error> {
258                ::postgresql::encode_copy_row(&[ #(#encoders),* ], buf)
259            }
260        }
261    })
262}
263
264/// Derives `json` / `jsonb` encoding and decoding for a type.
265///
266/// The type must also derive `serde::Serialize` (to be written as a parameter
267/// or `COPY` field) and `serde::Deserialize` (to be read from a result
268/// column). The generated impls map the whole value to a single `jsonb`
269/// column: `FromSql`, `ToSql`, `StaticType`, `FromRow` and `RowShape`.
270///
271/// Because it defines `FromRow` and `RowShape`, a type cannot derive both
272/// `Json` and `FromRow`; use `postgresql::Json<T>` for the JSON role in that
273/// case.
274#[proc_macro_derive(Json, attributes(postgresql, pg))]
275pub fn derive_json(input: TokenStream) -> TokenStream {
276    let input = parse_macro_input!(input as DeriveInput);
277    match expand_json(&input) {
278        Ok(tokens) => tokens.into(),
279        Err(err) => err.to_compile_error().into(),
280    }
281}
282
283fn expand_json(input: &DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
284    let name = &input.ident;
285    let (impl_generics, ty_generics, _) = input.generics.split_for_impl();
286    let self_ty = quote! { #name #ty_generics };
287
288    let deserialize: WherePredicate = syn::parse_quote! {
289        #self_ty: ::postgresql::__private::serde::de::DeserializeOwned
290    };
291    let serialize: WherePredicate = syn::parse_quote! {
292        #self_ty: ::postgresql::__private::serde::Serialize
293    };
294    let send_sync: WherePredicate = syn::parse_quote! {
295        #self_ty: ::core::marker::Send + ::core::marker::Sync
296    };
297
298    let from_sql_where = merge_where(&input.generics, [deserialize.clone()]);
299    let from_row_where = merge_where(&input.generics, [deserialize]);
300    let to_sql_where = merge_where(&input.generics, [serialize, send_sync]);
301    let plain_where = merge_where(&input.generics, []);
302
303    Ok(quote! {
304        impl #impl_generics ::postgresql::FromSql for #self_ty #from_sql_where {
305            fn from_sql(
306                oid: ::postgresql::types::Oid,
307                buf: &[u8],
308            ) -> ::core::result::Result<Self, ::postgresql::Error> {
309                ::postgresql::__private::json_from_sql(oid, buf)
310            }
311
312            fn accepts(oid: ::postgresql::types::Oid) -> bool {
313                ::postgresql::__private::json_accepts(oid)
314            }
315        }
316
317        impl #impl_generics ::postgresql::FromRow for #self_ty #from_row_where {
318            fn from_row(
319                row: &::postgresql::Row,
320            ) -> ::core::result::Result<Self, ::postgresql::Error> {
321                ::postgresql::Row::try_get_by_index(row, 0usize)
322            }
323        }
324
325        impl #impl_generics ::postgresql::ToSql for #self_ty #to_sql_where {
326            fn encode(
327                &self,
328                buf: &mut ::postgresql::__private::BytesMut,
329            ) -> ::core::result::Result<::postgresql::IsNull, ::postgresql::Error> {
330                ::postgresql::__private::json_encode(self, buf)
331            }
332
333            fn oid(&self) -> ::postgresql::types::Oid {
334                ::postgresql::__private::JSONBOID
335            }
336        }
337
338        impl #impl_generics ::postgresql::StaticType for #self_ty #plain_where {
339            const OID: ::postgresql::types::Oid = ::postgresql::__private::JSONBOID;
340        }
341
342        impl #impl_generics ::postgresql::RowShape for #self_ty #plain_where {
343            const SHAPE: &'static [::postgresql::ColumnSpec] = &[::postgresql::ColumnSpec {
344                name: "",
345                type_oid: ::postgresql::__private::JSONBOID,
346                nullable: false,
347                known: true,
348            }];
349        }
350    })
351}