fix: CopyResponse did not return correct column format (#633)

CopyResponse for BINARY copy statements did not return the correct format code for the individual columns. The COPY protocol will always use either only TEXT or only BINARY for each column, and the overall format is given in a separate field. The format per column is a precaution in the protocol to support a mix of TEXT and BINARY in potential future versions of the protocol. Most clients therefore ignore the values for the separate columns in a COPY statement. psycopg3 however does inspect these values and will fail if they contain an invalid value. PGAdapter would return a byte array containing only 1's for the column formats if the BINARY protocol was used. As the individual column formats are interpreted as int2 values, that meant that the client saw a format code value of 257 instead of 1.

Author: olavloite

src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java

@@ -241,8 +241,8 @@ public MutationWriter getMutationWriter() {
  }
 
  /** @return 0 for text/csv formatting and 1 for binary */
- public int getFormatCode() {
- return (parsedCopyStatement.format == Format.BINARY) ? 1 : 0;
+ public byte getFormatCode() {
+ return (parsedCopyStatement.format == Format.BINARY) ? (byte) 1 : (byte) 0;
  }
 
  private void verifyCopyColumns() {

src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiver.java

@@ -65,7 +65,7 @@ void handleCopy() throws Exception {
  this.connectionHandler.setActiveCopyStatement(copyStatement);
  new CopyInResponse(
  this.connectionHandler.getConnectionMetadata().getOutputStream(),
- copyStatement.getTableColumns().size(),
+ (short) copyStatement.getTableColumns().size(),
  copyStatement.getFormatCode())
  .send();
  ConnectionStatus initialConnectionStatus = this.connectionHandler.getStatus();

src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyInResponse.java

@@ -35,23 +35,25 @@ public class CopyInResponse extends WireOutput {
 
  protected static final char IDENTIFIER = 'G';
 
- private final int numColumns;
- private final int formatCode;
- private final byte[] columnFormat;
+ private final short numColumns;
+ private final byte formatCode;
+ private final short[] columnFormat;
 
- public CopyInResponse(DataOutputStream output, int numColumns, int formatCode) {
+ public CopyInResponse(DataOutputStream output, short numColumns, byte formatCode) {
  super(output, calculateLength(numColumns));
  this.numColumns = numColumns;
  this.formatCode = formatCode;
- columnFormat = new byte[COLUMN_NUM_LENGTH * this.numColumns];
- Arrays.fill(columnFormat, (byte) formatCode);
+ this.columnFormat = new short[this.numColumns];
+ Arrays.fill(this.columnFormat, formatCode);
  }
 
  @Override
  protected void sendPayload() throws IOException {
  this.outputStream.writeByte(this.formatCode);
  this.outputStream.writeShort(this.numColumns);
- this.outputStream.write(this.columnFormat);
+ for (int i = 0; i < this.numColumns; i++) {
+ this.outputStream.writeShort(this.columnFormat[i]);
+ }
  }
 
  @Override

src/test/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyInResponseTest.java

@@ -0,0 +1,70 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.cloud.spanner.pgadapter.wireoutput;
+
+import static org.junit.Assert.assertArrayEquals;
+
+import com.google.cloud.spanner.pgadapter.ProxyServer.DataFormat;
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class CopyInResponseTest {
+
+ @Test
+ public void testSendPayloadText() throws Exception {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ DataOutputStream output = new DataOutputStream(bytes);
+ short numColumns = 3;
+
+ CopyInResponse response =
+ new CopyInResponse(output, numColumns, (byte) DataFormat.POSTGRESQL_TEXT.getCode());
+ response.sendPayload();
+
+ assertArrayEquals(
+ new byte[] {
+ 0, // format
+ 0, 3, // numColumns
+ 0, 0, // format column 1
+ 0, 0, // format column 2
+ 0, 0, // format column 3
+ },
+ bytes.toByteArray());
+ }
+
+ @Test
+ public void testSendPayloadBinary() throws Exception {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ DataOutputStream output = new DataOutputStream(bytes);
+ short numColumns = 3;
+
+ CopyInResponse response =
+ new CopyInResponse(output, numColumns, (byte) DataFormat.POSTGRESQL_BINARY.getCode());
+ response.sendPayload();
+
+ assertArrayEquals(
+ new byte[] {
+ 1, // format
+ 0, 3, // numColumns
+ 0, 1, // format column 1
+ 0, 1, // format column 2
+ 0, 1, // format column 3
+ },
+ bytes.toByteArray());
+ }
+}