Missing ArgAbi and BoxRet trait implementations for datums (#1757)

After updating to version `0.12.0-beta.0` I can no longer use `JsonB` as
a `pg_extern` function argument, and I can no longer use `JsonB`,
`AnyArray`, or `AnyElement` as state type in `pg_aggregate`.
Unfortunately, pgrx lacked any tests for these types, so they were not
checked for compilation when adding ArgAbi and RetAbi.

Author: danbluhmhansen

pgrx-tests/src/tests/aggregate_tests.rs

@@ -143,6 +143,74 @@ impl Aggregate for DemoCustomState {
     }
 }
 
+struct FirstJson;
+
+#[pg_aggregate]
+impl Aggregate for FirstJson {
+    type State = pgrx::Json;
+    type Args = pgrx::name!(value, pgrx::Json);
+
+    #[pgrx(parallel_safe, immutable, strict)]
+    fn state(
+        current: Self::State,
+        _arg: Self::Args,
+        _fcinfo: pg_sys::FunctionCallInfo,
+    ) -> Self::State {
+        current
+    }
+}
+
+struct FirstJsonB;
+
+#[pg_aggregate]
+impl Aggregate for FirstJsonB {
+    type State = pgrx::JsonB;
+    type Args = pgrx::name!(value, pgrx::JsonB);
+
+    #[pgrx(parallel_safe, immutable, strict)]
+    fn state(
+        current: Self::State,
+        _arg: Self::Args,
+        _fcinfo: pg_sys::FunctionCallInfo,
+    ) -> Self::State {
+        current
+    }
+}
+
+struct FirstAnyArray;
+
+#[pg_aggregate]
+impl Aggregate for FirstAnyArray {
+    type State = pgrx::AnyArray;
+    type Args = pgrx::name!(value, pgrx::AnyArray);
+
+    #[pgrx(parallel_safe, immutable, strict)]
+    fn state(
+        current: Self::State,
+        _arg: Self::Args,
+        _fcinfo: pg_sys::FunctionCallInfo,
+    ) -> Self::State {
+        current
+    }
+}
+
+struct FirstAnyElement;
+
+#[pg_aggregate]
+impl Aggregate for FirstAnyElement {
+    type State = pgrx::AnyElement;
+    type Args = pgrx::name!(value, pgrx::AnyElement);
+
+    #[pgrx(parallel_safe, immutable, strict)]
+    fn state(
+        current: Self::State,
+        _arg: Self::Args,
+        _fcinfo: pg_sys::FunctionCallInfo,
+    ) -> Self::State {
+        current
+    }
+}
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pgrx::pg_schema]
 mod tests {
@@ -190,4 +258,60 @@ mod tests {
         );
         assert_eq!(retval, Ok(Some(4)));
     }
+
+    #[pg_test]
+    fn aggregate_first_json() -> Result<(), pgrx::spi::Error> {
+        let retval = Spi::get_one_with_args::<pgrx::Json>(
+            "SELECT FirstJson(value) FROM UNNEST(ARRAY [$1, $2]) as value;",
+            vec![
+                (
+                    PgBuiltInOids::JSONOID.oid(),
+                    pgrx::Json(serde_json::json!({ "foo": "one" })).into_datum(),
+                ),
+                (
+                    PgBuiltInOids::JSONOID.oid(),
+                    pgrx::Json(serde_json::json!({ "foo": "two" })).into_datum(),
+                ),
+            ],
+        )?
+        .map(|json| json.0);
+
+        assert_eq!(retval, Some(serde_json::json!({ "foo": "one" })));
+
+        Ok(())
+    }
+
+    #[pg_test]
+    fn aggregate_first_jsonb() -> Result<(), pgrx::spi::Error> {
+        let retval = Spi::get_one_with_args::<pgrx::JsonB>(
+            "SELECT FirstJsonB(value) FROM UNNEST(ARRAY [$1, $2]) as value;",
+            vec![
+                (
+                    PgBuiltInOids::JSONBOID.oid(),
+                    pgrx::JsonB(serde_json::json!({ "foo": "one" })).into_datum(),
+                ),
+                (
+                    PgBuiltInOids::JSONBOID.oid(),
+                    pgrx::JsonB(serde_json::json!({ "foo": "two" })).into_datum(),
+                ),
+            ],
+        )?
+        .map(|json| json.0);
+
+        assert_eq!(retval, Some(serde_json::json!({ "foo": "one" })));
+
+        Ok(())
+    }
+
+    #[pg_test]
+    fn aggregate_first_anyelement() -> Result<(), pgrx::spi::Error> {
+        let retval = Spi::get_one::<pgrx::AnyElement>(
+            "SELECT FirstAnyElement(value) FROM UNNEST(ARRAY [1, 2]) as value;",
+        )?
+        .map(|element| element.into_datum());
+
+        assert_eq!(retval, Some(1.into_datum()));
+
+        Ok(())
+    }
 }

pgrx-tests/src/tests/anyelement_tests.rs

@@ -0,0 +1,28 @@
+use pgrx::{prelude::*, AnyElement};
+
+#[pg_extern]
+fn anyelement_arg(element: AnyElement) -> AnyElement {
+    element
+}
+
+#[cfg(any(test, feature = "pg_test"))]
+#[pgrx::pg_schema]
+mod tests {
+    #[allow(unused_imports)]
+    use crate as pgrx_tests;
+
+    use pgrx::{prelude::*, AnyElement};
+
+    #[pg_test]
+    fn test_anyelement_arg() -> Result<(), pgrx::spi::Error> {
+        let element = Spi::get_one_with_args::<AnyElement>(
+            "SELECT anyelement_arg($1);",
+            vec![(PgBuiltInOids::ANYELEMENTOID.oid(), 123.into_datum())],
+        )?
+        .map(|e| e.datum());
+
+        assert_eq!(element, 123.into_datum());
+
+        Ok(())
+    }
+}

pgrx-tests/src/tests/anynumeric_tests.rs

@@ -0,0 +1,28 @@
+use pgrx::{prelude::*, AnyNumeric};
+
+#[pg_extern]
+fn anynumeric_arg(numeric: AnyNumeric) -> AnyNumeric {
+    numeric
+}
+
+#[cfg(any(test, feature = "pg_test"))]
+#[pgrx::pg_schema]
+mod tests {
+    #[allow(unused_imports)]
+    use crate as pgrx_tests;
+
+    use pgrx::{prelude::*, AnyNumeric};
+
+    #[pg_test]
+    fn test_anynumeric_arg() -> Result<(), pgrx::spi::Error> {
+        let numeric = Spi::get_one_with_args::<AnyNumeric>(
+            "SELECT anynumeric_arg($1);",
+            vec![(PgBuiltInOids::INT4OID.oid(), 123.into_datum())],
+        )?
+        .map(|n| n.normalize().to_string());
+
+        assert_eq!(numeric, Some("123".to_string()));
+
+        Ok(())
+    }
+}

pgrx-tests/src/tests/json_tests.rs

@@ -7,6 +7,19 @@
 //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::prelude::*;
+use pgrx::{Json, JsonB};
+
+#[pg_extern]
+fn json_arg(json: Json) -> Json {
+    json
+}
+
+#[pg_extern]
+fn jsonb_arg(json: JsonB) -> JsonB {
+    json
+}
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pgrx::pg_schema]
 mod tests {
@@ -61,4 +74,36 @@ mod tests {
         assert_eq!(user.last_name, "McBlahFace");
         Ok(())
     }
+
+    #[pg_test]
+    fn test_json_arg() -> Result<(), pgrx::spi::Error> {
+        let json = Spi::get_one_with_args::<Json>(
+            "SELECT json_arg($1);",
+            vec![(
+                PgBuiltInOids::JSONOID.oid(),
+                Json(serde_json::json!({ "foo": "bar" })).into_datum(),
+            )],
+        )?
+        .expect("json was null");
+
+        assert_eq!(json.0, serde_json::json!({ "foo": "bar" }));
+
+        Ok(())
+    }
+
+    #[pg_test]
+    fn test_jsonb_arg() -> Result<(), pgrx::spi::Error> {
+        let json = Spi::get_one_with_args::<JsonB>(
+            "SELECT jsonb_arg($1);",
+            vec![(
+                PgBuiltInOids::JSONBOID.oid(),
+                JsonB(serde_json::json!({ "foo": "bar" })).into_datum(),
+            )],
+        )?
+        .expect("json was null");
+
+        assert_eq!(json.0, serde_json::json!({ "foo": "bar" }));
+
+        Ok(())
+    }
 }

pgrx-tests/src/tests/mod.rs

@@ -9,6 +9,8 @@
 //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
 mod aggregate_tests;
 mod anyarray_tests;
+mod anyelement_tests;
+mod anynumeric_tests;
 mod array_tests;
 mod attributes_tests;
 mod bgworker_tests;

pgrx/src/callconv.rs

@@ -16,7 +16,7 @@ use crate::heap_tuple::PgHeapTuple;
 use crate::nullable::Nullable;
 use crate::{
     pg_return_null, pg_sys, AnyArray, AnyElement, AnyNumeric, Date, FromDatum, Inet, Internal,
-    Interval, IntoDatum, Json, PgBox, PgMemoryContexts, PgVarlena, Time, TimeWithTimeZone,
+    Interval, IntoDatum, Json, JsonB, PgBox, PgMemoryContexts, PgVarlena, Time, TimeWithTimeZone,
     Timestamp, TimestampWithTimeZone, UnboxDatum, Uuid,
 };
 use core::marker::PhantomData;
@@ -258,7 +258,7 @@ macro_rules! argue_from_datum {
 argue_from_datum! { 'fcx; i8, i16, i32, i64, f32, f64, bool, char, String, Vec<u8> }
 argue_from_datum! { 'fcx; Date, Interval, Time, TimeWithTimeZone, Timestamp, TimestampWithTimeZone }
 argue_from_datum! { 'fcx; AnyArray, AnyElement, AnyNumeric }
-argue_from_datum! { 'fcx; Inet, Internal, Json, Uuid }
+argue_from_datum! { 'fcx; Inet, Internal, Json, JsonB, Uuid }
 argue_from_datum! { 'fcx; pg_sys::Oid, pg_sys::Point, pg_sys::BOX  }
 argue_from_datum! { 'fcx; &'fcx str, &'fcx CStr, &'fcx [u8] }
 
@@ -505,7 +505,7 @@ macro_rules! impl_repackage_into_datum {
 
 impl_repackage_into_datum! {
     String, CString, Vec<u8>, char,
-    Json, Inet, Uuid, AnyNumeric, Internal,
+    Json, JsonB, Inet, Uuid, AnyNumeric, AnyArray, AnyElement, Internal,
     Date, Interval, Time, TimeWithTimeZone, Timestamp, TimestampWithTimeZone,
     pg_sys::Oid, pg_sys::BOX, pg_sys::Point
 }