Repository: GoogleCloudPlatform/pgadapter
Pull request #609 (2023-01-22 11:10:46)

fix: binary copy header should be included in first data message
The PostgreSQL wire protocol requires that the binary copy header is sent before any data. It does however not specify whether it should be in a separate CopyData message, or whether it should be included in the first actual CopyData message. PGAdapter used to send the header in a separate message. This is accepted by most clients. PostgreSQL sends the header as part of the first data message, and some clients (npgsql) assume that this will always be the case.

This fix therefore changes the behavior of PGAdapter so it will always include the header with the first actual CopyData message, instead of sending it in a separate message. This also means that if the copy operation does not actually copy any data at all, the header will be included in the trailer message. This is also consistent with real PostgreSQL.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#609](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (960ecb9) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/0c607c7c1bce48139f28688a5d7f1e202d839860?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (0c607c7) will **decrease** coverage by `0.06%`.
> The diff coverage is `87.87%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #609      +/-   ##
========================================================
- Coverage                 88.67%   88.62%   -0.06%     
- Complexity                 2136     2141       +5     
========================================================
  Files                       124      124              
  Lines                      7101     7118      +17     
  Branches                   1011     1023      +12     
========================================================
+ Hits                       6297     6308      +11     
- Misses                      567      570       +3     
- Partials                    237      240       +3     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `88.62% <87.87%> (-0.06%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...ud/spanner/pgadapter/metadata/OptionsMetadata.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL21ldGFkYXRhL09wdGlvbnNNZXRhZGF0YS5qYXZh) | `81.02% <ø> (ø)` | |
| [...e/cloud/spanner/pgadapter/utils/CsvCopyParser.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0NzdkNvcHlQYXJzZXIuamF2YQ==) | `98.33% <0.00%> (-1.67%)` | :arrow_down: |
| [...loud/spanner/pgadapter/utils/BinaryCopyParser.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0JpbmFyeUNvcHlQYXJzZXIuamF2YQ==) | `91.74% <33.33%> (-1.66%)` | :arrow_down: |
| [...oogle/cloud/spanner/pgadapter/utils/Converter.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0NvbnZlcnRlci5qYXZh) | `92.06% <90.90%> (-0.39%)` | :arrow_down: |
| [...ud/spanner/pgadapter/statements/CopyStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQ29weVN0YXRlbWVudC5qYXZh) | `94.84% <100.00%> (+0.02%)` | :arrow_up: |
| [.../spanner/pgadapter/statements/CopyToStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQ29weVRvU3RhdGVtZW50LmphdmE=) | `99.03% <100.00%> (-0.03%)` | :arrow_down: |
| [...spanner/pgadapter/wireoutput/CopyDataResponse.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVvdXRwdXQvQ29weURhdGFSZXNwb25zZS5qYXZh) | `90.00% <100.00%> (+0.41%)` | :arrow_up: |
| [...spanner/pgadapter/wireprotocol/ControlMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9Db250cm9sTWVzc2FnZS5qYXZh) | `86.97% <100.00%> (+0.12%)` | :arrow_up: |
| [...gle/cloud/spanner/pgadapter/ConnectionHandler.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL0Nvbm5lY3Rpb25IYW5kbGVyLmphdmE=) | `86.38% <0.00%> (-0.84%)` | :arrow_down: |
| ... and [1 more](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/609?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java
index 7c49c4bd..fe950b4e 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java
@@ -24,6 +24,7 @@ import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStateme
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.AutocommitDmlMode;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import com.google.cloud.spanner.pgadapter.ProxyServer.DataFormat;
 import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.error.SQLState;
@@ -80,7 +81,17 @@ public class CopyStatement extends IntermediatePortalStatement {
   public enum Format {
     CSV,
     TEXT,
-    BINARY,
+    BINARY {
+      @Override
+      public DataFormat getDataFormat() {
+        return DataFormat.POSTGRESQL_BINARY;
+      }
+    };
+
+    /** Returns the (default) data format that should be used for this copy format. */
+    public DataFormat getDataFormat() {
+      return DataFormat.POSTGRESQL_TEXT;
+    }
   }
 
   private static final String COLUMN_NAME = "column_name";
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyToStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyToStatement.java
index 4c4c0f95..75c6748c 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyToStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyToStatement.java
@@ -23,7 +23,6 @@ import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
-import com.google.cloud.spanner.pgadapter.ProxyServer.DataFormat;
 import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.error.SQLState;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
@@ -58,6 +57,7 @@ public class CopyToStatement extends IntermediatePortalStatement {
 
   private final ParsedCopyStatement parsedCopyStatement;
   private final CSVFormat csvFormat;
+  private boolean hasReturnedData;
 
   public CopyToStatement(
       ConnectionHandler connectionHandler,
@@ -220,22 +220,23 @@ public class CopyToStatement extends IntermediatePortalStatement {
 
   @Override
   public WireOutput[] createResultPrefix(ResultSet resultSet) {
-    return this.parsedCopyStatement.format == CopyStatement.Format.BINARY
-        ? new WireOutput[] {
-          new CopyOutResponse(
-              this.outputStream,
-              resultSet.getColumnCount(),
-              DataFormat.POSTGRESQL_BINARY.getCode()),
-          CopyDataResponse.createBinaryHeader(this.outputStream)
-        }
-        : new WireOutput[] {
-          new CopyOutResponse(
-              this.outputStream, resultSet.getColumnCount(), DataFormat.POSTGRESQL_TEXT.getCode())
-        };
+    return new WireOutput[] {
+      new CopyOutResponse(
+          this.outputStream,
+          resultSet.getColumnCount(),
+          this.parsedCopyStatement.format.getDataFormat().getCode())
+    };
   }
 
   @Override
   public CopyDataResponse createDataRowResponse(Converter converter) {
+    // Keep track of whether this COPY statement has returned at least one row. This is necessary to
+    // know whether we need to include the header in the current row and/or in the trailer.
+    // PostgreSQL includes the header in either the first data row or in the trailer if there are no
+    // rows. This is not specifically mentioned in the protocol description, but some clients assume
+    // this behavior. See
+    // https://github.com/npgsql/npgsql/blob/7f97dbad28c71b2202dd7bcccd05fc42a7de23c8/src/Npgsql/NpgsqlBinaryExporter.cs#L156
+    this.hasReturnedData = true;
     return parsedCopyStatement.format == CopyStatement.Format.BINARY
         ? createBinaryDataResponse(converter)
         : createDataResponse(converter.getResultSet());
@@ -245,7 +246,7 @@ public class CopyToStatement extends IntermediatePortalStatement {
   public WireOutput[] createResultSuffix() {
     return this.parsedCopyStatement.format == Format.BINARY
         ? new WireOutput[] {
-          CopyDataResponse.createBinaryTrailer(this.outputStream),
+          CopyDataResponse.createBinaryTrailer(this.outputStream, !hasReturnedData),
           new CopyDoneResponse(this.outputStream)
         }
         : new WireOutput[] {new CopyDoneResponse(this.outputStream)};
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/BinaryCopyParser.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/BinaryCopyParser.java
index d39e1bd7..9665d405 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/BinaryCopyParser.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/BinaryCopyParser.java
@@ -235,6 +235,14 @@ class BinaryCopyParser implements CopyInParser {
       return false;
     }
 
+    @Override
+    public boolean isNull(int columnIndex) {
+      Preconditions.checkArgument(
+          columnIndex >= 0 && columnIndex < numColumns(),
+          "columnIndex must be >= 0 && < numColumns");
+      return fields[columnIndex].data == null;
+    }
+
     @Override
     public Value getValue(Type type, String columnName) {
       // The binary copy format does not include any column name headers or any type information.
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/Converter.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/Converter.java
index 48211567..efebccb7 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/Converter.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/Converter.java
@@ -14,6 +14,9 @@
 
 package com.google.cloud.spanner.pgadapter.utils;
 
+import static com.google.cloud.spanner.pgadapter.statements.CopyToStatement.COPY_BINARY_HEADER;
+
+import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
@@ -38,6 +41,7 @@ import java.io.DataOutputStream;
 import java.io.IOException;
 
 /** Utility class for converting between generic PostgreSQL conversions. */
+@InternalApi
 public class Converter implements AutoCloseable {
   private final ByteArrayOutputStream buffer = new ByteArrayOutputStream(256);
   private final DataOutputStream outputStream = new DataOutputStream(buffer);
@@ -46,12 +50,15 @@ public class Converter implements AutoCloseable {
   private final OptionsMetadata options;
   private final ResultSet resultSet;
   private final SessionState sessionState;
+  private final boolean includeBinaryCopyHeaderInFirstRow;
+  private boolean firstRow = true;
 
   public Converter(
       IntermediateStatement statement,
       QueryMode mode,
       OptionsMetadata options,
-      ResultSet resultSet) {
+      ResultSet resultSet,
+      boolean includeBinaryCopyHeaderInFirstRow) {
     this.statement = statement;
     this.mode = mode;
     this.options = options;
@@ -62,6 +69,7 @@ public class Converter implements AutoCloseable {
             .getExtendedQueryProtocolHandler()
             .getBackendConnection()
             .getSessionState();
+    this.includeBinaryCopyHeaderInFirstRow = includeBinaryCopyHeaderInFirstRow;
   }
 
   @Override
@@ -83,6 +91,12 @@ public class Converter implements AutoCloseable {
               : DataFormat.POSTGRESQL_TEXT;
     }
     buffer.reset();
+    if (includeBinaryCopyHeaderInFirstRow && firstRow) {
+      outputStream.write(COPY_BINARY_HEADER);
+      outputStream.writeInt(0); // flags
+      outputStream.writeInt(0); // header extension area length
+    }
+    firstRow = false;
     outputStream.writeShort(resultSet.getColumnCount());
     for (int column_index = 0; /* column indices start at 0 */
         column_index < resultSet.getColumnCount();
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyRecord.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyRecord.java
index e6c252aa..e77f91cc 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyRecord.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyRecord.java
@@ -34,6 +34,9 @@ public interface CopyRecord {
    */
   boolean hasColumnNames();
 
+  /** Returns true if the value of the given column is null. */
+  boolean isNull(int columnIndex);
+
   /**
    * Returns the value of the given column as a Cloud Spanner {@link Value} of the given type. This
    * method is used by a COPY ... FROM ... operation to convert a value to the type of the column
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/CsvCopyParser.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/CsvCopyParser.java
index 9265faf6..882519b6 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/CsvCopyParser.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/CsvCopyParser.java
@@ -102,6 +102,11 @@ class CsvCopyParser implements CopyInParser {
       return this.hasHeader;
     }
 
+    @Override
+    public boolean isNull(int columnIndex) {
+      return record.get(columnIndex) == null;
+    }
+
     @Override
     public Value getValue(Type type, String columnName) throws SpannerException {
       String recordValue = record.get(columnName);
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyDataResponse.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyDataResponse.java
index c91250ba..871721d8 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyDataResponse.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyDataResponse.java
@@ -28,9 +28,9 @@ import java.nio.charset.StandardCharsets;
 public class CopyDataResponse extends WireOutput {
   @InternalApi
   public enum ResponseType {
-    HEADER,
     ROW,
     TRAILER,
+    TRAILER_WITH_HEADER,
   }
 
   private final ResponseType responseType;
@@ -39,16 +39,14 @@ public class CopyDataResponse extends WireOutput {
   private final char rowTerminator;
   private final Converter converter;
 
-  /** Creates a {@link CopyDataResponse} message containing the fixed binary COPY header. */
-  @InternalApi
-  public static CopyDataResponse createBinaryHeader(DataOutputStream output) {
-    return new CopyDataResponse(output, COPY_BINARY_HEADER.length + 8, ResponseType.HEADER);
-  }
-
   /** Creates a {@link CopyDataResponse} message containing the fixed binary COPY trailer. */
   @InternalApi
-  public static CopyDataResponse createBinaryTrailer(DataOutputStream output) {
-    return new CopyDataResponse(output, 2, ResponseType.TRAILER);
+  public static CopyDataResponse createBinaryTrailer(
+      DataOutputStream output, boolean includeHeader) {
+    return new CopyDataResponse(
+        output,
+        includeHeader ? 21 : 2,
+        includeHeader ? ResponseType.TRAILER_WITH_HEADER : ResponseType.TRAILER);
   }
 
   private CopyDataResponse(DataOutputStream output, int length, ResponseType responseType) {
@@ -78,6 +76,7 @@ public class CopyDataResponse extends WireOutput {
     this.converter = converter;
   }
 
+  @Override
   public void send(boolean flush) throws Exception {
     if (converter != null) {
       this.length = 4 + converter.convertResultSetRowToDataRowResponse();
@@ -91,10 +90,11 @@ public class CopyDataResponse extends WireOutput {
       this.outputStream.write(this.stringData.getBytes(StandardCharsets.UTF_8));
       this.outputStream.write(this.rowTerminator);
     } else if (this.format == DataFormat.POSTGRESQL_BINARY) {
-      if (this.responseType == ResponseType.HEADER) {
+      if (this.responseType == ResponseType.TRAILER_WITH_HEADER) {
         this.outputStream.write(COPY_BINARY_HEADER);
         this.outputStream.writeInt(0); // flags
         this.outputStream.writeInt(0); // header extension area length
+        this.outputStream.writeShort(-1);
       } else if (this.responseType == ResponseType.TRAILER) {
         this.outputStream.writeShort(-1);
       } else if (this.converter != null) {
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
index e233eb41..4aa61bdb 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
@@ -38,6 +38,7 @@ import com.google.cloud.spanner.pgadapter.error.SQLState;
 import com.google.cloud.spanner.pgadapter.error.Severity;
 import com.google.cloud.spanner.pgadapter.metadata.SendResultSetState;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.PartitionQueryResult;
+import com.google.cloud.spanner.pgadapter.statements.CopyToStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.cloud.spanner.pgadapter.utils.Converter;
 import com.google.cloud.spanner.pgadapter.wireoutput.CommandCompleteResponse;
@@ -444,7 +445,10 @@ public abstract class ControlMessage extends WireMessage {
               describedResult,
               mode,
               describedResult.getConnectionHandler().getServer().getOptions(),
-              resultSet);
+              resultSet,
+              includePrefix
+                  && describedResult instanceof CopyToStatement
+                  && ((CopyToStatement) describedResult).isBinary());
       this.batchReadOnlyTransaction = null;
       this.partition = null;
       this.maxRows = maxRows;
@@ -485,7 +489,8 @@ public abstract class ControlMessage extends WireMessage {
                     describedResult,
                     mode,
                     describedResult.getConnectionHandler().getServer().getOptions(),
-                    resultSet);
+                    resultSet,
+                    includePrefix);
             hasData = resultSet.next();
           } catch (Throwable t) {
             if (includePrefix) {
@@ -501,14 +506,15 @@ public abstract class ControlMessage extends WireMessage {
             for (WireOutput prefix : describedResult.createResultPrefix(resultSet)) {
               prefix.send(false);
             }
-            prefixSent.set(true);
           } catch (Throwable t) {
             prefixSent.setException(t);
             throw t;
           }
         }
         // Wait until the prefix (if any) has been sent.
-        prefixSent.get();
+        if (!includePrefix) {
+          prefixSent.get();
+        }
         long rows = 0L;
         while (hasData) {
           if (Thread.interrupted()) {
@@ -518,6 +524,10 @@ public abstract class ControlMessage extends WireMessage {
           synchronized (describedResult) {
             wireOutput.send(false);
           }
+          // The prefix will be included in the first data row.
+          if (includePrefix && rows == 0L) {
+            prefixSent.set(true);
+          }
           rows++;
           hasData = resultSet.next();
           if (rows % 1000 == 0) {


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/CopyOutMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/CopyOutMockServerTest.java
index fe7da284..41417504 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/CopyOutMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/CopyOutMockServerTest.java
@@ -487,6 +487,47 @@ public class CopyOutMockServerTest extends AbstractMockServerTest {
     assertEquals(Value.string("test"), record.getValue(Type.string(), 8));
   }
 
+  @Test
+  public void testCopyOutBinaryEmptyPsql() throws Exception {
+    assumeTrue("This test requires psql", isPsqlAvailable());
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of("select * from all_types"),
+            ALL_TYPES_RESULTSET.toBuilder().clearRows().build()));
+
+    ProcessBuilder builder = new ProcessBuilder();
+    builder.command(
+        "psql",
+        "-h",
+        (useDomainSocket ? "/tmp" : "localhost"),
+        "-p",
+        String.valueOf(pgServer.getLocalPort()),
+        "-c",
+        "copy all_types to stdout binary");
+    Process process = builder.start();
+    StringBuilder errorBuilder = new StringBuilder();
+    try (Scanner scanner = new Scanner(new InputStreamReader(process.getErrorStream()))) {
+      while (scanner.hasNextLine()) {
+        errorBuilder.append(scanner.nextLine()).append('\n');
+      }
+    }
+    PipedOutputStream pipedOutputStream = new PipedOutputStream();
+    PipedInputStream inputStream = new PipedInputStream(pipedOutputStream, 1 << 16);
+    SessionState sessionState = mock(SessionState.class);
+    CopyInParser copyParser =
+        CopyInParser.create(sessionState, Format.BINARY, null, inputStream, false);
+    int b;
+    while ((b = process.getInputStream().read()) != -1) {
+      pipedOutputStream.write(b);
+    }
+    int res = process.waitFor();
+    assertEquals("", errorBuilder.toString());
+    assertEquals(0, res);
+
+    Iterator<CopyRecord> iterator = copyParser.iterator();
+    assertFalse(iterator.hasNext());
+  }
+
   @Test
   public void testCopyOutNullsBinaryPsql() throws Exception {
     assumeTrue("This test requires psql", isPsqlAvailable());
@@ -564,6 +605,38 @@ public class CopyOutMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  @Test
+  public void testCopyOutPartitionedBinary() throws SQLException, IOException {
+    final int expectedRowCount = 100;
+    RandomResultSetGenerator randomResultSetGenerator =
+        new RandomResultSetGenerator(expectedRowCount, Dialect.POSTGRESQL);
+    com.google.spanner.v1.ResultSet resultSet = randomResultSetGenerator.generate();
+    mockSpanner.putStatementResult(
+        StatementResult.query(Statement.of("select * from random"), resultSet));
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
+      PipedOutputStream pipedOutputStream = new PipedOutputStream();
+      PipedInputStream inputStream = new PipedInputStream(pipedOutputStream, 1 << 20);
+      SessionState sessionState = mock(SessionState.class);
+      CopyInParser copyParser =
+          CopyInParser.create(sessionState, Format.BINARY, null, inputStream, false);
+      long rows = copyManager.copyOut("COPY random TO STDOUT (format binary)", pipedOutputStream);
+
+      assertEquals(expectedRowCount, rows);
+
+      Iterator<CopyRecord> iterator = copyParser.iterator();
+      int recordCount = 0;
+      while (iterator.hasNext()) {
+        recordCount++;
+        CopyRecord record = iterator.next();
+        int index = findIndex(resultSet, record);
+        assertNotEquals(String.format("Row %d not found: %s", recordCount, record), -1, index);
+      }
+      assertEquals(expectedRowCount, recordCount);
+    }
+  }
+
   @Test
   public void testCopyOutPartitioned_NonExistingTable() throws SQLException {
     StatusRuntimeException exception =
@@ -627,4 +700,45 @@ public class CopyOutMockServerTest extends AbstractMockServerTest {
     }
     return -1;
   }
+
+  static int findIndex(com.google.spanner.v1.ResultSet resultSet, CopyRecord record) {
+    for (int index = 0; index < resultSet.getRowsCount(); index++) {
+      boolean nullValuesEqual = true;
+      for (int colIndex = 0; colIndex < record.numColumns(); colIndex++) {
+        if (record.isNull(colIndex)
+            != resultSet.getRows(index).getValues(colIndex).hasNullValue()) {
+          nullValuesEqual = false;
+          break;
+        }
+      }
+      if (!nullValuesEqual) {
+        continue;
+      }
+
+      boolean valuesEqual = true;
+      for (int colIndex = 0; colIndex < record.numColumns(); colIndex++) {
+        if (!resultSet.getRows(index).getValues(colIndex).hasNullValue()) {
+          if (resultSet.getMetadata().getRowType().getFields(colIndex).getType().getCode()
+                  == TypeCode.STRING
+              && !record
+                  .getValue(Type.string(), colIndex)
+                  .getString()
+                  .equals(resultSet.getRows(index).getValues(colIndex).getStringValue())) {
+            valuesEqual = false;
+            break;
+          }
+          if (resultSet.getRows(index).getValues(colIndex).hasBoolValue()
+              && record.getValue(Type.bool(), colIndex).getBool()
+                  != resultSet.getRows(index).getValues(colIndex).getBoolValue()) {
+            valuesEqual = false;
+            break;
+          }
+        }
+      }
+      if (valuesEqual) {
+        return index;
+      }
+    }
+    return -1;
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ITPsqlTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ITPsqlTest.java
index fc4adcb2..39979910 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITPsqlTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITPsqlTest.java
@@ -394,8 +394,12 @@ public class ITPsqlTest implements IntegrationTest {
         output);
   }
 
+  /**
+   * This test copies data back and forth between PostgreSQL and Cloud Spanner and verifies that the
+   * contents are equal after the COPY operation in both directions.
+   */
   @Test
-  public void testCopyFromPostgreSQLToCloudSpanner() throws Exception {
+  public void testCopyBetweenPostgreSQLAndCloudSpanner() throws Exception {
     int numRows = 100;
 
     // Generate 99 random rows.
@@ -523,6 +527,103 @@ public class ITPsqlTest implements IntegrationTest {
     }
   }
 
+  /** This test verifies that we can copy an empty table between PostgreSQL and Cloud Spanner. */
+  @Test
+  public void testCopyEmptyTableBetweenCloudSpannerAndPostgreSQL() throws Exception {
+    // Remove all rows in the table in the local PostgreSQL database.
+    try (Connection connection =
+        DriverManager.getConnection(createJdbcUrlForLocalPg(), POSTGRES_USER, POSTGRES_PASSWORD)) {
+      connection.createStatement().executeUpdate("delete from all_types");
+    }
+
+    // Execute the COPY tests in both binary and text mode.
+    for (boolean binary : new boolean[] {false, true}) {
+      // Make sure the all_types table on Cloud Spanner is empty.
+      String databaseId = database.getId().getDatabase();
+      testEnv.write(databaseId, Collections.singleton(Mutation.delete("all_types", KeySet.all())));
+
+      // COPY the empty table to Cloud Spanner.
+      ProcessBuilder builder = new ProcessBuilder();
+      builder.command(
+          "bash",
+          "-c",
+          "psql"
+              + " -h "
+              + POSTGRES_HOST
+              + " -p "
+              + POSTGRES_PORT
+              + " -U "
+              + POSTGRES_USER
+              + " -d "
+              + POSTGRES_DATABASE
+              + " -c \"copy all_types to stdout"
+              + (binary ? " binary \" " : "\" ")
+              + "  | psql "
+              + " -h "
+              + (POSTGRES_HOST.startsWith("/") ? "/tmp" : testEnv.getPGAdapterHost())
+              + " -p "
+              + testEnv.getPGAdapterPort()
+              + " -d "
+              + database.getId().getDatabase()
+              + " -c \"copy all_types from stdin "
+              + (binary ? "binary" : "")
+              + ";\"\n");
+      setPgPassword(builder);
+      Process process = builder.start();
+      int res = process.waitFor();
+      assertEquals(0, res);
+
+      // Verify that still have 0 rows in Cloud Spanner.
+      try (Connection connection = DriverManager.getConnection(createJdbcUrlForPGAdapter())) {
+        try (ResultSet resultSet =
+            connection.createStatement().executeQuery("select count(*) from all_types")) {
+          assertTrue(resultSet.next());
+          assertEquals(0, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+      }
+
+      // COPY the empty table from Cloud Spanner to PostgreSQL.
+      ProcessBuilder copyToPostgresBuilder = new ProcessBuilder();
+      copyToPostgresBuilder.command(
+          "bash",
+          "-c",
+          "psql"
+              + " -c \"copy all_types to stdout"
+              + (binary ? " binary \" " : "\" ")
+              + " -h "
+              + (POSTGRES_HOST.startsWith("/") ? "/tmp" : testEnv.getPGAdapterHost())
+              + " -p "
+              + testEnv.getPGAdapterPort()
+              + " -d "
+              + database.getId().getDatabase()
+              + "  | psql "
+              + " -h "
+              + POSTGRES_HOST
+              + " -p "
+              + POSTGRES_PORT
+              + " -U "
+              + POSTGRES_USER
+              + " -d "
+              + POSTGRES_DATABASE
+              + " -c \"copy all_types from stdin "
+              + (binary ? "binary" : "")
+              + ";\"\n");
+      setPgPassword(copyToPostgresBuilder);
+      Process copyToPostgresProcess = copyToPostgresBuilder.start();
+      InputStream errorStream = copyToPostgresProcess.getErrorStream();
+      int copyToPostgresResult = copyToPostgresProcess.waitFor();
+      StringBuilder errors = new StringBuilder();
+      try (Scanner scanner = new Scanner(new InputStreamReader(errorStream))) {
+        while (scanner.hasNextLine()) {
+          errors.append(errors).append(scanner.nextLine()).append("\n");
+        }
+      }
+      assertEquals("", errors.toString());
+      assertEquals(0, copyToPostgresResult);
+    }
+  }
+
   @Test
   public void testTimestamptzParsing() throws Exception {
     final int numTests = 10;
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/PsqlMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/PsqlMockServerTest.java
index 8f0e4e93..8d177e93 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/PsqlMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/PsqlMockServerTest.java
@@ -14,16 +14,23 @@
 
 package com.google.cloud.spanner.pgadapter;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
 
 import com.google.cloud.spanner.MockSpannerServiceImpl;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.pgadapter.statements.CopyToStatement;
 import com.google.cloud.spanner.pgadapter.wireprotocol.SSLMessage;
 import com.google.common.collect.ImmutableList;
+import com.google.common.io.ByteStreams;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
 import java.io.BufferedReader;
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
 import java.io.File;
 import java.io.InputStreamReader;
 import java.util.List;
@@ -141,4 +148,91 @@ public class PsqlMockServerTest extends AbstractMockServerTest {
     assertEquals(
         1L, pgServer.getDebugMessages().stream().filter(m -> m instanceof SSLMessage).count());
   }
+
+  @Test
+  public void testCopyToBinaryPsql() throws Exception {
+    assumeTrue("This test requires psql to be installed", isPsqlAvailable());
+
+    ProcessBuilder builder = new ProcessBuilder();
+    String[] psqlCommand =
+        new String[] {
+          "psql",
+          "-h",
+          "localhost",
+          "-p",
+          String.valueOf(pgServer.getLocalPort()),
+          "-c",
+          "COPY (SELECT 1) TO STDOUT (FORMAT BINARY)"
+        };
+    builder.command(psqlCommand);
+    Process process = builder.start();
+    String errors;
+
+    try (BufferedReader errorReader =
+        new BufferedReader(new InputStreamReader(process.getErrorStream()))) {
+      errors = errorReader.lines().collect(Collectors.joining("\n"));
+    }
+    byte[] copyOutput = ByteStreams.toByteArray(process.getInputStream());
+
+    assertEquals("", errors);
+    int res = process.waitFor();
+    assertEquals(0, res);
+
+    ByteArrayOutputStream byteArrayOutputStream =
+        new ByteArrayOutputStream(CopyToStatement.COPY_BINARY_HEADER.length);
+    DataOutputStream expectedOutputStream = new DataOutputStream(byteArrayOutputStream);
+    expectedOutputStream.write(CopyToStatement.COPY_BINARY_HEADER);
+    expectedOutputStream.writeInt(0); // flags
+    expectedOutputStream.writeInt(0); // header extension area length
+    expectedOutputStream.writeShort(1); // Column count
+    expectedOutputStream.writeInt(8); // Value length
+    expectedOutputStream.writeLong(1L); // Column value
+    expectedOutputStream.writeShort(-1); // Column count == -1 means end of data.
+
+    assertArrayEquals(byteArrayOutputStream.toByteArray(), copyOutput);
+  }
+
+  @Test
+  public void testCopyToBinaryPsqlEmptyTable() throws Exception {
+    assumeTrue("This test requires psql to be installed", isPsqlAvailable());
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of("select * from all_types"),
+            ALL_TYPES_RESULTSET.toBuilder().clearRows().build()));
+
+    ProcessBuilder builder = new ProcessBuilder();
+    String[] psqlCommand =
+        new String[] {
+          "psql",
+          "-h",
+          "localhost",
+          "-p",
+          String.valueOf(pgServer.getLocalPort()),
+          "-c",
+          "COPY all_types TO STDOUT (FORMAT BINARY)"
+        };
+    builder.command(psqlCommand);
+    Process process = builder.start();
+    String errors;
+
+    try (BufferedReader errorReader =
+        new BufferedReader(new InputStreamReader(process.getErrorStream()))) {
+      errors = errorReader.lines().collect(Collectors.joining("\n"));
+    }
+    byte[] copyOutput = ByteStreams.toByteArray(process.getInputStream());
+
+    assertEquals("", errors);
+    int res = process.waitFor();
+    assertEquals(0, res);
+
+    ByteArrayOutputStream byteArrayOutputStream =
+        new ByteArrayOutputStream(CopyToStatement.COPY_BINARY_HEADER.length);
+    DataOutputStream expectedOutputStream = new DataOutputStream(byteArrayOutputStream);
+    expectedOutputStream.write(CopyToStatement.COPY_BINARY_HEADER);
+    expectedOutputStream.writeInt(0); // flags
+    expectedOutputStream.writeInt(0); // header extension area length
+    expectedOutputStream.writeShort(-1); // Column count == -1 means end of data.
+
+    assertArrayEquals(byteArrayOutputStream.toByteArray(), copyOutput);
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/BinaryCopyParserTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/BinaryCopyParserTest.java
index 3bb7ff7d..be8ba384 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/utils/BinaryCopyParserTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/BinaryCopyParserTest.java
@@ -361,4 +361,67 @@ public class BinaryCopyParserTest {
     assertThrows(IllegalArgumentException.class, () -> record.getValue(Type.string(), 10));
     assertThrows(SpannerException.class, () -> record.getValue(Type.numeric(), 0));
   }
+
+  @Test
+  public void testIsNull() throws IOException {
+    PipedOutputStream pipedOutputStream = new PipedOutputStream();
+    BinaryCopyParser parser = new BinaryCopyParser(new PipedInputStream(pipedOutputStream, 256));
+
+    DataOutputStream data = new DataOutputStream(pipedOutputStream);
+    data.write(COPY_BINARY_HEADER);
+    data.writeInt(0);
+    data.writeInt(0);
+
+    // Write a tuple with two non-null fields.
+    data.writeShort(2);
+    data.writeInt(8);
+    data.writeLong(1L);
+    data.writeInt(8);
+    data.writeLong(2L);
+
+    // Write a tuple with one non-null and one null field.
+    data.writeShort(2);
+    data.writeInt(8);
+    data.writeLong(1L);
+    data.writeInt(-1);
+
+    // Write a tuple with one null and one non-null field.
+    data.writeShort(2);
+    data.writeInt(-1);
+    data.writeInt(8);
+    data.writeLong(2L);
+
+    // Write a tuple with two non-null fields.
+    data.writeShort(2);
+    data.writeInt(-1);
+    data.writeInt(-1);
+
+    // Trailer.
+    data.writeShort(-1);
+
+    Iterator<CopyRecord> iterator = parser.iterator();
+    assertTrue(iterator.hasNext());
+    CopyRecord record1 = iterator.next();
+    assertFalse(record1.isNull(0));
+    assertFalse(record1.isNull(1));
+    assertThrows(IllegalArgumentException.class, () -> record1.isNull(-1));
+    assertThrows(IllegalArgumentException.class, () -> record1.isNull(2));
+
+    assertTrue(iterator.hasNext());
+    CopyRecord record2 = iterator.next();
+    assertFalse(record2.isNull(0));
+    assertTrue(record2.isNull(1));
+
+    assertTrue(iterator.hasNext());
+    CopyRecord record3 = iterator.next();
+    assertTrue(record3.isNull(0));
+    assertFalse(record3.isNull(1));
+
+    assertTrue(iterator.hasNext());
+    CopyRecord record4 = iterator.next();
+    assertTrue(record4.isNull(0));
+    assertTrue(record4.isNull(1));
+
+    assertFalse(iterator.hasNext());
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/CsvCopyParserTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/CsvCopyParserTest.java
index 15614b9f..3138ced5 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/utils/CsvCopyParserTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/CsvCopyParserTest.java
@@ -15,8 +15,10 @@
 package com.google.cloud.spanner.pgadapter.utils;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 
 import com.google.cloud.ByteArray;
@@ -33,6 +35,7 @@ import java.io.IOException;
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
 import java.nio.charset.StandardCharsets;
+import java.util.Iterator;
 import org.apache.commons.csv.CSVFormat;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -43,7 +46,6 @@ public class CsvCopyParserTest {
 
   @Test
   public void testCanCreateIterator() throws IOException {
-
     CsvCopyParser parser =
         new CsvCopyParser(
             mock(SessionState.class),
@@ -293,4 +295,60 @@ public class CsvCopyParserTest {
             CsvCopyRecord.getSpannerValue(
                 sessionState, Type.struct(StructField.of("f1", Type.string())), "value"));
   }
+
+  @Test
+  public void testIsNull() throws IOException {
+    PipedOutputStream outputStream = new PipedOutputStream();
+    DataOutputStream data = new DataOutputStream(outputStream);
+    new Thread(
+            () -> {
+              while (true) {
+                try {
+                  data.write("\"value1\"\t\"value2\"\n".getBytes(StandardCharsets.UTF_8));
+                  data.write("\"value1\"\t\\N\n".getBytes(StandardCharsets.UTF_8));
+                  data.write("\\N\t\"value2\"\n".getBytes(StandardCharsets.UTF_8));
+                  data.write("\\N\t\\N\n".getBytes(StandardCharsets.UTF_8));
+                  data.close();
+                  break;
+                } catch (IOException e) {
+                  if (e.getMessage().contains("Pipe not connected")) {
+                    Thread.yield();
+                  } else {
+                    throw new RuntimeException(e);
+                  }
+                }
+              }
+            })
+        .start();
+    CsvCopyParser parser =
+        new CsvCopyParser(
+            mock(SessionState.class),
+            CSVFormat.POSTGRESQL_TEXT,
+            new PipedInputStream(outputStream, 256),
+            false);
+    Iterator<CopyRecord> iterator = parser.iterator();
+
+    assertTrue(iterator.hasNext());
+    CopyRecord record = iterator.next();
+    assertFalse(record.isNull(0));
+    assertFalse(record.isNull(1));
+
+    assertTrue(iterator.hasNext());
+    record = iterator.next();
+    assertFalse(record.isNull(0));
+    assertTrue(record.isNull(1));
+
+    assertTrue(iterator.hasNext());
+    record = iterator.next();
+    assertTrue(record.isNull(0));
+    assertFalse(record.isNull(1));
+
+    assertTrue(iterator.hasNext());
+    record = iterator.next();
+    assertTrue(record.isNull(0));
+    assertTrue(record.isNull(1));
+
+    assertFalse(iterator.hasNext());
+    parser.close();
+  }
 }