From 01d1d24d435866446e24d2fea8058bab9b9a5a52 Mon Sep 17 00:00:00 2001
From: Yurii Rashkovskii <yrashk@gmail.com>
Date: Wed, 14 Dec 2022 17:25:54 -0800
Subject: [PATCH] Problem: cursors may or may not be readonly-friendly

However, we always assume them to be non-readonly.

Solution: split the API to `open_cursor` and `open_cursor_mut`

This way we can ensure we're squeezing out performance in the right
place, when feasible.
---
 pgx-tests/src/tests/spi_tests.rs | 21 +++++++++++++++++++++
 pgx/src/spi.rs                   | 24 +++++++++++++++++++++++-
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/pgx-tests/src/tests/spi_tests.rs b/pgx-tests/src/tests/spi_tests.rs
index 7444478f0..95d6df4f5 100644
--- a/pgx-tests/src/tests/spi_tests.rs
+++ b/pgx-tests/src/tests/spi_tests.rs
@@ -209,6 +209,27 @@ mod tests {
         });
     }
 
+    #[pg_test]
+    fn test_cursor_mut() {
+        Spi::execute(|mut client| {
+            client.update("CREATE TABLE tests.cursor_table (id int)", None, None);
+
+            let mut portal = client.open_cursor_mut(
+                "INSERT INTO tests.cursor_table (id) \
+            SELECT i FROM generate_series(1, 10) AS t(i) RETURNING id",
+                None,
+            );
+
+            fn sum_all(table: pgx::SpiTupleTable) -> i32 {
+                table.map(|r| r.by_ordinal(1).unwrap().value::<i32>().unwrap()).sum()
+            }
+            assert_eq!(sum_all(portal.fetch(3)), 1 + 2 + 3);
+            assert_eq!(sum_all(portal.fetch(3)), 4 + 5 + 6);
+            assert_eq!(sum_all(portal.fetch(3)), 7 + 8 + 9);
+            assert_eq!(sum_all(portal.fetch(3)), 10);
+        });
+    }
+
     #[pg_test]
     fn test_cursor_by_name() {
         let cursor_name = Spi::connect(|mut client| {
diff --git a/pgx/src/spi.rs b/pgx/src/spi.rs
index 62de0c45d..0f1c0338e 100644
--- a/pgx/src/spi.rs
+++ b/pgx/src/spi.rs
@@ -412,6 +412,28 @@ impl<'a> SpiClient<'a> {
         &self,
         query: &str,
         args: Option<Vec<(PgOid, Option<pg_sys::Datum>)>>,
+    ) -> SpiCursor {
+        self.open_cursor_impl(query, args)
+    }
+
+    /// Set up a cursor that will execute the specified update (mutating) query
+    ///
+    /// Rows may be then fetched using [`SpiCursor::fetch`].
+    ///
+    /// See [`SpiCursor`] docs for usage details.
+    pub fn open_cursor_mut(
+        &mut self,
+        query: &str,
+        args: Option<Vec<(PgOid, Option<pg_sys::Datum>)>>,
+    ) -> SpiCursor {
+        self.readonly = false;
+        self.open_cursor_impl(query, args)
+    }
+
+    fn open_cursor_impl(
+        &self,
+        query: &str,
+        args: Option<Vec<(PgOid, Option<pg_sys::Datum>)>>,
     ) -> SpiCursor {
         let src = std::ffi::CString::new(query).expect("query contained a null byte");
         let args = args.unwrap_or_default();
@@ -447,7 +469,7 @@ impl<'a> SpiClient<'a> {
                 argtypes.as_mut_ptr(),
                 datums.as_mut_ptr(),
                 nulls.as_ptr(),
-                false,
+                self.readonly,
                 0,
             )
         })