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::ToSqlConfig;
36use crate::enrich::{CodeEnrichment, ToEntityGraphTokens, ToRustCodeTokens};
37use crate::finfo::{finfo_v1_extern_c, finfo_v1_tokens};
38use crate::fmt::ErrHarder;
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") { Some(attr.meta.clone()) } else { None }
209        }) {
210            let Meta::NameValue(syn::MetaNameValue { ref value, .. }) = meta else { continue };
211            let syn::Expr::Lit(syn::ExprLit { lit: syn::Lit::Str(inner), .. }) = value else {
212                continue;
213            };
214            span.get_or_insert(value.span());
215            if !in_commented_sql_block && inner.value().trim() == "```pgrxsql" {
216                in_commented_sql_block = true;
217            } else if in_commented_sql_block && inner.value().trim() == "```" {
218                in_commented_sql_block = false;
219            } else if in_commented_sql_block {
220                let line = inner
221                    .value()
222                    .trim_start()
223                    .replace("@FUNCTION_NAME@", &(self.func.sig.ident.to_string() + "_wrapper"))
224                    + "\n";
225                retval.get_or_insert_with(String::default).push_str(&line);
226            }
227        }
228        retval.map(|s| syn::LitStr::new(s.as_ref(), span.unwrap()))
229    }
230
231    #[track_caller]
232    fn operator(func: &syn::ItemFn) -> syn::Result<Option<PgOperator>> {
233        let mut skel = Option::<PgOperator>::default();
234        for attr in &func.attrs {
235            let last_segment = attr.path().segments.last().unwrap();
236            match last_segment.ident.to_string().as_str() {
237                s @ "opname" => {
238                    // we've been accepting strings of tokens for a while now
239                    let attr = attr
240                        .parse_args::<PgrxOperatorOpName>()
241                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
242                    skel.get_or_insert_with(Default::default).opname.get_or_insert(attr);
243                }
244                s @ "commutator" => {
245                    let attr: PgrxOperatorAttributeWithIdent = attr
246                        .parse_args()
247                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
248
249                    skel.get_or_insert_with(Default::default).commutator.get_or_insert(attr);
250                }
251                s @ "negator" => {
252                    let attr: PgrxOperatorAttributeWithIdent = attr
253                        .parse_args()
254                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
255                    skel.get_or_insert_with(Default::default).negator.get_or_insert(attr);
256                }
257                s @ "join" => {
258                    let attr: PgrxOperatorAttributeWithIdent = attr
259                        .parse_args()
260                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
261
262                    skel.get_or_insert_with(Default::default).join.get_or_insert(attr);
263                }
264                s @ "restrict" => {
265                    let attr: PgrxOperatorAttributeWithIdent = attr
266                        .parse_args()
267                        .more_error(lazy_err!(attr.clone(), "bad parse of {s}?"))?;
268
269                    skel.get_or_insert_with(Default::default).restrict.get_or_insert(attr);
270                }
271                "hashes" => {
272                    skel.get_or_insert_with(Default::default).hashes = true;
273                }
274                "merges" => {
275                    skel.get_or_insert_with(Default::default).merges = true;
276                }
277                _ => (),
278            }
279        }
280        Ok(skel)
281    }
282
283    fn search_path(func: &syn::ItemFn) -> syn::Result<Option<SearchPathList>> {
284        func.attrs
285            .iter()
286            .find(|f| {
287                f.path()
288                    .segments
289                    .first()
290                    // FIXME: find out if we should be using synthetic spans, issue #1667
291                    .map(|f| f.ident == Ident::new("search_path", func.span()))
292                    .unwrap_or_default()
293            })
294            .map(|attr| attr.parse_args::<SearchPathList>())
295            .transpose()
296    }
297
298    fn inputs(func: &syn::ItemFn) -> syn::Result<Vec<PgExternArgument>> {
299        let mut args = Vec::default();
300        for input in &func.sig.inputs {
301            let arg = PgExternArgument::build(input.clone())?;
302            args.push(arg);
303        }
304        Ok(args)
305    }
306
307    fn entity_tokens(&self) -> TokenStream2 {
308        let ident = &self.func.sig.ident;
309        let name = self.name();
310        let unsafety = &self.func.sig.unsafety;
311        let schema = self.schema();
312        let schema_iter = schema.iter();
313        let extern_attrs = self
314            .attrs
315            .iter()
316            .map(|attr| attr.to_sql_entity_graph_tokens())
317            .collect::<Punctuated<_, Token![,]>>();
318        let search_path = self.search_path.clone().into_iter();
319        let inputs = &self.inputs;
320        let inputs_iter = inputs.iter().map(|v| v.entity_tokens());
321
322        let input_types = self.input_types.iter().cloned();
323
324        let returns = &self.returns;
325
326        let return_type = match &self.func.sig.output {
327            syn::ReturnType::Default => None,
328            syn::ReturnType::Type(arrow, ty) => {
329                let ty = ty.clone();
330                Some(syn::ReturnType::Type(*arrow, ty))
331            }
332        };
333
334        let operator = self.operator.clone().into_iter();
335        let cast = self.cast.clone().into_iter();
336        let to_sql_config = match self.overridden() {
337            None => self.to_sql_config.clone(),
338            Some(content) => ToSqlConfig { content: Some(content), ..self.to_sql_config.clone() },
339        };
340
341        let lifetimes = self
342            .func
343            .sig
344            .generics
345            .params
346            .iter()
347            .filter_map(|generic| match generic {
348                // syn::GenericParam::Type(_) => todo!(),
349                syn::GenericParam::Lifetime(lt) => Some(lt),
350                _ => None, // syn::GenericParam::Const(_) => todo!(),
351            })
352            .collect::<Vec<_>>();
353        let hrtb = if lifetimes.is_empty() { None } else { Some(quote! { for<#(#lifetimes),*> }) };
354
355        let sql_graph_entity_fn_name = format_ident!("__pgrx_internals_fn_{}", ident);
356        quote_spanned! { self.func.sig.span() =>
357            #[unsafe(no_mangle)]
358            #[doc(hidden)]
359            #[allow(unknown_lints, clippy::no_mangle_with_rust_abi)]
360            pub extern "Rust" fn  #sql_graph_entity_fn_name() -> ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity {
361                extern crate alloc;
362                #[allow(unused_imports)]
363                use alloc::{vec, vec::Vec};
364                type FunctionPointer = #hrtb #unsafety fn(#( #input_types ),*) #return_type;
365                let submission = ::pgrx::pgrx_sql_entity_graph::PgExternEntity {
366                    name: #name,
367                    unaliased_name: stringify!(#ident),
368                    module_path: core::module_path!(),
369                    full_path: concat!(core::module_path!(), "::", stringify!(#ident)),
370                    metadata: <FunctionPointer as ::pgrx::pgrx_sql_entity_graph::metadata::FunctionMetadata::<_>>::entity(),
371                    fn_args: vec![#(#inputs_iter),*],
372                    fn_return: #returns,
373                    #[allow(clippy::or_fun_call)]
374                    schema: None #( .unwrap_or_else(|| Some(#schema_iter)) )*,
375                    file: file!(),
376                    line: line!(),
377                    extern_attrs: vec![#extern_attrs],
378                    #[allow(clippy::or_fun_call)]
379                    search_path: None #( .unwrap_or_else(|| Some(vec![#search_path])) )*,
380                    #[allow(clippy::or_fun_call)]
381                    operator: None #( .unwrap_or_else(|| Some(#operator)) )*,
382                    cast: None #( .unwrap_or_else(|| Some(#cast)) )*,
383                    to_sql_config: #to_sql_config,
384                };
385                ::pgrx::pgrx_sql_entity_graph::SqlGraphEntity::Function(submission)
386            }
387        }
388    }
389
390    pub fn wrapper_func(&self) -> Result<syn::ItemFn, syn::Error> {
391        let signature = &self.func.sig;
392        let func_name = &signature.ident;
393        // we do this odd dance so we can pass the same ident to macros that don't know each other
394        let synthetic_ident_span = Span::mixed_site().located_at(signature.ident.span());
395        let fcinfo_ident = syn::Ident::new("fcinfo", synthetic_ident_span);
396        let mut lifetimes = signature
397            .generics
398            .lifetimes()
399            .cloned()
400            .collect::<syn::punctuated::Punctuated<_, Comma>>();
401        // we pick an arbitrary lifetime from the provided signature of the fn, if available,
402        // so lifetime-bound fn are easier to write with pgrx
403        let fc_lt = lifetimes
404            .first()
405            .map(|lt_p| lt_p.lifetime.clone())
406            .filter(|lt| lt.ident != "static")
407            .unwrap_or(syn::Lifetime::new("'fcx", Span::mixed_site()));
408        // we need the bare lifetime later, but we also jam it into the bounds if it is new
409        let fc_ltparam = syn::LifetimeParam::new(fc_lt.clone());
410        if lifetimes.first() != Some(&fc_ltparam) {
411            lifetimes.insert(0, fc_ltparam)
412        }
413
414        let args = &self.inputs;
415        // for unclear reasons the linker vomits if we don't do this
416        let arg_pats = args.iter().map(|v| format_ident!("{}_", &v.pat)).collect::<Vec<_>>();
417        let args_ident = proc_macro2::Ident::new("_args", Span::call_site());
418        let arg_fetches = arg_pats.iter().map(|pat| {
419                quote_spanned!{ pat.span() =>
420                    let #pat = #args_ident.next_arg_unchecked().unwrap_or_else(|| panic!("unboxing {} argument failed", stringify!(#pat)));
421                }
422            }
423        );
424
425        match &self.returns {
426            Returning::None
427            | Returning::Type(_)
428            | Returning::SetOf { .. }
429            | Returning::Iterated { .. } => {
430                let ret_ty = match &signature.output {
431                    syn::ReturnType::Default => syn::parse_quote! { () },
432                    syn::ReturnType::Type(_, ret_ty) => ret_ty.clone(),
433                };
434                let wrapper_code = quote_spanned! { self.func.block.span() =>
435                    fn _internal_wrapper<#lifetimes>(fcinfo: &mut ::pgrx::callconv::FcInfo<#fc_lt>) -> ::pgrx::datum::Datum<#fc_lt> {
436                        #[allow(unused_unsafe)]
437                        unsafe {
438                            let call_flow = <#ret_ty as ::pgrx::callconv::RetAbi>::check_and_prepare(fcinfo);
439                            let result = match call_flow {
440                                ::pgrx::callconv::CallCx::WrappedFn(mcx) => {
441                                    let mut mcx = ::pgrx::PgMemoryContexts::For(mcx);
442                                    let #args_ident = &mut fcinfo.args();
443                                    let call_result = mcx.switch_to(|_| {
444                                        #(#arg_fetches)*
445                                        #func_name( #(#arg_pats),* )
446                                    });
447                                    ::pgrx::callconv::RetAbi::to_ret(call_result)
448                                }
449                                ::pgrx::callconv::CallCx::RestoreCx => <#ret_ty as ::pgrx::callconv::RetAbi>::ret_from_fcx(fcinfo),
450                            };
451                            unsafe { <#ret_ty as ::pgrx::callconv::RetAbi>::box_ret_in(fcinfo, result) }
452                        }
453                    }
454                    // We preserve the invariants
455                    let datum = unsafe {
456                        ::pgrx::pg_sys::submodules::panic::pgrx_extern_c_guard(|| {
457                            let mut fcinfo = ::pgrx::callconv::FcInfo::from_ptr(#fcinfo_ident);
458                            _internal_wrapper(&mut fcinfo)
459                        })
460                    };
461                    datum.sans_lifetime()
462                };
463                finfo_v1_extern_c(&self.func, fcinfo_ident, wrapper_code)
464            }
465        }
466    }
467}
468
469trait LastIdent {
470    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment>;
471    #[inline]
472    fn last_ident_is(&self, id: &str) -> bool {
473        self.filter_last_ident(id).is_some()
474    }
475}
476
477impl LastIdent for syn::Type {
478    #[inline]
479    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment> {
480        let syn::Type::Path(syn::TypePath { path, .. }) = self else { return None };
481        path.filter_last_ident(id)
482    }
483}
484impl LastIdent for syn::TypePath {
485    #[inline]
486    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment> {
487        self.path.filter_last_ident(id)
488    }
489}
490impl LastIdent for syn::Path {
491    #[inline]
492    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment> {
493        self.segments.filter_last_ident(id)
494    }
495}
496impl<P> LastIdent for Punctuated<syn::PathSegment, P> {
497    #[inline]
498    fn filter_last_ident(&self, id: &str) -> Option<&syn::PathSegment> {
499        self.last().filter(|segment| segment.ident == id)
500    }
501}
502
503impl ToEntityGraphTokens for PgExtern {
504    fn to_entity_graph_tokens(&self) -> TokenStream2 {
505        self.entity_tokens()
506    }
507}
508
509impl ToRustCodeTokens for PgExtern {
510    fn to_rust_code_tokens(&self) -> TokenStream2 {
511        let original_func = &self.func;
512        let wrapper_func = self.wrapper_func().unwrap();
513        let finfo_tokens = finfo_v1_tokens(wrapper_func.sig.ident.clone()).unwrap();
514
515        quote_spanned! { self.func.sig.span() =>
516            #original_func
517            #wrapper_func
518            #finfo_tokens
519        }
520    }
521}
522
523impl Parse for CodeEnrichment<PgExtern> {
524    fn parse(input: ParseStream) -> Result<Self, syn::Error> {
525        let parser = Punctuated::<Attribute, Token![,]>::parse_terminated;
526        let punctuated_attrs = input.call(parser).ok().unwrap_or_default();
527        let attrs = punctuated_attrs.into_pairs().map(|pair| pair.into_value());
528        PgExtern::new(quote! {#(#attrs)*}, input.parse()?)
529    }
530}
pgrx_sql_entity_graph/pg_extern/mod.rs

pgrx_sql_entity_graph/pg_extern/
mod.rs
