From 5325b82277a9531704d88bce0e89c58fa8bac35e Mon Sep 17 00:00:00 2001 From: Lars Kanis Date: Wed, 26 Apr 2023 09:19:35 +0200 Subject: [PATCH] Fix regression in copy_data regarding binary format ... when not using an coder. To keep compatibility the handling if header and trailer data is now only done when a coder is given. Fixes #526 --- lib/pg/connection.rb | 26 ++++++++++++++++++++--- spec/pg/connection_spec.rb | 43 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 64 insertions(+), 5 deletions(-) diff --git a/lib/pg/connection.rb b/lib/pg/connection.rb index ada7d6758..9e4f71581 100644 --- a/lib/pg/connection.rb +++ b/lib/pg/connection.rb @@ -166,6 +166,14 @@ def inspect # conn.put_copy_data ['more', 'data', 'to', 'copy'] # end # + # Also PG::BinaryEncoder::CopyRow can be used to send data in binary format to the server. + # In this case copy_data generates the header and trailer data automatically: + # enco = PG::BinaryEncoder::CopyRow.new + # conn.copy_data "COPY my_table FROM STDIN (FORMAT binary)", enco do + # conn.put_copy_data ['some', 'data', 'to', 'copy'] + # conn.put_copy_data ['more', 'data', 'to', 'copy'] + # end + # # Example with CSV output format: # conn.copy_data "COPY my_table TO STDOUT CSV" do # while row=conn.get_copy_data @@ -187,6 +195,18 @@ def inspect # This receives all rows of +my_table+ as ruby array: # ["some", "data", "to", "copy"] # ["more", "data", "to", "copy"] + # + # Also PG::BinaryDecoder::CopyRow can be used to retrieve data in binary format from the server. + # In this case the header and trailer data is processed by the decoder and the remaining +nil+ from get_copy_data is processed by copy_data, so that binary data can be processes equally to text data: + # deco = PG::BinaryDecoder::CopyRow.new + # conn.copy_data "COPY my_table TO STDOUT (FORMAT binary)", deco do + # while row=conn.get_copy_data + # p row + # end + # end + # This receives all rows of +my_table+ as ruby array: + # ["some", "data", "to", "copy"] + # ["more", "data", "to", "copy"] def copy_data( sql, coder=nil ) raise PG::NotInBlockingMode.new("copy_data can not be used in nonblocking mode", connection: self) if nonblocking? @@ -195,7 +215,7 @@ def copy_data( sql, coder=nil ) case res.result_status when PGRES_COPY_IN begin - if res.binary_tuples == 1 + if coder && res.binary_tuples == 1 # Binary file header (11 byte signature, 32 bit flags and 32 bit extension length) put_copy_data(BinarySignature + ("\x00" * 8)) end @@ -219,7 +239,7 @@ def copy_data( sql, coder=nil ) begin self.encoder_for_put_copy_data = old_coder if coder - if res.binary_tuples == 1 + if coder && res.binary_tuples == 1 put_copy_data("\xFF\xFF") # Binary file trailer 16 bit "-1" end @@ -244,7 +264,7 @@ def copy_data( sql, coder=nil ) discard_results raise else - if res.binary_tuples == 1 + if coder && res.binary_tuples == 1 # there are two end markers in binary mode: file trailer and the final nil if get_copy_data discard_results diff --git a/spec/pg/connection_spec.rb b/spec/pg/connection_spec.rb index 2f647b61c..33fb02665 100644 --- a/spec/pg/connection_spec.rb +++ b/spec/pg/connection_spec.rb @@ -1212,11 +1212,12 @@ end describe "#copy_data" do - it "can process #copy_data output queries" do + it "can process #copy_data output queries in text format" do rows = [] res2 = @conn.copy_data( "COPY (SELECT 1 UNION ALL SELECT 2) TO STDOUT" ) do |res| expect( res.result_status ).to eq( PG::PGRES_COPY_OUT ) expect( res.nfields ).to eq( 1 ) + expect( res.binary_tuples ).to eq( 0 ) while row=@conn.get_copy_data rows << row end @@ -1226,6 +1227,21 @@ expect( @conn ).to still_be_usable end + it "can process #copy_data output queries in binary format" do + rows = [] + res2 = @conn.copy_data( "COPY (SELECT 1 UNION ALL SELECT 2) TO STDOUT (FORMAT binary)" ) do |res| + expect( res.result_status ).to eq( PG::PGRES_COPY_OUT ) + expect( res.nfields ).to eq( 1 ) + expect( res.binary_tuples ).to eq( 1 ) + while row=@conn.get_copy_data + rows << row + end + end + expect( rows ).to eq( ["PGCOPY\n\xFF\r\n\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x04\x00\x00\x00\x01".b, "\x00\x01\x00\x00\x00\x04\x00\x00\x00\x02".b, "\xFF\xFF".b] ) + expect( res2.result_status ).to eq( PG::PGRES_COMMAND_OK ) + expect( @conn ).to still_be_usable + end + it "can handle incomplete #copy_data output queries" do expect { @conn.copy_data( "COPY (SELECT 1 UNION ALL SELECT 2) TO STDOUT" ) do |res| @@ -1269,11 +1285,12 @@ expect( @conn ).to still_be_usable end - it "can process #copy_data input queries" do + it "can process #copy_data input queries in text format" do @conn.exec( "CREATE TEMP TABLE copytable (col1 TEXT)" ) res2 = @conn.copy_data( "COPY copytable FROM STDOUT" ) do |res| expect( res.result_status ).to eq( PG::PGRES_COPY_IN ) expect( res.nfields ).to eq( 1 ) + expect( res.binary_tuples ).to eq( 0 ) @conn.put_copy_data "1\n" @conn.put_copy_data "2\n" end @@ -1286,6 +1303,28 @@ @conn.exec( "DROP TABLE IF EXISTS copytable" ) end + it "can process #copy_data input queries in binary format" do + @conn.exec( "CREATE TEMP TABLE copytable (col1 TEXT)" ) + res2 = @conn.copy_data( "COPY copytable FROM STDOUT (FORMAT binary)" ) do |res| + expect( res.result_status ).to eq( PG::PGRES_COPY_IN ) + expect( res.nfields ).to eq( 1 ) + expect( res.binary_tuples ).to eq( 1 ) + # header and first record ("1") + @conn.put_copy_data "PGCOPY\n\xFF\r\n\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x31".b + # second record ("2") + @conn.put_copy_data "\x00\x01\x00\x00\x00\x01\x32".b + # trailer + @conn.put_copy_data "\xFF\xFF".b + end + expect( res2.result_status ).to eq( PG::PGRES_COMMAND_OK ) + + expect( @conn ).to still_be_usable + + 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 str = "abcd" * 2000 + "\n" @conn.exec( "CREATE TEMP TABLE copytable2 (col1 TEXT)" )