Improve copy_data error handling

lib/pg.rb

@@ -50,12 +50,6 @@ module PG
 		end
 	end
 
-
-	class NotAllCopyDataRetrieved < PG::Error
-	end
-	class NotInBlockingMode < PG::Error
-	end
-
 	# Get the PG library version.
 	#
 	# +include_buildnum+ is no longer used and any value passed will be ignored.

lib/pg/connection.rb

@@ -196,11 +196,19 @@ def copy_data( sql, coder=nil )
 				yield res
 			rescue Exception => err
 				errmsg = "%s while copy data: %s" % [ err.class.name, err.message ]
-				put_copy_end( errmsg )
-				get_result
-				raise
+				begin
+					put_copy_end( errmsg )
+				rescue PG::Error
+					# Ignore error in cleanup to avoid losing original exception
+				end
+				discard_results
+				raise err
 			else
-				put_copy_end
+				begin
+					put_copy_end
+				rescue PG::Error => err
+					raise PG::LostCopyState.new("#{err} (probably by executing another SQL query while running a COPY command)", connection: self)
+				end
 				get_last_result
 			ensure
 				self.encoder_for_put_copy_data = old_coder if coder
@@ -213,24 +221,17 @@ def copy_data( sql, coder=nil )
 					self.decoder_for_get_copy_data = coder
 				end
 				yield res
-			rescue Exception => err
+			rescue Exception
 				cancel
-				begin
-					while get_copy_data
-					end
-				rescue PG::Error
-					# Ignore error in cleanup to avoid losing original exception
-				end
-				while get_result
-				end
-				raise err
+				discard_results
+				raise
 			else
 				res = get_last_result
-				if !res || res.result_status != PGRES_COMMAND_OK
-					while get_copy_data
-					end
-					while get_result
-					end
+				if !res
+					discard_results
+					raise PG::LostCopyState.new("Lost COPY state (probably by executing another SQL query while running a COPY command)", connection: self)
+				elsif res.result_status != PGRES_COMMAND_OK
+					discard_results
 					raise PG::NotAllCopyDataRetrieved.new("Not all COPY data retrieved", connection: self)
 				end
 				res

lib/pg/exceptions.rb

@@ -14,5 +14,12 @@ def initialize(msg=nil, connection: nil, result: nil)
 		end
 	end
 
+	class NotAllCopyDataRetrieved < PG::Error
+	end
+	class LostCopyState < PG::Error
+	end
+	class NotInBlockingMode < PG::Error
+	end
+
 end # module PG
 

spec/pg/connection_spec.rb

@@ -1166,6 +1166,7 @@
 		@conn.sync_put_copy_end
 		res = @conn.get_last_result
 		expect( res.result_status ).to eq( PG::PGRES_COMMAND_OK )
+		@conn.exec( "DROP TABLE IF EXISTS copytable2" )
 	end
 
 	describe "#copy_data" do
@@ -1240,6 +1241,7 @@
 
 			res = @conn.exec( "SELECT * FROM copytable ORDER BY col1" )
 			expect( res.values ).to eq( [["1"], ["2"]] )
+			@conn.exec( "DROP TABLE IF EXISTS copytable" )
 		end
 
 		it "can process #copy_data input queries with lots of data" do
@@ -1256,6 +1258,7 @@
 			expect( res.values ).to eq( [["1000"]] )
 			res = @conn.exec( "SELECT * FROM copytable2 LIMIT 1" )
 			expect( res.values ).to eq( [[str.chomp]] )
+			@conn.exec( "DROP TABLE IF EXISTS copytable2" )
 		end
 
 		it "can handle client errors in #copy_data for input" do
@@ -1270,6 +1273,7 @@
 			end
 
 			expect( @conn ).to still_be_usable
+			@conn.exec( "DROP TABLE IF EXISTS copytable" )
 		end
 
 		it "can handle server errors in #copy_data for input" do
@@ -1283,30 +1287,43 @@
 				}.to raise_error(PG::Error, /invalid input syntax for .*integer/){|err| expect(err).to have_attributes(connection: @conn) }
 			end
 			expect( @conn ).to still_be_usable
+			@conn.exec( "DROP TABLE IF EXISTS copytable" )
 		end
 
-		it "gracefully handle SQL statements while in #copy_data for input" do
+		it "doesn't lose client error when #copy_data can not be finished" do
 			@conn.exec "ROLLBACK"
 			@conn.transaction do
 				@conn.exec( "CREATE TEMP TABLE copytable (col1 INT)" )
 				expect {
 					@conn.copy_data( "COPY copytable FROM STDOUT" ) do |res|
-						@conn.exec "SELECT 1"
+						@conn.discard_results # end copy state so that put_copy_end fails in copy_data
+						raise "boom"
 					end
-				}.to raise_error(PG::Error, /no COPY in progress/){|err| expect(err).to have_attributes(connection: @conn) }
+				}.to raise_error(RuntimeError, "boom")
 			end
 			expect( @conn ).to still_be_usable
+			@conn.exec( "DROP TABLE IF EXISTS copytable" )
+		end
+
+		it "gracefully handle SQL statements while in #copy_data for input" do
+			@conn.exec "ROLLBACK"
+			@conn.exec( "CREATE TEMP TABLE copytable (col1 INT)" )
+			expect {
+				@conn.copy_data( "COPY copytable FROM STDOUT" ) do |res|
+					@conn.exec "SELECT 1"
+				end
+			}.to raise_error(PG::LostCopyState, /another SQL query/){|err| expect(err).to have_attributes(connection: @conn) }
+			expect( @conn ).to still_be_usable
+			@conn.exec( "DROP TABLE copytable" )
 		end
 
 		it "gracefully handle SQL statements while in #copy_data for output" do
 			@conn.exec "ROLLBACK"
-			@conn.transaction do
-				expect {
-					@conn.copy_data( "COPY (VALUES(1), (2)) TO STDOUT" ) do |res|
-						@conn.exec "SELECT 3"
-					end
-				}.to raise_error(PG::Error, /no COPY in progress/){|err| expect(err).to have_attributes(connection: @conn) }
-			end
+			expect {
+				@conn.copy_data( "COPY (VALUES(1), (2)) TO STDOUT" ) do |res|
+					@conn.exec "SELECT 3"
+				end
+			}.to raise_error(PG::LostCopyState, /another SQL query/){|err| expect(err).to have_attributes(connection: @conn) }
 			expect( @conn ).to still_be_usable
 		end