From 3e9c1438a01b7eeea3d3402da0372b3624163188 Mon Sep 17 00:00:00 2001
From: Lars Kanis <lars.kanis@sincnovation.com>
Date: Wed, 6 Mar 2024 11:08:08 +0100
Subject: [PATCH] Add PG::RollbackTransaction as an option to exit
 conn.transaction

---
 lib/pg/connection.rb       |  5 +++++
 lib/pg/exceptions.rb       |  6 ++++++
 spec/pg/connection_spec.rb | 18 ++++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/lib/pg/connection.rb b/lib/pg/connection.rb
index 47647697c..daf3991c7 100644
--- a/lib/pg/connection.rb
+++ b/lib/pg/connection.rb
@@ -306,6 +306,11 @@ def transaction
 		rollback = false
 		exec "BEGIN"
 		yield(self)
+	rescue PG::RollbackTransaction
+		rollback = true
+		cancel if transaction_status == PG::PQTRANS_ACTIVE
+		block
+		exec "ROLLBACK"
 	rescue Exception
 		rollback = true
 		cancel if transaction_status == PG::PQTRANS_ACTIVE
diff --git a/lib/pg/exceptions.rb b/lib/pg/exceptions.rb
index c50a9aa52..7152a268c 100644
--- a/lib/pg/exceptions.rb
+++ b/lib/pg/exceptions.rb
@@ -21,5 +21,11 @@ class LostCopyState < PG::Error
 	class NotInBlockingMode < PG::Error
 	end
 
+	# PG::Connection#transaction uses this exception to distinguish a deliberate rollback from other exceptional situations.
+	# Normally, raising an exception will cause the .transaction method to rollback the database transaction and pass on the exception.
+	# But if you raise an PG::RollbackTransaction exception, then the database transaction will be rolled back, without passing on the exception.
+	class RollbackTransaction < StandardError
+	end
+
 end # module PG
 
diff --git a/spec/pg/connection_spec.rb b/spec/pg/connection_spec.rb
index caac2e50a..1ea6d4be1 100644
--- a/spec/pg/connection_spec.rb
+++ b/spec/pg/connection_spec.rb
@@ -920,6 +920,24 @@
 			end
 		end
 
+		it "rolls back a transaction if a PG::RollbackTransaction exception is raised" do
+			# abort the per-example transaction so we can test our own
+			@conn.exec( 'ROLLBACK' )
+			@conn.exec( "CREATE TABLE pie ( flavor TEXT )" )
+
+			begin
+				@conn.transaction do
+					@conn.exec( "INSERT INTO pie VALUES ('rhubarb'), ('cherry'), ('schizophrenia')" )
+					raise PG::RollbackTransaction
+				end
+
+				res = @conn.exec( "SELECT * FROM pie" )
+				expect( res.ntuples ).to eq( 0 )
+			ensure
+				@conn.exec( "DROP TABLE pie" )
+			end
+		end
+
 		it "commits even if the block includes an early break/return" do
 			# abort the per-example transaction so we can test our own
 			@conn.exec( 'ROLLBACK' )