Repository: GoogleCloudPlatform/pgadapter
Pull request #615 (2023-01-25 14:49:17)

fix: empty partitions could skip binary copy header
Binary COPY TO STDOUT statements could fail if it used a partitioned query that contained one or more empty partitions. The binary copy header would be included for the first partition, but if that partition happened to return 0 rows, the copy statement would never send a binary copy header. This again would cause all other partitions never to be processed and sent, because they would wait for the first partition to send the header before starting to send data.

This is now fixed by:
1. Include the binary header with the first partition that actually sends data.
2. Include the binary header with the trailer if no partition sent any data.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/615?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#615](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/615?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (d92d731) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/2fbf89e6a6b3ba0b66f126abf019e386e9276d4c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (2fbf89e) will **increase** coverage by `0.11%`.
> The diff coverage is `92.10%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #615      +/-   ##
========================================================
+ Coverage                 88.71%   88.82%   +0.11%     
- Complexity                 2141     2148       +7     
========================================================
  Files                       124      124              
  Lines                      7116     7116              
  Branches                   1024     1019       -5     
========================================================
+ Hits                       6313     6321       +8     
+ Misses                      567      558       -9     
- Partials                    236      237       +1     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `88.82% <92.10%> (+0.11%)` | :arrow_up: |

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/615?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...spanner/pgadapter/wireprotocol/ControlMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/615?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9Db250cm9sTWVzc2FnZS5qYXZh) | `89.85% <87.50%> (+3.28%)` | :arrow_up: |
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/615?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `93.57% <100.00%> (+0.04%)` | :arrow_up: |
| [.../spanner/pgadapter/statements/CopyToStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/615?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQ29weVRvU3RhdGVtZW50LmphdmE=) | `100.00% <100.00%> (ø)` | |
| [...oogle/cloud/spanner/pgadapter/utils/Converter.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/615?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0NvbnZlcnRlci5qYXZh) | `93.65% <100.00%> (+0.31%)` | :arrow_up: |

: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/BackendConnection.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
index d3ba9aa3..a5cdf21d 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
@@ -372,7 +372,7 @@ public class BackendConnection {
     }
   }
 
-  private static final int MAX_PARTITIONS =
+  public static final int MAX_PARTITIONS =
       Math.max(16, 2 * Runtime.getRuntime().availableProcessors());
 
   private final class CopyOut extends BufferedStatement<StatementResult> {
@@ -401,9 +401,14 @@ public class BackendConnection {
               // No need for the extra complexity of a partitioned query.
               result.set(spannerConnection.execute(statement));
             } else {
+              // Get the metadata of the query, so we can include that in the result.
+              ResultSet metadataResultSet =
+                  spannerConnection.analyzeQuery(statement, QueryAnalyzeMode.PLAN);
               result.set(
                   new PartitionQueryResult(
-                      batchReadOnlyTransaction.getBatchTransactionId(), partitions));
+                      batchReadOnlyTransaction.getBatchTransactionId(),
+                      partitions,
+                      metadataResultSet));
             }
           } catch (SpannerException spannerException) {
             // The query might not be suitable for partitioning. Just try with a normal query.
@@ -1262,10 +1267,15 @@ public class BackendConnection {
   public static final class PartitionQueryResult implements StatementResult {
     private final BatchTransactionId batchTransactionId;
     private final List<Partition> partitions;
+    private final ResultSet metadataResultSet;
 
-    public PartitionQueryResult(BatchTransactionId batchTransactionId, List<Partition> partitions) {
+    public PartitionQueryResult(
+        BatchTransactionId batchTransactionId,
+        List<Partition> partitions,
+        ResultSet metadataResultSet) {
       this.batchTransactionId = batchTransactionId;
       this.partitions = partitions;
+      this.metadataResultSet = metadataResultSet;
     }
 
     public BatchTransactionId getBatchTransactionId() {
@@ -1276,6 +1286,10 @@ public class BackendConnection {
       return partitions;
     }
 
+    public ResultSet getMetadataResultSet() {
+      return metadataResultSet;
+    }
+
     @Override
     public ResultType getResultType() {
       return ResultType.RESULT_SET;
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 75c6748c..876e8862 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
@@ -41,6 +41,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.Futures;
 import java.util.List;
 import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 import org.apache.commons.csv.CSVFormat;
 import org.apache.commons.csv.QuoteMode;
@@ -57,7 +58,7 @@ public class CopyToStatement extends IntermediatePortalStatement {
 
   private final ParsedCopyStatement parsedCopyStatement;
   private final CSVFormat csvFormat;
-  private boolean hasReturnedData;
+  private final AtomicBoolean hasReturnedData = new AtomicBoolean(false);
 
   public CopyToStatement(
       ConnectionHandler connectionHandler,
@@ -236,7 +237,9 @@ public class CopyToStatement extends IntermediatePortalStatement {
     // 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;
+    if (parsedCopyStatement.format == Format.BINARY && !hasReturnedData.getAndSet(true)) {
+      converter = converter.includeBinaryCopyHeader();
+    }
     return parsedCopyStatement.format == CopyStatement.Format.BINARY
         ? createBinaryDataResponse(converter)
         : createDataResponse(converter.getResultSet());
@@ -246,7 +249,7 @@ public class CopyToStatement extends IntermediatePortalStatement {
   public WireOutput[] createResultSuffix() {
     return this.parsedCopyStatement.format == Format.BINARY
         ? new WireOutput[] {
-          CopyDataResponse.createBinaryTrailer(this.outputStream, !hasReturnedData),
+          CopyDataResponse.createBinaryTrailer(this.outputStream, !hasReturnedData.get()),
           new CopyDoneResponse(this.outputStream)
         }
         : new WireOutput[] {new CopyDoneResponse(this.outputStream)};
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 efebccb7..37857bd9 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
@@ -50,7 +50,7 @@ public class Converter implements AutoCloseable {
   private final OptionsMetadata options;
   private final ResultSet resultSet;
   private final SessionState sessionState;
-  private final boolean includeBinaryCopyHeaderInFirstRow;
+  private boolean includeBinaryCopyHeaderInFirstRow;
   private boolean firstRow = true;
 
   public Converter(
@@ -72,6 +72,15 @@ public class Converter implements AutoCloseable {
     this.includeBinaryCopyHeaderInFirstRow = includeBinaryCopyHeaderInFirstRow;
   }
 
+  public Converter includeBinaryCopyHeader() {
+    this.includeBinaryCopyHeaderInFirstRow = true;
+    return this;
+  }
+
+  public boolean isIncludeBinaryCopyHeaderInFirstRow() {
+    return this.includeBinaryCopyHeaderInFirstRow;
+  }
+
   @Override
   public void close() throws Exception {
     buffer.close();
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 4aa61bdb..e97f269a 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
@@ -52,7 +52,6 @@ import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
-import com.google.common.util.concurrent.SettableFuture;
 import io.grpc.Context;
 import io.grpc.MethodDescriptor;
 import java.io.DataInputStream;
@@ -60,6 +59,7 @@ import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
 import java.util.logging.Level;
@@ -313,6 +313,7 @@ public abstract class ControlMessage extends WireMessage {
     if (statementResult instanceof PartitionQueryResult) {
       hasData = false;
       PartitionQueryResult partitionQueryResult = (PartitionQueryResult) statementResult;
+      sendPrefix(describedResult, ((PartitionQueryResult) statementResult).getMetadataResultSet());
       rows =
           sendPartitionedQuery(
               describedResult,
@@ -322,17 +323,28 @@ public abstract class ControlMessage extends WireMessage {
     } else {
       hasData = describedResult.isHasMoreData();
       ResultSet resultSet = describedResult.getStatementResult().getResultSet();
+      sendPrefix(describedResult, resultSet);
       SendResultSetRunnable runnable =
-          SendResultSetRunnable.forResultSet(
-              describedResult, resultSet, maxRows, mode, SettableFuture.create(), hasData);
+          SendResultSetRunnable.forResultSet(describedResult, resultSet, maxRows, mode, hasData);
       rows = runnable.call();
       hasData = runnable.hasData;
     }
 
+    sendSuffix(describedResult);
+    return new SendResultSetState(describedResult.getCommandTag(), rows, hasData);
+  }
+
+  private void sendPrefix(IntermediateStatement describedResult, ResultSet resultSet)
+      throws Exception {
+    for (WireOutput prefix : describedResult.createResultPrefix(resultSet)) {
+      prefix.send(false);
+    }
+  }
+
+  private void sendSuffix(IntermediateStatement describedResult) throws Exception {
     for (WireOutput suffix : describedResult.createResultSuffix()) {
       suffix.send(false);
     }
-    return new SendResultSetState(describedResult.getCommandTag(), rows, hasData);
   }
 
   long sendPartitionedQuery(
@@ -345,7 +357,6 @@ public abstract class ControlMessage extends WireMessage {
             Executors.newFixedThreadPool(
                 Math.min(8 * Runtime.getRuntime().availableProcessors(), partitions.size())));
     List<ListenableFuture<Long>> futures = new ArrayList<>(partitions.size());
-    SettableFuture<Boolean> prefixSent = SettableFuture.create();
     Connection spannerConnection = connection.getSpannerConnection();
     Spanner spanner = ConnectionOptionsHelper.getSpanner(spannerConnection);
     BatchClient batchClient = spanner.getBatchClient(connection.getDatabaseId());
@@ -362,6 +373,10 @@ public abstract class ControlMessage extends WireMessage {
                     return GrpcCallContext.createDefault().withTimeout(Duration.ofHours(24L));
                   }
                 });
+    CountDownLatch binaryCopyHeaderSentLatch =
+        describedResult instanceof CopyToStatement && ((CopyToStatement) describedResult).isBinary()
+            ? new CountDownLatch(1)
+            : new CountDownLatch(0);
     for (int i = 0; i < partitions.size(); i++) {
       futures.add(
           executorService.submit(
@@ -371,8 +386,7 @@ public abstract class ControlMessage extends WireMessage {
                       batchReadOnlyTransaction,
                       partitions.get(i),
                       mode,
-                      i == 0,
-                      prefixSent))));
+                      binaryCopyHeaderSentLatch))));
     }
     executorService.shutdown();
     try {
@@ -404,8 +418,7 @@ public abstract class ControlMessage extends WireMessage {
     private final Partition partition;
     private final long maxRows;
     private final QueryMode mode;
-    private final boolean includePrefix;
-    private final SettableFuture<Boolean> prefixSent;
+    private final CountDownLatch binaryCopyHeaderSentLatch;
     private boolean hasData;
 
     static SendResultSetRunnable forResultSet(
@@ -413,10 +426,8 @@ public abstract class ControlMessage extends WireMessage {
         ResultSet resultSet,
         long maxRows,
         QueryMode mode,
-        SettableFuture<Boolean> prefixSent,
         boolean hasData) {
-      return new SendResultSetRunnable(
-          describedResult, resultSet, maxRows, mode, true, prefixSent, hasData);
+      return new SendResultSetRunnable(describedResult, resultSet, maxRows, mode, true, hasData);
     }
 
     static SendResultSetRunnable forPartition(
@@ -424,10 +435,9 @@ public abstract class ControlMessage extends WireMessage {
         BatchReadOnlyTransaction batchReadOnlyTransaction,
         Partition partition,
         QueryMode mode,
-        boolean includePrefix,
-        SettableFuture<Boolean> prefixSent) {
+        CountDownLatch binaryCopyHeaderSentLatch) {
       return new SendResultSetRunnable(
-          describedResult, batchReadOnlyTransaction, partition, mode, includePrefix, prefixSent);
+          describedResult, batchReadOnlyTransaction, partition, mode, binaryCopyHeaderSentLatch);
     }
 
     private SendResultSetRunnable(
@@ -436,7 +446,6 @@ public abstract class ControlMessage extends WireMessage {
         long maxRows,
         QueryMode mode,
         boolean includePrefix,
-        SettableFuture<Boolean> prefixSent,
         boolean hasData) {
       this.describedResult = describedResult;
       this.resultSet = resultSet;
@@ -453,8 +462,7 @@ public abstract class ControlMessage extends WireMessage {
       this.partition = null;
       this.maxRows = maxRows;
       this.mode = mode;
-      this.includePrefix = includePrefix;
-      this.prefixSent = prefixSent;
+      this.binaryCopyHeaderSentLatch = new CountDownLatch(0);
       this.hasData = hasData;
     }
 
@@ -463,16 +471,14 @@ public abstract class ControlMessage extends WireMessage {
         BatchReadOnlyTransaction batchReadOnlyTransaction,
         Partition partition,
         QueryMode mode,
-        boolean includePrefix,
-        SettableFuture<Boolean> prefixSent) {
+        CountDownLatch binaryCopyHeaderSentLatch) {
       this.describedResult = describedResult;
       this.resultSet = null;
       this.batchReadOnlyTransaction = batchReadOnlyTransaction;
       this.partition = partition;
       this.maxRows = 0L;
       this.mode = mode;
-      this.includePrefix = includePrefix;
-      this.prefixSent = prefixSent;
+      this.binaryCopyHeaderSentLatch = binaryCopyHeaderSentLatch;
       this.hasData = false;
     }
 
@@ -482,52 +488,30 @@ public abstract class ControlMessage extends WireMessage {
         if (resultSet == null && batchReadOnlyTransaction != null && partition != null) {
           // Note: It is OK not to close this result set, as the underlying transaction and session
           // will be cleaned up at a later moment.
-          try {
-            resultSet = batchReadOnlyTransaction.execute(partition);
-            converter =
-                new Converter(
-                    describedResult,
-                    mode,
-                    describedResult.getConnectionHandler().getServer().getOptions(),
-                    resultSet,
-                    includePrefix);
-            hasData = resultSet.next();
-          } catch (Throwable t) {
-            if (includePrefix) {
-              synchronized (describedResult) {
-                prefixSent.setException(t);
-              }
-            }
-            throw t;
-          }
-        }
-        if (includePrefix) {
-          try {
-            for (WireOutput prefix : describedResult.createResultPrefix(resultSet)) {
-              prefix.send(false);
-            }
-          } catch (Throwable t) {
-            prefixSent.setException(t);
-            throw t;
-          }
-        }
-        // Wait until the prefix (if any) has been sent.
-        if (!includePrefix) {
-          prefixSent.get();
+          resultSet = batchReadOnlyTransaction.execute(partition);
+          hasData = resultSet.next();
+          converter =
+              new Converter(
+                  describedResult,
+                  mode,
+                  describedResult.getConnectionHandler().getServer().getOptions(),
+                  resultSet,
+                  false);
         }
         long rows = 0L;
         while (hasData) {
-          if (Thread.interrupted()) {
-            throw PGExceptionFactory.newQueryCancelledException();
-          }
           WireOutput wireOutput = describedResult.createDataRowResponse(converter);
+          if (!converter.isIncludeBinaryCopyHeaderInFirstRow()) {
+            binaryCopyHeaderSentLatch.await();
+          }
           synchronized (describedResult) {
             wireOutput.send(false);
           }
-          // The prefix will be included in the first data row.
-          if (includePrefix && rows == 0L) {
-            prefixSent.set(true);
+          binaryCopyHeaderSentLatch.countDown();
+          if (Thread.interrupted()) {
+            throw PGExceptionFactory.newQueryCancelledException();
           }
+
           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 41417504..d151f407 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/CopyOutMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/CopyOutMockServerTest.java
@@ -35,6 +35,7 @@ import com.google.cloud.spanner.Value;
 import com.google.cloud.spanner.connection.RandomResultSetGenerator;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.session.SessionState;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
 import com.google.cloud.spanner.pgadapter.statements.CopyStatement.Format;
 import com.google.cloud.spanner.pgadapter.utils.CopyInParser;
 import com.google.cloud.spanner.pgadapter.utils.CopyRecord;
@@ -577,63 +578,65 @@ public class CopyOutMockServerTest extends AbstractMockServerTest {
 
   @Test
   public void testCopyOutPartitioned() 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));
-      StringWriter writer = new StringWriter();
-      long rows = copyManager.copyOut("COPY random TO STDOUT", writer);
-
-      assertEquals(expectedRowCount, rows);
-
-      try (Scanner scanner = new Scanner(writer.toString())) {
-        int lineCount = 0;
-        while (scanner.hasNextLine()) {
-          lineCount++;
-          String line = scanner.nextLine();
-          String[] columns = line.split("\t");
-          int index = findIndex(resultSet, columns);
-          assertNotEquals(String.format("Row %d not found: %s", lineCount, line), -1, index);
+    for (int expectedRowCount : new int[] {0, 1, 2, 3, 5, BackendConnection.MAX_PARTITIONS, 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));
+        StringWriter writer = new StringWriter();
+        long rows = copyManager.copyOut("COPY random TO STDOUT", writer);
+
+        assertEquals(expectedRowCount, rows);
+
+        try (Scanner scanner = new Scanner(writer.toString())) {
+          int lineCount = 0;
+          while (scanner.hasNextLine()) {
+            lineCount++;
+            String line = scanner.nextLine();
+            String[] columns = line.split("\t");
+            int index = findIndex(resultSet, columns);
+            assertNotEquals(String.format("Row %d not found: %s", lineCount, line), -1, index);
+          }
+          assertEquals(expectedRowCount, lineCount);
         }
-        assertEquals(expectedRowCount, lineCount);
       }
     }
   }
 
   @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);
+    for (int expectedRowCount : new int[] {0, 1, 2, 3, 5, BackendConnection.MAX_PARTITIONS, 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);
       }
-      assertEquals(expectedRowCount, recordCount);
     }
   }
 
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 39979910..095dd56f 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITPsqlTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITPsqlTest.java
@@ -51,6 +51,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.Random;
 import java.util.Scanner;
+import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -68,6 +69,8 @@ import org.junit.runners.JUnit4;
 @Category(IntegrationTest.class)
 @RunWith(JUnit4.class)
 public class ITPsqlTest implements IntegrationTest {
+  private static final Logger logger = Logger.getLogger(ITPsqlTest.class.getName());
+
   private static final PgAdapterTestEnv testEnv = new PgAdapterTestEnv();
   private static Database database;
   private static boolean allAssumptionsPassed = false;
@@ -402,6 +405,7 @@ public class ITPsqlTest implements IntegrationTest {
   public void testCopyBetweenPostgreSQLAndCloudSpanner() throws Exception {
     int numRows = 100;
 
+    logger.info("Copying initial data to PG");
     // Generate 99 random rows.
     copyRandomRowsToPostgreSQL(numRows - 1);
     // Also add one row with all nulls to ensure that nulls are also copied correctly.
@@ -414,6 +418,7 @@ public class ITPsqlTest implements IntegrationTest {
       }
     }
 
+    logger.info("Verifying initial data in PG");
     // Verify that we have 100 rows in PostgreSQL.
     try (Connection connection =
         DriverManager.getConnection(createJdbcUrlForLocalPg(), POSTGRES_USER, POSTGRES_PASSWORD)) {
@@ -427,10 +432,12 @@ public class ITPsqlTest implements IntegrationTest {
 
     // Execute the COPY tests in both binary and text mode.
     for (boolean binary : new boolean[] {false, true}) {
+      logger.info("Testing binary: " + binary);
       // 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())));
 
+      logger.info("Copying rows to Spanner");
       // COPY the rows to Cloud Spanner.
       ProcessBuilder builder = new ProcessBuilder();
       builder.command(
@@ -462,6 +469,7 @@ public class ITPsqlTest implements IntegrationTest {
       int res = process.waitFor();
       assertEquals(0, res);
 
+      logger.info("Verifying data in Spanner");
       // Verify that we now also have 100 rows in Spanner.
       try (Connection connection = DriverManager.getConnection(createJdbcUrlForPGAdapter())) {
         try (ResultSet resultSet =
@@ -472,9 +480,11 @@ public class ITPsqlTest implements IntegrationTest {
         }
       }
 
+      logger.info("Comparing table contents");
       // Verify that the rows in both databases are equal.
       compareTableContents();
 
+      logger.info("Deleting all data in PG");
       // Remove all rows in the table in the local PostgreSQL database and then copy everything from
       // Cloud Spanner to PostgreSQL.
       try (Connection connection =
@@ -483,6 +493,7 @@ public class ITPsqlTest implements IntegrationTest {
         assertEquals(numRows, connection.createStatement().executeUpdate("delete from all_types"));
       }
 
+      logger.info("Copying rows to PG");
       // COPY the rows from Cloud Spanner to PostgreSQL.
       ProcessBuilder copyToPostgresBuilder = new ProcessBuilder();
       copyToPostgresBuilder.command(
@@ -522,6 +533,7 @@ public class ITPsqlTest implements IntegrationTest {
       assertEquals("", errors.toString());
       assertEquals(0, copyToPostgresResult);
 
+      logger.info("Compare table contents");
       // Compare table contents again.
       compareTableContents();
     }
@@ -530,6 +542,7 @@ 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 {
+    logger.info("Deleting all data in PG");
     // Remove all rows in the table in the local PostgreSQL database.
     try (Connection connection =
         DriverManager.getConnection(createJdbcUrlForLocalPg(), POSTGRES_USER, POSTGRES_PASSWORD)) {
@@ -538,10 +551,12 @@ public class ITPsqlTest implements IntegrationTest {
 
     // Execute the COPY tests in both binary and text mode.
     for (boolean binary : new boolean[] {false, true}) {
+      logger.info("Testing binary: " + binary);
       // 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())));
 
+      logger.info("Copy empty table to CS");
       // COPY the empty table to Cloud Spanner.
       ProcessBuilder builder = new ProcessBuilder();
       builder.command(
@@ -573,6 +588,7 @@ public class ITPsqlTest implements IntegrationTest {
       int res = process.waitFor();
       assertEquals(0, res);
 
+      logger.info("Verify that CS is empty");
       // Verify that still have 0 rows in Cloud Spanner.
       try (Connection connection = DriverManager.getConnection(createJdbcUrlForPGAdapter())) {
         try (ResultSet resultSet =
@@ -583,6 +599,7 @@ public class ITPsqlTest implements IntegrationTest {
         }
       }
 
+      logger.info("Copy empty table to PG");
       // COPY the empty table from Cloud Spanner to PostgreSQL.
       ProcessBuilder copyToPostgresBuilder = new ProcessBuilder();
       copyToPostgresBuilder.command(