pgrx_sql_entity_graph/pg_extern/
mod.rs

1//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC.
2//LICENSE
3//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc.
4//LICENSE
5//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. <contact@pgcentral.org>
6//LICENSE
7//LICENSE All rights reserved.
8//LICENSE
9//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
10/*!
11
12`#[pg_extern]` related macro expansion for Rust to SQL translation
13
14> Like all of the [`sql_entity_graph`][crate] APIs, this is considered **internal**
15> to the `pgrx` framework and very subject to change between versions. While you may use this, please do it with caution.
16
17*/
18mod argument;
19mod attribute;
20mod cast;
21pub mod entity;
22mod operator;
23mod returning;
24mod search_path;
25
26pub use argument::PgExternArgument;
27pub(crate) use attribute::Attribute;
28pub use cast::PgCast;
29pub use operator::PgOperator;
30pub use returning::NameMacro;
31use syn::token::Comma;
32
33use self::returning::Returning;
34use super::UsedType;
35use crate::enrich::{CodeEnrichment, ToEntityGraphTokens, ToRustCodeTokens};
36use crate::finfo::{finfo_v1_extern_c, finfo_v1_tokens};
37use crate::fmt::ErrHarder;
38use crate::ToSqlConfig;
39use operator::{PgrxOperatorAttributeWithIdent, PgrxOperatorOpName};
40use search_path::SearchPathList;
41
42use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
43use quote::quote;
44use syn::parse::{Parse, ParseStream, Parser};
45use syn::punctuated::Punctuated;
46use syn::spanned::Spanned;
47use syn::{Meta, Token};
48
49macro_rules! quote_spanned {
50    ($span:expr=> $($expansion:tt)*) => {
51        {
52            let synthetic = Span::mixed_site();
53            let synthetic = synthetic.located_at($span);
54            quote::quote_spanned! {synthetic=> $($expansion)* }
55        }
56    };
57}
58
59macro_rules! format_ident {
60    ($s:literal, $e:expr) => {{
61        let mut synthetic = $e.clone();
62        synthetic.set_span(Span::call_site().located_at($e.span()));
63        quote::format_ident!($s, synthetic)
64    }};
65}
66
67/// A parsed `#[pg_extern]` item.
68///
69/// It should be used with [`syn::parse::Parse`] functions.
70///
71/// Using [`quote::ToTokens`] will output the declaration for a [`PgExternEntity`][crate::PgExternEntity].
72///
73/// ```rust
74/// use syn::{Macro, parse::Parse, parse_quote, parse};
75/// use quote::{quote, ToTokens};
76/// use pgrx_sql_entity_graph::PgExtern;
77///
78/// # fn main() -> eyre::Result<()> {
79/// use pgrx_sql_entity_graph::CodeEnrichment;
80/// let parsed: CodeEnrichment<PgExtern> = parse_quote! {
81///     fn example(x: Option<str>) -> Option<&'a str> {
82///         unimplemented!()
83///     }
84/// };
85/// let sql_graph_entity_tokens = parsed.to_token_stream();
86/// # Ok(())
87/// # }
88/// ```
89#[derive(Debug, Clone)]
90pub struct PgExtern {
91    attrs: Vec<Attribute>,
92    func: syn::ItemFn,
93    to_sql_config: ToSqlConfig,
94    operator: Option<PgOperator>,
95    cast: Option<PgCast>,
96    search_path: Option<SearchPathList>,
97    inputs: Vec<PgExternArgument>,
98    input_types: Vec<syn::Type>,
99    returns: Returning,
100}
101
102impl PgExtern {
103    #[track_caller]
104    pub fn new(attr: TokenStream2, item: TokenStream2) -> Result<CodeEnrichment<Self>, syn::Error> {
105        let mut attrs = Vec::new();
106        let mut to_sql_config: Option<ToSqlConfig> = None;
107
108        let parser = Punctuated::<Attribute, Token![,]>::parse_terminated;
109        let attr_ts = attr.clone();
110        let punctuated_attrs = parser
111            .parse2(attr)
112            .more_error(lazy_err!(attr_ts, "failed parsing pg_extern arguments"))?;
113        for pair in punctuated_attrs.into_pairs() {
114            match pair.into_value() {
115                Attribute::Sql(config) => to_sql_config = to_sql_config.or(Some(config)),
116                attr => attrs.push(attr),
117            }
118        }
119
120        let mut to_sql_config = to_sql_config.unwrap_or_default();
121
122        let func = syn::parse2::<syn::ItemFn>(item)?;
123
124        if let Some(ref mut content) = to_sql_config.content {
125            let value = content.value();
126            // FIXME: find out if we should be using synthetic spans, issue #1667
127            let span = content.span();
128            let updated_value =
129                value.replace("@FUNCTION_NAME@", &(func.sig.ident.to_string() + "_wrapper")) + "\n";
130            *content = syn::LitStr::new(&updated_value, span);
131        }
132
133        if !to_sql_config.overrides_default() {
134            crate::ident_is_acceptable_to_postgres(&func.sig.ident)?;
135        }
136        let operator = Self::operator(&func)?;
137        let search_path = Self::search_path(&func)?;
138        let inputs = Self::inputs(&func)?;
139        let input_types = Self::input_types(&func)?;
140        let returns = Returning::try_from(&func.sig.output)?;
141        Ok(CodeEnrichment(Self {
142            attrs,
143            func,
144            to_sql_config,
145            operator,
146            cast: None,
147            search_path,
148            inputs,
149            input_types,
150            returns,
151        }))
152    }
153
154    /// Returns a new instance of this `PgExtern` with `cast` overwritten to `pg_cast`.
155    pub fn as_cast(&self, pg_cast: PgCast) -> PgExtern {
156        let mut result = self.clone();
157        result.cast = Some(pg_cast);
158        result
159    }
160
161    #[track_caller]
162    fn input_types(func: &syn::ItemFn) -> syn::Result<Vec<syn::Type>> {
163        func.sig
164            .inputs
165            .iter()
166            .filter_map(|v| -> Option<syn::Result<syn::Type>> {
167                match v {
168                    syn::FnArg::Receiver(_) => None,
169                    syn::FnArg::Typed(pat_ty) => {
170                        let ty = match UsedType::new(*pat_ty.ty.clone()) {
171                            Ok(v) => v.resolved_ty,
172                            Err(e) => return Some(Err(e)),
173                        };
174                        Some(Ok(ty))
175                    }
176                }
177            })
178            .collect()
179    }
180
181    fn name(&self) -> String {
182        self.attrs
183            .iter()
184            .find_map(|a| match a {
185                Attribute::Name(name) => Some(name.value()),
186                _ => None,
187            })
188            .unwrap_or_else(|| self.func.sig.ident.to_string())
189    }
190
191    fn schema(&self) -> Option<String> {
192        self.attrs.iter().find_map(|a| match a {
193            Attribute::Schema(name) => Some(name.value()),
194            _ => None,
195        })
196    }
197
198    pub fn extern_attrs(&self) -> &[Attribute] {
199        self.attrs.as_slice()
200    }
201
202    #[track_caller]
203    fn overridden(&self) -> Option<syn::LitStr> {
204        let mut span = None;
205        let mut retval = None;
206        let mut in_commented_sql_block = false;
207        for meta in self.func.attrs.iter().filter_map(|attr| {
208            if attr.meta.path().is_ident("doc") {
209                Some(attr.meta.clone())
210            } else {
211                None
212            }
213        }) {
214            let Meta::NameValue(syn::MetaNameValue { ref value, .. }) = meta else { continue };
215            let syn::Expr::Lit(syn::ExprLit { lit: syn::Lit::Str(inner), .. }) = value else {
216                continue;
217            };
218            span.get_or_insert(value.span());
219            if !in_commented_sql_block && inner.value().trim() == "```pgrxsql" {
220                in_commented_sql_block = true;
221            } else if in_commented_sql_block && inner.value().trim() == "```" {
222                in_commented_sql_block = false;
223            } else if in_commented_sql_block {
224                let line = inner
225                    .value()
226                    .trim_start()
227                    .replace("@FUNCTION_NAME@", &(self.func.sig.ident.to_string() + "_wrapper"))
228                    + "\n";
229                retval.get_or_insert_with(String::default).push_str(&line);
230            }
231        }
232        retval.map(|s| syn::LitStr::new(s.as_ref(), span.unwrap()))
233    }
234
235    #[track_caller]
236    fn operator(func: &syn::ItemFn) -> syn::Result<Option<PgOperator>> {
237        let mut skel = Option::<PgOperator>::default();
238        for attr in &func.attrs {
239            let last_segment = attr.path().segments.last().unwrap();
240            match last_segment.ident.to_string().as_str() {
241                s @ "opname" => {
242                    // we've been accepting strings of tokens for a while now
243                    let attr = attr
244                        .parse_args::<PgrxOperatorOpName>()
245                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
246                    skel.get_or_insert_with(Default::default).opname.get_or_insert(attr);
247                }
248                s @ "commutator" => {
249                    let attr: PgrxOperatorAttributeWithIdent = attr
250                        .parse_args()
251                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
252
253                    skel.get_or_insert_with(Default::default).commutator.get_or_insert(attr);
254                }
255                s @ "negator" => {
256                    let attr: PgrxOperatorAttributeWithIdent = attr
257                        .parse_args()
258                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
259                    skel.get_or_insert_with(Default::default).negator.get_or_insert(attr);
260                }
261                s @ "join" => {
262                    let attr: PgrxOperatorAttributeWithIdent = attr
263                        .parse_args()
264                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
265
266                    skel.get_or_insert_with(Default::default).join.get_or_insert(attr);
267                }
268                s @ "restrict" => {
269                    let attr: PgrxOperatorAttributeWithIdent = attr
270                        .parse_args()
271                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
272
273                    skel.get_or_insert_with(Default::default).restrict.get_or_insert(attr);
274                }
275                "hashes" => {
276                    skel.get_or_insert_with(Default::default).hashes = true;
277                }
278                "merges" => {
279                    skel.get_or_insert_with(Default::default).merges = true;
280                }
281                _ => (),
282            }
283        }
284        Ok(skel)
285    }
286
287    fn search_path(func: &syn::ItemFn) -> syn::Result<Option<SearchPathList>> {
288        func.attrs
289            .iter()
290            .find(|f| {
291                f.path()
292                    .segments
293                    .first()
294                    // FIXME: find out if we should be using synthetic spans, issue #1667
295                    .map(|f| f.ident == Ident::new("search_path", func.span()))
296                    .unwrap_or_default()
297            })
298            .map(|attr| attr.parse_args::<SearchPathList>())
299            .transpose()
300    }
301
302    fn inputs(func: &syn::ItemFn) -> syn::Result<Vec<PgExternArgument>> {
303        let mut args = Vec::default();
304        for input in &func.sig.inputs {
305            let arg = PgExternArgument::build(input.clone())?;
306            args.push(arg);
307        }
308        Ok(args)
309    }
310
311    fn entity_tokens(&self) -> TokenStream2 {
312        let ident = &self.func.sig.ident;
313        let name = self.name();
314        let unsafety = &self.func.sig.unsafety;
315        let schema = self.schema();
316        let schema_iter = schema.iter();
317        let extern_attrs = self
318            .attrs
319            .iter()
320            .map(|attr| attr.to_sql_entity_graph_tokens())
321            .collect::<Punctuated<_, Token![,]>>();
322        let search_path = self.search_path.clone().into_iter();
323        let inputs = &self.inputs;
324        let inputs_iter = inputs.iter().map(|v| v.entity_tokens());
325
326        let input_types = self.input_types.iter().cloned();
327
328        let returns = &self.returns;
329
330        let return_type = match &self.func.sig.output {
331            syn::ReturnType::Default => None,
332            syn::ReturnType::Type(arrow, ty) => {
333                let ty = ty.clone();
334                Some(syn::ReturnType::Type(*arrow, ty))
335            }
336        };
337
338        let operator = self.operator.clone().into_iter();
339        let cast = self.cast.clone().into_iter();
340        let to_sql_config = match self.overridden() {
341            None => self.to_sql_config.clone(),
342            Some(content) => ToSqlConfig { content: Some(content), ..self.to_sql_config.clone() },
343        };
344
345        let lifetimes = self
346            .func
347            .sig
348            .generics
349            .params
350            .iter()
351            .filter_map(|generic| match generic {
352                // syn::GenericParam::Type(_) => todo!(),
353                syn::GenericParam::Lifetime(lt) => Some(lt),
354                _ => None, // syn::GenericParam::Const(_) => todo!(),
355            })
356            .collect::<Vec<_>>();
357        let hrtb = if lifetimes.is_empty() { None } else { Some(quote! { for<#(#lifetimes),*> }) };
358
359        let sql_graph_entity_fn_name = format_ident!("__pgrx_internals_fn_{}", ident);
360        quote_spanned! { self.func.sig.span() =>
361            #[no_mangle]
362            #[doc(hidden)]
363            #[allow(unknown_lints, clippy::no_mangle_with_rust_abi)]
364            pub extern "Rust" fn  #sql_graph_entity_fn_name() -> ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity {
365                extern crate alloc;
366                #[allow(unused_imports)]
367                use alloc::{vec, vec::Vec};
368                type FunctionPointer = #hrtb #unsafety fn(#( #input_types ),*) #return_type;
369                let submission = ::pgrx::pgrx_sql_entity_graph::PgExternEntity {
370                    name: #name,
371                    unaliased_name: stringify!(#ident),
372                    module_path: core::module_path!(),
373                    full_path: concat!(core::module_path!(), "::", stringify!(#ident)),
374                    metadata: <FunctionPointer as ::pgrx::pgrx_sql_entity_graph::metadata::FunctionMetadata::<_>>::entity(),
375                    fn_args: vec![#(#inputs_iter),*],
376                    fn_return: #returns,
377                    #[allow(clippy::or_fun_call)]
378                    schema: None #( .unwrap_or_else(|| Some(#schema_iter)) )*,
379                    file: file!(),
380                    line: line!(),
381                    extern_attrs: vec![#extern_attrs],
382                    #[allow(clippy::or_fun_call)]
383                    search_path: None #( .unwrap_or_else(|| Some(vec![#search_path])) )*,
384                    #[allow(clippy::or_fun_call)]
385                    operator: None #( .unwrap_or_else(|| Some(#operator)) )*,
386                    cast: None #( .unwrap_or_else(|| Some(#cast)) )*,
387                    to_sql_config: #to_sql_config,
388                };
389                ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity::Function(submission)
390            }
391        }
392    }
393
394    pub fn wrapper_func(&self) -> Result<syn::ItemFn, syn::Error> {
395        let signature = &self.func.sig;
396        let func_name = &signature.ident;
397        // we do this odd dance so we can pass the same ident to macros that don't know each other
398        let synthetic_ident_span = Span::mixed_site().located_at(signature.ident.span());
399        let fcinfo_ident = syn::Ident::new("fcinfo", synthetic_ident_span);
400        let mut lifetimes = signature
401            .generics
402            .lifetimes()
403            .cloned()
404            .collect::<syn::punctuated::Punctuated<_, Comma>>();
405        // we pick an arbitrary lifetime from the provided signature of the fn, if available,
406        // so lifetime-bound fn are easier to write with pgrx
407        let fc_lt = lifetimes
408            .first()
409            .map(|lt_p| lt_p.lifetime.clone())
410            .filter(|lt| lt.ident != "static")
411            .unwrap_or(syn::Lifetime::new("'fcx", Span::mixed_site()));
412        // we need the bare lifetime later, but we also jam it into the bounds if it is new
413        let fc_ltparam = syn::LifetimeParam::new(fc_lt.clone());
414        if lifetimes.first() != Some(&fc_ltparam) {
415            lifetimes.insert(0, fc_ltparam)
416        }
417
418        let args = &self.inputs;
419        // for unclear reasons the linker vomits if we don't do this
420        let arg_pats = args.iter().map(|v| format_ident!("{}_", &v.pat)).collect::<Vec<_>>();
421        let args_ident = proc_macro2::Ident::new("_args", Span::call_site());
422        let arg_fetches = arg_pats.iter().map(|pat| {
423                quote_spanned!{ pat.span() =>
424                    let #pat = #args_ident.next_arg_unchecked().unwrap_or_else(|| panic!("unboxing {} argument failed", stringify!(#pat)));
425                }
426            }
427        );
428
429        match &self.returns {
430            Returning::None
431            | Returning::Type(_)
432            | Returning::SetOf { .. }
433            | Returning::Iterated { .. } => {
434                let ret_ty = match &signature.output {
435                    syn::ReturnType::Default => syn::parse_quote! { () },
436                    syn::ReturnType::Type(_, ret_ty) => ret_ty.clone(),
437                };
438                let wrapper_code = quote_spanned! { self.func.block.span() =>
439                    fn _internal_wrapper<#lifetimes>(fcinfo: &mut ::pgrx::callconv::FcInfo<#fc_lt>) -> ::pgrx::datum::Datum<#fc_lt> {
440                        #[allow(unused_unsafe)]
441                        unsafe {
442                            let call_flow = <#ret_ty as ::pgrx::callconv::RetAbi>::check_and_prepare(fcinfo);
443                            let result = match call_flow {
444                                ::pgrx::callconv::CallCx::WrappedFn(mcx) => {
445                                    let mut mcx = ::pgrx::PgMemoryContexts::For(mcx);
446                                    let #args_ident = &mut fcinfo.args();
447                                    let call_result = mcx.switch_to(|_| {
448                                        #(#arg_fetches)*
449                                        #func_name( #(#arg_pats),* )
450                                    });
451                                    ::pgrx::callconv::RetAbi::to_ret(call_result)
452                                }
453                                ::pgrx::callconv::CallCx::RestoreCx => <#ret_ty as ::pgrx::callconv::RetAbi>::ret_from_fcx(fcinfo),
454                            };
455                            unsafe { <#ret_ty as ::pgrx::callconv::RetAbi>::box_ret_in(fcinfo, result) }
456                        }
457                    }
458                    // We preserve the invariants
459                    let datum = unsafe {
460                        ::pgrx::pg_sys::submodules::panic::pgrx_extern_c_guard(|| {
461                            let mut fcinfo = ::pgrx::callconv::FcInfo::from_ptr(#fcinfo_ident);
462                            _internal_wrapper(&mut fcinfo)
463                        })
464                    };
465                    datum.sans_lifetime()
466                };
467                finfo_v1_extern_c(&self.func, fcinfo_ident, wrapper_code)
468            }
469        }
470    }
471}
472
473trait LastIdent {
474    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment>;
475    #[inline]
476    fn last_ident_is(&self, id: &str) -> bool {
477        self.filter_last_ident(id).is_some()
478    }
479}
480
481impl LastIdent for syn::Type {
482    #[inline]
483    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment> {
484        let syn::Type::Path(syn::TypePath { path, .. }) = self else { return None };
485        path.filter_last_ident(id)
486    }
487}
488impl LastIdent for syn::TypePath {
489    #[inline]
490    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment> {
491        self.path.filter_last_ident(id)
492    }
493}
494impl LastIdent for syn::Path {
495    #[inline]
496    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment> {
497        self.segments.filter_last_ident(id)
498    }
499}
500impl<P> LastIdent for Punctuated<syn::PathSegment, P> {
501    #[inline]
502    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment> {
503        self.last().filter(|segment| segment.ident == id)
504    }
505}
506
507impl ToEntityGraphTokens for PgExtern {
508    fn to_entity_graph_tokens(&self) -> TokenStream2 {
509        self.entity_tokens()
510    }
511}
512
513impl ToRustCodeTokens for PgExtern {
514    fn to_rust_code_tokens(&self) -> TokenStream2 {
515        let original_func = &self.func;
516        let wrapper_func = self.wrapper_func().unwrap();
517        let finfo_tokens = finfo_v1_tokens(wrapper_func.sig.ident.clone()).unwrap();
518
519        quote_spanned! { self.func.sig.span() =>
520            #original_func
521            #wrapper_func
522            #finfo_tokens
523        }
524    }
525}
526
527impl Parse for CodeEnrichment<PgExtern> {
528    fn parse(input: ParseStream) -> Result<Self, syn::Error> {
529        let parser = Punctuated::<Attribute, Token![,]>::parse_terminated;
530        let punctuated_attrs = input.call(parser).ok().unwrap_or_default();
531        let attrs = punctuated_attrs.into_pairs().map(|pair| pair.into_value());
532        PgExtern::new(quote! {#(#attrs)*}, input.parse()?)
533    }
534}
pgrx_sql_entity_graph/pg_extern/mod.rs

pgrx_sql_entity_graph/pg_extern/
mod.rs
