postgresql_macros/
postgresql_macros.rs

1//! Compile-time checked query macros for the `postgresql` crate.
2//!
3//! `query!` and `query_as!` build a query and, when `DATABASE_URL` is set at
4//! build time (and `STDX_POSTGRESQL_CHECK` is not `false`), connect to the
5//! database, describe the statement and check that the parameters match the
6//! server-inferred types. When the SQL is not a string literal, or checking is
7//! disabled, the same code compiles and is decoded at runtime.
8
9#![forbid(unsafe_code)]
10#![warn(missing_docs)]
11
12mod describe;
13
14use proc_macro::TokenStream;
15use quote::quote;
16use syn::{
17    Expr, Lit, Token, Type,
18    parse::{Parse, ParseStream},
19    parse_macro_input,
20    punctuated::Punctuated,
21};
22
23struct QueryInput {
24    sql: Expr,
25    args: Punctuated<Expr, Token![,]>,
26}
27
28impl Parse for QueryInput {
29    fn parse(input: ParseStream) -> syn::Result<Self> {
30        let sql: Expr = input.parse()?;
31        let args = parse_trailing_args(input)?;
32        Ok(QueryInput {
33            sql,
34            args,
35        })
36    }
37}
38
39struct QueryAsInput {
40    ty: Type,
41    sql: Expr,
42    args: Punctuated<Expr, Token![,]>,
43}
44
45impl Parse for QueryAsInput {
46    fn parse(input: ParseStream) -> syn::Result<Self> {
47        let ty: Type = input.parse()?;
48        input.parse::<Token![,]>()?;
49        let sql: Expr = input.parse()?;
50        let args = parse_trailing_args(input)?;
51        Ok(QueryAsInput {
52            ty,
53            sql,
54            args,
55        })
56    }
57}
58
59fn parse_trailing_args(input: ParseStream) -> syn::Result<Punctuated<Expr, Token![,]>> {
60    if input.is_empty() {
61        return Ok(Punctuated::new());
62    }
63    input.parse::<Token![,]>()?;
64    if input.is_empty() {
65        return Ok(Punctuated::new());
66    }
67    Punctuated::parse_terminated(input)
68}
69
70/// Builds an untyped query. See the crate documentation.
71#[proc_macro]
72pub fn query(input: TokenStream) -> TokenStream {
73    let input = parse_macro_input!(input as QueryInput);
74    match expand(input.sql, input.args, None) {
75        Ok(tokens) => tokens.into(),
76        Err(err) => err.to_compile_error().into(),
77    }
78}
79
80/// Builds a query decoded into a named type. See the crate documentation.
81#[proc_macro]
82pub fn query_as(input: TokenStream) -> TokenStream {
83    let input = parse_macro_input!(input as QueryAsInput);
84    match expand(input.sql, input.args, Some(input.ty)) {
85        Ok(tokens) => tokens.into(),
86        Err(err) => err.to_compile_error().into(),
87    }
88}
89
90fn checking_enabled() -> Option<String> {
91    if let Ok(flag) = std::env::var("STDX_POSTGRESQL_CHECK") {
92        if flag.eq_ignore_ascii_case("false") {
93            return None;
94        }
95    }
96    std::env::var("DATABASE_URL").ok().filter(|url| !url.is_empty())
97}
98
99fn sql_literal(sql: &Expr) -> Option<String> {
100    if let Expr::Lit(lit) = sql {
101        if let Lit::Str(s) = &lit.lit {
102            return Some(s.value());
103        }
104    }
105    None
106}
107
108fn expand(sql: Expr, args: Punctuated<Expr, Token![,]>, target: Option<Type>) -> syn::Result<proc_macro2::TokenStream> {
109    let described = match (checking_enabled(), sql_literal(&sql)) {
110        (Some(url), Some(text)) => Some(
111            describe::describe(&url, &text).map_err(|e| syn::Error::new_spanned(&sql, format!("postgresql: {e}")))?,
112        ),
113        _ => None,
114    };
115
116    if let Some(described) = &described {
117        if described.param_oids.len() != args.len() {
118            return Err(syn::Error::new_spanned(
119                &sql,
120                format!(
121                    "postgresql: query expects {} parameter(s) but {} were provided",
122                    described.param_oids.len(),
123                    args.len()
124                ),
125            ));
126        }
127    }
128
129    let params: Vec<proc_macro2::TokenStream> = args
130        .iter()
131        .enumerate()
132        .map(|(index, arg)| match &described {
133            Some(described) => {
134                let oid = described.param_oids[index];
135                quote! {
136                    ::postgresql::__private::param::<#oid, _>(&(#arg)) as &dyn ::postgresql::ToSql
137                }
138            }
139            None => quote! {
140                &(#arg) as &dyn ::postgresql::ToSql
141            },
142        })
143        .collect();
144
145    let params_slice = quote! { &[ #(#params),* ] };
146
147    let tokens = match target {
148        Some(ty) => {
149            let base = quote! {
150                ::postgresql::QueryAs::<#ty>::new(#sql, #params_slice)
151            };
152            match &described {
153                Some(described) => {
154                    let specs = described.columns.iter().map(|column| {
155                        let name = column.name.as_str();
156                        let oid = column.type_oid;
157                        let nullable = column.nullable;
158                        let known = column.nullable_known;
159                        quote! {
160                            ::postgresql::ColumnSpec {
161                                name: #name,
162                                type_oid: #oid,
163                                nullable: #nullable,
164                                known: #known,
165                            },
166                        }
167                    });
168                    quote! {
169                        (
170                            const {
171                                ::postgresql::__private::verify_shape(
172                                    <#ty as ::postgresql::RowShape>::SHAPE,
173                                    &[ #(#specs)* ],
174                                )
175                            },
176                            #base,
177                        ).1
178                    }
179                }
180                None => base,
181            }
182        }
183        None => quote! {
184            ::postgresql::Query::new(#sql, #params_slice)
185        },
186    };
187    Ok(tokens)
188}
189
190struct CopyInput {
191    conn: Expr,
192    ty: Type,
193    sql: Expr,
194}
195
196impl Parse for CopyInput {
197    fn parse(input: ParseStream) -> syn::Result<Self> {
198        let conn: Expr = input.parse()?;
199        input.parse::<Token![,]>()?;
200        let ty: Type = input.parse()?;
201        input.parse::<Token![,]>()?;
202        let sql: Expr = input.parse()?;
203        Ok(CopyInput {
204            conn,
205            ty,
206            sql,
207        })
208    }
209}
210
211/// Builds a compile-time checked binary `COPY ... FROM STDIN`.
212///
213/// Usage: `copy_in!(conn, RowType, "COPY t (a, b) FROM STDIN BINARY").await?`.
214#[proc_macro]
215pub fn copy_in(input: TokenStream) -> TokenStream {
216    let input = parse_macro_input!(input as CopyInput);
217    match expand_copy(input, false) {
218        Ok(tokens) => tokens.into(),
219        Err(err) => err.to_compile_error().into(),
220    }
221}
222
223/// Builds a compile-time checked binary `COPY ... TO STDOUT`.
224///
225/// Usage: `copy_out!(conn, RowType, "COPY t TO STDOUT BINARY").await?`.
226#[proc_macro]
227pub fn copy_out(input: TokenStream) -> TokenStream {
228    let input = parse_macro_input!(input as CopyInput);
229    match expand_copy(input, true) {
230        Ok(tokens) => tokens.into(),
231        Err(err) => err.to_compile_error().into(),
232    }
233}
234
235fn expand_copy(input: CopyInput, out: bool) -> syn::Result<proc_macro2::TokenStream> {
236    let columns = match (checking_enabled(), sql_literal(&input.sql)) {
237        (Some(url), Some(text)) => Some(
238            describe::copy_columns(&url, &text)
239                .map_err(|e| syn::Error::new_spanned(&input.sql, format!("postgresql: {e}")))?,
240        ),
241        _ => None,
242    };
243
244    let ty = &input.ty;
245    let conn = &input.conn;
246    let sql = &input.sql;
247
248    let check = columns.as_ref().map(|columns| {
249        let specs = columns.iter().map(|column| {
250            let name = column.name.as_str();
251            let oid = column.type_oid;
252            let nullable = column.nullable;
253            let known = column.nullable_known;
254            quote! {
255                ::postgresql::ColumnSpec {
256                    name: #name,
257                    type_oid: #oid,
258                    nullable: #nullable,
259                    known: #known,
260                },
261            }
262        });
263        quote! {
264            const {
265                ::postgresql::__private::verify_shape(
266                    <#ty as ::postgresql::RowShape>::SHAPE,
267                    &[ #(#specs)* ],
268                )
269            }
270        }
271    });
272
273    let call = if out {
274        let oids: Vec<u32> = columns
275            .as_ref()
276            .map(|columns| columns.iter().map(|c| c.type_oid).collect())
277            .unwrap_or_default();
278        quote! {
279            ::postgresql::__private::copy_out::<#ty>(&(#conn), #sql, &[ #(#oids),* ])
280        }
281    } else {
282        quote! {
283            ::postgresql::__private::copy_in::<#ty>(&(#conn), #sql)
284        }
285    };
286
287    Ok(match check {
288        Some(check) => quote! { ( #check, #call ).1 },
289        None => call,
290    })
291}
postgresql_macros/postgresql_macros.rs

postgresql_macros/
postgresql_macros.rs
