pgdog_macros/
lib.rs

1//! Macros used by PgDog plugins.
2//!
3//! Required and exported by the `pgdog-plugin` crate. You don't have to add this crate separately.
4//!
5use proc_macro::TokenStream;
6use quote::quote;
7use syn::{ItemFn, parse_macro_input};
8
9/// Generates required methods for PgDog to run at plugin load time.
10///
11/// ### Methods
12///
13/// * `pgdog_rustc_version`: Returns the version of the Rust compiler used to build the plugin.
14/// * `pgdog_pg_query_version`: Returns the version of the pg_query library used by the plugin.
15/// * `pgdog_plugin_version`: Returns the version of the plugin itself, taken from Cargo.toml.
16///
17#[proc_macro]
18pub fn plugin(_input: TokenStream) -> TokenStream {
19    let expanded = quote! {
20        #[unsafe(no_mangle)]
21        pub unsafe extern "C" fn pgdog_rustc_version(output: *mut pgdog_plugin::PdStr) {
22            let version = pgdog_plugin::comp::rustc_version();
23            unsafe {
24                *output = version;
25            }
26        }
27
28        #[unsafe(no_mangle)]
29        pub unsafe extern "C" fn pgdog_pg_query_version(output: *mut pgdog_plugin::PdStr) {
30            let version: pgdog_plugin::PdStr = option_env!("PGDOG_PGQUERY_VERSION")
31                .unwrap_or_default()
32                .into();
33            unsafe {
34                *output = version;
35            }
36        }
37
38        #[unsafe(no_mangle)]
39        pub unsafe extern "C" fn pgdog_plugin_version(output: *mut pgdog_plugin::PdStr) {
40            let version: pgdog_plugin::PdStr = env!("CARGO_PKG_VERSION").into();
41            unsafe {
42                *output = version;
43            }
44        }
45    };
46    TokenStream::from(expanded)
47}
48
49/// Generate the `pgdog_init` method that's executed at plugin load time.
50#[proc_macro_attribute]
51pub fn init(_attr: TokenStream, item: TokenStream) -> TokenStream {
52    let input_fn = parse_macro_input!(item as ItemFn);
53    let fn_name = &input_fn.sig.ident;
54
55    let expanded = quote! {
56
57        #[unsafe(no_mangle)]
58        pub extern "C" fn pgdog_init() {
59            #input_fn
60
61            #fn_name();
62        }
63    };
64
65    TokenStream::from(expanded)
66}
67
68/// Generate the `pgdog_fini` method that runs at PgDog shutdown.
69#[proc_macro_attribute]
70pub fn fini(_attr: TokenStream, item: TokenStream) -> TokenStream {
71    let input_fn = parse_macro_input!(item as ItemFn);
72    let fn_name = &input_fn.sig.ident;
73
74    let expanded = quote! {
75        #[unsafe(no_mangle)]
76        pub extern "C" fn pgdog_fini() {
77            #input_fn
78
79            #fn_name();
80        }
81    };
82
83    TokenStream::from(expanded)
84}
85
86/// Generates the `pgdog_route` method for routing queries.
87#[proc_macro_attribute]
88pub fn route(_attr: TokenStream, item: TokenStream) -> TokenStream {
89    let input_fn = parse_macro_input!(item as ItemFn);
90    let fn_name = &input_fn.sig.ident;
91    let fn_inputs = &input_fn.sig.inputs;
92
93    // Extract the first parameter name and type for the pgdog_route function signature
94    let (first_param_name, _) = fn_inputs
95        .iter()
96        .filter_map(|input| {
97            if let syn::FnArg::Typed(pat_type) = input {
98                if let syn::Pat::Ident(pat_ident) = &*pat_type.pat {
99                    Some((pat_ident.ident.clone(), pat_type.ty.clone()))
100                } else {
101                    None
102                }
103            } else {
104                None
105            }
106        })
107        .next()
108        .expect("Route function must have at least one named parameter");
109
110    let expanded = quote! {
111        #[unsafe(no_mangle)]
112        pub unsafe extern "C" fn pgdog_route(#first_param_name: pgdog_plugin::PdRouterContext, output: *mut pgdog_plugin::PdRoute) {
113            #input_fn
114
115            let pgdog_context: pgdog_plugin::Context = #first_param_name.into();
116            let route: pgdog_plugin::PdRoute = #fn_name(pgdog_context).into();
117            unsafe {
118                *output = route;
119            }
120        }
121    };
122
123    TokenStream::from(expanded)
124}
pgdog_macros/lib.rs

pgdog_macros/
lib.rs
