pg_extern_attr/
lib.rs

1// Copyright 2018 Benjamin Fry <benjaminfry@me.com>
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8#![recursion_limit = "128"]
9
10extern crate proc_macro;
11extern crate proc_macro2;
12#[macro_use]
13extern crate syn;
14#[macro_use]
15extern crate quote;
16
17use proc_macro2::{Ident, Span, TokenStream};
18use quote::ToTokens;
19use syn::punctuated::Punctuated;
20use syn::token::Comma;
21
22fn create_function_params(num_args: usize) -> TokenStream {
23    let mut tokens = TokenStream::new();
24
25    for i in 0..num_args {
26        let arg_name = Ident::new(&format!("arg_{}", i), Span::call_site());
27
28        tokens.extend(quote!(
29            #arg_name,
30        ));
31    }
32
33    tokens
34}
35
36fn extract_arg_data(inputs: &Punctuated<syn::FnArg, Comma>) -> TokenStream {
37    let mut get_args_stream = TokenStream::new();
38
39    for (i, arg) in inputs.iter().enumerate() {
40        let arg_type: &syn::Type = match *arg {
41            syn::FnArg::SelfRef(_) | syn::FnArg::SelfValue(_) => {
42                panic!("self functions not supported")
43            }
44            syn::FnArg::Inferred(_) => panic!("inferred function parameters not supported"),
45            syn::FnArg::Captured(ref captured) => &captured.ty,
46            syn::FnArg::Ignored(ref ty) => ty,
47        };
48
49        let arg_name = Ident::new(&format!("arg_{}", i), Span::call_site());
50        let arg_error = format!("unsupported function argument type for {}", arg_name);
51
52        let get_arg = quote!(
53            let #arg_name: #arg_type = pg_extend::pg_datum::TryFromPgDatum::try_from(
54                pg_extend::pg_datum::PgDatum::from_raw(
55                    args[#i],
56                    args_null[#i]
57                ),
58            )
59            .expect(#arg_error);
60        );
61
62        get_args_stream.extend(get_arg);
63    }
64
65    get_args_stream
66}
67
68fn impl_info_for_fn(item: &syn::Item) -> TokenStream {
69    let func = if let syn::Item::Fn(func) = item {
70        func
71    } else {
72        panic!("annotation only supported on functions");
73    };
74
75    let func_name = &func.ident;
76    let func_decl = &func.decl;
77
78    if func_decl.variadic.is_some() {
79        panic!("variadic functions (...) not supported")
80    }
81
82    //let generics = &func_decl.generics;
83    let inputs = &func_decl.inputs;
84    //let result = &func_decl.output;
85    //let func_block = &func.block;
86
87    // declare the function
88    let mut function = item.clone().into_token_stream();
89
90    let func_wrapper_name = syn::Ident::new(&format!("pg_{}", func_name), Span::call_site());
91    let func_info_name = syn::Ident::new(
92        &format!("pg_finfo_{}", func_wrapper_name),
93        Span::call_site(),
94    );
95
96    // create the postgres info
97    let func_info = quote!(
98        #[no_mangle]
99        pub extern "C" fn #func_info_name () -> &'static pg_extend::pg_sys::Pg_finfo_record {
100            const my_finfo: pg_extend::pg_sys::Pg_finfo_record = pg_extend::pg_sys::Pg_finfo_record { api_version: 1 };
101            &my_finfo
102        }
103    );
104
105    // join the function information in
106    function.extend(func_info);
107
108    let get_args_from_datums = extract_arg_data(inputs);
109    let func_params = create_function_params(inputs.len());
110
111    // wrap the original function in a pg_wrapper function
112    let func_wrapper = quote!(
113        #[no_mangle]
114        pub extern "C" fn #func_wrapper_name (func_call_info: pg_extend::pg_sys::FunctionCallInfo) -> pg_extend::pg_sys::Datum {
115            use std::panic;
116
117            let func_info: &mut pg_extend::pg_sys::FunctionCallInfoData = unsafe {
118                func_call_info
119                    .as_mut()
120                    .expect("func_call_info was unexpectedly NULL")
121            };
122
123            // guard the Postgres process against the panic, and give us an oportunity to cleanup
124            let panic_result = panic::catch_unwind(|| {
125                // extract the argument list
126                let (args, args_null) = pg_extend::get_args(func_info);
127
128                // arbitrary Datum conversions occur here, and could panic
129                //   so this is inside the catch unwind
130                #get_args_from_datums
131
132                // this is the meat of the function call into the extension code
133                let result = #func_name(#func_params);
134
135                // arbitrary Rust code could panic, so this is guarded
136                pg_extend::pg_datum::PgDatum::from(result)
137            });
138
139            // see if we caught a panic
140            match panic_result {
141                Ok(result) => {
142                    // in addition to the null case, we should handle result types probably
143                    func_info.isnull = result.is_null();
144
145                    // return the datum
146                    result.into_datum()
147                }
148                Err(err) => {
149                    // ensure the return value is null
150                    func_info.isnull = true;
151
152                    // TODO: anything else to cean up before resuming the panic?
153                    panic::resume_unwind(err)
154                }
155            }
156        }
157    );
158
159    function.extend(func_wrapper);
160    function
161}
162
163/// An attribute macro for wrapping Rust functions with boiler plate for defining and
164///   calling conventions between Postgres and Rust.
165///
166///  This mimics the C macro for defining functions
167///
168/// ```c
169/// #define PG_FUNCTION_INFO_V1(funcname) \
170/// extern Datum funcname(PG_FUNCTION_ARGS); \
171/// extern PGDLLEXPORT const Pg_finfo_record * CppConcat(pg_finfo_,funcname)(void); \
172/// const Pg_finfo_record * \
173/// CppConcat(pg_finfo_,funcname) (void) \
174/// { \
175///     static const Pg_finfo_record my_finfo = { 1 }; \
176///     return &my_finfo; \
177/// } \
178/// ```
179///
180/// # Returns
181///
182/// The result of this macro will be to produce a new function wrapping the one annotated but prepended with
183/// `pg_` to distinquish them and also declares a function for Postgres to get the Function information;
184///
185/// For example: if the signature `fn add_one(value: i32) -> i32` is annotated, two functions will be produced,
186///  the wrapper function with a signature of:
187///
188/// ```rust,no_run
189///  #[no_mangle]
190///  pub extern "C" fn pg_add_one(func_call_info: pg_sys::FunctionCallInfo) -> pg_sys::Datum
191/// # {
192/// # unimplemented!()
193/// # }
194/// ```
195///
196/// and the info function with a signature of:
197///
198/// ```rust,no_run
199/// #[no_mangle]
200/// pub extern "C" fn pg_finfo_pg_add_one() -> &'static Pg_finfo_record
201/// # {
202/// # unimplemented!()
203/// # }
204/// ```
205///
206#[proc_macro_attribute]
207pub fn pg_extern(
208    _attr: proc_macro::TokenStream,
209    item: proc_macro::TokenStream,
210) -> proc_macro::TokenStream {
211    // get a usable token stream
212    let ast: syn::Item = parse_macro_input!(item as syn::Item);
213
214    // Build the impl
215    let expanded: TokenStream = impl_info_for_fn(&ast);
216
217    // Return the generated impl
218    proc_macro::TokenStream::from(expanded)
219}
pg_extern_attr/lib.rs

pg_extern_attr/
lib.rs
