spi: Add support for specifying mutating mode to PreparedStatement

pgrx-tests/src/tests/spi_tests.rs

@@ -17,6 +17,7 @@ mod tests {
 
     use pgrx::prelude::*;
     use pgrx::spi;
+    use pgrx::spi::Query;
 
     #[pg_test(error = "syntax error at or near \"THIS\"")]
     fn test_spi_failure() -> Result<(), spi::Error> {
@@ -419,6 +420,26 @@ mod tests {
         })
     }
 
+    #[pg_test(error = "CREATE TABLE is not allowed in a non-volatile function")]
+    fn test_execute_prepared_statement_in_readonly() -> Result<(), spi::Error> {
+        Spi::connect(|client| {
+            let stmt = client.prepare("CREATE TABLE a ()", None)?;
+            // This is supposed to run in read-only
+            stmt.execute(&client, Some(1), None)?;
+            Ok(())
+        })
+    }
+
+    #[pg_test]
+    fn test_execute_prepared_statement_in_readwrite() -> Result<(), spi::Error> {
+        Spi::connect(|client| {
+            let stmt = client.prepare_mut("CREATE TABLE a ()", None)?;
+            // This is supposed to run in read-write
+            stmt.execute(&client, Some(1), None)?;
+            Ok(())
+        })
+    }
+
     #[pg_test]
     fn test_spi_select_sees_update() -> spi::Result<()> {
         let with_select = Spi::connect(|mut client| {

pgrx/src/spi/client.rs

@@ -20,6 +20,28 @@ impl<'conn> SpiClient<'conn> {
         &self,
         query: &str,
         args: Option<Vec<PgOid>>,
+    ) -> SpiResult<PreparedStatement<'conn>> {
+        self.make_prepare_statement(query, args, false)
+    }
+
+    /// Prepares a mutating statement that is valid for the lifetime of the client
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the supplied `query` string contained a NULL byte
+    pub fn prepare_mut(
+        &self,
+        query: &str,
+        args: Option<Vec<PgOid>>,
+    ) -> SpiResult<PreparedStatement<'conn>> {
+        self.make_prepare_statement(query, args, true)
+    }
+
+    fn make_prepare_statement(
+        &self,
+        query: &str,
+        args: Option<Vec<PgOid>>,
+        mutating: bool,
     ) -> SpiResult<PreparedStatement<'conn>> {
         let src = CString::new(query).expect("query contained a null byte");
         let args = args.unwrap_or_default();
@@ -43,6 +65,7 @@ impl<'conn> SpiClient<'conn> {
                 .unwrap()
             })?,
             __marker: PhantomData,
+            mutating,
         })
     }
 

pgrx/src/spi/query.rs

@@ -148,6 +148,7 @@ impl<'conn> Query<'conn> for &str {
 pub struct PreparedStatement<'conn> {
     pub(super) plan: NonNull<pg_sys::_SPI_plan>,
     pub(super) __marker: PhantomData<&'conn ()>,
+    pub(super) mutating: bool,
 }
 
 /// Static lifetime-bound prepared statement
@@ -214,7 +215,11 @@ impl<'conn> PreparedStatement<'conn> {
         unsafe {
             pg_sys::SPI_keepplan(self.plan.as_ptr());
         }
-        OwnedPreparedStatement(PreparedStatement { __marker: PhantomData, plan: self.plan })
+        OwnedPreparedStatement(PreparedStatement {
+            __marker: PhantomData,
+            plan: self.plan,
+            mutating: self.mutating,
+        })
     }
 }
 
@@ -248,7 +253,7 @@ impl<'conn: 'stmt, 'stmt> Query<'conn> for &'stmt PreparedStatement<'conn> {
                 self.plan.as_ptr(),
                 datums.as_mut_ptr(),
                 nulls.as_mut_ptr(),
-                Spi::is_xact_still_immutable(),
+                !self.mutating && Spi::is_xact_still_immutable(),
                 limit.unwrap_or(0),
             )
         };
@@ -269,7 +274,7 @@ impl<'conn: 'stmt, 'stmt> Query<'conn> for &'stmt PreparedStatement<'conn> {
                 self.plan.as_ptr(),
                 datums.as_mut_ptr(),
                 nulls.as_ptr(),
-                Spi::is_xact_still_immutable(),
+                !self.mutating && Spi::is_xact_still_immutable(),
             ))
         };
         SpiCursor { ptr, __marker: PhantomData }