From eda75496a08dbf5b1077694a9d1d3b5100d6aead Mon Sep 17 00:00:00 2001
From: Yoh Deadfall <yoh.deadfall@hotmail.com>
Date: Tue, 26 Nov 2024 02:19:06 +0300
Subject: [PATCH] Added `oids_of!` macro (#1879)

That's an improvement for SPI making it less cumbersome. Just to feel
the difference look at the tests. Who would really want to write
something like `PgBuiltInOids::INT4OID.oid()` instead of just telling a
Rust type?
---
 pgrx-tests/src/tests/spi_tests.rs | 19 ++++++-------------
 pgrx/src/datum/mod.rs             | 27 +++++++++++++++++++++++++++
 pgrx/src/prelude.rs               | 13 ++++++++-----
 3 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/pgrx-tests/src/tests/spi_tests.rs b/pgrx-tests/src/tests/spi_tests.rs
index 05b16ee616..2df12a2fe8 100644
--- a/pgrx-tests/src/tests/spi_tests.rs
+++ b/pgrx-tests/src/tests/spi_tests.rs
@@ -17,8 +17,7 @@ mod tests {
     use std::error::Error;
 
     use pgrx::prelude::*;
-    use pgrx::spi;
-    use pgrx::spi::Query;
+    use pgrx::spi::{self, Query};
 
     #[pg_test(error = "syntax error at or near \"THIS\"")]
     fn test_spi_failure() -> Result<(), spi::Error> {
@@ -253,10 +252,8 @@ mod tests {
                 None,
                 &[],
             )?;
-            let prepared = client.prepare(
-                "SELECT * FROM tests.cursor_table WHERE id = $1",
-                &[PgBuiltInOids::INT4OID.oid()],
-            )?;
+            let prepared =
+                client.prepare("SELECT * FROM tests.cursor_table WHERE id = $1", &oids_of![i32])?;
             client.open_cursor(&prepared, args);
             unreachable!();
         })
@@ -376,8 +373,7 @@ mod tests {
     #[pg_test]
     fn test_prepared_statement() -> Result<(), spi::Error> {
         let rc = Spi::connect(|client| {
-            let prepared =
-                client.prepare("SELECT $1", &[PgOid::BuiltIn(PgBuiltInOids::INT4OID)])?;
+            let prepared = client.prepare("SELECT $1", &oids_of![i32])?;
             client.select(&prepared, None, &[42.into()])?.first().get::<i32>(1)
         })?;
 
@@ -388,8 +384,7 @@ mod tests {
     #[pg_test]
     fn test_prepared_statement_argument_mismatch() {
         let err = Spi::connect(|client| {
-            let prepared =
-                client.prepare("SELECT $1", &[PgOid::BuiltIn(PgBuiltInOids::INT4OID)])?;
+            let prepared = client.prepare("SELECT $1", &oids_of![i32])?;
             client.select(&prepared, None, &[]).map(|_| ())
         })
         .unwrap_err();
@@ -403,9 +398,7 @@ mod tests {
     #[pg_test]
     fn test_owned_prepared_statement() -> Result<(), spi::Error> {
         let prepared = Spi::connect(|client| {
-            Ok::<_, spi::Error>(
-                client.prepare("SELECT $1", &[PgOid::BuiltIn(PgBuiltInOids::INT4OID)])?.keep(),
-            )
+            Ok::<_, spi::Error>(client.prepare("SELECT $1", &oids_of![i32])?.keep())
         })?;
         let rc = Spi::connect(|client| {
             client.select(&prepared, None, &[42.into()])?.first().get::<i32>(1)
diff --git a/pgrx/src/datum/mod.rs b/pgrx/src/datum/mod.rs
index cc5a7ba3eb..4c1b854ed0 100644
--- a/pgrx/src/datum/mod.rs
+++ b/pgrx/src/datum/mod.rs
@@ -204,3 +204,30 @@ impl<'src, T: IntoDatum> From<T> for DatumWithOid<'src> {
 /// A tagging trait to indicate a user type is also meant to be used by Postgres
 /// Implemented automatically by `#[derive(PostgresType)]`
 pub trait PostgresType {}
+
+/// Creates an array of [`pg_sys::Oid`] with the OID of each provided type
+///
+/// # Examples
+///
+/// ```
+/// use pgrx::{oids_of, datum::IntoDatum};
+///
+/// let oids = oids_of![i32, f64];
+/// assert_eq!(oids[0], i32::type_oid().into());
+/// assert_eq!(oids[1], f64::type_oid().into());
+///
+/// // the usual conversions or coercions are available
+/// let oid_vec = oids_of![i8, i16].to_vec();
+/// let no_oid = &oids_of![];
+/// assert_eq!(no_oid.len(), 0);
+/// ```
+#[macro_export]
+macro_rules! oids_of {
+    () =>(
+        // avoid coercions to an ambiguously-typed array or slice
+        [$crate::pg_sys::PgOid::Invalid; 0]
+    );
+    ($($t:path),+ $(,)?) => (
+        [$($crate::pg_sys::PgOid::from(<$t>::type_oid())),*]
+    );
+}
diff --git a/pgrx/src/prelude.rs b/pgrx/src/prelude.rs
index 44e2a67878..fa0fec391c 100644
--- a/pgrx/src/prelude.rs
+++ b/pgrx/src/prelude.rs
@@ -29,12 +29,15 @@ pub use crate::pgbox::{AllocatedByPostgres, AllocatedByRust, PgBox, WhoAllocated
 // These could be factored into a temporal type module that could be easily imported for code which works with them.
 // However, reexporting them seems fine for now.
 
-pub use crate::datum::{
-    datetime_support::*, AnyNumeric, Array, ArraySliceError, Date, FromDatum, Interval, IntoDatum,
-    Numeric, PgVarlena, PostgresType, Range, RangeBound, RangeSubType, Time, TimeWithTimeZone,
-    Timestamp, TimestampWithTimeZone, VariadicArray,
-};
 pub use crate::inoutfuncs::{InOutFuncs, PgVarlenaInOutFuncs};
+pub use crate::{
+    datum::{
+        datetime_support::*, AnyNumeric, Array, ArraySliceError, Date, FromDatum, Interval,
+        IntoDatum, Numeric, PgVarlena, PostgresType, Range, RangeBound, RangeSubType, Time,
+        TimeWithTimeZone, Timestamp, TimestampWithTimeZone, VariadicArray,
+    },
+    oids_of,
+};
 
 // Trigger support
 pub use crate::trigger_support::{