Unbox types properly from arrays (#1476)

Introducing `Datum<'src>` allows using it with HRTB fns in schema generation to persist lifetimes all the way through `#[pg_extern]`'s wrapper-fn infrastructure without compromising the integrity of lifetime annotations. This is immediately implemented to prevent many types from exceeding the lifetime of the Array they are unboxed from, unless they are already "self-owned" and thus 'static. Unfortunately, this does not fix all instances using `pg_getarg`, and those simply are left unsolved for now in some cases. I reorganized the UI tests to make it easier to notice these problems and when they get solved.
pgcentralfoundation · Jan 17, 2024 · 4fa87bd · 4fa87bd
1 parent a562516
commit 4fa87bd
Show file tree

Hide file tree

Showing 48 changed files with 2,427 additions and 760 deletions.
diff --git a/pgrx-examples/bad_ideas/src/lib.rs b/pgrx-examples/bad_ideas/src/lib.rs
@@ -55,7 +55,7 @@ fn warning(s: &str) -> bool {
 #[pg_extern]
 fn exec<'a>(
     command: &'a str,
-    args: default!(Vec<Option<&'a str>>, "ARRAY[]::text[]"),
+    args: default!(Vec<Option<String>>, "ARRAY[]::text[]"),
 ) -> TableIterator<'static, (name!(status, Option<i32>), name!(stdout, String))> {
     let mut command = &mut Command::new(command);
 

diff --git a/pgrx-macros/src/lib.rs b/pgrx-macros/src/lib.rs
@@ -618,6 +618,7 @@ pub fn postgres_enum(input: TokenStream) -> TokenStream {
 fn impl_postgres_enum(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
     let mut stream = proc_macro2::TokenStream::new();
     let sql_graph_entity_ast = ast.clone();
+    let generics = &ast.generics;
     let enum_ident = &ast.ident;
     let enum_name = enum_ident.to_string();
 
@@ -660,6 +661,14 @@ fn impl_postgres_enum(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
             }
         }
 
+        unsafe impl #generics ::pgrx::datum::UnboxDatum for #enum_ident #generics {
+            type As<'dat> = #enum_ident #generics where Self: 'dat;
+            #[inline]
+            unsafe fn unbox<'dat>(d: ::pgrx::datum::Datum<'dat>) -> Self::As<'dat> where Self: 'dat {
+                Self::from_datum(::core::mem::transmute(d), false).unwrap()
+            }
+        }
+
         impl ::pgrx::datum::IntoDatum for #enum_ident {
             #[inline]
             fn into_datum(self) -> Option<::pgrx::pg_sys::Datum> {
@@ -808,6 +817,13 @@ fn impl_postgres_type(ast: DeriveInput) -> syn::Result<proc_macro2::TokenStream>
                         }
                     }
                 }
+
+                unsafe impl #generics ::pgrx::datum::UnboxDatum for #name #generics {
+                    type As<'dat> = Self where Self: 'dat;
+                    unsafe fn unbox<'dat>(datum: ::pgrx::datum::Datum<'dat>) -> Self::As<'dat> where Self: 'dat {
+                        <Self as ::pgrx::datum::FromDatum>::from_datum(::core::mem::transmute(datum), false).unwrap()
+                    }
+                }
             }
         )
     }

diff --git a/pgrx-macros/src/rewriter.rs b/pgrx-macros/src/rewriter.rs
@@ -102,7 +102,7 @@ pub fn item_fn_without_rewrite(mut func: ItemFn) -> syn::Result<proc_macro2::Tok
             #[allow(unused_unsafe)]
             unsafe {
                 // NB: this is purposely not spelled `::pgrx` as pgrx itself uses #[pg_guard]
-                pgrx::pg_sys::submodules::panic::pgrx_extern_c_guard( || #body )
+                pgrx::pg_sys::submodules::panic::pgrx_extern_c_guard(move || #body )
             }
         }
     })

diff --git a/pgrx-sql-entity-graph/src/lifetimes.rs b/pgrx-sql-entity-graph/src/lifetimes.rs
@@ -7,76 +7,6 @@
 //LICENSE All rights reserved.
 //LICENSE
 //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
-use crate::NameMacro;
-use proc_macro2::TokenStream;
-
-pub fn staticize_lifetimes_in_type_path(value: syn::TypePath) -> syn::TypePath {
-    let mut ty = syn::Type::Path(value);
-    staticize_lifetimes(&mut ty);
-    match ty {
-        syn::Type::Path(type_path) => type_path,
-
-        // shouldn't happen
-        _ => panic!("not a TypePath"),
-    }
-}
-
-pub fn staticize_lifetimes(value: &mut syn::Type) {
-    match value {
-        syn::Type::Path(syn::TypePath { path: syn::Path { segments, .. }, .. }) => segments
-            .iter_mut()
-            .filter_map(|segment| match &mut segment.arguments {
-                syn::PathArguments::AngleBracketed(bracketed) => Some(bracketed),
-                _ => None,
-            })
-            .flat_map(|bracketed| &mut bracketed.args)
-            .for_each(|arg| match arg {
-                // rename lifetimes to the static lifetime so the TypeIds match.
-                syn::GenericArgument::Lifetime(lifetime) => {
-                    lifetime.ident = syn::Ident::new("static", lifetime.ident.span());
-                }
-                // recurse
-                syn::GenericArgument::Type(ty) => staticize_lifetimes(ty),
-                syn::GenericArgument::Binding(binding) => staticize_lifetimes(&mut binding.ty),
-                syn::GenericArgument::Constraint(constraint) => {
-                    constraint.bounds.iter_mut().for_each(|bound| {
-                        if let syn::TypeParamBound::Lifetime(lifetime) = bound {
-                            lifetime.ident = syn::Ident::new("static", lifetime.ident.span())
-                        }
-                    })
-                }
-                // nothing to do otherwise
-                _ => {}
-            }),
-
-        syn::Type::Reference(type_ref) => {
-            if let Some(lifetime) = &mut type_ref.lifetime {
-                lifetime.ident = syn::Ident::new("static", lifetime.ident.span());
-            } else {
-                type_ref.lifetime = Some(syn::parse_quote!('static));
-            }
-        }
-
-        syn::Type::Tuple(type_tuple) => type_tuple.elems.iter_mut().for_each(staticize_lifetimes),
-
-        syn::Type::Macro(syn::TypeMacro { mac })
-            if mac.path.segments.last().is_some_and(|seg| seg.ident == "name") =>
-        {
-            // We don't particularly care what the identifier is, so we parse a
-            // raw TokenStream.  Specifically, it's okay for the identifier String,
-            // which we end up using as a Postgres column name, to be nearly any
-            // string, which can include Rust reserved words such as "type" or "match"
-            let Ok(out) = mac.parse_body::<NameMacro>() else { return };
-            let Ok(ident) = syn::parse_str::<TokenStream>(&out.ident) else { return };
-            let mut ty = out.used_ty.resolved_ty;
-
-            // rewrite the name!() macro's type so that it has a static lifetime, if any
-            staticize_lifetimes(&mut ty);
-            *mac = syn::parse_quote! {::pgrx::name!(#ident, #ty)};
-        }
-        _ => {}
-    }
-}
 
 pub fn anonymize_lifetimes_in_type_path(value: syn::TypePath) -> syn::TypePath {
     let mut ty = syn::Type::Path(value);

diff --git a/pgrx-sql-entity-graph/src/metadata/function_metadata.rs b/pgrx-sql-entity-graph/src/metadata/function_metadata.rs
@@ -42,61 +42,39 @@ pub trait FunctionMetadata<A> {
     fn entity(&self) -> FunctionMetadataEntity;
 }
 
-impl<R> FunctionMetadata<()> for fn() -> R
-where
-    R: SqlTranslatable,
-{
-    fn entity(&self) -> FunctionMetadataEntity {
-        FunctionMetadataEntity {
-            arguments: vec![],
-            retval: {
-                let marker: PhantomData<R> = PhantomData;
-                marker.entity()
-            },
-            path: self.path(),
-        }
-    }
-}
-
-impl<R> FunctionMetadata<()> for unsafe fn() -> R
-where
-    R: SqlTranslatable,
-{
-    fn entity(&self) -> FunctionMetadataEntity {
-        FunctionMetadataEntity {
-            arguments: vec![],
-            retval: {
-                let marker: PhantomData<R> = PhantomData;
-                marker.entity()
-            },
-            path: self.path(),
-        }
-    }
-}
-
 macro_rules! impl_fn {
-    ($($T:ident),* $(,)?) => {
-        impl<$($T: SqlTranslatable,)* R: SqlTranslatable> FunctionMetadata<($($T,)*)> for fn($($T,)*) -> R {
+    ($($A:ident),* $(,)?) => {
+        impl<$($A,)* R, F> FunctionMetadata<($($A,)*)> for F
+        where
+            $($A: SqlTranslatable,)*
+            R: SqlTranslatable,
+            F: FnMut($($A,)*) -> R,
+        {
             fn entity(&self) -> FunctionMetadataEntity {
                 FunctionMetadataEntity {
-                    arguments: vec![$(PhantomData::<$T>.entity()),+],
+                    arguments: vec![$(PhantomData::<$A>.entity()),*],
                     retval: PhantomData::<R>.entity(),
                     path: self.path(),
                 }
             }
         }
-        impl<$($T: SqlTranslatable,)* R: SqlTranslatable> FunctionMetadata<($($T,)*)> for unsafe fn($($T,)*) -> R {
+        impl<$($A,)* R> FunctionMetadata<($($A,)*)> for unsafe fn($($A,)*) -> R
+        where
+            $($A: SqlTranslatable,)*
+            R: SqlTranslatable,
+        {
             fn entity(&self) -> FunctionMetadataEntity {
                 FunctionMetadataEntity {
-                    arguments: vec![$(PhantomData::<$T>.entity()),+],
+                    arguments: vec![$(PhantomData::<$A>.entity()),*],
                     retval: PhantomData::<R>.entity(),
                     path: self.path(),
                 }
             }
         }
     };
 }
-// empty tuples are above
+
+impl_fn!();
 impl_fn!(T0);
 impl_fn!(T0, T1);
 impl_fn!(T0, T1, T2);

diff --git a/pgrx-sql-entity-graph/src/pg_extern/mod.rs b/pgrx-sql-entity-graph/src/pg_extern/mod.rs
@@ -34,7 +34,6 @@ use search_path::SearchPathList;
 use crate::enrich::CodeEnrichment;
 use crate::enrich::ToEntityGraphTokens;
 use crate::enrich::ToRustCodeTokens;
-use crate::lifetimes::staticize_lifetimes;
 use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
 use quote::{quote, quote_spanned, ToTokens};
 use syn::parse::{Parse, ParseStream, Parser};
@@ -133,12 +132,11 @@ impl PgExtern {
                 match v {
                     syn::FnArg::Receiver(_) => None,
                     syn::FnArg::Typed(pat_ty) => {
-                        let mut static_ty = match UsedType::new(*pat_ty.ty.clone()) {
+                        let ty = match UsedType::new(*pat_ty.ty.clone()) {
                             Ok(v) => v.resolved_ty,
                             Err(e) => return Some(Err(e)),
                         };
-                        staticize_lifetimes(&mut static_ty);
-                        Some(Ok(static_ty))
+                        Some(Ok(ty))
                     }
                 }
             })
@@ -275,9 +273,8 @@ impl PgExtern {
         let return_type = match &self.func.sig.output {
             syn::ReturnType::Default => None,
             syn::ReturnType::Type(arrow, ty) => {
-                let mut static_ty = ty.clone();
-                staticize_lifetimes(&mut static_ty);
-                Some(syn::ReturnType::Type(*arrow, static_ty))
+                let ty = ty.clone();
+                Some(syn::ReturnType::Type(*arrow, ty))
             }
         };
 
@@ -287,6 +284,20 @@ impl PgExtern {
             Some(content) => ToSqlConfig { content: Some(content), ..self.to_sql_config.clone() },
         };
 
+        let lifetimes = self
+            .func
+            .sig
+            .generics
+            .params
+            .iter()
+            .filter_map(|generic| match generic {
+                // syn::GenericParam::Type(_) => todo!(),
+                syn::GenericParam::Lifetime(lt) => Some(lt),
+                _ => None, // syn::GenericParam::Const(_) => todo!(),
+            })
+            .collect::<Vec<_>>();
+        let hrtb = if lifetimes.is_empty() { None } else { Some(quote! { for<#(#lifetimes),*> }) };
+
         let sql_graph_entity_fn_name =
             syn::Ident::new(&format!("__pgrx_internals_fn_{}", ident), Span::call_site());
         quote_spanned! { self.func.sig.span() =>
@@ -297,7 +308,7 @@ impl PgExtern {
                 extern crate alloc;
                 #[allow(unused_imports)]
                 use alloc::{vec, vec::Vec};
-                type FunctionPointer = #unsafety fn(#( #input_types ),*) #return_type;
+                type FunctionPointer = #hrtb #unsafety fn (#( #input_types ),*) #return_type;
                 let metadata: FunctionPointer = #ident;
                 let submission = ::pgrx::pgrx_sql_entity_graph::PgExternEntity {
                     name: #name,

diff --git a/pgrx-sql-entity-graph/src/postgres_type/mod.rs b/pgrx-sql-entity-graph/src/postgres_type/mod.rs
@@ -21,7 +21,7 @@ use crate::enrich::{ToEntityGraphTokens, ToRustCodeTokens};
 use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
 use quote::quote;
 use syn::parse::{Parse, ParseStream};
-use syn::{DeriveInput, Generics, ItemStruct};
+use syn::{DeriveInput, Generics, ItemStruct, Lifetime, LifetimeDef};
 
 use crate::{CodeEnrichment, ToSqlConfig};
 
@@ -103,37 +103,24 @@ impl PostgresTypeDerive {
 impl ToEntityGraphTokens for PostgresTypeDerive {
     fn to_entity_graph_tokens(&self) -> TokenStream2 {
         let name = &self.name;
-        let mut static_generics = self.generics.clone();
-        static_generics.params = static_generics
-            .params
-            .clone()
-            .into_iter()
-            .flat_map(|param| match param {
-                item @ syn::GenericParam::Type(_) | item @ syn::GenericParam::Const(_) => {
-                    Some(item)
-                }
-                syn::GenericParam::Lifetime(mut lifetime) => {
-                    lifetime.lifetime.ident = Ident::new("static", Span::call_site());
-                    Some(syn::GenericParam::Lifetime(lifetime))
-                }
-            })
-            .collect();
-        let mut staticless_generics = self.generics.clone();
-        staticless_generics.params = static_generics
+        let generics = self.generics.clone();
+        let (impl_generics, ty_generics, where_clauses) = generics.split_for_impl();
+
+        // We need some generics we can use inside a fn without a lifetime for qualified paths.
+        let mut anon_generics = generics.clone();
+        anon_generics.params = anon_generics
             .params
-            .clone()
             .into_iter()
             .flat_map(|param| match param {
                 item @ syn::GenericParam::Type(_) | item @ syn::GenericParam::Const(_) => {
                     Some(item)
                 }
-                syn::GenericParam::Lifetime(_) => None,
+                syn::GenericParam::Lifetime(lt_def) => Some(syn::GenericParam::Lifetime(
+                    LifetimeDef::new(Lifetime::new("'_", lt_def.lifetime.span())),
+                )),
             })
             .collect();
-        let (staticless_impl_generics, _staticless_ty_generics, _staticless_where_clauses) =
-            staticless_generics.split_for_impl();
-        let (_static_impl_generics, static_ty_generics, static_where_clauses) =
-            static_generics.split_for_impl();
+        let (_, anon_ty_gen, _) = anon_generics.split_for_impl();
 
         let in_fn = &self.in_fn;
         let out_fn = &self.out_fn;
@@ -144,7 +131,7 @@ impl ToEntityGraphTokens for PostgresTypeDerive {
         let to_sql_config = &self.to_sql_config;
 
         quote! {
-            unsafe impl #staticless_impl_generics ::pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable for #name #static_ty_generics #static_where_clauses {
+            unsafe impl #impl_generics ::pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable for #name #ty_generics #where_clauses {
                 fn argument_sql() -> core::result::Result<::pgrx::pgrx_sql_entity_graph::metadata::SqlMapping, ::pgrx::pgrx_sql_entity_graph::metadata::ArgumentError> {
                     Ok(::pgrx::pgrx_sql_entity_graph::metadata::SqlMapping::As(String::from(stringify!(#name))))
                 }
@@ -166,19 +153,19 @@ impl ToEntityGraphTokens for PostgresTypeDerive {
                 use ::pgrx::datum::WithTypeIds;
 
                 let mut mappings = Default::default();
-                <#name #static_ty_generics as ::pgrx::datum::WithTypeIds>::register_with_refs(
+                <#name #anon_ty_gen as ::pgrx::datum::WithTypeIds>::register_with_refs(
                     &mut mappings,
                     stringify!(#name).to_string()
                 );
-                ::pgrx::datum::WithSizedTypeIds::<#name #static_ty_generics>::register_sized_with_refs(
+                ::pgrx::datum::WithSizedTypeIds::<#name #anon_ty_gen>::register_sized_with_refs(
                     &mut mappings,
                     stringify!(#name).to_string()
                 );
-                ::pgrx::datum::WithArrayTypeIds::<#name #static_ty_generics>::register_array_with_refs(
+                ::pgrx::datum::WithArrayTypeIds::<#name #anon_ty_gen>::register_array_with_refs(
                     &mut mappings,
                     stringify!(#name).to_string()
                 );
-                ::pgrx::datum::WithVarlenaTypeIds::<#name #static_ty_generics>::register_varlena_with_refs(
+                ::pgrx::datum::WithVarlenaTypeIds::<#name #anon_ty_gen>::register_varlena_with_refs(
                     &mut mappings,
                     stringify!(#name).to_string()
                 );
@@ -187,7 +174,7 @@ impl ToEntityGraphTokens for PostgresTypeDerive {
                     file: file!(),
                     line: line!(),
                     module_path: module_path!(),
-                    full_path: core::any::type_name::<#name #static_ty_generics>(),
+                    full_path: core::any::type_name::<#name #anon_ty_gen>(),
                     mappings: mappings.into_iter().collect(),
                     in_fn: stringify!(#in_fn),
                     in_fn_module_path: {