Include type when getting argument (#740)

* Include type when getting argument

* Test that oid of AnyElement is correct
* Only get argument type when needed
* Use separate trait function when OID needed
* from_datum_with_typid -> from_polymorphic_datum

Co-authored-by: Jubilee <46493976+workingjubilee@users.noreply.github.com>

pgx-tests/src/tests/pg_extern_tests.rs

@@ -96,4 +96,17 @@ mod tests {
                 .expect("failed to get SPI result");
         assert_eq!(create_result, 42);
     }
+
+    #[pg_extern]
+    fn anyele_type(x: pgx::AnyElement) -> i32 {
+        x.oid() as i32
+    }
+
+    #[pg_test]
+    fn test_anyele_type() {
+        let interval_type =
+            Spi::get_one::<i32>(r#"SELECT tests."anyele_type"('5 hours'::interval)"#)
+                .expect("failed to get SPI result");
+        assert_eq!(interval_type as u32, pg_sys::INTERVALOID);
+    }
 }

pgx/src/datum/anyarray.rs

@@ -34,11 +34,22 @@ impl AnyArray {
 }
 
 impl FromDatum for AnyArray {
+    const GET_TYPOID: bool = true;
+
     #[inline]
     unsafe fn from_datum(
         datum: pg_sys::Datum,
         is_null: bool,
         typoid: pg_sys::Oid,
+    ) -> Option<AnyArray> {
+        FromDatum::from_polymorphic_datum(datum, is_null, typoid)
+    }
+
+    #[inline]
+    unsafe fn from_polymorphic_datum(
+        datum: pg_sys::Datum,
+        is_null: bool,
+        typoid: pg_sys::Oid,
     ) -> Option<AnyArray> {
         if is_null {
             None

pgx/src/datum/anyelement.rs

@@ -34,11 +34,22 @@ impl AnyElement {
 }
 
 impl FromDatum for AnyElement {
+    const GET_TYPOID: bool = true;
+
     #[inline]
     unsafe fn from_datum(
         datum: pg_sys::Datum,
         is_null: bool,
         typoid: pg_sys::Oid,
+    ) -> Option<AnyElement> {
+        FromDatum::from_polymorphic_datum(datum, is_null, typoid)
+    }
+
+    #[inline]
+    unsafe fn from_polymorphic_datum(
+        datum: pg_sys::Datum,
+        is_null: bool,
+        typoid: pg_sys::Oid,
     ) -> Option<AnyElement> {
         if is_null {
             None

pgx/src/datum/from.rs

@@ -39,6 +39,9 @@ pub enum TryFromDatumError {
 /// If implementing this, also implement `IntoDatum` for the reverse
 /// conversion.
 pub trait FromDatum {
+    /// Should a type OID be fetched when calling `from_datum`?
+    const GET_TYPOID: bool = false;
+
     /// ## Safety
     ///
     /// This method is inherently unsafe as the `datum` argument can represent an arbitrary
@@ -54,6 +57,23 @@ pub trait FromDatum {
     where
         Self: Sized;
 
+    /// Like `from_datum` for instantiating polymorphic types
+    /// which require preserving the dynamic type metadata.
+    ///
+    /// ## Safety
+    ///
+    /// Same caveats as `FromDatum::from_datum(...)`.
+    unsafe fn from_polymorphic_datum(
+        datum: pg_sys::Datum,
+        is_null: bool,
+        typoid: pg_sys::Oid,
+    ) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        FromDatum::from_datum(datum, is_null, typoid)
+    }
+
     /// Default implementation switched to the specified memory context and then simply calls
     /// `FromDatum::from_datum(...)` from within that context.
     ///

pgx/src/fcinfo.rs

@@ -95,8 +95,11 @@ mod pg_10_11 {
         let datum = unsafe { fcinfo.as_ref() }.unwrap().arg[num];
         let isnull = pg_arg_is_null(fcinfo, num);
         unsafe {
-            let typid = pg_sys::InvalidOid;
-            T::from_datum(datum, isnull, typid)
+            if T::GET_TYPOID {
+                T::from_polymorphic_datum(datum, isnull, super::pg_getarg_type(fcinfo, num))
+            } else {
+                T::from_datum(datum, isnull, pg_sys::InvalidOid)
+            }
         }
     }
 
@@ -134,8 +137,15 @@ mod pg_12_13_14 {
     pub fn pg_getarg<T: FromDatum>(fcinfo: pg_sys::FunctionCallInfo, num: usize) -> Option<T> {
         let datum = get_nullable_datum(fcinfo, num);
         unsafe {
-            let typid = pg_sys::InvalidOid;
-            T::from_datum(datum.value, datum.isnull, typid)
+            if T::GET_TYPOID {
+                T::from_polymorphic_datum(
+                    datum.value,
+                    datum.isnull,
+                    super::pg_getarg_type(fcinfo, num),
+                )
+            } else {
+                T::from_datum(datum.value, datum.isnull, pg_sys::InvalidOid)
+            }
         }
     }