pgx_utils/sql_entity_graph/pg_extern/
mod.rs

1/*
2Portions Copyright 2019-2021 ZomboDB, LLC.
3Portions Copyright 2021-2022 Technology Concepts & Design, Inc. <support@tcdi.com>
4
5All rights reserved.
6
7Use of this source code is governed by the MIT license that can be found in the LICENSE file.
8*/
9
10/*!
11
12`#[pg_extern]` related macro expansion for Rust to SQL translation
13
14> Like all of the [`sql_entity_graph`][crate::sql_entity_graph] APIs, this is considered **internal**
15to the `pgx` framework and very subject to change between versions. While you may use this, please do it with caution.
16
17*/
18mod argument;
19mod attribute;
20pub mod entity;
21mod operator;
22mod returning;
23mod search_path;
24
25pub use argument::PgExternArgument;
26pub use operator::PgOperator;
27pub use returning::NameMacro;
28
29use crate::sql_entity_graph::ToSqlConfig;
30use crate::staticize_lifetimes;
31use attribute::Attribute;
32use operator::{PgxOperatorAttributeWithIdent, PgxOperatorOpName};
33use search_path::SearchPathList;
34
35use eyre::WrapErr;
36use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
37use quote::{quote, quote_spanned, ToTokens, TokenStreamExt};
38use syn::parse::{Parse, ParseStream, Parser};
39use syn::punctuated::Punctuated;
40use syn::spanned::Spanned;
41use syn::{Meta, Token};
42
43use self::returning::Returning;
44
45use super::UsedType;
46
47/// A parsed `#[pg_extern]` item.
48///
49/// It should be used with [`syn::parse::Parse`] functions.
50///
51/// Using [`quote::ToTokens`] will output the declaration for a [`PgExternEntity`][crate::sql_entity_graph::PgExternEntity].
52///
53/// ```rust
54/// use syn::{Macro, parse::Parse, parse_quote, parse};
55/// use quote::{quote, ToTokens};
56/// use pgx_utils::sql_entity_graph::PgExtern;
57///
58/// # fn main() -> eyre::Result<()> {
59/// let parsed: PgExtern = parse_quote! {
60///     fn example(x: Option<str>) -> Option<&'a str> {
61///         unimplemented!()
62///     }
63/// };
64/// let sql_graph_entity_tokens = parsed.to_token_stream();
65/// # Ok(())
66/// # }
67/// ```
68#[derive(Debug, Clone)]
69pub struct PgExtern {
70    attrs: Vec<Attribute>,
71    func: syn::ItemFn,
72    to_sql_config: ToSqlConfig,
73}
74
75impl PgExtern {
76    pub fn new(attr: TokenStream2, item: TokenStream2) -> Result<Self, syn::Error> {
77        let mut attrs = Vec::new();
78        let mut to_sql_config: Option<ToSqlConfig> = None;
79
80        let parser = Punctuated::<Attribute, Token![,]>::parse_terminated;
81        let punctuated_attrs = parser.parse2(attr)?;
82        for pair in punctuated_attrs.into_pairs() {
83            match pair.into_value() {
84                Attribute::Sql(config) => {
85                    to_sql_config.get_or_insert(config);
86                }
87                attr => {
88                    attrs.push(attr);
89                }
90            }
91        }
92
93        let mut to_sql_config = to_sql_config.unwrap_or_default();
94
95        let func = syn::parse2::<syn::ItemFn>(item)?;
96
97        if let Some(ref mut content) = to_sql_config.content {
98            let value = content.value();
99            let updated_value = value
100                .replace("@FUNCTION_NAME@", &*(func.sig.ident.to_string() + "_wrapper"))
101                + "\n";
102            *content = syn::LitStr::new(&updated_value, Span::call_site());
103        }
104
105        if !to_sql_config.overrides_default() {
106            crate::ident_is_acceptable_to_postgres(&func.sig.ident)?;
107        }
108
109        Ok(Self { attrs, func, to_sql_config })
110    }
111
112    fn name(&self) -> String {
113        self.attrs
114            .iter()
115            .find_map(|a| match a {
116                Attribute::Name(name) => Some(name.value()),
117                _ => None,
118            })
119            .unwrap_or_else(|| self.func.sig.ident.to_string())
120    }
121
122    fn schema(&self) -> Option<String> {
123        self.attrs.iter().find_map(|a| match a {
124            Attribute::Schema(name) => Some(name.value()),
125            _ => None,
126        })
127    }
128
129    pub fn extern_attrs(&self) -> &[Attribute] {
130        self.attrs.as_slice()
131    }
132
133    fn overridden(&self) -> Option<syn::LitStr> {
134        let mut span = None;
135        let mut retval = None;
136        let mut in_commented_sql_block = false;
137        for attr in &self.func.attrs {
138            let meta = attr.parse_meta().ok();
139            if let Some(meta) = meta {
140                if meta.path().is_ident("doc") {
141                    let content = match meta {
142                        Meta::Path(_) | Meta::List(_) => continue,
143                        Meta::NameValue(mnv) => mnv,
144                    };
145                    if let syn::Lit::Str(ref inner) = content.lit {
146                        span.get_or_insert(content.lit.span());
147                        if !in_commented_sql_block && inner.value().trim() == "```pgxsql" {
148                            in_commented_sql_block = true;
149                        } else if in_commented_sql_block && inner.value().trim() == "```" {
150                            in_commented_sql_block = false;
151                        } else if in_commented_sql_block {
152                            let sql = retval.get_or_insert_with(String::default);
153                            let line = inner.value().trim_start().replace(
154                                "@FUNCTION_NAME@",
155                                &*(self.func.sig.ident.to_string() + "_wrapper"),
156                            ) + "\n";
157                            sql.push_str(&*line);
158                        }
159                    }
160                }
161            }
162        }
163        retval.map(|s| syn::LitStr::new(s.as_ref(), span.unwrap()))
164    }
165
166    fn operator(&self) -> Option<PgOperator> {
167        let mut skel = Option::<PgOperator>::default();
168        for attr in &self.func.attrs {
169            let last_segment = attr.path.segments.last().unwrap();
170            match last_segment.ident.to_string().as_str() {
171                "opname" => {
172                    let attr: PgxOperatorOpName = syn::parse2(attr.tokens.clone())
173                        .expect(&format!("Unable to parse {:?}", &attr.tokens));
174                    skel.get_or_insert_with(Default::default).opname.get_or_insert(attr);
175                }
176                "commutator" => {
177                    let attr: PgxOperatorAttributeWithIdent = syn::parse2(attr.tokens.clone())
178                        .expect(&format!("Unable to parse {:?}", &attr.tokens));
179                    skel.get_or_insert_with(Default::default).commutator.get_or_insert(attr);
180                }
181                "negator" => {
182                    let attr: PgxOperatorAttributeWithIdent = syn::parse2(attr.tokens.clone())
183                        .expect(&format!("Unable to parse {:?}", &attr.tokens));
184                    skel.get_or_insert_with(Default::default).negator.get_or_insert(attr);
185                }
186                "join" => {
187                    let attr: PgxOperatorAttributeWithIdent = syn::parse2(attr.tokens.clone())
188                        .expect(&format!("Unable to parse {:?}", &attr.tokens));
189                    skel.get_or_insert_with(Default::default).join.get_or_insert(attr);
190                }
191                "restrict" => {
192                    let attr: PgxOperatorAttributeWithIdent = syn::parse2(attr.tokens.clone())
193                        .expect(&format!("Unable to parse {:?}", &attr.tokens));
194                    skel.get_or_insert_with(Default::default).restrict.get_or_insert(attr);
195                }
196                "hashes" => {
197                    skel.get_or_insert_with(Default::default).hashes = true;
198                }
199                "merges" => {
200                    skel.get_or_insert_with(Default::default).merges = true;
201                }
202                _ => (),
203            }
204        }
205        skel
206    }
207
208    fn search_path(&self) -> Option<SearchPathList> {
209        self.func
210            .attrs
211            .iter()
212            .find(|f| {
213                f.path
214                    .segments
215                    .first()
216                    .map(|f| f.ident == Ident::new("search_path", Span::call_site()))
217                    .unwrap_or_default()
218            })
219            .and_then(|attr| Some(attr.parse_args::<SearchPathList>().unwrap()))
220    }
221
222    fn inputs(&self) -> eyre::Result<Vec<PgExternArgument>> {
223        let mut args = Vec::default();
224        for input in &self.func.sig.inputs {
225            let arg = PgExternArgument::build(input.clone())
226                .wrap_err_with(|| format!("Could not map {:?}", input))?;
227            args.push(arg);
228        }
229        Ok(args)
230    }
231
232    fn returns(&self) -> Result<Returning, syn::Error> {
233        Returning::try_from(&self.func.sig.output)
234    }
235
236    fn entity_tokens(&self) -> TokenStream2 {
237        let ident = &self.func.sig.ident;
238        let name = self.name();
239        let unsafety = &self.func.sig.unsafety;
240        let schema = self.schema();
241        let schema_iter = schema.iter();
242        let extern_attrs = self
243            .attrs
244            .iter()
245            .map(|attr| attr.to_sql_entity_graph_tokens())
246            .collect::<Punctuated<_, Token![,]>>();
247        let search_path = self.search_path().into_iter();
248        let inputs = self.inputs().unwrap();
249        let inputs_iter = inputs.iter().map(|v| v.entity_tokens());
250
251        let input_types = self.func.sig.inputs.iter().filter_map(|v| match v {
252            syn::FnArg::Receiver(_) => None,
253            syn::FnArg::Typed(pat_ty) => {
254                let static_ty = pat_ty.ty.clone();
255                let mut static_ty = UsedType::new(*static_ty).unwrap().resolved_ty;
256                staticize_lifetimes(&mut static_ty);
257                Some(static_ty)
258            }
259        });
260
261        let returns = match self.returns() {
262            Ok(returns) => returns,
263            Err(e) => {
264                let msg = e.to_string();
265                return quote! {
266                    std::compile_error!(#msg);
267                };
268            }
269        };
270
271        let return_type = match &self.func.sig.output {
272            syn::ReturnType::Default => None,
273            syn::ReturnType::Type(arrow, ty) => {
274                let mut static_ty = ty.clone();
275                staticize_lifetimes(&mut static_ty);
276                Some(syn::ReturnType::Type(*arrow, static_ty))
277            }
278        };
279
280        let operator = self.operator().into_iter();
281        let to_sql_config = match self.overridden() {
282            None => self.to_sql_config.clone(),
283            Some(content) => {
284                let mut config = self.to_sql_config.clone();
285                config.content = Some(content);
286                config
287            }
288        };
289
290        let sql_graph_entity_fn_name =
291            syn::Ident::new(&format!("__pgx_internals_fn_{}", ident), Span::call_site());
292        quote_spanned! { self.func.sig.span() =>
293            #[no_mangle]
294            #[doc(hidden)]
295            pub extern "Rust" fn  #sql_graph_entity_fn_name() -> ::pgx::utils::sql_entity_graph::SqlGraphEntity {
296                extern crate alloc;
297                #[allow(unused_imports)]
298                use alloc::{vec, vec::Vec};
299                type FunctionPointer = #unsafety fn(#( #input_types ),*) #return_type;
300                let metadata: FunctionPointer = #ident;
301                let submission = ::pgx::utils::sql_entity_graph::PgExternEntity {
302                    name: #name,
303                    unaliased_name: stringify!(#ident),
304                    module_path: core::module_path!(),
305                    full_path: concat!(core::module_path!(), "::", stringify!(#ident)),
306                    metadata: pgx::utils::sql_entity_graph::metadata::FunctionMetadata::entity(&metadata),
307                    fn_args: vec![#(#inputs_iter),*],
308                    fn_return: #returns,
309                    #[allow(clippy::or_fun_call)]
310                    schema: None #( .unwrap_or_else(|| Some(#schema_iter)) )*,
311                    file: file!(),
312                    line: line!(),
313                    extern_attrs: vec![#extern_attrs],
314                    #[allow(clippy::or_fun_call)]
315                    search_path: None #( .unwrap_or_else(|| Some(vec![#search_path])) )*,
316                    #[allow(clippy::or_fun_call)]
317                    operator: None #( .unwrap_or_else(|| Some(#operator)) )*,
318                    to_sql_config: #to_sql_config,
319                };
320                ::pgx::utils::sql_entity_graph::SqlGraphEntity::Function(submission)
321            }
322        }
323    }
324
325    fn finfo_tokens(&self) -> TokenStream2 {
326        let finfo_name = syn::Ident::new(
327            &format!("pg_finfo_{}_wrapper", self.func.sig.ident),
328            Span::call_site(),
329        );
330        quote_spanned! { self.func.sig.span() =>
331            #[no_mangle]
332            #[doc(hidden)]
333            pub extern "C" fn #finfo_name() -> &'static pg_sys::Pg_finfo_record {
334                const V1_API: pg_sys::Pg_finfo_record = pg_sys::Pg_finfo_record { api_version: 1 };
335                &V1_API
336            }
337        }
338    }
339
340    pub fn wrapper_func(&self) -> TokenStream2 {
341        let func_name = &self.func.sig.ident;
342        let func_name_wrapper = Ident::new(
343            &format!("{}_wrapper", &self.func.sig.ident.to_string()),
344            self.func.sig.ident.span(),
345        );
346        let func_generics = &self.func.sig.generics;
347        let is_raw = self.extern_attrs().contains(&Attribute::Raw);
348        // We use a `_` prefix to make functions with no args more satisfied during linting.
349        let fcinfo_ident = syn::Ident::new("_fcinfo", self.func.sig.ident.span());
350
351        let args = self.inputs().unwrap();
352        let arg_pats = args
353            .iter()
354            .map(|v| syn::Ident::new(&format!("{}_", &v.pat), self.func.sig.span()))
355            .collect::<Vec<_>>();
356        let arg_fetches = args.iter().enumerate().map(|(idx, arg)| {
357            let pat = &arg_pats[idx];
358            let resolved_ty = &arg.used_ty.resolved_ty;
359            if arg.used_ty.resolved_ty.to_token_stream().to_string() == quote!(pgx::pg_sys::FunctionCallInfo).to_token_stream().to_string()
360            || arg.used_ty.resolved_ty.to_token_stream().to_string() == quote!(pg_sys::FunctionCallInfo).to_token_stream().to_string() {
361                quote_spanned! {pat.span()=>
362                    let #pat = #fcinfo_ident;
363                }
364            } else if arg.used_ty.resolved_ty.to_token_stream().to_string() == quote!(()).to_token_stream().to_string() {
365                quote_spanned! {pat.span()=>
366                    debug_assert!(pgx::pg_getarg::<()>(#fcinfo_ident, #idx).is_none(), "A `()` argument should always recieve `NULL`");
367                    let #pat = ();
368                }
369            } else {
370                match (is_raw, &arg.used_ty.optional) {
371                    (true, None) | (true, Some(_)) => quote_spanned! { pat.span() =>
372                        let #pat = pgx::pg_getarg_datum_raw(#fcinfo_ident, #idx) as #resolved_ty;
373                    },
374                    (false, None) => quote_spanned! { pat.span() =>
375                        let #pat = pgx::pg_getarg::<#resolved_ty>(#fcinfo_ident, #idx).unwrap_or_else(|| panic!("{} is null", stringify!{#pat}));
376                    },
377                    (false, Some(inner)) => quote_spanned! { pat.span() =>
378                        let #pat = pgx::pg_getarg::<#inner>(#fcinfo_ident, #idx);
379                    },
380                }
381            }
382        });
383
384        match self.returns().unwrap() {
385            Returning::None => quote_spanned! { self.func.sig.span() =>
386                    #[no_mangle]
387                    #[doc(hidden)]
388                    #[pg_guard]
389                    pub unsafe extern "C" fn #func_name_wrapper #func_generics(#fcinfo_ident: ::pgx::pg_sys::FunctionCallInfo) {
390                        #(
391                            #arg_fetches
392                        )*
393
394                        #[allow(unused_unsafe)] // unwrapped fn might be unsafe
395                        unsafe { #func_name(#(#arg_pats),*) }
396                    }
397            },
398            Returning::Type(retval_ty) => {
399                let result_ident = syn::Ident::new("result", self.func.sig.span());
400                let retval_transform = if retval_ty.resolved_ty == syn::parse_quote!(()) {
401                    quote_spanned! { self.func.sig.output.span() =>
402                       pgx::pg_return_void()
403                    }
404                } else if retval_ty.resolved_ty == syn::parse_quote!(pg_sys::Datum)
405                    || retval_ty.resolved_ty == syn::parse_quote!(pgx::pg_sys::Datum)
406                {
407                    quote_spanned! { self.func.sig.output.span() =>
408                       #result_ident
409                    }
410                } else if retval_ty.optional.is_some() {
411                    quote_spanned! { self.func.sig.output.span() =>
412                        match #result_ident {
413                            Some(result) => {
414                                pgx::datum::IntoDatum::into_datum(result).unwrap_or_else(|| panic!("returned Option<T> was NULL"))
415                            },
416                            None => pgx::pg_return_null(#fcinfo_ident)
417                        }
418                    }
419                } else {
420                    quote_spanned! { self.func.sig.output.span() =>
421                        pgx::datum::IntoDatum::into_datum(#result_ident).unwrap_or_else(|| panic!("returned Datum was NULL"))
422                    }
423                };
424
425                quote_spanned! { self.func.sig.span() =>
426                    #[no_mangle]
427                    #[doc(hidden)]
428                    #[pg_guard]
429                    pub unsafe extern "C" fn #func_name_wrapper #func_generics(#fcinfo_ident: ::pgx::pg_sys::FunctionCallInfo) -> ::pgx::pg_sys::Datum {
430                        #(
431                            #arg_fetches
432                        )*
433
434                        #[allow(unused_unsafe)] // unwrapped fn might be unsafe
435                        let #result_ident = unsafe { #func_name(#(#arg_pats),*) };
436
437                        #retval_transform
438                    }
439                }
440            }
441            Returning::SetOf {
442                ty: retval_ty,
443                optional,
444            } => {
445                let result_ident = syn::Ident::new("result", self.func.sig.span());
446                let retval_ty_resolved = retval_ty.original_ty;
447                let result_handler = if optional {
448                    // don't need unsafe annotations because of the larger unsafe block coming up
449                    quote_spanned! { self.func.sig.span() =>
450                        #func_name(#(#arg_pats),*)
451                    }
452                } else {
453                    quote_spanned! { self.func.sig.span() =>
454                        Some(#func_name(#(#arg_pats),*))
455                    }
456                };
457
458                quote_spanned! { self.func.sig.span() =>
459                    #[no_mangle]
460                    #[doc(hidden)]
461                    #[pg_guard]
462                    #[warn(unsafe_op_in_unsafe_fn)]
463                    pub unsafe extern "C" fn #func_name_wrapper #func_generics(#fcinfo_ident: ::pgx::pg_sys::FunctionCallInfo) -> ::pgx::pg_sys::Datum {
464                        struct IteratorHolder<'__pgx_internal_lifetime, T: std::panic::UnwindSafe + std::panic::RefUnwindSafe> {
465                            iter: *mut SetOfIterator<'__pgx_internal_lifetime, T>,
466                        }
467
468                        let mut funcctx: ::pgx::PgBox<pg_sys::FuncCallContext>;
469                        let mut iterator_holder: ::pgx::PgBox<IteratorHolder<#retval_ty_resolved>>;
470
471                        unsafe {
472                            if ::pgx::srf_is_first_call(#fcinfo_ident) {
473                                funcctx = pgx::srf_first_call_init(#fcinfo_ident);
474                                funcctx.user_fctx = pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).palloc_struct::<IteratorHolder<#retval_ty_resolved>>() as *mut ::core::ffi::c_void;
475                                iterator_holder = pgx::PgBox::from_pg(funcctx.user_fctx as *mut IteratorHolder<#retval_ty_resolved>);
476
477                                // function arguments need to be "fetched" while in the function call's
478                                // multi-call-memory-context to ensure that any detoasted datums will
479                                // live long enough for the SRF to use them over each call
480                                let #result_ident = match pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).switch_to(|_| {
481                                    #( #arg_fetches )*
482                                    #result_handler
483                                }) {
484                                    Some(result) => result,
485                                    None => {
486                                        pgx::srf_return_done(#fcinfo_ident, &mut funcctx);
487                                        return pgx::pg_return_null(#fcinfo_ident)
488                                    }
489                                };
490
491                                iterator_holder.iter = pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).leak_trivial_alloc(result);
492                            }
493
494                            funcctx = pgx::srf_per_call_setup(#fcinfo_ident);
495                            iterator_holder = pgx::PgBox::from_pg(funcctx.user_fctx as *mut IteratorHolder<#retval_ty_resolved>);
496                        }
497
498                        // SAFETY: should have been set up correctly on this or previous call
499                        let mut iter = unsafe { Box::from_raw(iterator_holder.iter) };
500                        match iter.next() {
501                            Some(result) => {
502                                // we need to leak the boxed iterator so that it's not freed by Rust and we can
503                                // continue to use it
504                                Box::leak(iter);
505
506                                // SAFETY: what is an srf if it does not return?
507                                unsafe { pgx::srf_return_next(#fcinfo_ident, &mut funcctx) };
508                                match pgx::datum::IntoDatum::into_datum(result) {
509                                    Some(datum) => datum,
510                                    None => pgx::pg_return_null(#fcinfo_ident),
511                                }
512                            },
513                            None => {
514                                // leak the iterator here too, even tho we're done, b/c our MemoryContextCallback
515                                // function is going to properly drop it for us
516                                Box::leak(iter);
517
518                                // SAFETY: seem to be finished
519                                unsafe { pgx::srf_return_done(#fcinfo_ident, &mut funcctx) };
520                                pgx::pg_return_null(#fcinfo_ident)
521                            },
522                        }
523                    }
524                }
525            }
526            Returning::Iterated {
527                tys: retval_tys,
528                optional,
529            } => {
530                let result_ident = syn::Ident::new("result", self.func.sig.span());
531                let funcctx_ident = syn::Ident::new("funcctx", self.func.sig.span());
532                let retval_tys_resolved = retval_tys.iter().map(|v| &v.used_ty.resolved_ty);
533                let retval_tys_tuple = quote! { (#(#retval_tys_resolved,)*) };
534
535                let retval_tuple_indexes = (0..retval_tys.len()).map(syn::Index::from);
536                let retval_tuple_len = retval_tuple_indexes.len();
537                let create_heap_tuple = quote! {
538                    let mut datums: [::pgx::pg_sys::Datum; #retval_tuple_len] = [::pgx::pg_sys::Datum::from(0); #retval_tuple_len];
539                    let mut nulls: [bool; #retval_tuple_len] = [false; #retval_tuple_len];
540
541                    #(
542                        let datum = pgx::datum::IntoDatum::into_datum(result.#retval_tuple_indexes);
543                        match datum {
544                            Some(datum) => { datums[#retval_tuple_indexes] = datum.into(); },
545                            None => { nulls[#retval_tuple_indexes] = true; }
546                        }
547                    )*
548
549                    // SAFETY: just went to considerable trouble to make sure these are well-formed for a tuple
550                    let heap_tuple = unsafe { pgx::pg_sys::heap_form_tuple(#funcctx_ident.tuple_desc, datums.as_mut_ptr(), nulls.as_mut_ptr()) };
551                };
552
553                let result_handler = if optional {
554                    // don't need unsafe annotations because of the larger unsafe block coming up
555                    quote_spanned! { self.func.sig.span() =>
556                        #func_name(#(#arg_pats),*)
557                    }
558                } else {
559                    quote_spanned! { self.func.sig.span() =>
560                        Some(#func_name(#(#arg_pats),*))
561                    }
562                };
563
564                quote_spanned! { self.func.sig.span() =>
565                    #[no_mangle]
566                    #[doc(hidden)]
567                    #[pg_guard]
568                    #[warn(unsafe_op_in_unsafe_fn)]
569                    pub unsafe extern "C" fn #func_name_wrapper #func_generics(#fcinfo_ident: ::pgx::pg_sys::FunctionCallInfo) -> ::pgx::pg_sys::Datum {
570                        struct IteratorHolder<'__pgx_internal_lifetime, T: std::panic::UnwindSafe + std::panic::RefUnwindSafe> {
571                            iter: *mut ::pgx::iter::TableIterator<'__pgx_internal_lifetime, T>,
572                        }
573
574                        let mut funcctx: pgx::PgBox<pg_sys::FuncCallContext>;
575                        let mut iterator_holder: pgx::PgBox<IteratorHolder<#retval_tys_tuple>>;
576
577                        unsafe {
578                            if ::pgx::srf_is_first_call(#fcinfo_ident) {
579                                funcctx = ::pgx::srf_first_call_init(#fcinfo_ident);
580                                funcctx.user_fctx = pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).palloc_struct::<IteratorHolder<#retval_tys_tuple>>() as *mut ::core::ffi::c_void;
581                                funcctx.tuple_desc = pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).switch_to(|_| {
582                                    let mut tupdesc: *mut pgx::pg_sys::TupleDescData = std::ptr::null_mut();
583
584                                    /* Build a tuple descriptor for our result type */
585                                    if pgx::pg_sys::get_call_result_type(#fcinfo_ident, std::ptr::null_mut(), &mut tupdesc) != pgx::pg_sys::TypeFuncClass_TYPEFUNC_COMPOSITE {
586                                        pgx::error!("return type must be a row type");
587                                    }
588
589                                    pgx::pg_sys::BlessTupleDesc(tupdesc)
590                                });
591                                iterator_holder = pgx::PgBox::from_pg(funcctx.user_fctx as *mut IteratorHolder<#retval_tys_tuple>);
592
593                                // function arguments need to be "fetched" while in the function call's
594                                // multi-call-memory-context to ensure that any detoasted datums will
595                                // live long enough for the SRF to use them over each call
596                                let #result_ident = match pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).switch_to(|_| {
597                                    #( #arg_fetches )*
598                                    #result_handler
599                                }) {
600                                    Some(result) => result,
601                                    None => {
602                                        pgx::srf_return_done(#fcinfo_ident, &mut funcctx);
603                                        return pgx::pg_return_null(#fcinfo_ident)
604                                    }
605                                };
606
607                                iterator_holder.iter = pgx::PgMemoryContexts::For(funcctx.multi_call_memory_ctx).leak_and_drop_on_delete(result);
608                            }
609
610                            funcctx = pgx::srf_per_call_setup(#fcinfo_ident);
611                            iterator_holder = pgx::PgBox::from_pg(funcctx.user_fctx as *mut IteratorHolder<#retval_tys_tuple>);
612                        }
613
614                        // SAFETY: should have been set up correctly on this or previous call
615                        let mut iter = unsafe { Box::from_raw(iterator_holder.iter) };
616                        match iter.next() {
617                            Some(result) => {
618                                // we need to leak the boxed iterator so that it's not freed by rust and we can
619                                // continue to use it
620                                Box::leak(iter);
621
622                                #create_heap_tuple
623
624                                let datum = pgx::heap_tuple_get_datum(heap_tuple);
625                                // SAFETY: what is an srf if it does not return?
626                                unsafe { pgx::srf_return_next(#fcinfo_ident, &mut funcctx) };
627                                datum
628                            },
629                            None => {
630                                // leak the iterator here too, even tho we're done, b/c our MemoryContextCallback
631                                // function is going to properly drop it for us
632                                Box::leak(iter);
633
634                                // SAFETY: seem to be finished
635                                unsafe { pgx::srf_return_done(#fcinfo_ident, &mut funcctx) };
636                                pgx::pg_return_null(#fcinfo_ident)
637                            },
638                        }
639                    }
640                }
641            }
642            // /// We don't actually do this since triggers have their own cool kids zone
643            // Returning::Trigger => {
644            //     quote_spanned! { self.func.sig.span() =>
645            //         #[no_mangle]
646            //         #[doc(hidden)]
647            //         #[pg_guard]
648            //         pub unsafe extern "C" fn #func_name_wrapper #func_generics(#fcinfo_ident: ::pgx::pg_sys::FunctionCallInfo) -> ::pgx::pg_sys::Datum {
649            //             let maybe_pg_trigger = unsafe { ::pgx::trigger_support::PgTrigger::from_fcinfo(#fcinfo_ident) };
650            //             let pg_trigger = maybe_pg_trigger.expect("PgTrigger::from_fcinfo failed");
651            //             let trigger_fn_result: Result<
652            //                 ::pgx::heap_tuple::PgHeapTuple<'_, _>,
653            //                 _,
654            //             > = #func_name(&pg_trigger);
655
656            //             let trigger_retval = trigger_fn_result.expect("Trigger function panic");
657            //             match trigger_retval.into_trigger_datum() {
658            //                 None => ::pgx::pg_return_null(fcinfo),
659            //                 Some(datum) => datum,
660            //             }
661            //         }
662            //     }
663            // },
664        }
665    }
666}
667
668impl ToTokens for PgExtern {
669    fn to_tokens(&self, tokens: &mut TokenStream2) {
670        let original_func = &self.func;
671        let wrapper_func = self.wrapper_func();
672        let entity_func = self.entity_tokens();
673        let finfo_tokens = self.finfo_tokens();
674
675        let expansion = quote_spanned! { self.func.sig.span() =>
676            #original_func
677
678            #wrapper_func
679
680            #entity_func
681
682            #finfo_tokens
683        };
684        tokens.append_all(expansion);
685    }
686}
687
688impl Parse for PgExtern {
689    fn parse(input: ParseStream) -> Result<Self, syn::Error> {
690        let mut attrs = Vec::new();
691        let mut to_sql_config: Option<ToSqlConfig> = None;
692
693        let parser = Punctuated::<Attribute, Token![,]>::parse_terminated;
694        let punctuated_attrs = input.call(parser).ok().unwrap_or_default();
695        for pair in punctuated_attrs.into_pairs() {
696            match pair.into_value() {
697                Attribute::Sql(config) => {
698                    to_sql_config.get_or_insert(config);
699                }
700                attr => {
701                    attrs.push(attr);
702                }
703            }
704        }
705
706        let to_sql_config = to_sql_config.unwrap_or_default();
707
708        let func: syn::ItemFn = input.parse()?;
709
710        if !to_sql_config.overrides_default() {
711            crate::ident_is_acceptable_to_postgres(&func.sig.ident)?;
712        }
713
714        Ok(Self { attrs, func, to_sql_config })
715    }
716}
pgx_utils/sql_entity_graph/pg_extern/mod.rs

pgx_utils/sql_entity_graph/pg_extern/
mod.rs
