From 1563e7313b50bd2c25b3540eb369098da192c074 Mon Sep 17 00:00:00 2001 From: Lars Kanis Date: Fri, 24 Feb 2023 11:29:20 +0100 Subject: [PATCH] Improve copy_data error handling Make sure an error in put_copy_end doesn't lose the original exception. Use a dedicated error class PG::LostCopyState for errors due to another query while COPYing and mention that it's probably due to another query. Previously the "no COPY in progress" PG::Error was less specific. Remove transactions around the "another query" specs, previously the put_copy_data direction had a double fault due to another query while transaction is in error state. Use discard_result which we have now and which does essentially what previous error handling did. Cleanup temporary tables after running. Move definition of error classes to exception.rb, where they better fit into. --- lib/pg.rb | 6 ------ lib/pg/connection.rb | 39 +++++++++++++++++++------------------- lib/pg/exceptions.rb | 7 +++++++ spec/pg/connection_spec.rb | 37 ++++++++++++++++++++++++++---------- 4 files changed, 54 insertions(+), 35 deletions(-) diff --git a/lib/pg.rb b/lib/pg.rb index a8375b98d..1cd16e7fa 100644 --- a/lib/pg.rb +++ b/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. diff --git a/lib/pg/connection.rb b/lib/pg/connection.rb index 39db3502a..67695fb4e 100644 --- a/lib/pg/connection.rb +++ b/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 diff --git a/lib/pg/exceptions.rb b/lib/pg/exceptions.rb index 8940ce7c3..c50a9aa52 100644 --- a/lib/pg/exceptions.rb +++ b/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 diff --git a/spec/pg/connection_spec.rb b/spec/pg/connection_spec.rb index b7e892f59..34d3c64fe 100644 --- a/spec/pg/connection_spec.rb +++ b/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