diff --git a/pgrx-examples/custom_types/src/hexint.rs b/pgrx-examples/custom_types/src/hexint.rs index 492548edf..8a5f73fc2 100644 --- a/pgrx-examples/custom_types/src/hexint.rs +++ b/pgrx-examples/custom_types/src/hexint.rs @@ -7,6 +7,7 @@ //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 pgrx::callconv::BoxRet; use pgrx::pg_sys::{Datum, Oid}; use pgrx::pgrx_sql_entity_graph::metadata::{ ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable, @@ -93,6 +94,12 @@ impl IntoDatum for HexInt { } } +unsafe impl BoxRet for HexInt { + unsafe fn box_in_fcinfo(self, _fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + Datum::from(self.value) + } +} + /// Input function for `HexInt`. Parses any valid "radix(16)" text string, with or without a leading /// `0x` (or `0X`) into a `HexInt` type. Parse errors are returned and handled by pgrx /// diff --git a/pgrx-examples/spi/src/lib.rs b/pgrx-examples/spi/src/lib.rs index 0d7632129..8032ce78c 100644 --- a/pgrx-examples/spi/src/lib.rs +++ b/pgrx-examples/spi/src/lib.rs @@ -30,13 +30,7 @@ INSERT INTO spi_example (title) VALUES ('I like pudding'); #[pg_extern] fn spi_return_query() -> Result< - TableIterator< - 'static, - ( - name!(oid, Result, pgrx::spi::Error>), - name!(name, Result, pgrx::spi::Error>), - ), - >, + TableIterator<'static, (name!(oid, Option), name!(name, Option))>, spi::Error, > { #[cfg(feature = "pg12")] @@ -51,10 +45,10 @@ fn spi_return_query() -> Result< let query = "SELECT oid, relname::text || '-pg16' FROM pg_class"; Spi::connect(|client| { - Ok(client + client .select(query, None, None)? - .map(|row| (row["oid"].value(), row[2].value())) - .collect::>()) + .map(|row| Ok((row["oid"].value()?, row[2].value()?))) + .collect::, _>>() }) .map(TableIterator::new) } diff --git a/pgrx-examples/spi_srf/src/lib.rs b/pgrx-examples/spi_srf/src/lib.rs index c99a50e15..f1f2a9a17 100644 --- a/pgrx-examples/spi_srf/src/lib.rs +++ b/pgrx-examples/spi_srf/src/lib.rs @@ -38,9 +38,9 @@ fn calculate_human_years() -> Result< TableIterator< 'static, ( - name!(dog_name, Result, pgrx::spi::Error>), + name!(dog_name, Option), name!(dog_age, i32), - name!(dog_breed, Result, pgrx::spi::Error>), + name!(dog_breed, Option), name!(human_age, i32), ), >, @@ -62,7 +62,7 @@ fn calculate_human_years() -> Result< let dog_age = row["dog_age"].value::()?.expect("dog_age was null"); let dog_breed = row["dog_breed"].value::(); let human_age = dog_age * 7; - results.push((dog_name, dog_age, dog_breed, human_age)); + results.push((dog_name?, dog_age, dog_breed?, human_age)); } Ok(TableIterator::new(results)) @@ -76,9 +76,9 @@ fn filter_by_breed( TableIterator< 'static, ( - name!(dog_name, Result, pgrx::spi::Error>), - name!(dog_age, Result, pgrx::spi::Error>), - name!(dog_breed, Result, pgrx::spi::Error>), + name!(dog_name, Option), + name!(dog_age, Option), + name!(dog_breed, Option), ), >, spi::Error, @@ -95,9 +95,11 @@ fn filter_by_breed( let tup_table = client.select(query, None, Some(args))?; let filtered = tup_table - .map(|row| (row["dog_name"].value(), row["dog_age"].value(), row["dog_breed"].value())) - .collect::>(); - Ok(TableIterator::new(filtered)) + .map(|row| { + Ok((row["dog_name"].value()?, row["dog_age"].value()?, row["dog_breed"].value()?)) + }) + .collect::, _>>(); + filtered.map(|v| TableIterator::new(v)) }) } @@ -107,19 +109,17 @@ mod tests { use crate::calculate_human_years; use pgrx::prelude::*; - #[rustfmt::skip] #[pg_test] fn test_calculate_human_years() -> Result<(), pgrx::spi::Error> { - let mut results: Vec<(Result, _>, i32, Result, _>, i32)> = - Vec::new(); - - results.push((Ok(Some("Fido".to_string())), 3, Ok(Some("Labrador".to_string())), 21)); - results.push((Ok(Some("Spot".to_string())), 5, Ok(Some("Poodle".to_string())), 35)); - results.push((Ok(Some("Rover".to_string())), 7, Ok(Some("Golden Retriever".to_string())), 49)); - results.push((Ok(Some("Snoopy".to_string())), 9, Ok(Some("Beagle".to_string())), 63)); - results.push((Ok(Some("Lassie".to_string())), 11, Ok(Some("Collie".to_string())), 77)); - results.push((Ok(Some("Scooby".to_string())), 13, Ok(Some("Great Dane".to_string())), 91)); - results.push((Ok(Some("Moomba".to_string())), 15, Ok(Some("Labrador".to_string())), 105)); + let mut results = Vec::new(); + + results.push((Some("Fido".to_string()), 3, Some("Labrador".to_string()), 21)); + results.push((Some("Spot".to_string()), 5, Some("Poodle".to_string()), 35)); + results.push((Some("Rover".to_string()), 7, Some("Golden Retriever".to_string()), 49)); + results.push((Some("Snoopy".to_string()), 9, Some("Beagle".to_string()), 63)); + results.push((Some("Lassie".to_string()), 11, Some("Collie".to_string()), 77)); + results.push((Some("Scooby".to_string()), 13, Some("Great Dane".to_string()), 91)); + results.push((Some("Moomba".to_string()), 15, Some("Labrador".to_string()), 105)); let func_results = calculate_human_years()?; for (expected, actual) in results.iter().zip(func_results) { diff --git a/pgrx-examples/srf/src/lib.rs b/pgrx-examples/srf/src/lib.rs index 036666a00..3f95fefb1 100644 --- a/pgrx-examples/srf/src/lib.rs +++ b/pgrx-examples/srf/src/lib.rs @@ -34,10 +34,10 @@ fn random_values(num_rows: i32) -> TableIterator<'static, (name!(index, i32), na #[pg_extern] fn result_table() -> Result< - Option<::pgrx::iter::TableIterator<'static, (name!(a, Option), name!(b, Option))>>, + ::pgrx::iter::TableIterator<'static, (name!(a, Option), name!(b, Option))>, Box, > { - Ok(Some(TableIterator::new(vec![(Some(1), Some(2))]))) + Ok(TableIterator::new(vec![(Some(1), Some(2))])) } #[pg_extern] diff --git a/pgrx-macros/src/lib.rs b/pgrx-macros/src/lib.rs index d3f411148..29f2f080a 100644 --- a/pgrx-macros/src/lib.rs +++ b/pgrx-macros/src/lib.rs @@ -704,6 +704,12 @@ fn impl_postgres_enum(ast: DeriveInput) -> syn::Result } } + + unsafe impl ::pgrx::callconv::BoxRet for #enum_ident { + unsafe fn box_in_fcinfo(self, fcinfo: ::pgrx::pg_sys::FunctionCallInfo) -> ::pgrx::pg_sys::Datum { + ::pgrx::datum::IntoDatum::into_datum(self).unwrap() + } + } }); let sql_graph_entity_item = PostgresEnum::from_derive_input(sql_graph_entity_ast)?; @@ -808,6 +814,15 @@ fn impl_postgres_type(ast: DeriveInput) -> syn::Result } } + unsafe impl #generics ::pgrx::callconv::BoxRet for #name #generics { + unsafe fn box_in_fcinfo(self, fcinfo: ::pgrx::pg_sys::FunctionCallInfo) -> ::pgrx::pg_sys::Datum { + match ::pgrx::datum::IntoDatum::into_datum(self) { + None => ::pgrx::fcinfo::pg_return_null(fcinfo), + Some(datum) => datum, + } + } + } + impl #generics ::pgrx::datum::FromDatum for #name #generics { unsafe fn from_polymorphic_datum( datum: ::pgrx::pg_sys::Datum, diff --git a/pgrx-pg-sys/src/lib.rs b/pgrx-pg-sys/src/lib.rs index a6e0a7fb8..c66b809a3 100644 --- a/pgrx-pg-sys/src/lib.rs +++ b/pgrx-pg-sys/src/lib.rs @@ -17,7 +17,6 @@ #![allow(non_upper_case_globals)] #![allow(improper_ctypes)] #![allow(clippy::unneeded_field_pattern)] -#![cfg_attr(nightly, feature(strict_provenance))] #[cfg( // no features at all will cause problems diff --git a/pgrx-sql-entity-graph/src/pg_extern/entity/returning.rs b/pgrx-sql-entity-graph/src/pg_extern/entity/returning.rs index e12652f44..5f74725f8 100644 --- a/pgrx-sql-entity-graph/src/pg_extern/entity/returning.rs +++ b/pgrx-sql-entity-graph/src/pg_extern/entity/returning.rs @@ -20,19 +20,9 @@ use crate::UsedTypeEntity; #[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)] pub enum PgExternReturnEntity { None, - Type { - ty: UsedTypeEntity, - }, - SetOf { - ty: UsedTypeEntity, - is_option: bool, /* Eg `Option>` */ - is_result: bool, /* Eg `Result, E>` */ - }, - Iterated { - tys: Vec, - is_option: bool, /* Eg `Option>` */ - is_result: bool, /* Eg `Result, E>` */ - }, + Type { ty: UsedTypeEntity }, + SetOf { ty: UsedTypeEntity }, + Iterated { tys: Vec }, Trigger, } diff --git a/pgrx-sql-entity-graph/src/pg_extern/mod.rs b/pgrx-sql-entity-graph/src/pg_extern/mod.rs index bc9580720..8c73ab4e4 100644 --- a/pgrx-sql-entity-graph/src/pg_extern/mod.rs +++ b/pgrx-sql-entity-graph/src/pg_extern/mod.rs @@ -38,7 +38,7 @@ use crate::ToSqlConfig; use operator::{PgrxOperatorAttributeWithIdent, PgrxOperatorOpName}; use search_path::SearchPathList; -use proc_macro2::{Ident, Span, TokenStream as TokenStream2}; +use proc_macro2::{Ident, TokenStream as TokenStream2}; use quote::{format_ident, quote, quote_spanned}; use syn::parse::{Parse, ParseStream, Parser}; use syn::punctuated::Punctuated; @@ -406,141 +406,33 @@ impl PgExtern { } }); - // Iterators require fancy handling for their retvals - let emit_result_handler = |span: Span, optional: bool, result: bool| { - let mut ret_expr = quote! { #func_name(#(#arg_pats),*) }; - if result { - // If it's a result, we need to report it. - ret_expr = quote! { #ret_expr.unwrap_or_report() }; - } - if !optional { - // If it's not already an option, we need to wrap it. - ret_expr = quote! { Some(#ret_expr) }; - } - let import = result.then(|| quote! { use ::pgrx::pg_sys::panic::ErrorReportable; }); - quote_spanned! { span => - #import - #ret_expr - } - }; - match &self.returns { - Returning::None => { - let fn_contents = quote! { - #(#arg_fetches)* - #[allow(unused_unsafe)] - unsafe { #func_name(#(#arg_pats),*) }; - // -> () means always returning the zero Datum - ::pgrx::pg_sys::Datum::from(0) - }; - finfo_v1_extern_c(&self.func, fcinfo_ident, fn_contents) - } - Returning::Type(retval_ty) => { - let result_ident = syn::Ident::new("result", self.func.sig.span()); - let retval_transform = if retval_ty.resolved_ty == syn::parse_quote!(()) { - quote_spanned! { self.func.sig.output.span() => - unsafe { ::pgrx::fcinfo::pg_return_void() } - } - } else if retval_ty.result && retval_ty.optional.is_some() { - // returning `Result>` - quote_spanned! { self.func.sig.output.span() => - match ::pgrx::datum::IntoDatum::into_datum(#result_ident) { - Some(datum) => datum, - None => unsafe { ::pgrx::fcinfo::pg_return_null(#fcinfo_ident) }, - } - } - } else if retval_ty.result { - // returning Result - quote_spanned! { self.func.sig.output.span() => - ::pgrx::datum::IntoDatum::into_datum(#result_ident).unwrap_or_else(|| panic!("returned Datum was NULL")) - } - } else if retval_ty.resolved_ty.last_ident_is("Datum") { - // As before, we can just throw this in because it must typecheck - quote_spanned! { self.func.sig.output.span() => - #result_ident - } - } else if retval_ty.optional.is_some() { - quote_spanned! { self.func.sig.output.span() => - match #result_ident { - Some(result) => { - ::pgrx::datum::IntoDatum::into_datum(result).unwrap_or_else(|| panic!("returned Option was NULL")) - }, - None => unsafe { ::pgrx::fcinfo::pg_return_null(#fcinfo_ident) } - } - } - } else { - quote_spanned! { self.func.sig.output.span() => - ::pgrx::datum::IntoDatum::into_datum(#result_ident).unwrap_or_else(|| panic!("returned Datum was NULL")) - } - }; - - let fn_contents = quote! { - #(#arg_fetches)* - - #[allow(unused_unsafe)] // unwrapped fn might be unsafe - let #result_ident = unsafe { #func_name(#(#arg_pats),*) }; - - #retval_transform + Returning::None + | Returning::Type(_) + | Returning::SetOf { .. } + | Returning::Iterated { .. } => { + let ret_ty = match &self.func.sig.output { + syn::ReturnType::Default => syn::parse_quote! { () }, + syn::ReturnType::Type(_, ret_ty) => ret_ty.clone(), }; - finfo_v1_extern_c(&self.func, fcinfo_ident, fn_contents) - } - Returning::SetOf { ty: _retval_ty, is_option, is_result } => { - let result_handler = - emit_result_handler(self.func.sig.span(), *is_option, *is_result); - let setof_closure = quote! { + let wrapper_code = quote_spanned! { self.func.block.span() => #[allow(unused_unsafe)] unsafe { - // SAFETY: the caller has asserted that `fcinfo` is a valid FunctionCallInfo pointer, allocated by Postgres - // with all its fields properly setup. Unless the user is calling this wrapper function directly, this - // will always be the case - ::pgrx::iter::SetOfIterator::srf_next(#fcinfo_ident, || { - #( #arg_fetches )* - #result_handler - }) - } - }; - finfo_v1_extern_c(&self.func, fcinfo_ident, setof_closure) - } - Returning::Iterated { tys: retval_tys, is_option, is_result } => { - let result_handler = - emit_result_handler(self.func.sig.span(), *is_option, *is_result); - - let iter_closure = if retval_tys.len() == 1 { - // Postgres considers functions returning a 1-field table (`RETURNS TABLE (T)`) to be - // a function that `RETURNS SETOF T`. So we write a different wrapper implementation - // that transparently transforms the `TableIterator` returned by the user into a `SetOfIterator` - quote! { - #[allow(unused_unsafe)] - unsafe { - // SAFETY: the caller has asserted that `fcinfo` is a valid FunctionCallInfo pointer, allocated by Postgres - // with all its fields properly setup. Unless the user is calling this wrapper function directly, this - // will always be the case - ::pgrx::iter::SetOfIterator::srf_next(#fcinfo_ident, || { - #( #arg_fetches )* - let table_iterator = { #result_handler }; - - // we need to convert the 1-field `TableIterator` provided by the user - // into a SetOfIterator in order to properly handle the case of `RETURNS TABLE (T)`, - // which is a table that returns only 1 field. - table_iterator.map(|i| ::pgrx::iter::SetOfIterator::new(i.into_iter().map(|(v,)| v))) - }) - } - } - } else { - quote! { - #[allow(unused_unsafe)] - unsafe { - // SAFETY: the caller has asserted that `fcinfo` is a valid FunctionCallInfo pointer, allocated by Postgres - // with all its fields properly setup. Unless the user is calling this wrapper function directly, this - // will always be the case - ::pgrx::iter::TableIterator::srf_next(#fcinfo_ident, || { - #( #arg_fetches )* - #result_handler - }) - } + let fcinfo = #fcinfo_ident; + let result = match <#ret_ty as ::pgrx::callconv::RetAbi>::check_fcinfo_and_prepare(fcinfo) { + ::pgrx::callconv::CallCx::WrappedFn(mcx) => { + let mut mcx = ::pgrx::PgMemoryContexts::For(mcx); + ::pgrx::callconv::RetAbi::to_ret(mcx.switch_to(|_| { + #(#arg_fetches)* + #func_name( #(#arg_pats),* ) + })) + } + ::pgrx::callconv::CallCx::RestoreCx => <#ret_ty as ::pgrx::callconv::RetAbi>::ret_from_fcinfo_fcx(fcinfo), + }; + unsafe { <#ret_ty as ::pgrx::callconv::RetAbi>::box_ret_in_fcinfo(fcinfo, result) } } }; - finfo_v1_extern_c(&self.func, fcinfo_ident, iter_closure) + finfo_v1_extern_c(&self.func, fcinfo_ident, wrapper_code) } } } diff --git a/pgrx-sql-entity-graph/src/pg_extern/returning.rs b/pgrx-sql-entity-graph/src/pg_extern/returning.rs index e535ea3db..b2f984056 100644 --- a/pgrx-sql-entity-graph/src/pg_extern/returning.rs +++ b/pgrx-sql-entity-graph/src/pg_extern/returning.rs @@ -35,8 +35,8 @@ pub struct ReturningIteratedItem { pub enum Returning { None, Type(UsedType), - SetOf { ty: UsedType, is_option: bool, is_result: bool }, - Iterated { tys: Vec, is_option: bool, is_result: bool }, + SetOf { ty: UsedType }, + Iterated { tys: Vec }, // /// Technically we don't ever create this, single triggers have their own macro. // Trigger, } @@ -61,7 +61,7 @@ impl Returning { match &mut *ty { syn::Type::Path(typepath) => { - let mut is_option = typepath.last_ident_is("Option"); + let is_option = typepath.last_ident_is("Option"); let is_result = typepath.last_ident_is("Result"); let mut is_setof_iter = typepath.last_ident_is("SetOfIterator"); let mut is_table_iter = typepath.last_ident_is("TableIterator"); @@ -116,7 +116,6 @@ impl Returning { "where's the generic args?", )); }; - is_option = true; segments = this_path.path.segments.clone(); // recurse deeper } else { if segments.last_ident_is("SetOfIterator") { @@ -163,7 +162,7 @@ impl Returning { )) } }; - Ok(Returning::SetOf { ty: used_ty, is_option, is_result }) + Ok(Returning::SetOf { ty: used_ty }) } else if is_table_iter { let last_path_segment = segments.last_mut().unwrap(); let mut iterated_items = vec![]; @@ -246,7 +245,7 @@ impl Returning { )) } }; - Ok(Returning::Iterated { tys: iterated_items, is_option, is_result }) + Ok(Returning::Iterated { tys: iterated_items }) } else { let used_ty = UsedType::new(syn::Type::Path(typepath.clone()))?; Ok(Returning::Type(used_ty)) @@ -302,17 +301,15 @@ impl ToTokens for Returning { } } } - Returning::SetOf { ty: used_ty, is_option, is_result } => { + Returning::SetOf { ty: used_ty } => { let used_ty_entity_tokens = used_ty.entity_tokens(); quote! { ::pgrx::pgrx_sql_entity_graph::PgExternReturnEntity::SetOf { ty: #used_ty_entity_tokens, - is_option: #is_option, - is_result: #is_result - } + } } } - Returning::Iterated { tys: items, is_option, is_result } => { + Returning::Iterated { tys: items } => { let quoted_items = items .iter() .map(|ReturningIteratedItem { used_ty, name }| { @@ -331,8 +328,6 @@ impl ToTokens for Returning { tys: vec![ #(#quoted_items),* ], - is_option: #is_option, - is_result: #is_result } } } diff --git a/pgrx-tests/src/tests/srf_tests.rs b/pgrx-tests/src/tests/srf_tests.rs index 1b671701a..0bb74a8c6 100644 --- a/pgrx-tests/src/tests/srf_tests.rs +++ b/pgrx-tests/src/tests/srf_tests.rs @@ -27,33 +27,33 @@ fn example_composite_set() -> TableIterator<'static, (name!(idx, i32), name!(val } #[pg_extern] -fn return_some_iterator( -) -> Option> { - Some(TableIterator::new( +fn return_table_iterator( +) -> TableIterator<'static, (name!(idx, i32), name!(some_value, &'static str))> { + TableIterator::new( vec!["a", "b", "c"].into_iter().enumerate().map(|(idx, value)| ((idx + 1) as i32, value)), - )) + ) } #[pg_extern] -fn return_none_iterator( -) -> Option> { - None +fn return_empty_iterator( +) -> TableIterator<'static, (name!(idx, i32), name!(some_value, &'static str))> { + TableIterator::empty() } #[pg_extern] -fn return_some_setof_iterator() -> Option> { - Some(SetOfIterator::new(vec![1, 2, 3].into_iter())) +fn return_setof_iterator() -> SetOfIterator<'static, i32> { + SetOfIterator::new(vec![1, 2, 3].into_iter()) } #[pg_extern] -fn return_none_setof_iterator() -> Option> { - None +fn return_empty_setof_iterator() -> SetOfIterator<'static, i32> { + SetOfIterator::empty() } #[pg_extern] -fn return_none_result_setof_iterator( -) -> Result>, Box> { - Ok(None) +fn return_empty_result_setof_iterator( +) -> Result, Box> { + Ok(SetOfIterator::empty()) } // TODO: We don't yet support returning Result> because the code generator @@ -82,31 +82,31 @@ fn split_table_with_borrow<'a>( #[pg_extern] fn result_table_1() -> Result< - Option<::pgrx::iter::TableIterator<'static, (name!(a, Option), name!(b, Option))>>, + ::pgrx::iter::TableIterator<'static, (name!(a, Option), name!(b, Option))>, Box, > { - Ok(Some(TableIterator::new(vec![(Some(1), Some(2))]))) + Ok(TableIterator::new(vec![(Some(1), Some(2))])) } #[pg_extern] fn result_table_2() -> Result< - Option), name!(b, Option))>>, + TableIterator<'static, (name!(a, Option), name!(b, Option))>, Box, > { - Ok(Some(TableIterator::new(vec![(Some(1), Some(2))]))) + Ok(TableIterator::new(vec![(Some(1), Some(2))])) } #[pg_extern] fn result_table_3() -> Result< - Option), name!(b, Option))>>, + TableIterator<'static, (name!(a, Option), name!(b, Option))>, Box, > { - Ok(Some(TableIterator::new(vec![(Some(1), Some(2))]))) + Ok(TableIterator::new(vec![(Some(1), Some(2))])) } #[pg_extern] fn result_table_4_err() -> Result< - Option), name!(b, Option))>>, + TableIterator<'static, (name!(a, Option), name!(b, Option))>, Box, > { Err("oh no")? @@ -114,10 +114,10 @@ fn result_table_4_err() -> Result< #[pg_extern] fn result_table_5_none() -> Result< - Option), name!(b, Option))>>, + TableIterator<'static, (name!(a, Option), name!(b, Option))>, Box, > { - Ok(None) + Ok(TableIterator::empty()) } #[pg_extern] @@ -126,8 +126,8 @@ fn one_col() -> TableIterator<'static, (name!(a, i32),)> { } #[pg_extern] -fn one_col_option() -> Option> { - Some(TableIterator::once((42,))) +fn one_col_option() -> TableIterator<'static, (name!(a, i32),)> { + TableIterator::once((42,)) } #[pg_extern] @@ -138,8 +138,8 @@ fn one_col_result() -> Result, Box Result>, Box> { - Ok(Some(TableIterator::once((42,)))) +) -> Result, Box> { + Ok(TableIterator::once((42,))) } #[cfg(any(test, feature = "pg_test"))] @@ -198,9 +198,9 @@ mod tests { } #[pg_test] - fn test_return_some_iterator() { + fn test_return_table_iterator() { let cnt = Spi::connect(|client| { - let table = client.select("SELECT * from return_some_iterator();", None, None)?; + let table = client.select("SELECT * from return_table_iterator();", None, None)?; Ok::<_, spi::Error>(table.len() as i64) }); @@ -209,9 +209,9 @@ mod tests { } #[pg_test] - fn test_return_none_iterator() { + fn test_return_empty_iterator() { let cnt = Spi::connect(|client| { - let table = client.select("SELECT * from return_none_iterator();", None, None)?; + let table = client.select("SELECT * from return_empty_iterator();", None, None)?; Ok::<_, spi::Error>(table.len() as i64) }); @@ -220,9 +220,9 @@ mod tests { } #[pg_test] - fn test_return_some_setof_iterator() { + fn test_return_setof_iterator() { let cnt = Spi::connect(|client| { - let table = client.select("SELECT * from return_some_setof_iterator();", None, None)?; + let table = client.select("SELECT * from return_setof_iterator();", None, None)?; Ok::<_, spi::Error>(table.len() as i64) }); @@ -231,9 +231,10 @@ mod tests { } #[pg_test] - fn test_return_none_setof_iterator() { + fn test_return_empty_setof_iterator() { let cnt = Spi::connect(|client| { - let table = client.select("SELECT * from return_none_setof_iterator();", None, None)?; + let table = + client.select("SELECT * from return_empty_setof_iterator();", None, None)?; Ok::<_, spi::Error>(table.len() as i64) }); @@ -279,22 +280,33 @@ mod tests { #[pg_test(error = "column \"cause_an_error\" does not exist")] pub fn spi_in_iterator( - ) -> TableIterator<'static, (name!(id, i32), name!(relname, Result, spi::Error>))> + ) -> Result))>, spi::Error> { let oids = vec![1213, 1214, 1232, 1233, 1247, 1249, 1255]; - - TableIterator::new(oids.into_iter().map(|oid| { - (oid, Spi::get_one(&format!("SELECT CAUSE_AN_ERROR FROM pg_class WHERE oid = {oid}"))) - })) + let result = oids + .into_iter() + .map(|oid| { + Ok(( + oid, + Spi::get_one(&format!( + "SELECT CAUSE_AN_ERROR FROM pg_class WHERE oid = {oid}" + ))?, + )) + }) + .collect::, _>>(); + result.map(|v| TableIterator::new(v)) } #[pg_test(error = "column \"cause_an_error\" does not exist")] - pub fn spi_in_setof() -> SetOfIterator<'static, Result, spi::Error>> { + pub fn spi_in_setof() -> Result>, spi::Error> { let oids = vec![1213, 1214, 1232, 1233, 1247, 1249, 1255]; - - SetOfIterator::new(oids.into_iter().map(|oid| { - Spi::get_one(&format!("SELECT CAUSE_AN_ERROR FROM pg_class WHERE oid = {oid}")) - })) + let result = oids + .into_iter() + .map(|oid| { + Spi::get_one(&format!("SELECT CAUSE_AN_ERROR FROM pg_class WHERE oid = {oid}")) + }) + .collect::>, _>>(); + result.map(SetOfIterator::new) } #[pg_test] diff --git a/pgrx-tests/tests/compile-fail/eq-for-postgres_hash.stderr b/pgrx-tests/tests/compile-fail/eq-for-postgres_hash.stderr index 1e05137b8..d649617d0 100644 --- a/pgrx-tests/tests/compile-fail/eq-for-postgres_hash.stderr +++ b/pgrx-tests/tests/compile-fail/eq-for-postgres_hash.stderr @@ -39,7 +39,7 @@ note: required by a bound in `brokentype_hash` | ^^^^^^^^^^^^ required by this bound in `brokentype_hash` 5 | pub struct BrokenType { | ---------- required by a bound in this function - = note: this error originates in the attribute macro `::pgrx::pgrx_macros::pg_extern` which comes from the expansion of the derive macro `PostgresHash` (in Nightly builds, run with -Z macro-backtrace for more info) + = note: this error originates in the derive macro `PostgresHash` (in Nightly builds, run with -Z macro-backtrace for more info) help: consider annotating `BrokenType` with `#[derive(Hash)]` | 5 + #[derive(Hash)] @@ -59,7 +59,7 @@ note: required by a bound in `brokentype_hash` | ^^^^^^^^^^^^ required by this bound in `brokentype_hash` 5 | pub struct BrokenType { | ---------- required by a bound in this function - = note: this error originates in the attribute macro `::pgrx::pgrx_macros::pg_extern` which comes from the expansion of the derive macro `PostgresHash` (in Nightly builds, run with -Z macro-backtrace for more info) + = note: this error originates in the derive macro `PostgresHash` (in Nightly builds, run with -Z macro-backtrace for more info) help: consider annotating `BrokenType` with `#[derive(Eq)]` | 5 + #[derive(Eq)] diff --git a/pgrx-tests/tests/compile-fail/total-eq-for-postgres_eq.stderr b/pgrx-tests/tests/compile-fail/total-eq-for-postgres_eq.stderr index ec2a80321..05a918704 100644 --- a/pgrx-tests/tests/compile-fail/total-eq-for-postgres_eq.stderr +++ b/pgrx-tests/tests/compile-fail/total-eq-for-postgres_eq.stderr @@ -25,7 +25,7 @@ note: required by a bound in `brokentype_eq` | ^^^^^^^^^^ required by this bound in `brokentype_eq` 5 | pub struct BrokenType { | ---------- required by a bound in this function - = note: this error originates in the attribute macro `::pgrx::pgrx_macros::pg_operator` which comes from the expansion of the derive macro `PostgresEq` (in Nightly builds, run with -Z macro-backtrace for more info) + = note: this error originates in the derive macro `PostgresEq` (in Nightly builds, run with -Z macro-backtrace for more info) help: consider annotating `BrokenType` with `#[derive(Eq)]` | 5 + #[derive(Eq)] diff --git a/pgrx/src/callconv.rs b/pgrx/src/callconv.rs index 5b2169a7f..8e49124b1 100644 --- a/pgrx/src/callconv.rs +++ b/pgrx/src/callconv.rs @@ -8,129 +8,278 @@ //LICENSE //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. #![doc(hidden)] +#![deny(unsafe_op_in_unsafe_fn)] //! Helper implementations for returning sets and tables from `#[pg_extern]`-style functions -use crate::iter::{SetOfIterator, TableIterator}; +use crate::heap_tuple::PgHeapTuple; use crate::{ - pg_return_null, pg_sys, srf_is_first_call, srf_return_done, srf_return_next, IntoDatum, - IntoHeapTuple, PgMemoryContexts, + pg_return_null, pg_sys, AnyNumeric, Date, Inet, Internal, Interval, IntoDatum, Json, PgBox, + PgVarlena, Time, TimeWithTimeZone, Timestamp, TimestampWithTimeZone, Uuid, }; -use core::ops::ControlFlow; -use core::ptr; - -impl<'a, T: IntoDatum> SetOfIterator<'a, T> { - #[doc(hidden)] - pub unsafe fn srf_next( - fcinfo: pg_sys::FunctionCallInfo, - wrapped_fn: impl FnOnce() -> Option>, - ) -> pg_sys::Datum { - if fcx_needs_setup(fcinfo) { - let fcx = deref_fcx(fcinfo); - // first off, ask the user's function to do the needful and return Option> - let setof_iterator = srf_memcx(fcx).switch_to(|_| wrapped_fn()); - if let ControlFlow::Break(datum) = finish_srf_init(setof_iterator, fcinfo) { - return datum; - } - } +use std::ffi::{CStr, CString}; - let fcx = deref_fcx(fcinfo); - // SAFETY: fcx.user_fctx was set earlier, immediately before or in a prior call - let setof_iterator = &mut *(*fcx).user_fctx.cast::>(); +/// How to return a value from Rust to Postgres +/// +/// This bound is necessary to distinguish things which can be returned from a `#[pg_extern] fn`. +/// This bound is not accurately described by IntoDatum or similar traits, as value conversions are +/// handled in a special way at function return boundaries, and may require mutating multiple fields +/// behind the FunctionCallInfo. The most exceptional case are set-returning functions, which +/// require special handling for the fcinfo and also for certain inner types. +/// +/// This trait is exposed to external code so macro-generated wrapper fn may expand to calls to it. +/// The number of invariants implementers must uphold is unlikely to be adequately documented. +/// Prefer to use RetAbi as a trait bound instead of implementing it, or even calling it, yourself. +pub unsafe trait RetAbi: Sized { + /// Type returned to Postgres + type Item: Sized; + /// Driver for complex returns + type Ret; - match setof_iterator.next() { - Some(datum) => { - srf_return_next(fcinfo, fcx); - datum.into_datum().unwrap_or_else(|| pg_return_null(fcinfo)) - } - None => empty_srf(fcinfo), - } + /// Initialize the FunctionCallInfo for returns + /// + /// The implementer must pick the correct memory context for the wrapped fn's allocations. + /// # Safety + /// Requires a valid FunctionCallInfo. + unsafe fn check_fcinfo_and_prepare(_fcinfo: pg_sys::FunctionCallInfo) -> CallCx { + CallCx::WrappedFn(unsafe { pg_sys::CurrentMemoryContext }) + } + + /// answer what kind and how many returns happen from this type + fn to_ret(self) -> Self::Ret; + + /// box the return value + /// # Safety + /// must be called with a valid fcinfo + unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum; + + /// Multi-call types want to be in the fcinfo so they can be restored + /// # Safety + /// must be called with a valid fcinfo + unsafe fn move_into_fcinfo_fcx(self, _fcinfo: pg_sys::FunctionCallInfo); + + /// Other types want to add metadata to the fcinfo + /// # Safety + /// must be called with a valid fcinfo + unsafe fn fill_fcinfo_fcx(&self, _fcinfo: pg_sys::FunctionCallInfo); + + /// for multi-call types, how to restore them from the multi-call context + /// + /// for all others: panic + /// # Safety + /// must be called with a valid fcinfo + unsafe fn ret_from_fcinfo_fcx(_fcinfo: pg_sys::FunctionCallInfo) -> Self::Ret { + unimplemented!() } + + /// must be called with a valid fcinfo + unsafe fn finish_call_fcinfo(_fcinfo: pg_sys::FunctionCallInfo) {} +} + +/// A simplified blanket RetAbi +pub unsafe trait BoxRet: Sized { + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum; } -impl<'a, T: IntoHeapTuple> TableIterator<'a, T> { - #[doc(hidden)] - pub unsafe fn srf_next( - fcinfo: pg_sys::FunctionCallInfo, - wrapped_fn: impl FnOnce() -> Option>, - ) -> pg_sys::Datum { - if fcx_needs_setup(fcinfo) { - let fcx = deref_fcx(fcinfo); - let table_iterator = srf_memcx(fcx).switch_to(|_| { - // first off, ask the user's function to do the needful and return Option> - let table_iterator = wrapped_fn(); - - // Build a tuple descriptor for our result type - let mut tupdesc = ptr::null_mut(); - let ty_class = pg_sys::get_call_result_type(fcinfo, ptr::null_mut(), &mut tupdesc); - if ty_class != pg_sys::TypeFuncClass_TYPEFUNC_COMPOSITE { - pg_sys::error!("return type must be a row type"); - } - pg_sys::BlessTupleDesc(tupdesc); - (*fcx).tuple_desc = tupdesc; - - table_iterator - }); - - if let ControlFlow::Break(datum) = finish_srf_init(table_iterator, fcinfo) { - return datum; +unsafe impl RetAbi for T +where + T: BoxRet, +{ + type Item = Self; + type Ret = Self; + + fn to_ret(self) -> Self::Ret { + self + } + + unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum { + unsafe { ret.box_in_fcinfo(fcinfo) } + } + + unsafe fn check_fcinfo_and_prepare(_fcinfo: pg_sys::FunctionCallInfo) -> CallCx { + CallCx::WrappedFn(unsafe { pg_sys::CurrentMemoryContext }) + } + + unsafe fn fill_fcinfo_fcx(&self, _fcinfo: pg_sys::FunctionCallInfo) {} + unsafe fn move_into_fcinfo_fcx(self, _fcinfo: pg_sys::FunctionCallInfo) {} + unsafe fn ret_from_fcinfo_fcx(_fcinfo: pg_sys::FunctionCallInfo) -> Self::Ret { + unimplemented!() + } + unsafe fn finish_call_fcinfo(_fcinfo: pg_sys::FunctionCallInfo) {} +} + +/// Control flow for RetAbi +pub enum CallCx { + RestoreCx, + WrappedFn(pg_sys::MemoryContext), +} + +unsafe impl BoxRet for Option +where + T: BoxRet, +{ + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + unsafe { + match self { + None => pg_return_null(fcinfo), + Some(value) => value.box_in_fcinfo(fcinfo), } } + } +} - let fcx = deref_fcx(fcinfo); - // SAFETY: fcx.user_fctx was set earlier, immediately before or in a prior call - let table_iterator = &mut *(*fcx).user_fctx.cast::>(); +unsafe impl RetAbi for Result +where + T: RetAbi, + T::Item: RetAbi, + E: core::any::Any + core::fmt::Display, +{ + type Item = T::Item; + type Ret = T::Ret; - match table_iterator.next() { - Some(tuple) => { - let heap_tuple = tuple.into_heap_tuple((*fcx).tuple_desc); - srf_return_next(fcinfo, fcx); - pg_sys::HeapTupleHeaderGetDatum((*heap_tuple).t_data) - } - None => empty_srf(fcinfo), + unsafe fn check_fcinfo_and_prepare(fcinfo: pg_sys::FunctionCallInfo) -> CallCx { + unsafe { T::check_fcinfo_and_prepare(fcinfo) } + } + + fn to_ret(self) -> Self::Ret { + let value = pg_sys::panic::ErrorReportable::unwrap_or_report(self); + value.to_ret() + } + + unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum { + unsafe { T::box_ret_in_fcinfo(fcinfo, ret) } + } + + unsafe fn fill_fcinfo_fcx(&self, fcinfo: pg_sys::FunctionCallInfo) { + match self { + Ok(value) => unsafe { value.fill_fcinfo_fcx(fcinfo) }, + Err(_) => (), + } + } + + unsafe fn move_into_fcinfo_fcx(self, fcinfo: pg_sys::FunctionCallInfo) { + match self { + Ok(value) => unsafe { value.move_into_fcinfo_fcx(fcinfo) }, + Err(_) => (), } } + + unsafe fn ret_from_fcinfo_fcx(fcinfo: pg_sys::FunctionCallInfo) -> Self::Ret { + unsafe { T::ret_from_fcinfo_fcx(fcinfo) } + } + + unsafe fn finish_call_fcinfo(fcinfo: pg_sys::FunctionCallInfo) { + unsafe { T::finish_call_fcinfo(fcinfo) } + } } -fn fcx_needs_setup(fcinfo: pg_sys::FunctionCallInfo) -> bool { - let need = unsafe { srf_is_first_call(fcinfo) }; - if need { - unsafe { pg_sys::init_MultiFuncCall(fcinfo) }; +macro_rules! return_packaging_for_primitives { + ($($scalar:ty),*) => { + $(unsafe impl BoxRet for $scalar { + unsafe fn box_in_fcinfo(self, _fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + $crate::pg_sys::Datum::from(self) + } + })* } - need } -fn empty_srf(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { - unsafe { - let fcx = deref_fcx(fcinfo); - srf_return_done(fcinfo, fcx); - pg_return_null(fcinfo) +return_packaging_for_primitives!(i8, i16, i32, i64, bool); + +unsafe impl BoxRet for () { + unsafe fn box_in_fcinfo(self, _fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + pg_sys::Datum::from(0) } } -/// "per_MultiFuncCall" but no FFI cost -fn deref_fcx(fcinfo: pg_sys::FunctionCallInfo) -> *mut pg_sys::FuncCallContext { - unsafe { (*(*fcinfo).flinfo).fn_extra.cast() } +unsafe impl BoxRet for f32 { + unsafe fn box_in_fcinfo(self, _fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + pg_sys::Datum::from(self.to_bits()) + } } -fn srf_memcx(fcx: *mut pg_sys::FuncCallContext) -> PgMemoryContexts { - unsafe { PgMemoryContexts::For((*fcx).multi_call_memory_ctx) } +unsafe impl BoxRet for f64 { + unsafe fn box_in_fcinfo(self, _fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + pg_sys::Datum::from(self.to_bits()) + } } -fn finish_srf_init( - arg: Option, - fcinfo: pg_sys::FunctionCallInfo, -) -> ControlFlow { - match arg { - // nothing to iterate? - None => ControlFlow::Break(empty_srf(fcinfo)), - // must be saved for the next call by leaking it into the multi-call memory context - Some(value) => { - let fcx = deref_fcx(fcinfo); - unsafe { - let ptr = srf_memcx(fcx).leak_and_drop_on_delete(value); - // it's the first call so we need to finish setting up fcx - (*fcx).user_fctx = ptr.cast(); - } - ControlFlow::Continue(()) - } +unsafe impl<'a> BoxRet for &'a [u8] { + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) + } +} + +unsafe impl<'a> BoxRet for &'a str { + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) + } +} + +unsafe impl<'a> BoxRet for &'a CStr { + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) + } +} + +macro_rules! impl_repackage_into_datum { + ($($boxable:ty),*) => { + $(unsafe impl BoxRet for $boxable { + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) + } + })* + }; +} + +impl_repackage_into_datum! { + String, CString, Vec, char, + Json, Inet, Uuid, AnyNumeric, Internal, + Date, Interval, Time, TimeWithTimeZone, Timestamp, TimestampWithTimeZone, + pg_sys::Oid, pg_sys::BOX, pg_sys::Point +} + +unsafe impl BoxRet for crate::Numeric { + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) + } +} + +unsafe impl BoxRet for crate::Range +where + T: IntoDatum + crate::RangeSubType, +{ + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) + } +} + +unsafe impl BoxRet for Vec +where + T: IntoDatum, +{ + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) + } +} + +unsafe impl BoxRet for PgVarlena { + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) + } +} + +unsafe impl<'mcx, A> BoxRet for PgHeapTuple<'mcx, A> +where + A: crate::WhoAllocated, +{ + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) + } +} + +unsafe impl BoxRet for PgBox +where + A: crate::WhoAllocated, +{ + unsafe fn box_in_fcinfo(self, fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + self.into_datum().unwrap_or_else(|| unsafe { pg_return_null(fcinfo) }) } } diff --git a/pgrx/src/iter.rs b/pgrx/src/iter.rs index 38b0bb16e..84924f82e 100644 --- a/pgrx/src/iter.rs +++ b/pgrx/src/iter.rs @@ -7,10 +7,14 @@ //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. -#![allow(clippy::vec_init_then_push)] -use std::iter; +#![deny(unsafe_op_in_unsafe_fn)] +use core::{iter, ptr}; -use crate::{pg_sys, IntoDatum, IntoHeapTuple}; +use crate::callconv::{BoxRet, CallCx, RetAbi}; +use crate::fcinfo::{pg_return_null, srf_is_first_call, srf_return_done, srf_return_next}; +use crate::ptr::PointerExt; +use crate::{pg_sys, IntoDatum, IntoHeapTuple, PgMemoryContexts}; +use pgrx_pg_sys::TypeFuncClass_TYPEFUNC_COMPOSITE; use pgrx_sql_entity_graph::metadata::{ ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable, }; @@ -45,13 +49,16 @@ use pgrx_sql_entity_graph::metadata::{ /// SetOfIterator::new(input.split_whitespace()) /// } /// ``` -pub struct SetOfIterator<'a, T> { - iter: Box + 'a>, -} +#[repr(transparent)] +pub struct SetOfIterator<'a, T>( + // Postgres uses the same ABI for `returns setof` and 1-column `returns table` + TableIterator<'a, (T,)>, +); impl<'a, T: 'a> SetOfIterator<'a, T> { pub fn new(iter: impl IntoIterator + 'a) -> Self { - Self { iter: Box::new(iter.into_iter()) } + // Forward the impl using remapping to minimize `unsafe` code and keep them in sync + Self(TableIterator::new(iter.into_iter().map(|c| (c,)))) } pub fn empty() -> Self { @@ -68,10 +75,11 @@ impl<'a, T> Iterator for SetOfIterator<'a, T> { #[inline] fn next(&mut self) -> Option { - self.iter.next() + self.0.next().map(|(val,)| val) } } +/// `SetOfIterator<'_, T>` differs from `TableIterator<'a, (T,)>` in generated SQL unsafe impl<'a, T> SqlTranslatable for SetOfIterator<'a, T> where T: SqlTranslatable, @@ -131,15 +139,12 @@ where /// TableIterator::new(input.split_whitespace().enumerate().map(|(n, w)| (n as i32, w))) /// } /// ``` -pub struct TableIterator<'a, T> { - iter: Box + 'a>, +pub struct TableIterator<'a, Row> { + iter: Box + 'a>, } -impl<'a, T> TableIterator<'a, T> -where - T: IntoHeapTuple + 'a, -{ - pub fn new(iter: impl IntoIterator + 'a) -> Self { +impl<'a, Row: 'a> TableIterator<'a, Row> { + pub fn new(iter: impl IntoIterator + 'a) -> Self { Self { iter: Box::new(iter.into_iter()) } } @@ -147,13 +152,13 @@ where Self::new(iter::empty()) } - pub fn once(value: T) -> Self { + pub fn once(value: Row) -> Self { Self::new(iter::once(value)) } } -impl<'a, T> Iterator for TableIterator<'a, T> { - type Item = T; +impl<'a, Row> Iterator for TableIterator<'a, Row> { + type Item = Row; #[inline] fn next(&mut self) -> Option { @@ -161,6 +166,204 @@ impl<'a, T> Iterator for TableIterator<'a, T> { } } +unsafe impl<'iter, C> SqlTranslatable for TableIterator<'iter, (C,)> +where + C: SqlTranslatable + 'iter, +{ + fn argument_sql() -> Result { + Err(ArgumentError::Table) + } + fn return_sql() -> Result { + let vec = vec![C::return_sql().and_then(|sql| match sql { + Returns::One(sql) => Ok(sql), + Returns::SetOf(_) => Err(ReturnsError::TableContainingSetOf), + Returns::Table(_) => Err(ReturnsError::NestedTable), + })?]; + Ok(Returns::Table(vec)) + } +} + +unsafe impl<'a, T> RetAbi for SetOfIterator<'a, T> +where + T: BoxRet, +{ + type Item = ::Item; + type Ret = IterRet; + + unsafe fn check_fcinfo_and_prepare(fcinfo: pg_sys::FunctionCallInfo) -> CallCx { + unsafe { TableIterator::<(T,)>::check_fcinfo_and_prepare(fcinfo) } + } + + fn to_ret(self) -> Self::Ret { + let mut iter = self; + IterRet(match iter.next() { + None => Step::Done, + Some(value) => Step::Init(iter, value), + }) + } + + unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum { + let ret = match ret.0 { + Step::Done => Step::Done, + Step::Once(value) => Step::Once((value,)), + Step::Init(iter, value) => Step::Init(iter.0, (value,)), + }; + unsafe { TableIterator::<(T,)>::box_ret_in_fcinfo(fcinfo, IterRet(ret)) } + } + + unsafe fn fill_fcinfo_fcx(&self, _fcinfo: pg_sys::FunctionCallInfo) {} + + unsafe fn move_into_fcinfo_fcx(self, fcinfo: pg_sys::FunctionCallInfo) { + unsafe { self.0.move_into_fcinfo_fcx(fcinfo) } + } + + unsafe fn ret_from_fcinfo_fcx(fcinfo: pg_sys::FunctionCallInfo) -> Self::Ret { + let step = match unsafe { TableIterator::<(T,)>::ret_from_fcinfo_fcx(fcinfo).0 } { + Step::Done => Step::Done, + Step::Once((item,)) => Step::Once(item), + Step::Init(iter, (value,)) => Step::Init(Self(iter), value), + }; + IterRet(step) + } + + unsafe fn finish_call_fcinfo(fcinfo: pg_sys::FunctionCallInfo) { + unsafe { TableIterator::<(T,)>::finish_call_fcinfo(fcinfo) } + } +} + +unsafe impl<'a, Row> RetAbi for TableIterator<'a, Row> +where + Row: RetAbi, +{ + type Item = ::Item; + type Ret = IterRet; + + unsafe fn check_fcinfo_and_prepare(fcinfo: pg_sys::FunctionCallInfo) -> CallCx { + unsafe { + if srf_is_first_call(fcinfo) { + let fn_call_cx = pg_sys::init_MultiFuncCall(fcinfo); + CallCx::WrappedFn((*fn_call_cx).multi_call_memory_ctx) + } else { + CallCx::RestoreCx + } + } + } + + fn to_ret(self) -> Self::Ret { + let mut iter = self; + IterRet(match iter.next() { + None => Step::Done, + Some(value) => Step::Init(iter, value), + }) + } + + unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum { + let value = unsafe { + match ret.0 { + Step::Done => return empty_srf(fcinfo), + Step::Once(value) => value, + Step::Init(iter, value) => { + // Move the iterator in and don't worry about putting it back + iter.move_into_fcinfo_fcx(fcinfo); + value.fill_fcinfo_fcx(fcinfo); + value + } + } + }; + + unsafe { + let fcx = deref_fcx(fcinfo); + srf_return_next(fcinfo, fcx); + ::box_ret_in_fcinfo(fcinfo, value.to_ret()) + } + } + + unsafe fn fill_fcinfo_fcx(&self, _fcinfo: pg_sys::FunctionCallInfo) {} + + unsafe fn move_into_fcinfo_fcx(self, fcinfo: pg_sys::FunctionCallInfo) { + unsafe { + let fcx = deref_fcx(fcinfo); + let ptr = srf_memcx(fcx).leak_and_drop_on_delete(self); + // it's the first call so we need to finish setting up fcx + (*fcx).user_fctx = ptr.cast(); + } + } + + unsafe fn ret_from_fcinfo_fcx(fcinfo: pg_sys::FunctionCallInfo) -> Self::Ret { + // SAFETY: fcx.user_fctx was set earlier, immediately before or in a prior call + let iter = unsafe { + let fcx = deref_fcx(fcinfo); + &mut *(*fcx).user_fctx.cast::>() + }; + IterRet(match iter.next() { + None => Step::Done, + Some(value) => Step::Once(value), + }) + } + + unsafe fn finish_call_fcinfo(fcinfo: pg_sys::FunctionCallInfo) { + unsafe { + let fcx = deref_fcx(fcinfo); + srf_return_done(fcinfo, fcx) + } + } +} + +/// How iterators are returned +pub struct IterRet(Step); + +/// ValuePerCall SRF steps +enum Step { + Done, + Once(T::Item), + Init(T, T::Item), +} + +pub(crate) unsafe fn empty_srf(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum { + unsafe { + let fcx = deref_fcx(fcinfo); + srf_return_done(fcinfo, fcx); + pg_return_null(fcinfo) + } +} + +/// "per_MultiFuncCall" but no FFI cost +pub(crate) unsafe fn deref_fcx(fcinfo: pg_sys::FunctionCallInfo) -> *mut pg_sys::FuncCallContext { + unsafe { (*(*fcinfo).flinfo).fn_extra.cast() } +} + +pub(crate) unsafe fn srf_memcx(fcx: *mut pg_sys::FuncCallContext) -> PgMemoryContexts { + unsafe { PgMemoryContexts::For((*fcx).multi_call_memory_ctx) } +} + +/// Return ABI for single-column tuples +/// +/// Postgres extended SQL with `returns setof $ty` before SQL had `returns table($($ident $ty),*)`. +/// Due to this history, single-column `returns table` reuses the internals of `returns setof`. +/// This lets them simply return a simple Datum instead of handling a TupleDesc and HeapTuple, but +/// means we need to have this distinct impl, as the return type is not `TYPEFUNC_COMPOSITE`! +/// Fortunately, RetAbi lets `TableIterator<'a, Tup>` handle this by calling ``. +unsafe impl RetAbi for (C,) +where + C: BoxRet, // so we support TableIterator<'a, (Option,)> as well +{ + type Item = C; + type Ret = C; + + fn to_ret(self) -> Self::Ret { + self.0 + } + + /// Returning a "row" of only one column is identical to returning that single type + unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum { + unsafe { C::box_ret_in_fcinfo(fcinfo, ret.to_ret()) } + } + + unsafe fn fill_fcinfo_fcx(&self, _fcinfo: pg_sys::FunctionCallInfo) {} + + unsafe fn move_into_fcinfo_fcx(self, _fcinfo: pg_sys::FunctionCallInfo) {} +} + macro_rules! impl_table_iter { ($($C:ident),* $(,)?) => { unsafe impl<'iter, $($C,)*> SqlTranslatable for TableIterator<'iter, ($($C,)*)> @@ -173,12 +376,11 @@ macro_rules! impl_table_iter { fn return_sql() -> Result { let vec = vec![ $( - match $C::return_sql() { - Ok(Returns::One(sql)) => sql, - Ok(Returns::SetOf(_)) => return Err(ReturnsError::TableContainingSetOf), - Ok(Returns::Table(_)) => return Err(ReturnsError::NestedTable), - err => return err, - }, + $C::return_sql().and_then(|sql| match sql { + Returns::One(sql) => Ok(sql), + Returns::SetOf(_) => Err(ReturnsError::TableContainingSetOf), + Returns::Table(_) => Err(ReturnsError::NestedTable), + })?, )* ]; Ok(Returns::Table(vec)) @@ -194,7 +396,6 @@ macro_rules! impl_table_iter { let mut nulls = datums.map(|option| option.is_none()); let mut datums = datums.map(|option| option.unwrap_or(pg_sys::Datum::from(0))); - unsafe { // SAFETY: Caller has asserted that `tupdesc` is valid, and we just went // through a little bit of effort to setup properly sized arrays for @@ -204,10 +405,48 @@ macro_rules! impl_table_iter { } } + unsafe impl<$($C),*> RetAbi for ($($C,)*) + where + $($C: BoxRet,)* + Self: IntoHeapTuple, + { + type Item = Self; + type Ret = Self; + + fn to_ret(self) -> Self::Ret { + self + } + + unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum { + unsafe { + let fcx = deref_fcx(fcinfo); + let heap_tuple = ret.into_heap_tuple((*fcx).tuple_desc); + pg_sys::HeapTupleHeaderGetDatum((*heap_tuple).t_data) + } + } + + unsafe fn move_into_fcinfo_fcx(self, _fcinfo: pg_sys::FunctionCallInfo) {} + + unsafe fn fill_fcinfo_fcx(&self, fcinfo: pg_sys::FunctionCallInfo) { + // Pure side effect, leave the value in place. + unsafe { + let fcx = deref_fcx(fcinfo); + srf_memcx(fcx).switch_to(|_| { + let mut tupdesc = ptr::null_mut(); + let mut oid = pg_sys::Oid::default(); + let ty_class = pg_sys::get_call_result_type(fcinfo, &mut oid, &mut tupdesc); + if tupdesc.is_non_null() && ty_class == TypeFuncClass_TYPEFUNC_COMPOSITE { + pg_sys::BlessTupleDesc(tupdesc); + (*fcx).tuple_desc = tupdesc; + } + }); + } + } + } + } } -impl_table_iter!(T0); impl_table_iter!(T0, T1); impl_table_iter!(T0, T1, T2); impl_table_iter!(T0, T1, T2, T3);