From 6abbdffe8d8a7dd1900e021b0e419c16ab48893e Mon Sep 17 00:00:00 2001
From: "Eric B. Ridge" <eebbrr@gmail.com>
Date: Tue, 27 Dec 2022 11:05:24 -0500
Subject: [PATCH] fix issue #983

A readonly SPI session can now see the database modifications of a previous non-readonly SPI session.  This necessitates copying the active snapshot and incrementing its command id when a non-readonly SPI session is done.
---
 pgx-tests/src/tests/spi_tests.rs | 42 ++++++++++++++++++++++++++++++++
 pgx/src/spi.rs                   | 14 +++++++++++
 2 files changed, 56 insertions(+)

diff --git a/pgx-tests/src/tests/spi_tests.rs b/pgx-tests/src/tests/spi_tests.rs
index b8a81beba..24837325e 100644
--- a/pgx-tests/src/tests/spi_tests.rs
+++ b/pgx-tests/src/tests/spi_tests.rs
@@ -417,4 +417,46 @@ mod tests {
             Ok(())
         })
     }
+
+    #[pg_test]
+    fn test_spi_select_sees_update() -> spi::Result<()> {
+        let with_select = Spi::connect(|mut client| {
+            client.update("CREATE TABLE asd(id int)", None, None)?;
+            client.update("INSERT INTO asd(id) VALUES (1)", None, None)?;
+            client.select("SELECT COUNT(*) FROM asd", None, None)?.first().get_one::<i64>()
+        })?;
+        let with_get_one = Spi::get_one::<i64>("SELECT COUNT(*) FROM asd")?;
+
+        assert_eq!(with_select, with_get_one);
+        Ok(())
+    }
+
+    #[pg_test]
+    fn test_spi_select_sees_run() -> spi::Result<()> {
+        Spi::run("CREATE TABLE asd(id int)")?;
+        Spi::run("INSERT INTO asd(id) VALUES (1)")?;
+        let with_select = Spi::connect(|client| {
+            client.select("SELECT COUNT(*) FROM asd", None, None)?.first().get_one::<i64>()
+        })?;
+        let with_get_one = Spi::get_one::<i64>("SELECT COUNT(*) FROM asd")?;
+
+        assert_eq!(with_select, with_get_one);
+        Ok(())
+    }
+
+    #[pg_test]
+    fn test_spi_select_sees_update_in_other_session() -> spi::Result<()> {
+        Spi::connect::<spi::Result<()>, _>(|mut client| {
+            client.update("CREATE TABLE asd(id int)", None, None)?;
+            client.update("INSERT INTO asd(id) VALUES (1)", None, None)?;
+            Ok(())
+        })?;
+        let with_select = Spi::connect(|client| {
+            client.select("SELECT COUNT(*) FROM asd", None, None)?.first().get_one::<i64>()
+        })?;
+        let with_get_one = Spi::get_one::<i64>("SELECT COUNT(*) FROM asd")?;
+
+        assert_eq!(with_select, with_get_one);
+        Ok(())
+    }
 }
diff --git a/pgx/src/spi.rs b/pgx/src/spi.rs
index 6b81f32f0..4612c6518 100644
--- a/pgx/src/spi.rs
+++ b/pgx/src/spi.rs
@@ -161,6 +161,20 @@ pub struct SpiClient<'conn> {
     __marker: PhantomData<&'conn SpiConnection>,
 }
 
+impl<'conn> Drop for SpiClient<'conn> {
+    fn drop(&mut self) {
+        if self.readonly == false {
+            // Assume the user actually modified the database during this SpiClient's lifetime
+            // and push a new snapshot with an updated CommandId.  This ensures that subsequent
+            // `readonly == true` SpiClients will see the changes made by this one
+            unsafe {
+                pg_sys::PushCopiedSnapshot(pg_sys::GetActiveSnapshot());
+                pg_sys::UpdateActiveSnapshotCommandId();
+            }
+        }
+    }
+}
+
 /// a struct to manage our SPI connection lifetime
 struct SpiConnection(PhantomData<*mut ()>);