Merge pull request #527 from larskanis/fix-526

Fix regression in copy_data regarding binary format

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 processed 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

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)" )