do some cleanups around PgHeapTuple to make them easier to use (#1153)

Also addresses a safety issue by marking a function `unsafe`.

Author: eeeebbbbrrrr

pgrx-tests/src/tests/heap_tuple.rs

@@ -600,16 +600,18 @@ mod sql_generator_tests {
     #[pg_extern]
     fn generate_lots_of_dogs() -> SetOfIterator<'static, pgrx::composite_type!("Dog")> {
         let tuple_desc =
-            pgrx::PgTupleDesc::for_composite_type("Dog").expect("Coudln't find TestType");
+            pgrx::PgTupleDesc::for_composite_type("Dog").expect("Couldn't find TestType");
 
         let tuples: Vec<PgHeapTuple<'_, AllocatedByRust>> = (0..10_000)
             .into_iter()
             .map(move |i| {
                 let datums: Vec<Option<pg_sys::Datum>> =
                     vec!["good boy".into_datum(), i.into_datum()];
 
-                PgHeapTuple::from_datums(tuple_desc.clone(), datums)
-                    .expect("couldn't get heap tuple")
+                unsafe {
+                    PgHeapTuple::from_datums(tuple_desc.clone(), datums)
+                        .expect("couldn't get heap tuple")
+                }
             })
             .collect();
 

pgrx/src/heap_tuple.rs

@@ -9,6 +9,8 @@ use crate::{
     FromDatum, IntoDatum, PgBox, PgMemoryContexts, PgTupleDesc, TriggerTuple, TryFromDatumError,
     WhoAllocated,
 };
+use pgrx_pg_sys::errcodes::PgSqlErrorCode;
+use pgrx_pg_sys::PgTryBuilder;
 use pgrx_sql_entity_graph::metadata::{
     ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable,
 };
@@ -22,6 +24,12 @@ pub enum PgHeapTupleError {
 
     #[error("The specified composite type, {0}, does not exist")]
     NoSuchType(String),
+
+    #[error("The specified composite type, {0}, does not exist")]
+    NoSuchTypeOid(pg_sys::Oid),
+
+    #[error("Oid `{0}` is not a composite type")]
+    NotACompositeType(pg_sys::Oid),
 }
 
 /// A [`PgHeapTuple`] is a lightweight wrapper around Postgres' [`pg_sys::HeapTuple`] object and a [`PgTupleDesc`].
@@ -199,28 +207,52 @@ impl<'a> PgHeapTuple<'a, AllocatedByRust> {
     ) -> Result<PgHeapTuple<'a, AllocatedByRust>, PgHeapTupleError> {
         let tuple_desc = PgTupleDesc::for_composite_type(type_name)
             .ok_or_else(|| PgHeapTupleError::NoSuchType(type_name.to_string()))?;
-        let natts = tuple_desc.len();
-        unsafe {
-            let datums =
-                pg_sys::palloc0(natts * std::mem::size_of::<pg_sys::Datum>()) as *mut pg_sys::Datum;
-            let mut is_null = (0..natts).map(|_| true).collect::<Vec<_>>();
 
-            let heap_tuple =
-                pg_sys::heap_form_tuple(tuple_desc.as_ptr(), datums, is_null.as_mut_ptr());
+        Self::new_composite_type_by_oid(tuple_desc.oid())
+    }
 
-            Ok(PgHeapTuple {
-                tuple: PgBox::<pg_sys::HeapTupleData, AllocatedByRust>::from_rust(heap_tuple),
-                tupdesc: tuple_desc,
-            })
-        }
+    pub fn new_composite_type_by_oid(
+        typoid: pg_sys::Oid,
+    ) -> Result<PgHeapTuple<'a, AllocatedByRust>, PgHeapTupleError> {
+        PgTryBuilder::new(|| {
+            let tuple_desc = PgTupleDesc::for_composite_type_by_oid(typoid)
+                .ok_or_else(|| PgHeapTupleError::NotACompositeType(typoid))?;
+            let natts = tuple_desc.len();
+
+            unsafe {
+                let datums = pg_sys::palloc0(natts * std::mem::size_of::<pg_sys::Datum>())
+                    as *mut pg_sys::Datum;
+                let mut is_null = vec![true; natts];
+
+                let heap_tuple =
+                    pg_sys::heap_form_tuple(tuple_desc.as_ptr(), datums, is_null.as_mut_ptr());
+
+                Ok(PgHeapTuple {
+                    tuple: PgBox::<pg_sys::HeapTupleData, AllocatedByRust>::from_rust(heap_tuple),
+                    tupdesc: tuple_desc,
+                })
+            }
+        })
+        .catch_when(PgSqlErrorCode::ERRCODE_WRONG_OBJECT_TYPE, |_| {
+            Err(PgHeapTupleError::NotACompositeType(typoid))
+        })
+        .catch_when(PgSqlErrorCode::ERRCODE_UNDEFINED_OBJECT, |_| {
+            Err(PgHeapTupleError::NoSuchTypeOid(typoid))
+        })
+        .execute()
     }
 
     /// Create a new [PgHeapTuple] from a [PgTupleDesc] from an iterator of Datums.
     ///
     /// ## Errors
     /// - [PgHeapTupleError::IncorrectAttributeCount] if the number of items in the iterator
     /// does not match the number of attributes in the [PgTupleDesc].
-    pub fn from_datums<I: IntoIterator<Item = Option<pg_sys::Datum>>>(
+    ///
+    /// # Safety
+    ///
+    /// This function is unsafe as we cannot guarantee the provided [`pg_sys::Datum`]s are valid
+    /// as the specified [`PgTupleDesc`] might expect
+    pub unsafe fn from_datums<I: IntoIterator<Item = Option<pg_sys::Datum>>>(
         tupdesc: PgTupleDesc<'a>,
         datums: I,
     ) -> Result<PgHeapTuple<'a, AllocatedByRust>, PgHeapTupleError> {

pgrx/src/tupdesc.rs

@@ -156,7 +156,7 @@ impl<'a> PgTupleDesc<'a> {
     ```rust,no_run
     use pgrx::{prelude::*, PgTupleDesc};
 
-    Spi::run("CREATE TYPE Dog AS (name text, age int);");
+    Spi::run("CREATE TYPE Dog AS (name text, age int);").unwrap();
     let tuple_desc = PgTupleDesc::for_composite_type("Dog").unwrap();
     let natts = tuple_desc.len();
 
@@ -179,12 +179,20 @@ impl<'a> PgTupleDesc<'a> {
             let mut typmod = 0;
             pg_sys::parseTypeString(name.as_pg_cstr(), &mut typoid, &mut typmod, true);
 
+            Self::for_composite_type_by_oid(typoid)
+        }
+    }
+
+    /// Similar to [`PgTupleDesc::for_composite_type()`] but using the type's [`pg_sys::Oid`] instead
+    /// of its name.
+    pub fn for_composite_type_by_oid(typoid: pg_sys::Oid) -> Option<PgTupleDesc<'a>> {
+        unsafe {
             if typoid == pg_sys::InvalidOid {
                 return None;
             }
 
             // It's important to make a copy of the tupledesc: https://www.postgresql.org/message-id/flat/24471.1136768659%40sss.pgh.pa.us
-            let tuple_desc = pg_sys::lookup_rowtype_tupdesc_copy(typoid, typmod);
+            let tuple_desc = pg_sys::lookup_rowtype_tupdesc_copy(typoid, -1);
 
             Some(PgTupleDesc::from_pg_copy(tuple_desc))
         }