From 3e4e457fda09a5008e56a3c99381693bbc24de2b Mon Sep 17 00:00:00 2001
From: Lars Kanis <lars@greiz-reinsdorf.de>
Date: Wed, 18 May 2022 20:49:46 +0200
Subject: [PATCH] Add block-call semantics to PG::Connection.new

This was an oversight, when implementing the async interface of pg-1.3.0.

Fixes #454
---
 lib/pg.rb                  |  8 ++++----
 lib/pg/connection.rb       | 12 ++++++++++--
 spec/pg/connection_spec.rb | 11 ++++++-----
 spec/pg_spec.rb            | 10 ++++++++++
 4 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/lib/pg.rb b/lib/pg.rb
index adcc4094a..0e6a00c4b 100644
--- a/lib/pg.rb
+++ b/lib/pg.rb
@@ -59,14 +59,14 @@ class NotInBlockingMode < PG::Error
 	# Get the PG library version.
 	#
 	# +include_buildnum+ is no longer used and any value passed will be ignored.
-	def self::version_string( include_buildnum=nil )
-		return "%s %s" % [ self.name, VERSION ]
+	def self.version_string( include_buildnum=nil )
+		"%s %s" % [ self.name, VERSION ]
 	end
 
 
 	### Convenience alias for PG::Connection.new.
-	def self::connect( *args, **kwargs )
-		return PG::Connection.new( *args, **kwargs )
+	def self.connect( *args, **kwargs, &block )
+		Connection.new( *args, **kwargs, &block )
 	end
 
 
diff --git a/lib/pg/connection.rb b/lib/pg/connection.rb
index 8d97ef1e5..375058a7f 100644
--- a/lib/pg/connection.rb
+++ b/lib/pg/connection.rb
@@ -549,6 +549,7 @@ def encrypt_password( password, username, algorithm=nil )
 	def reset
 		reset_start
 		async_connect_or_reset(:reset_poll)
+		self
 	end
 	alias async_reset reset
 
@@ -646,8 +647,6 @@ def cancel
 		sync_setnonblocking(true)
 		self.flush_data = true
 		set_default_encoding
-
-		self
 	end
 
 	class << self
@@ -709,6 +708,15 @@ def new(*args, **kwargs)
 			raise(PG::ConnectionBad, conn.error_message) if conn.status == PG::CONNECTION_BAD
 
 			conn.send(:async_connect_or_reset, :connect_poll)
+
+			if block_given?
+				begin
+					yield conn
+				ensure
+					conn.finish
+				end
+			end
+			conn
 		end
 		alias async_connect new
 		alias connect new
diff --git a/spec/pg/connection_spec.rb b/spec/pg/connection_spec.rb
index d8e905dce..ddcdd2ae3 100644
--- a/spec/pg/connection_spec.rb
+++ b/spec/pg/connection_spec.rb
@@ -315,10 +315,10 @@
 			klass = Class.new(described_class) do
 				alias execute exec
 			end
-			conn = klass.send(meth, @conninfo)
-			expect( conn ).to be_a_kind_of( klass )
-			expect( conn.execute("SELECT 1") ).to be_a_kind_of( PG::Result )
-			conn.close
+			klass.send(meth, @conninfo) do |conn|
+				expect( conn ).to be_a_kind_of( klass )
+				expect( conn.execute("SELECT 1") ).to be_a_kind_of( PG::Result )
+			end
 		end
 	end
 
@@ -1436,7 +1436,8 @@ def interrupt_thread(exc=nil)
 		# Close the two pipe file descriptors, so that the file descriptor of
 		# newly established connection is probably distinct from the previous one.
 		ios.each(&:close)
-		conn.reset
+		res = conn.reset
+		expect( res ).to eq( conn )
 
 		# The new connection should work even when the file descriptor has changed.
 		expect( conn.exec("SELECT 1").values ).to eq([["1"]])
diff --git a/spec/pg_spec.rb b/spec/pg_spec.rb
index 8e3f25120..a95439a68 100644
--- a/spec/pg_spec.rb
+++ b/spec/pg_spec.rb
@@ -62,4 +62,14 @@
 		expect( c ).to be_a_kind_of( PG::Connection )
 		c.close
 	end
+
+	it "can #connect with block" do
+		res = PG.connect(@conninfo) do |c|
+			res = c.exec "SELECT 5"
+			expect( res.values ).to eq( [["5"]] )
+		end
+
+		expect( res ).to be_a_kind_of(PG::Connection)
+		expect( res ).to be_finished
+	end
 end