Repository: GoogleCloudPlatform/pgadapter
Pull request #217 (2022-06-20 05:02:10)

feat: support COPY statement in a batch of sql statements
Adds support for `COPY` statements in query messages with multiple SQL statements. This also adds support for `COPY` statements in the extended query protocol, but there does not seem to be any clients that actually use that. This change removes the special handling of `COPY` statements and incorporates it into the existing generic handling of other statements. This means that the actual execution of the statement is delayed and moved to the `BackendConnection`. One major difference with other statements though, is that a `COPY` statement requires an auto-flush of the current pipeline of messages, as the `COPY` operation will start a sub-protocol and require messages to be sent back and forth.


codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#217](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (73257f9) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/9ec8e67dadbec2e46fa5e5ebdd76e869fe09077d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (9ec8e67) will **decrease** coverage by `0.51%`.
> The diff coverage is `85.47%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #217      +/-   ##
========================================================
- Coverage                 79.73%   79.22%   -0.52%     
- Complexity                 1041     1062      +21     
========================================================
  Files                        91       92       +1     
  Lines                      3543     3595      +52     
  Branches                    423      433      +10     
========================================================
+ Hits                       2825     2848      +23     
- Misses                      553      577      +24     
- Partials                    165      170       +5     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `79.22% <85.47%> (-0.52%)` | :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/217?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...panner/pgadapter/wireprotocol/DescribeMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9EZXNjcmliZU1lc3NhZ2UuamF2YQ==) | `83.87% <ø> (ø)` | |
| [...d/spanner/pgadapter/wireprotocol/FlushMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9GbHVzaE1lc3NhZ2UuamF2YQ==) | `57.14% <ø> (-12.86%)` | :arrow_down: |
| [...ud/spanner/pgadapter/wireprotocol/SkipMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9Ta2lwTWVzc2FnZS5qYXZh) | `70.00% <70.00%> (ø)` | |
| [...panner/pgadapter/wireprotocol/CopyFailMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9Db3B5RmFpbE1lc3NhZ2UuamF2YQ==) | `70.00% <71.42%> (-5.00%)` | :arrow_down: |
| [...spanner/pgadapter/wireprotocol/ControlMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9Db250cm9sTWVzc2FnZS5qYXZh) | `86.17% <75.00%> (-8.77%)` | :arrow_down: |
| [...panner/pgadapter/wireprotocol/CopyDataMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9Db3B5RGF0YU1lc3NhZ2UuamF2YQ==) | `68.18% <75.00%> (-10.08%)` | :arrow_down: |
| [...ud/spanner/pgadapter/statements/CopyStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQ29weVN0YXRlbWVudC5qYXZh) | `72.54% <82.75%> (+2.09%)` | :arrow_up: |
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `91.59% <84.21%> (-0.83%)` | :arrow_down: |
| [...gle/cloud/spanner/pgadapter/ConnectionHandler.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL0Nvbm5lY3Rpb25IYW5kbGVyLmphdmE=) | `76.23% <90.90%> (-0.30%)` | :arrow_down: |
| [...er/pgadapter/statements/IntermediateStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvSW50ZXJtZWRpYXRlU3RhdGVtZW50LmphdmE=) | `95.06% <94.73%> (+0.61%)` | :arrow_up: |
| ... and [12 more](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform). Last update [9ec8e67...73257f9](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/217?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform).


Patch:
diff --git a/.ci/evaluate-with-psql.sh b/.ci/evaluate-with-psql.sh
index 5a14e708..b0238cbe 100644
--- a/.ci/evaluate-with-psql.sh
+++ b/.ci/evaluate-with-psql.sh
@@ -131,6 +131,17 @@ echo "------Test \"COPY FROM STDIN\"------"
 17	17	17
 \.
 EOF
-echo "SELECT * FROM users;" | /usr/lib/postgresql/"${PSQL_VERSION}"/bin/psql -h localhost -p 4242 -d "${GOOGLE_CLOUD_DATABASE_WITH_VERSION}" > .ci/e2e-result/copy-from-stdin.txt
+echo "SELECT * FROM users order by id;" | /usr/lib/postgresql/"${PSQL_VERSION}"/bin/psql -h localhost -p 4242 -d "${GOOGLE_CLOUD_DATABASE_WITH_VERSION}" > .ci/e2e-result/copy-from-stdin.txt
 diff -i -w -s .ci/e2e-result/copy-from-stdin.txt .ci/e2e-expected/copy-from-stdin.txt
-RETURN_CODE=$((${RETURN_CODE}||$?))
\ No newline at end of file
+RETURN_CODE=$((${RETURN_CODE}||$?))
+
+echo "------Test \"COPY FROM STDIN in batch\"------"
+/usr/lib/postgresql/"${PSQL_VERSION}"/bin/psql -h localhost -p 4242 -d "${GOOGLE_CLOUD_DATABASE_WITH_VERSION}" -c "select 0; COPY users FROM STDIN; COPY users FROM STDIN; select 1;" <<EOF
+18	18	18
+\.
+19	19	19
+\.
+EOF
+echo "SELECT * FROM users order by id;" | /usr/lib/postgresql/"${PSQL_VERSION}"/bin/psql -h localhost -p 4242 -d "${GOOGLE_CLOUD_DATABASE_WITH_VERSION}" > .ci/e2e-result/copy-from-stdin-batch.txt
+diff -i -w -s .ci/e2e-result/copy-from-stdin-batch.txt .ci/e2e-expected/copy-from-stdin-batch.txt
+RETURN_CODE=$((${RETURN_CODE}||$?))
diff --git a/README.md b/README.md
index 254dc014..f3fa1062 100644
--- a/README.md
+++ b/README.md
@@ -293,7 +293,7 @@ Other `psql` meta-commands are __not__ supported.
 Spanner database. `COPY` operations are atomic by default, but the standard transaction limits of
 Cloud Spanner apply to these transactions. That means that at most 20,000 mutations can be included
 in one `COPY` operation. `COPY` can also be executed in non-atomic mode by executing the statement
-`SET AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing the copy operation.
+`SET SPANNER.AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing the copy operation.
 
 Although only `STDIN` is supported, export files can still be imported using `COPY` by piping files
 into `psql`. See the examples below.
@@ -315,7 +315,7 @@ create table numbers (number bigint not null primary key, name varchar);
 ```
 
 ```shell
-cat numbers.txt | psql -h localhost -d test-db -c "set autocommit_dml_mode='partitioned_non_atomic'" -c "copy numbers from stdin;"
+cat numbers.txt | psql -h localhost -d test-db -c "set spanner.autocommit_dml_mode='partitioned_non_atomic'; copy numbers from stdin;"
 ```
 
 The above operation will automatically split the data over multiple transactions if the file
@@ -348,8 +348,7 @@ Larger datasets require that the Cloud Spanner database is set to `PARTITIONED_N
 psql -h localhost -p 5432 -d my-local-db \
   -c "copy (select i, to_char(i, 'fm000') from generate_series(1, 1000000) s(i)) to stdout" \
   | psql -h localhost -p 5433 -d my-spanner-db \
-  -c "set autocommit_dml_mode='partitioned_non_atomic'" \
-  -c "copy numbers from stdin;"
+  -c "set spanner.autocommit_dml_mode='partitioned_non_atomic'; copy numbers from stdin;"
 ```
 
 ## Limitations
diff --git a/pom.xml b/pom.xml
index 9108a513..8a580a7d 100644
--- a/pom.xml
+++ b/pom.xml
@@ -302,6 +302,7 @@
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-surefire-plugin</artifactId>
         <configuration>
+          <forkedProcessExitTimeoutInSeconds>3</forkedProcessExitTimeoutInSeconds>
           <forkedProcessTimeoutInSeconds>3600</forkedProcessTimeoutInSeconds>
           <argLine>${surefire.jacoco.args}</argLine>
           <excludedGroups>${excludedTests}</excludedGroups>
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
index abd0eb43..f98da5a4 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -98,6 +98,7 @@ public class ConnectionHandler extends Thread {
   private static final AtomicInteger incrementingConnectionId = new AtomicInteger(0);
   private ConnectionMetadata connectionMetadata;
   private WireMessage message;
+  private int invalidMessagesCount;
   private Connection spannerConnection;
   private DatabaseId databaseId;
   private WellKnownClient wellKnownClient;
@@ -251,15 +252,7 @@ public class ConnectionHandler extends Thread {
           }
         }
         while (this.status != ConnectionStatus.TERMINATED) {
-          try {
-            message.nextHandler();
-            message.send();
-          } catch (IllegalArgumentException | IllegalStateException | EOFException fatalException) {
-            this.handleError(output, fatalException);
-            this.status = ConnectionStatus.TERMINATED;
-          } catch (Exception e) {
-            this.handleError(output, e);
-          }
+          handleMessages();
         }
       } catch (Exception e) {
         this.handleError(output, e);
@@ -293,6 +286,23 @@ public class ConnectionHandler extends Thread {
     }
   }
 
+  /**
+   * Reads and handles wire-protocol messages. This method is normally only called from this {@link
+   * ConnectionHandler}, but certain sub-protocols such as the COPY protocol also need to process
+   * messages in line.
+   */
+  public void handleMessages() throws Exception {
+    try {
+      message.nextHandler();
+      message.send();
+    } catch (IllegalArgumentException | IllegalStateException | EOFException fatalException) {
+      this.handleError(getConnectionMetadata().getOutputStream(), fatalException);
+      this.status = ConnectionStatus.TERMINATED;
+    } catch (Exception e) {
+      this.handleError(getConnectionMetadata().getOutputStream(), e);
+    }
+  }
+
   /** Called when a Terminate message is received. This closes this {@link ConnectionHandler}. */
   public void handleTerminate() {
     closeAllPortals();
@@ -307,42 +317,41 @@ public class ConnectionHandler extends Thread {
    * shutting down while the connection is still active.
    */
   void terminate() {
-    if (this.status != ConnectionStatus.TERMINATED) {
-      handleTerminate();
-      try {
-        if (!socket.isClosed()) {
-          socket.close();
-        }
-      } catch (IOException exception) {
-        logger.log(
-            Level.WARNING,
-            exception,
-            () ->
-                String.format(
-                    "Failed to close connection handler with ID %s: %s", getName(), exception));
+    handleTerminate();
+    try {
+      if (!socket.isClosed()) {
+        socket.close();
       }
+    } catch (IOException exception) {
+      logger.log(
+          Level.WARNING,
+          exception,
+          () ->
+              String.format(
+                  "Failed to close connection handler with ID %s: %s", getName(), exception));
     }
   }
 
   /**
    * Takes an Exception Object and relates its results to the user within the client.
    *
-   * @param e The exception to be related.
+   * @param exception The exception to be related.
    * @throws IOException if there is some issue in the sending of the error messages.
    */
-  private void handleError(DataOutputStream output, Exception e) throws Exception {
+  private void handleError(DataOutputStream output, Exception exception) throws Exception {
     logger.log(
         Level.WARNING,
-        e,
-        () -> String.format("Exception on connection handler with ID %s: %s", getName(), e));
+        exception,
+        () ->
+            String.format("Exception on connection handler with ID %s: %s", getName(), exception));
     if (this.status == ConnectionStatus.TERMINATED) {
-      new ErrorResponse(output, e, ErrorResponse.State.InternalError, Severity.FATAL).send();
+      new ErrorResponse(output, exception, ErrorResponse.State.InternalError, Severity.FATAL)
+          .send();
       new TerminateResponse(output).send();
     } else if (this.status == ConnectionStatus.COPY_IN) {
-      new ErrorResponse(output, e, ErrorResponse.State.InternalError).send();
+      new ErrorResponse(output, exception, ErrorResponse.State.InternalError).send();
     } else {
-      this.status = ConnectionStatus.AUTHENTICATED;
-      new ErrorResponse(output, e, ErrorResponse.State.InternalError).send();
+      new ErrorResponse(output, exception, ErrorResponse.State.InternalError).send();
       new ReadyResponse(output, ReadyResponse.Status.IDLE).send();
     }
   }
@@ -497,6 +506,31 @@ public class ConnectionHandler extends Thread {
     this.message = this.server.recordMessage(message);
   }
 
+  /**
+   * Returns the number of invalid messages that this connection has received in a row. This can for
+   * example happen if a client has initiated a COPY operation and the copy operation fails on the
+   * server. The server will then respond with an error response, but if the client fails to read
+   * that message and continues to send copy data messages, the server could get flooded. This value
+   * is used to detect such a situation and breaks the connection if too many invalid messages in a
+   * row are received.
+   */
+  public int getInvalidMessageCount() {
+    return this.invalidMessagesCount;
+  }
+
+  /** Increases the number of invalid messages that was received in a row by 1. */
+  public void increaseInvalidMessageCount() {
+    this.invalidMessagesCount++;
+  }
+
+  /**
+   * Clears the number of invalid messages that was received. This is called whenever a valid
+   * message is encountered.
+   */
+  public void clearInvalidMessageCount() {
+    this.invalidMessagesCount = 0;
+  }
+
   public ConnectionMetadata getConnectionMetadata() {
     return connectionMetadata;
   }
@@ -530,6 +564,8 @@ public class ConnectionHandler extends Thread {
     UNAUTHENTICATED,
     AUTHENTICATED,
     COPY_IN,
+    COPY_DONE,
+    COPY_FAILED,
     TERMINATED,
   }
 
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 d233d994..7eaddfa9 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
@@ -14,6 +14,8 @@
 
 package com.google.cloud.spanner.pgadapter.statements;
 
+import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.COPY;
+
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.ResultSet;
@@ -28,15 +30,25 @@ import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.DdlTransactionMode;
 import com.google.cloud.spanner.pgadapter.statements.local.LocalStatement;
+import com.google.cloud.spanner.pgadapter.utils.CopyDataReceiver;
+import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
+import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse.Status;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+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 java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.LinkedList;
 import java.util.Objects;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.stream.Collectors;
 
@@ -154,6 +166,57 @@ public class BackendConnection {
     }
   }
 
+  /**
+   * This statement represents a COPY table FROM STDIN statement. This has no one-on-one mapping
+   * with a Cloud Spanner SQL statement and is therefore executed using a custom {@link
+   * MutationWriter}. As the COPY implementation uses mutations instead of DML, it has slightly
+   * different transaction semantics than in real PostgreSQL. A COPY operation will by default be
+   * atomic, but can be configured to behave non-atomically for large batches. Also, if a COPY
+   * operation is executed in a transaction (both implicit and explicit), it will commit the
+   * transaction when the COPY operation is done. This is required to flush the mutations to the
+   * database.
+   */
+  private final class Copy extends BufferedStatement<StatementResult> {
+    private final CopyDataReceiver copyDataReceiver;
+    private final MutationWriter mutationWriter;
+    private final ListeningExecutorService executor;
+
+    Copy(
+        ParsedStatement parsedStatement,
+        Statement statement,
+        CopyDataReceiver copyDataReceiver,
+        MutationWriter mutationWriter,
+        ExecutorService executor) {
+      super(parsedStatement, statement);
+      this.copyDataReceiver = copyDataReceiver;
+      this.mutationWriter = mutationWriter;
+      this.executor = MoreExecutors.listeningDecorator(executor);
+    }
+
+    @Override
+    void execute() {
+      try {
+        checkConnectionState();
+        // Execute the MutationWriter and the CopyDataReceiver both asynchronously and wait for both
+        // to finish before continuing with the next statement. This ensures that all statements are
+        // applied in sequential order.
+        ListenableFuture<StatementResult> statementResultFuture = executor.submit(mutationWriter);
+        ListenableFuture<Void> copyDataReceiverFuture = executor.submit(copyDataReceiver);
+        this.result.setFuture(statementResultFuture);
+        Futures.allAsList(copyDataReceiverFuture, statementResultFuture).get();
+      } catch (ExecutionException executionException) {
+        result.setException(executionException.getCause());
+        throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
+      } catch (InterruptedException interruptedException) {
+        result.setException(SpannerExceptionFactory.propagateInterrupt(interruptedException));
+        throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
+      } catch (Exception exception) {
+        result.setException(exception);
+        throw exception;
+      }
+    }
+  }
+
   private static final ImmutableMap<String, LocalStatement> EMPTY_LOCAL_STATEMENTS =
       ImmutableMap.of();
   private static final StatementResult NO_RESULT = new NoResult();
@@ -205,6 +268,22 @@ public class BackendConnection {
     return execute.result;
   }
 
+  /**
+   * Buffers the given COPY operation for execution on the backend connection when the next
+   * flush/sync message is received. The returned future will contain the result of the COPY
+   * operation when execution has finished.
+   */
+  public Future<StatementResult> executeCopy(
+      ParsedStatement parsedStatement,
+      Statement statement,
+      CopyDataReceiver copyDataReceiver,
+      MutationWriter mutationWriter,
+      ExecutorService executor) {
+    Copy copy = new Copy(parsedStatement, statement, copyDataReceiver, mutationWriter, executor);
+    bufferedStatements.add(copy);
+    return copy.result;
+  }
+
   /** Flushes the buffered statements to Spanner. */
   void flush() {
     flush(false);
@@ -298,7 +377,7 @@ public class BackendConnection {
     // We need to start an implicit transaction.
     // Check if a read-only transaction suffices.
     spannerConnection.beginTransaction();
-    if (isSync && !hasDmlStatementsAfter(index)) {
+    if (isSync && !hasDmlOrCopyStatementsAfter(index)) {
       spannerConnection.setTransactionMode(
           com.google.cloud.spanner.connection.TransactionMode.READ_ONLY_TRANSACTION);
     }
@@ -431,14 +510,20 @@ public class BackendConnection {
     return isBegin(index) || isCommit(index) || isRollback(index);
   }
 
-  private boolean hasDmlStatementsAfter(int index) {
+  private boolean hasOnlyDmlStatementsAfter(int index) {
     return bufferedStatements.subList(index, bufferedStatements.size()).stream()
-        .anyMatch(statement -> statement.parsedStatement.getType() == StatementType.UPDATE);
+        .allMatch(statement -> statement.parsedStatement.getType() == StatementType.UPDATE);
   }
 
-  private boolean hasOnlyDmlStatementsAfter(int index) {
+  @VisibleForTesting
+  boolean hasDmlOrCopyStatementsAfter(int index) {
     return bufferedStatements.subList(index, bufferedStatements.size()).stream()
-        .allMatch(statement -> statement.parsedStatement.getType() == StatementType.UPDATE);
+        .anyMatch(
+            statement ->
+                statement.parsedStatement.getType() == StatementType.UPDATE
+                    || statement.parsedStatement.getType() == StatementType.UNKNOWN
+                        && StatementParser.isCommand(
+                            COPY, statement.parsedStatement.getSqlWithoutComments()));
   }
 
   private int getStatementCount() {
@@ -450,11 +535,11 @@ public class BackendConnection {
   }
 
   private boolean canBeBatchedTogether(StatementType statementType1, StatementType statementType2) {
-    if (Objects.equals(statementType1, StatementType.QUERY)
-        || Objects.equals(statementType1, StatementType.CLIENT_SIDE)) {
-      return false;
+    if (Objects.equals(statementType1, StatementType.DDL)
+        || Objects.equals(statementType2, StatementType.UPDATE)) {
+      return Objects.equals(statementType1, statementType2);
     }
-    return Objects.equals(statementType1, statementType2);
+    return false;
   }
 
   /**
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 f2a2375f..c211427e 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
@@ -30,6 +30,7 @@ import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser;
 import com.google.cloud.spanner.pgadapter.parsers.copy.TokenMgrError;
+import com.google.cloud.spanner.pgadapter.utils.CopyDataReceiver;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter.CopyTransactionMode;
 import com.google.common.base.Strings;
@@ -38,17 +39,19 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.apache.commons.csv.CSVFormat;
 
+/**
+ * {@link CopyStatement} models a `COPY table FROM STDIN` statement. The same class is used both as
+ * an {@link IntermediatePreparedStatement} and {@link IntermediatePortalStatement}, as COPY does
+ * not support any statement parameters, which means that there is no difference between the two.
+ */
 @InternalApi
-public class CopyStatement extends IntermediateStatement {
-
+public class CopyStatement extends IntermediatePortalStatement {
   private static final String COLUMN_NAME = "column_name";
   private static final String DATA_TYPE = "data_type";
 
@@ -59,22 +62,43 @@ public class CopyStatement extends IntermediateStatement {
   private Map<String, TypeCode> tableColumns;
   private int indexedColumnsCount;
   private MutationWriter mutationWriter;
-  private Future<Long> updateCount;
-  private final ExecutorService executor = Executors.newSingleThreadExecutor();
+  // We need two threads to execute a CopyStatement:
+  // 1. A thread to receive the CopyData messages from the client.
+  // 2. A thread to write the received data to Cloud Spanner.
+  // The two are kept separate to allow each to execute at the maximum speed that it can.
+  // The link between them is a piped output/input stream that ensures that backpressure is applied
+  // to the client if the client is sending data at a higher speed than that the writer can handle.
+  private final ExecutorService executor = Executors.newFixedThreadPool(2);
 
   public CopyStatement(
       ConnectionHandler connectionHandler,
       OptionsMetadata options,
+      String name,
       ParsedStatement parsedStatement,
       Statement originalStatement) {
-    super(connectionHandler, options, parsedStatement, originalStatement);
+    super(connectionHandler, options, name, parsedStatement, originalStatement);
   }
 
-  public Exception getException() {
+  @Override
+  public boolean hasException() {
+    return this.exception != null;
+  }
+
+  @Override
+  public SpannerException getException() {
     // Do not clear exceptions on a CopyStatement.
     return this.exception;
   }
 
+  @Override
+  public long getUpdateCount() {
+    // COPY statements continue to execute while the server continues to receive a stream of
+    // CopyData messages AFTER we have received a flush/sync message. We therefore need to block
+    // while waiting for the update count to be available once the server has received a CopyDone
+    // or CopyFailed message.
+    return getUpdateCount(ResultNotReadyBehavior.BLOCK);
+  }
+
   @Override
   public void close() throws Exception {
     if (this.mutationWriter != null) {
@@ -84,16 +108,6 @@ public class CopyStatement extends IntermediateStatement {
     super.close();
   }
 
-  public long getUpdateCount() {
-    try {
-      return updateCount.get();
-    } catch (ExecutionException e) {
-      throw SpannerExceptionFactory.asSpannerException(e.getCause());
-    } catch (InterruptedException e) {
-      throw SpannerExceptionFactory.propagateInterrupt(e);
-    }
-  }
-
   @Override
   public StatementType getStatementType() {
     return StatementType.UPDATE;
@@ -243,7 +257,7 @@ public class CopyStatement extends IntermediateStatement {
             .bind("p1")
             .to(getTableName())
             .build();
-    try (ResultSet result = connection.executeQuery(statement)) {
+    try (ResultSet result = connection.getDatabaseClient().singleUse().executeQuery(statement)) {
       while (result.next()) {
         String columnName = result.getString(COLUMN_NAME);
         TypeCode type = parsePostgreSQLDataType(result.getString(DATA_TYPE));
@@ -282,7 +296,7 @@ public class CopyStatement extends IntermediateStatement {
       paramIndex++;
     }
     Statement statement = builder.build();
-    try (ResultSet resultSet = connection.executeQuery(statement)) {
+    try (ResultSet resultSet = connection.getDatabaseClient().singleUse().executeQuery(statement)) {
       if (resultSet.next()) {
         return (int) resultSet.getLong(0);
       }
@@ -290,13 +304,24 @@ public class CopyStatement extends IntermediateStatement {
     return 0;
   }
 
+  @Override
+  public IntermediatePortalStatement bind(
+      String name,
+      byte[][] parameters,
+      List<Short> parameterFormatCodes,
+      List<Short> resultFormatCodes) {
+    // COPY does not support binding any parameters, so we just return the same statement.
+    return this;
+  }
+
   @Override
   public void handleExecutionException(SpannerException exception) {
     executor.shutdownNow();
     super.handleExecutionException(exception);
   }
 
-  public void execute() {
+  @Override
+  public void executeAsync(BackendConnection backendConnection) {
     this.executed = true;
     try {
       parseCopyStatement();
@@ -311,7 +336,13 @@ public class CopyStatement extends IntermediateStatement {
               indexedColumnsCount,
               getParserFormat(),
               hasHeader());
-      updateCount = executor.submit(mutationWriter);
+      setFutureStatementResult(
+          backendConnection.executeCopy(
+              parsedStatement,
+              statement,
+              new CopyDataReceiver(this, this.connectionHandler),
+              mutationWriter,
+              executor));
     } catch (Exception e) {
       SpannerException spannerException = SpannerExceptionFactory.asSpannerException(e);
       handleExecutionException(spannerException);
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java
index 3ea5e098..d823b095 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java
@@ -15,6 +15,7 @@
 package com.google.cloud.spanner.pgadapter.statements;
 
 import static com.google.cloud.spanner.pgadapter.utils.StatementParser.splitStatements;
+import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.COPY;
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.Dialect;
@@ -26,10 +27,12 @@ import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.commands.Command;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
+import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.cloud.spanner.pgadapter.wireprotocol.BindMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage.ManuallyCreatedToken;
 import com.google.cloud.spanner.pgadapter.wireprotocol.DescribeMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage;
+import com.google.cloud.spanner.pgadapter.wireprotocol.FlushMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.SyncMessage;
 import com.google.common.annotations.VisibleForTesting;
@@ -76,10 +79,20 @@ public class SimpleQueryStatement {
         // The original statement was replaced.
         originalStatement = Statement.of(parsedStatement.getSqlWithoutComments());
       }
+      // We need to flush the entire pipeline if we encounter a COPY statement, as COPY statements
+      // require additional messages to be sent back and forth, and this ensures that we get
+      // everything in the correct order.
+      boolean isCopy = StatementParser.isCommand(COPY, parsedStatement.getSqlWithoutComments());
+      if (isCopy) {
+        new FlushMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN).send();
+      }
       new ParseMessage(connectionHandler, parsedStatement, originalStatement).send();
       new BindMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN).send();
       new DescribeMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN).send();
       new ExecuteMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN).send();
+      if (isCopy) {
+        new FlushMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN).send();
+      }
     }
     new SyncMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN).send();
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiver.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiver.java
new file mode 100644
index 00000000..a349b2e5
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiver.java
@@ -0,0 +1,92 @@
+// Copyright 2022 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.utils;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
+import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
+import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement.ResultNotReadyBehavior;
+import com.google.cloud.spanner.pgadapter.wireoutput.CopyInResponse;
+import com.google.common.annotations.VisibleForTesting;
+import java.util.concurrent.Callable;
+
+/**
+ * This message receiver is activated when the COPY sub-protocol is activated by a COPY statement.
+ * This handler will run and receive all incoming messages until the state of the connection changes
+ * from COPY_IN to either COPY_DONE or COPY_FAILED.
+ */
+@InternalApi
+public class CopyDataReceiver implements Callable<Void> {
+  private final CopyStatement copyStatement;
+  private final ConnectionHandler connectionHandler;
+
+  public CopyDataReceiver(CopyStatement copyStatement, ConnectionHandler connectionHandler) {
+    this.copyStatement = copyStatement;
+    this.connectionHandler = connectionHandler;
+  }
+
+  @Override
+  public Void call() throws Exception {
+    handleCopy();
+    return null;
+  }
+
+  /**
+   * Sends a {@link CopyInResponse} to the client and then waits for copy data, done and fail
+   * messages. The incoming data messages are fed into the {@link MutationWriter} that is associated
+   * with this {@link CopyStatement}. The method blocks until it sees a {@link
+   * com.google.cloud.spanner.pgadapter.wireprotocol.CopyDoneMessage} or {@link
+   * com.google.cloud.spanner.pgadapter.wireprotocol.CopyFailMessage}, or until the {@link
+   * MutationWriter} changes the status of the connection to a non-COPY_IN status, for example as a
+   * result of an error while copying the data to Cloud Spanner.
+   */
+  @VisibleForTesting
+  void handleCopy() throws Exception {
+    if (copyStatement.hasException()) {
+      throw copyStatement.getException();
+    } else {
+      this.connectionHandler.addActiveStatement(copyStatement);
+      new CopyInResponse(
+              this.connectionHandler.getConnectionMetadata().getOutputStream(),
+              copyStatement.getTableColumns().size(),
+              copyStatement.getFormatCode())
+          .send();
+      ConnectionStatus initialConnectionStatus = this.connectionHandler.getStatus();
+      try {
+        this.connectionHandler.setStatus(ConnectionStatus.COPY_IN);
+        // Loop here until COPY_IN mode has finished.
+        while (this.connectionHandler.getStatus() == ConnectionStatus.COPY_IN) {
+          this.connectionHandler.handleMessages();
+        }
+        if (copyStatement.hasException(ResultNotReadyBehavior.BLOCK)
+            || this.connectionHandler.getStatus() == ConnectionStatus.COPY_FAILED) {
+          if (copyStatement.hasException(ResultNotReadyBehavior.BLOCK)) {
+            throw copyStatement.getException();
+          } else {
+            throw SpannerExceptionFactory.newSpannerException(
+                ErrorCode.INTERNAL, "Copy failed with unknown reason");
+          }
+        }
+      } finally {
+        this.connectionHandler.removeActiveStatement(copyStatement);
+        this.copyStatement.close();
+        this.connectionHandler.setStatus(initialConnectionStatus);
+      }
+    }
+  }
+}
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java
index b4e6b878..986990e9 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/MutationWriter.java
@@ -28,6 +28,8 @@ import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Value;
 import com.google.cloud.spanner.connection.Connection;
+import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection.UpdateCount;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
@@ -50,6 +52,7 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingDeque;
@@ -66,7 +69,7 @@ import org.apache.commons.csv.CSVRecord;
 import org.postgresql.jdbc.TimestampUtils;
 
 @InternalApi
-public class MutationWriter implements Callable<Long>, Closeable {
+public class MutationWriter implements Callable<StatementResult>, Closeable {
   public enum CopyTransactionMode {
     /**
      * 'Normal' auto-commit mode. The entire COPY operation is atomic. If the number of mutations
@@ -129,7 +132,9 @@ public class MutationWriter implements Callable<Long>, Closeable {
   private final CSVParser parser;
   private PrintWriter errorFileWriter;
   private final PipedOutputStream payload = new PipedOutputStream();
+  private final AtomicBoolean commit = new AtomicBoolean(false);
   private final AtomicBoolean rollback = new AtomicBoolean(false);
+  private final CountDownLatch closedLatch = new CountDownLatch(1);
   private final ListeningExecutorService executorService =
       MoreExecutors.listeningDecorator(Executors.newCachedThreadPool());
 
@@ -199,14 +204,23 @@ public class MutationWriter implements Callable<Long>, Closeable {
     try {
       this.payload.write(payload);
     } catch (IOException e) {
-      SpannerException spannerException =
-          SpannerExceptionFactory.newSpannerException(
-              ErrorCode.INTERNAL, "Could not write copy data to buffer", e);
-      logger.log(Level.SEVERE, spannerException.getMessage(), spannerException);
-      throw spannerException;
+      // Ignore the exception if the executor has already been shutdown. That means that an error
+      // occurred that ended the COPY operation while we were writing data to the buffer.
+      if (!executorService.isShutdown()) {
+        SpannerException spannerException =
+            SpannerExceptionFactory.newSpannerException(
+                ErrorCode.INTERNAL, "Could not write copy data to buffer", e);
+        logger.log(Level.SEVERE, spannerException.getMessage(), spannerException);
+        throw spannerException;
+      }
     }
   }
 
+  /** Indicate that this mutation writer should commit. */
+  public void commit() {
+    this.commit.set(true);
+  }
+
   /**
    * Indicate that this mutation writer should be rolled back. This will not rollback any changes
    * that have already been committed if the mutation writer is running in {@link
@@ -219,10 +233,11 @@ public class MutationWriter implements Callable<Long>, Closeable {
   @Override
   public void close() throws IOException {
     this.payload.close();
+    this.closedLatch.countDown();
   }
 
   @Override
-  public Long call() throws Exception {
+  public StatementResult call() throws Exception {
     // This LinkedBlockingDeque holds a reference to all transactions that are currently active. The
     // max capacity of this deque is what ensures that we never have more than maxParallelism
     // transactions running at the same time. We could also achieve that by using a thread pool with
@@ -304,9 +319,15 @@ public class MutationWriter implements Callable<Long>, Closeable {
       // Write any remaining mutations in the buffer.
       if (!rollback.get() && !mutations.isEmpty()) {
         if (transactionMode == CopyTransactionMode.Explicit) {
-          connection.bufferedWrite(mutations);
+          connection.write(mutations);
         } else {
-          allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
+          // Wait until we have received a CopyDone message before writing the remaining data to
+          // Spanner. If we are in a non-atomic transaction, there might already be data that have
+          // been written to Spanner.
+          closedLatch.await();
+          if (commit.get()) {
+            allCommitFutures.add(writeToSpannerAsync(activeCommitFutures, mutations));
+          }
         }
       }
       // Wait for all commits to finish. We do this even if something went wrong, as it ensures two
@@ -339,7 +360,7 @@ public class MutationWriter implements Callable<Long>, Closeable {
       this.payload.close();
       this.parser.close();
     }
-    return rowCount;
+    return new UpdateCount(rowCount);
   }
 
   private long addMutationAndMaybeFlushTransaction(
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyInResponse.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyInResponse.java
index e3b9b646..588ea9c5 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyInResponse.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/CopyInResponse.java
@@ -52,7 +52,6 @@ public class CopyInResponse extends WireOutput {
     this.outputStream.writeByte(this.formatCode);
     this.outputStream.writeShort(this.numColumns);
     this.outputStream.write(this.columnFormat);
-    this.outputStream.flush();
   }
 
   @Override
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 26f64cad..99251230 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
@@ -25,6 +25,7 @@ import com.google.cloud.spanner.pgadapter.wireoutput.CommandCompleteResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.DataRowResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.EmptyQueryResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse;
+import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse.Severity;
 import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse.State;
 import com.google.cloud.spanner.pgadapter.wireoutput.PortalSuspendedResponse;
 import com.google.common.base.Preconditions;
@@ -40,6 +41,9 @@ import java.util.List;
  */
 @InternalApi
 public abstract class ControlMessage extends WireMessage {
+  /** Maximum number of invalid messages in a row allowed before we terminate the connection. */
+  static final int MAX_INVALID_MESSAGE_COUNT = 50;
+
   /**
    * Token that is used to mark {@link ControlMessage}s that are manually created to execute a
    * {@link QueryMessage}.
@@ -73,50 +77,84 @@ public abstract class ControlMessage extends WireMessage {
    * @throws Exception If construction or reading fails.
    */
   public static ControlMessage create(ConnectionHandler connection) throws Exception {
+    boolean validMessage = true;
     char nextMsg = (char) connection.getConnectionMetadata().getInputStream().readUnsignedByte();
-    if (connection.getStatus() == ConnectionStatus.COPY_IN) {
-      switch (nextMsg) {
-        case CopyDoneMessage.IDENTIFIER:
-          return new CopyDoneMessage(connection);
-        case CopyDataMessage.IDENTIFIER:
-          return new CopyDataMessage(connection);
-        case CopyFailMessage.IDENTIFIER:
-          return new CopyFailMessage(connection);
-        default:
-          // Drop the connection if we receive an invalid message to prevent further CopyData
-          // messages from coming in. There is no error handling in the COPY protocol, and some
-          // clients will blindly continue to send data and never check any possible responses from
-          // the server during a (large) copy operation, so this is the safest option.
-          connection.setStatus(ConnectionStatus.TERMINATED);
-          throw new IllegalStateException(
-              String.format(
-                  "Expected CopyData ('d'), CopyDone ('c') or CopyFail ('f') messages, got: '%c'",
-                  nextMsg));
+    try {
+      if (connection.getStatus() == ConnectionStatus.COPY_IN) {
+        switch (nextMsg) {
+          case CopyDoneMessage.IDENTIFIER:
+            return new CopyDoneMessage(connection);
+          case CopyDataMessage.IDENTIFIER:
+            return new CopyDataMessage(connection);
+          case CopyFailMessage.IDENTIFIER:
+            return new CopyFailMessage(connection);
+          case SyncMessage.IDENTIFIER:
+          case FlushMessage.IDENTIFIER:
+            // Skip sync/flush in COPY_IN. This is consistent with real PostgreSQL which also does
+            // this to accommodate clients that do not check what type of statement they sent in an
+            // ExecuteMessage, and instead always blindly send a flush/sync after each execute.
+            return new SkipMessage(connection);
+          default:
+            // Skip other unexpected messages and throw an exception to fail the copy operation.
+            validMessage = false;
+            new SkipMessage(connection);
+            throw new IllegalStateException(
+                String.format(
+                    "Expected CopyData ('d'), CopyDone ('c') or CopyFail ('f') messages, got: '%c'",
+                    nextMsg));
+        }
+      } else {
+        switch (nextMsg) {
+          case QueryMessage.IDENTIFIER:
+            return new QueryMessage(connection);
+          case ParseMessage.IDENTIFIER:
+            return new ParseMessage(connection);
+          case BindMessage.IDENTIFIER:
+            return new BindMessage(connection);
+          case DescribeMessage.IDENTIFIER:
+            return new DescribeMessage(connection);
+          case ExecuteMessage.IDENTIFIER:
+            return new ExecuteMessage(connection);
+          case CloseMessage.IDENTIFIER:
+            return new CloseMessage(connection);
+          case SyncMessage.IDENTIFIER:
+            return new SyncMessage(connection);
+          case TerminateMessage.IDENTIFIER:
+            return new TerminateMessage(connection);
+          case FunctionCallMessage.IDENTIFIER:
+            return new FunctionCallMessage(connection);
+          case FlushMessage.IDENTIFIER:
+            return new FlushMessage(connection);
+          case CopyDoneMessage.IDENTIFIER:
+          case CopyDataMessage.IDENTIFIER:
+          case CopyFailMessage.IDENTIFIER:
+            // Silently skip COPY messages in non-COPY mode. This is consistent with the PG wire
+            // protocol. If we continue to receive COPY messages while in non-COPY mode, we'll
+            // terminate the connection to prevent the server from being flooded with invalid
+            // messages.
+            validMessage = false;
+            return new SkipMessage(connection);
+          default:
+            throw new IllegalStateException(String.format("Unknown message: %c", nextMsg));
+        }
       }
-    } else {
-      switch (nextMsg) {
-        case QueryMessage.IDENTIFIER:
-          return new QueryMessage(connection);
-        case ParseMessage.IDENTIFIER:
-          return new ParseMessage(connection);
-        case BindMessage.IDENTIFIER:
-          return new BindMessage(connection);
-        case DescribeMessage.IDENTIFIER:
-          return new DescribeMessage(connection);
-        case ExecuteMessage.IDENTIFIER:
-          return new ExecuteMessage(connection);
-        case CloseMessage.IDENTIFIER:
-          return new CloseMessage(connection);
-        case SyncMessage.IDENTIFIER:
-          return new SyncMessage(connection);
-        case TerminateMessage.IDENTIFIER:
-          return new TerminateMessage(connection);
-        case FunctionCallMessage.IDENTIFIER:
-          return new FunctionCallMessage(connection);
-        case FlushMessage.IDENTIFIER:
-          return new FlushMessage(connection);
-        default:
-          throw new IllegalStateException(String.format("Unknown message: %c", nextMsg));
+    } finally {
+      if (validMessage) {
+        connection.clearInvalidMessageCount();
+      } else {
+        connection.increaseInvalidMessageCount();
+        if (connection.getInvalidMessageCount() > MAX_INVALID_MESSAGE_COUNT) {
+          new ErrorResponse(
+                  connection.getConnectionMetadata().getOutputStream(),
+                  new IllegalStateException(
+                      String.format(
+                          "Received %d invalid/unexpected messages. Last received message: '%c'",
+                          connection.getInvalidMessageCount(), nextMsg)),
+                  State.InternalError,
+                  Severity.FATAL)
+              .send();
+          connection.setStatus(ConnectionStatus.TERMINATED);
+        }
       }
     }
   }
@@ -204,6 +242,7 @@ public abstract class ControlMessage extends WireMessage {
         command += ("INSERT".equals(command) ? " 0 " : " ") + statement.getUpdateCount();
         new CommandCompleteResponse(this.outputStream, command).send(false);
         break;
+      case UNKNOWN:
       default:
         throw new IllegalStateException("Unknown statement type: " + statement.getStatement());
     }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDataMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDataMessage.java
index cd1d7c24..25613a83 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDataMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDataMessage.java
@@ -17,7 +17,6 @@ package com.google.cloud.spanner.pgadapter.wireprotocol;
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
-import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
 import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import java.text.MessageFormat;
@@ -48,7 +47,10 @@ public class CopyDataMessage extends ControlMessage {
   @Override
   protected void sendPayload() throws Exception {
     if (statement == null) {
-      throw new IllegalStateException("This connection is no (longer) in COPY IN mode.");
+      // Do not handle this message if there is no CopyStatement available anymore. This means that
+      // the copy operation failed and stopped while the client was still sending data to the
+      // server.
+      return;
     }
     // If backend error occurred during copy-in mode, drop any subsequent CopyData messages.
     MutationWriter mutationWriter = this.statement.getMutationWriter();
@@ -59,12 +61,6 @@ public class CopyDataMessage extends ControlMessage {
         statement.handleExecutionException(exception);
         throw exception;
       }
-    } else {
-      // If we get another CopyData message after we have already failed, then we drop the
-      // connection to ensure that the client is not overflowing us with messages. Some clients do
-      // not check for error messages during a COPY operation.
-      connection.setStatus(ConnectionStatus.TERMINATED);
-      throw statement.getException();
     }
   }
 
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDoneMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDoneMessage.java
index 4c013730..7c2d4e1a 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDoneMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDoneMessage.java
@@ -17,11 +17,7 @@ package com.google.cloud.spanner.pgadapter.wireprotocol;
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
-import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
-import com.google.cloud.spanner.pgadapter.statements.BackendConnection.UpdateCount;
 import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
-import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
-import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
 import java.text.MessageFormat;
 
 /**
@@ -44,26 +40,13 @@ public class CopyDoneMessage extends ControlMessage {
   protected void sendPayload() throws Exception {
     // If backend error occurred during copy-in mode, drop any subsequent CopyDone messages.
     if (this.statement != null) {
-      MutationWriter mutationWriter = this.statement.getMutationWriter();
+      statement.getMutationWriter().commit();
       statement.close();
-      if (!statement.hasException()) {
-        try {
-          long rowCount = this.statement.getUpdateCount();
-          statement.setStatementResult(
-              new UpdateCount(rowCount)); // Set the row count of number of rows copied.
-          this.sendSpannerResult(this.statement, QueryMode.SIMPLE, 0L);
-          this.outputStream.flush();
-        } catch (Exception e) {
-          // Spanner returned an error when trying to commit the batch of mutations.
-          this.connection.setStatus(ConnectionStatus.AUTHENTICATED);
-          this.connection.removeActiveStatement(this.statement);
-          throw e;
-        }
-      }
     }
-    this.connection.setStatus(ConnectionStatus.AUTHENTICATED);
-    this.connection.removeActiveStatement(this.statement);
-    new ReadyResponse(this.outputStream, ReadyResponse.Status.IDLE).send();
+    // Clear the COPY_IN status to indicate that we finished successfully. This will cause the
+    // inline handling of incoming (copy) messages to stop and the server to resume normal
+    // operation.
+    this.connection.setStatus(ConnectionStatus.COPY_DONE);
   }
 
   @Override
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyFailMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyFailMessage.java
index cdb9b9c6..0c55b5d4 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyFailMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyFailMessage.java
@@ -15,13 +15,13 @@
 package com.google.cloud.spanner.pgadapter.wireprotocol;
 
 import com.google.api.core.InternalApi;
+import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
 import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
+import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
-import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse;
-import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse.State;
-import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
 import java.text.MessageFormat;
 
 /**
@@ -32,33 +32,35 @@ import java.text.MessageFormat;
  */
 @InternalApi
 public class CopyFailMessage extends ControlMessage {
-
   protected static final char IDENTIFIER = 'f';
-  private CopyStatement statement;
 
-  private String errorMessage;
+  private final CopyStatement statement;
+  private final String errorMessage;
 
   public CopyFailMessage(ConnectionHandler connection) throws Exception {
     super(connection);
     this.errorMessage = this.readAll();
-    this.statement = (CopyStatement) connection.getActiveStatement();
+    IntermediateStatement activeStatement = connection.getActiveStatement();
+    if (activeStatement instanceof CopyStatement) {
+      this.statement = (CopyStatement) activeStatement;
+    } else {
+      this.statement = null;
+    }
   }
 
   @Override
   protected void sendPayload() throws Exception {
-    // If backend error occurred during copy-in mode, drop any subsequent CopyFail messages.
     if (this.statement != null) {
       MutationWriter mutationWriter = this.statement.getMutationWriter();
       mutationWriter.rollback();
       statement.close();
-      if (!statement.hasException()) {
-        new ErrorResponse(this.outputStream, new Exception(this.errorMessage), State.IOError)
-            .send();
-      }
+      this.statement.handleExecutionException(
+          SpannerExceptionFactory.newSpannerException(ErrorCode.CANCELLED, this.errorMessage));
     }
-    this.connection.setStatus(ConnectionStatus.AUTHENTICATED);
-    this.connection.removeActiveStatement(this.statement);
-    new ReadyResponse(this.outputStream, ReadyResponse.Status.IDLE).send();
+    // Clear the COPY_IN status to indicate that we finished unsuccessfully. This will cause the
+    // inline handling of incoming (copy) messages to stop and the server to return an error message
+    // and to resume normal operation.
+    this.connection.setStatus(ConnectionStatus.COPY_FAILED);
   }
 
   @Override
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/DescribeMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/DescribeMessage.java
index a3b4834d..880cbf02 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/DescribeMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/DescribeMessage.java
@@ -21,7 +21,7 @@ import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.metadata.DescribePortalMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeStatementMetadata;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
-import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
+import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
 import com.google.cloud.spanner.pgadapter.wireoutput.NoDataResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ParameterDescriptionResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.RowDescriptionResponse;
@@ -38,7 +38,7 @@ public class DescribeMessage extends AbstractQueryProtocolMessage {
 
   private final PreparedType type;
   private final String name;
-  private final IntermediateStatement statement;
+  private final IntermediatePreparedStatement statement;
   private Future<DescribePortalMetadata> describePortalMetadata;
 
   public DescribeMessage(ConnectionHandler connection) throws Exception {
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ExecuteMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ExecuteMessage.java
index 7fb148da..3fb6f623 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ExecuteMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ExecuteMessage.java
@@ -17,7 +17,7 @@ package com.google.cloud.spanner.pgadapter.wireprotocol;
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
-import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
+import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
 import java.text.MessageFormat;
 
 /** Executes a portal. */
@@ -27,7 +27,7 @@ public class ExecuteMessage extends AbstractQueryProtocolMessage {
 
   private final String name;
   private final int maxRows;
-  private final IntermediateStatement statement;
+  private final IntermediatePreparedStatement statement;
 
   public ExecuteMessage(ConnectionHandler connection) throws Exception {
     super(connection);
@@ -94,7 +94,6 @@ public class ExecuteMessage extends AbstractQueryProtocolMessage {
     } else {
       try {
         this.sendSpannerResult(this.statement, this.queryMode, this.maxRows);
-        this.outputStream.flush();
       } catch (Exception exception) {
         handleError(exception);
       }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/FlushMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/FlushMessage.java
index 8aeab56b..a40a5bc4 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/FlushMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/FlushMessage.java
@@ -16,13 +16,11 @@ package com.google.cloud.spanner.pgadapter.wireprotocol;
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
-import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
-import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse.Status;
 import java.text.MessageFormat;
 
 /**
- * Handles a flush command from the user. JDBC does not require this step, so server-side this is a
- * noop.
+ * Handles a flush command from the user. This flushes any pending changes to Cloud Spanner and then
+ * sends any pending messages to the client.
  */
 @InternalApi
 public class FlushMessage extends ControlMessage {
@@ -33,13 +31,13 @@ public class FlushMessage extends ControlMessage {
     super(connection);
   }
 
+  public FlushMessage(ConnectionHandler connection, ManuallyCreatedToken manuallyCreatedToken) {
+    super(connection, 4, manuallyCreatedToken);
+  }
+
   @Override
   protected void sendPayload() throws Exception {
     connection.getExtendedQueryProtocolHandler().flush();
-    boolean inTransaction = connection.getSpannerConnection().isInTransaction();
-    new ReadyResponse(
-            this.outputStream, inTransaction ? Status.TRANSACTION : ReadyResponse.Status.IDLE)
-        .send();
   }
 
   @Override
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
index 635ee581..4f279c0d 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
@@ -14,6 +14,8 @@
 
 package com.google.cloud.spanner.pgadapter.wireprotocol;
 
+import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.COPY;
+
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.Statement;
@@ -21,7 +23,9 @@ import com.google.cloud.spanner.connection.AbstractStatementParser;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
+import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
+import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.cloud.spanner.pgadapter.wireoutput.ParseCompleteResponse;
 import com.google.common.base.Strings;
 import java.text.MessageFormat;
@@ -48,13 +52,7 @@ public class ParseMessage extends AbstractQueryProtocolMessage {
       parameterDataTypes[i] = this.inputStream.readInt();
     }
     this.statement =
-        new IntermediatePreparedStatement(
-            connection,
-            connection.getServer().getOptions(),
-            name,
-            parsedStatement,
-            originalStatement);
-    this.statement.setParameterDataTypes(this.parameterDataTypes);
+        createStatement(connection, name, parsedStatement, originalStatement, parameterDataTypes);
   }
 
   /**
@@ -69,13 +67,33 @@ public class ParseMessage extends AbstractQueryProtocolMessage {
     this.name = "";
     this.parameterDataTypes = new int[0];
     this.statement =
-        new IntermediatePreparedStatement(
-            connection,
-            connection.getServer().getOptions(),
-            name,
-            parsedStatement,
-            originalStatement);
-    this.statement.setParameterDataTypes(this.parameterDataTypes);
+        createStatement(connection, name, parsedStatement, originalStatement, parameterDataTypes);
+  }
+
+  static IntermediatePreparedStatement createStatement(
+      ConnectionHandler connectionHandler,
+      String name,
+      ParsedStatement parsedStatement,
+      Statement originalStatement,
+      int[] parameterDataTypes) {
+    if (StatementParser.isCommand(COPY, parsedStatement.getSqlWithoutComments())) {
+      return new CopyStatement(
+          connectionHandler,
+          connectionHandler.getServer().getOptions(),
+          name,
+          parsedStatement,
+          originalStatement);
+    } else {
+      IntermediatePreparedStatement statement =
+          new IntermediatePreparedStatement(
+              connectionHandler,
+              connectionHandler.getServer().getOptions(),
+              name,
+              parsedStatement,
+              originalStatement);
+      statement.setParameterDataTypes(parameterDataTypes);
+      return statement;
+    }
   }
 
   @Override
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java
index 91280cc1..7ee21079 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java
@@ -15,62 +15,31 @@
 package com.google.cloud.spanner.pgadapter.wireprotocol;
 
 import com.google.api.core.InternalApi;
-import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.Statement;
-import com.google.cloud.spanner.connection.AbstractStatementParser;
-import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
-import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
-import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
 import com.google.cloud.spanner.pgadapter.statements.SimpleQueryStatement;
-import com.google.cloud.spanner.pgadapter.utils.StatementParser;
-import com.google.cloud.spanner.pgadapter.wireoutput.CopyInResponse;
-import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
 import java.text.MessageFormat;
 
 /** Executes a simple statement. */
 @InternalApi
 public class QueryMessage extends ControlMessage {
-  private static final AbstractStatementParser PARSER =
-      AbstractStatementParser.getInstance(Dialect.POSTGRESQL);
   protected static final char IDENTIFIER = 'Q';
   public static final String COPY = "COPY";
 
   private final Statement originalStatement;
-  private final CopyStatement copyStatement;
   private final SimpleQueryStatement simpleQueryStatement;
-  private final boolean isCopy;
 
   public QueryMessage(ConnectionHandler connection) throws Exception {
     super(connection);
     this.originalStatement = Statement.of(this.readAll());
-    ParsedStatement parsedStatement = PARSER.parse(originalStatement);
-    this.isCopy = StatementParser.isCommand(COPY, parsedStatement.getSqlWithoutComments());
-    if (isCopy) {
-      this.copyStatement =
-          new CopyStatement(
-              connection,
-              connection.getServer().getOptions(),
-              parsedStatement,
-              this.originalStatement);
-      this.simpleQueryStatement = null;
-      this.connection.addActiveStatement(this.copyStatement);
-    } else {
-      this.simpleQueryStatement =
-          new SimpleQueryStatement(
-              connection.getServer().getOptions(), this.originalStatement, this.connection);
-      this.copyStatement = null;
-    }
+    this.simpleQueryStatement =
+        new SimpleQueryStatement(
+            connection.getServer().getOptions(), this.originalStatement, this.connection);
   }
 
   @Override
   protected void sendPayload() throws Exception {
-    if (this.isCopy) {
-      this.copyStatement.execute();
-      handleCopy();
-    } else {
-      this.simpleQueryStatement.execute();
-    }
+    this.simpleQueryStatement.execute();
   }
 
   @Override
@@ -96,25 +65,4 @@ public class QueryMessage extends ControlMessage {
   public SimpleQueryStatement getSimpleQueryStatement() {
     return this.simpleQueryStatement;
   }
-
-  private void handleCopy() throws Exception {
-    if (this.copyStatement.hasException()) {
-      handleError(this.copyStatement.getException());
-      new ReadyResponse(
-              this.outputStream,
-              connection
-                  .getExtendedQueryProtocolHandler()
-                  .getBackendConnection()
-                  .getConnectionState()
-                  .getReadyResponseStatus())
-          .send();
-    } else {
-      new CopyInResponse(
-              this.outputStream,
-              copyStatement.getTableColumns().size(),
-              copyStatement.getFormatCode())
-          .send();
-      this.connection.setStatus(ConnectionStatus.COPY_IN);
-    }
-  }
 }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SkipMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SkipMessage.java
new file mode 100644
index 00000000..f4b3d8ae
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SkipMessage.java
@@ -0,0 +1,56 @@
+// Copyright 2022 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.wireprotocol;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import java.io.IOException;
+import java.text.MessageFormat;
+
+/**
+ * {@link SkipMessage} is a pseudo wire-protocol message that is used to read and skip messages that
+ * we receive that we do not want at that time. This applies to COPY messages during normal
+ * operation, and sync/flush messages during COPY operation.
+ */
+@InternalApi
+public class SkipMessage extends ControlMessage {
+  public SkipMessage(ConnectionHandler connectionHandler) throws IOException {
+    super(connectionHandler);
+    int skipLength = this.length - 4;
+    int skipped = 0;
+    // Read and skip bytes until we have reached the total message length.
+    while (skipped < skipLength) {
+      skipped += connection.getConnectionMetadata().getInputStream().skip(skipLength - skipped);
+    }
+  }
+
+  @Override
+  protected void sendPayload() throws Exception {}
+
+  @Override
+  protected String getMessageName() {
+    return "Skip";
+  }
+
+  @Override
+  protected String getPayloadString() {
+    return new MessageFormat("Length: {0}").format(new Object[] {this.length});
+  }
+
+  @Override
+  protected String getIdentifier() {
+    return "";
+  }
+}


Test patch:
diff --git a/.ci/e2e-expected/copy-from-stdin-batch.txt b/.ci/e2e-expected/copy-from-stdin-batch.txt
new file mode 100644
index 00000000..0da509b7
--- /dev/null
+++ b/.ci/e2e-expected/copy-from-stdin-batch.txt
@@ -0,0 +1,22 @@
+  id | age | name
+-----+-----+--------
+   1 |   1 | John
+   2 |  20 | Joe
+   3 |  23 | Jack
+   7 |   7 | 7
+   8 |   8 | 8
+   9 |   9 | 9
+  10 |  10 | 10
+  11 |  11 | 11
+  12 |  12 | 12
+  13 |  13 | 13
+  14 |  14 | 14
+  15 |  15 | 15
+  16 |  16 | 16
+  17 |  17 | 17
+  18 |  18 | 18
+  19 |  19 | 19
+  99 |  99 | person
+ 100 | 100 | test
+(18 rows)
+
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
index 6f4deda2..4b477774 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
@@ -41,6 +41,19 @@ import java.util.concurrent.Future;
  */
 @InternalApi
 public class IntermediateStatement {
+  /**
+   * Indicates whether an attempt to get the result of a statement should block or fail if the
+   * result is not yet available. Normal SQL commands that can be executed directly on Cloud Spanner
+   * should always have their results available when a sync/flush message is received. COPY
+   * statements do not have that, as they require additional messages after a flush/sync has been
+   * received. Attempts to get the result of a COPY statement should therefore block until it is
+   * available, which is after a CopyDone or CopyFail message has been received.
+   */
+  public enum ResultNotReadyBehavior {
+    FAIL,
+    BLOCK;
+  }
+
   protected static final PostgreSQLStatementParser PARSER =
       (PostgreSQLStatementParser) AbstractStatementParser.getInstance(Dialect.POSTGRESQL);
 
@@ -122,11 +135,24 @@ public class IntermediateStatement {
 
   /**
    * @return The number of items that were modified by this execution for DML. 0 for DDL and -1 for
-   *     QUERY.
+   *     QUERY. Fails if the result is not yet available.
    */
   public long getUpdateCount() {
-    initFutureResult();
-    switch (this.parsedStatement.getType()) {
+    return getUpdateCount(ResultNotReadyBehavior.FAIL);
+  }
+
+  /**
+   * @return The number of items that were modified by this execution for DML. 0 for DDL and -1 for
+   *     QUERY. Will block or fail depending on the given {@link ResultNotReadyBehavior} if the
+   *     result is not yet available.
+   */
+  public long getUpdateCount(ResultNotReadyBehavior resultNotReadyBehavior) {
+    initFutureResult(resultNotReadyBehavior);
+    if (hasException()) {
+      throw getException();
+    }
+    // Note: getStatementType() returns UPDATE for COPY statements.
+    switch (getStatementType()) {
       case QUERY:
         return -1L;
       case UPDATE:
@@ -139,9 +165,20 @@ public class IntermediateStatement {
     }
   }
 
-  /** @return True if at some point in execution, and exception was thrown. */
+  /**
+   * @return True if at some point in execution an exception was thrown. Fails if execution has not
+   *     yet finished.
+   */
   public boolean hasException() {
-    initFutureResult();
+    return hasException(ResultNotReadyBehavior.FAIL);
+  }
+
+  /**
+   * @return True if at some point in execution an exception was thrown. Fails or blocks depending
+   *     on the given {@link ResultNotReadyBehavior} if execution has not yet finished.
+   */
+  public boolean hasException(ResultNotReadyBehavior resultNotReadyBehavior) {
+    initFutureResult(resultNotReadyBehavior);
     return this.exception != null;
   }
 
@@ -162,25 +199,30 @@ public class IntermediateStatement {
     return this.parsedStatement.getSqlWithoutComments();
   }
 
-  private void initFutureResult() {
+  private void initFutureResult(ResultNotReadyBehavior resultNotReadyBehavior) {
     if (this.futureStatementResult != null) {
-      if (!this.futureStatementResult.isDone()) {
+      if (resultNotReadyBehavior == ResultNotReadyBehavior.FAIL
+          && !this.futureStatementResult.isDone()) {
         throw new IllegalStateException("Statement result cannot be retrieved before flush/sync");
       }
       try {
         setStatementResult(this.futureStatementResult.get());
       } catch (ExecutionException executionException) {
-        this.exception = SpannerExceptionFactory.asSpannerException(executionException.getCause());
+        setException(SpannerExceptionFactory.asSpannerException(executionException.getCause()));
       } catch (InterruptedException interruptedException) {
-        this.exception = SpannerExceptionFactory.propagateInterrupt(interruptedException);
+        setException(SpannerExceptionFactory.propagateInterrupt(interruptedException));
       } finally {
         this.futureStatementResult = null;
       }
     }
   }
 
+  /**
+   * Returns the result of this statement as a {@link StatementResult}. Fails if the result is not
+   * yet available.
+   */
   public StatementResult getStatementResult() {
-    initFutureResult();
+    initFutureResult(ResultNotReadyBehavior.FAIL);
     return this.statementResult;
   }
 
@@ -206,10 +248,18 @@ public class IntermediateStatement {
     return this.parsedStatement.getSqlWithoutComments();
   }
 
-  public Exception getException() {
-    Exception e = this.exception;
-    this.exception = null;
-    return e;
+  /** Returns any execution exception registered for this statement. */
+  public SpannerException getException() {
+    return this.exception;
+  }
+
+  void setException(SpannerException exception) {
+    // Do not override any exception that has already been registered. COPY statements can receive
+    // multiple errors as they execute asynchronously while receiving a stream of data from the
+    // client. We always return the first exception that we encounter.
+    if (this.exception == null) {
+      this.exception = exception;
+    }
   }
 
   /**
@@ -218,7 +268,7 @@ public class IntermediateStatement {
    * @param exception The exception to store.
    */
   protected void handleExecutionException(SpannerException exception) {
-    this.exception = exception;
+    setException(exception);
     this.hasMoreData = false;
   }
 
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
index 0e5e65e0..e05ca7a1 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
@@ -14,11 +14,20 @@
 
 package com.google.cloud.spanner.pgadapter;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
+import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
+import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import java.io.EOFException;
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.Socket;
@@ -74,4 +83,60 @@ public class ConnectionHandlerTest {
     connection.terminate();
     verify(socket).close();
   }
+
+  @Test
+  public void testHandleMessages_NonFatalException() throws Exception {
+    ProxyServer server = mock(ProxyServer.class);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+    DataOutputStream dataOutputStream = new DataOutputStream(new ByteArrayOutputStream());
+    ConnectionMetadata connectionMetadata = mock(ConnectionMetadata.class);
+    when(connectionMetadata.getOutputStream()).thenReturn(dataOutputStream);
+    WireMessage message = mock(WireMessage.class);
+    when(server.recordMessage(message)).thenReturn(message);
+    doThrow(
+            SpannerExceptionFactory.newSpannerException(
+                ErrorCode.FAILED_PRECONDITION, "non-fatal test exception"))
+        .when(message)
+        .send();
+
+    ConnectionHandler connection =
+        new ConnectionHandler(server, socket) {
+          public ConnectionMetadata getConnectionMetadata() {
+            return connectionMetadata;
+          }
+        };
+
+    connection.setMessageState(message);
+    connection.handleMessages();
+
+    assertEquals(ConnectionStatus.UNAUTHENTICATED, connection.getStatus());
+  }
+
+  @Test
+  public void testHandleMessages_FatalException() throws Exception {
+    ProxyServer server = mock(ProxyServer.class);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+    DataOutputStream dataOutputStream = new DataOutputStream(new ByteArrayOutputStream());
+    ConnectionMetadata connectionMetadata = mock(ConnectionMetadata.class);
+    when(connectionMetadata.getOutputStream()).thenReturn(dataOutputStream);
+    WireMessage message = mock(WireMessage.class);
+    when(server.recordMessage(message)).thenReturn(message);
+    doThrow(new EOFException("fatal test exception")).when(message).send();
+
+    ConnectionHandler connection =
+        new ConnectionHandler(server, socket) {
+          public ConnectionMetadata getConnectionMetadata() {
+            return connectionMetadata;
+          }
+        };
+
+    connection.setMessageState(message);
+    connection.handleMessages();
+
+    assertEquals(ConnectionStatus.TERMINATED, connection.getStatus());
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java
index 639a281a..a66bb9b1 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java
@@ -15,6 +15,8 @@
 package com.google.cloud.spanner.pgadapter;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
@@ -25,6 +27,8 @@ import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.TextFormat;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage;
+import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage.ManuallyCreatedToken;
+import com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.DataInputStream;
@@ -99,4 +103,21 @@ public final class ControlMessageTest {
     assertEquals(numOfBytes, outputReader.read(bytes, 0, numOfBytes));
     assertEquals(resultMessage, new String(bytes, UTF8));
   }
+
+  @Test
+  public void testUnknownStatementTypeThrowsError() {
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    ExecuteMessage executeMessage =
+        new ExecuteMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN);
+    IntermediateStatement intermediateStatement = mock(IntermediateStatement.class);
+    when(intermediateStatement.getCommandTag()).thenReturn("PARSE");
+    when(intermediateStatement.getStatementType()).thenReturn(StatementType.UNKNOWN);
+    when(intermediateStatement.getStatement()).thenReturn("parse foo from bar");
+
+    IllegalStateException exception =
+        assertThrows(
+            IllegalStateException.class,
+            () -> executeMessage.sendSpannerResult(intermediateStatement, QueryMode.SIMPLE, 0L));
+    assertEquals("Unknown statement type: parse foo from bar", exception.getMessage());
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
index 482a65a0..53a3a519 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
@@ -18,6 +18,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeTrue;
 
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
@@ -33,8 +34,11 @@ import com.google.spanner.v1.StructType.Field;
 import com.google.spanner.v1.Type;
 import com.google.spanner.v1.TypeCode;
 import io.grpc.Status;
+import java.io.BufferedReader;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.io.InputStreamReader;
+import java.io.OutputStreamWriter;
 import java.io.StringReader;
 import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
@@ -42,6 +46,7 @@ import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.util.List;
+import java.util.stream.Collectors;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -54,8 +59,8 @@ import org.postgresql.copy.CopyManager;
 import org.postgresql.core.BaseConnection;
 import org.postgresql.core.PGStream;
 import org.postgresql.core.QueryExecutorBase;
-import org.postgresql.core.v3.CopyOperationImpl;
 import org.postgresql.core.v3.QueryExecutorImpl;
+import org.postgresql.jdbc.PgConnection;
 
 @RunWith(Parameterized.class)
 public class CopyInMockServerTest extends AbstractMockServerTest {
@@ -79,14 +84,20 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
    * mode for queries and DML statements.
    */
   private String createUrl() {
+    return createUrl("extended");
+  }
+
+  private String createUrl(String queryMode) {
     if (useDomainSocket) {
       return String.format(
           "jdbc:postgresql://localhost/?"
               + "socketFactory=org.newsclub.net.unix.AFUNIXSocketFactory$FactoryArg"
-              + "&socketFactoryArg=/tmp/.s.PGSQL.%d",
-          pgServer.getLocalPort());
+              + "&socketFactoryArg=/tmp/.s.PGSQL.%d"
+              + "&preferQueryMode=%s",
+          pgServer.getLocalPort(), queryMode);
     }
-    return String.format("jdbc:postgresql://localhost:%d/", pgServer.getLocalPort());
+    return String.format(
+        "jdbc:postgresql://localhost:%d/?preferQueryMode=%s", pgServer.getLocalPort(), queryMode);
   }
 
   @Test
@@ -115,6 +126,63 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
     assertEquals(3, mutation.getInsert().getValuesCount());
   }
 
+  @Test
+  public void testCopyInWithExplicitTransaction() throws SQLException, IOException {
+    setupCopyInformationSchemaResults();
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      connection.setAutoCommit(false);
+
+      CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
+      copyManager.copyIn("COPY users FROM STDIN;", new StringReader("5\t5\t5\n6\t6\t6\n7\t7\t7\n"));
+
+      // Verify that we can use the connection for normal queries.
+      try (ResultSet resultSet = connection.createStatement().executeQuery("SELECT 1")) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+
+      // Verify that the transaction has not yet been committed.
+      assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+
+      connection.commit();
+    }
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+    CommitRequest commitRequest = commitRequests.get(0);
+    assertEquals(1, commitRequest.getMutationsCount());
+
+    Mutation mutation = commitRequest.getMutations(0);
+    assertEquals(OperationCase.INSERT, mutation.getOperationCase());
+    assertEquals(3, mutation.getInsert().getValuesCount());
+  }
+
+  @Test
+  public void testCopyInWithExplicitTransaction_Rollback() throws SQLException, IOException {
+    setupCopyInformationSchemaResults();
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      connection.setAutoCommit(false);
+
+      CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
+      copyManager.copyIn("COPY users FROM STDIN;", new StringReader("5\t5\t5\n6\t6\t6\n7\t7\t7\n"));
+
+      // Verify that we can use the connection for normal queries.
+      try (ResultSet resultSet = connection.createStatement().executeQuery("SELECT 1")) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+
+      connection.rollback();
+    }
+
+    // Verify that the COPY operation was not committed.
+    assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+  }
+
   @Test
   public void testCopyIn_Small() throws SQLException, IOException {
     setupCopyInformationSchemaResults();
@@ -131,7 +199,7 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
   }
 
   @Test
-  public void testCopyIn_Large_FailsWhenAtomic() throws SQLException {
+  public void testCopyIn_Large_FailsWhenAtomic() throws Exception {
     setupCopyInformationSchemaResults();
 
     try (Connection connection = DriverManager.getConnection(createUrl())) {
@@ -144,12 +212,14 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
                   copyManager.copyIn(
                       "copy all_types from stdin;",
                       new FileInputStream("./src/test/resources/all_types_data.txt")));
-      // The JDBC driver CopyManager takes the COPY protocol quite literally, and as the COPY
-      // protocol does not include any error handling, the JDBC driver will just send all data to
-      // the server and ignore any error messages the server might send during the copy operation.
-      // PGAdapter therefore drops the connection if it continues to receive CopyData messages after
-      // it sent back an error message.
-      assertTrue(exception.getMessage().contains("Database connection failed"));
+      assertTrue(
+          exception.getMessage(),
+          exception
+                  .getMessage()
+                  .contains("FAILED_PRECONDITION: Record count: 2001 has exceeded the limit: 2000.")
+              || exception
+                  .getMessage()
+                  .contains("Database connection failed when canceling copy operation"));
     }
   }
 
@@ -187,6 +257,7 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
                   copyManager.copyIn(
                       "COPY users FROM STDIN;", new StringReader("5\t5\t5\n6\t6\t6\n7\t7\t7\n")));
       assertTrue(
+          exception.getMessage(),
           exception.getMessage().contains("io.grpc.StatusRuntimeException: INVALID_ARGUMENT"));
     }
 
@@ -222,7 +293,7 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
   }
 
   @Test
-  public void testCopyInWithInvalidRow() throws SQLException {
+  public void testCopyInWithInvalidRow() throws SQLException, IOException {
     setupCopyInformationSchemaResults();
 
     try (Connection connection = DriverManager.getConnection(createUrl())) {
@@ -255,7 +326,9 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
               () ->
                   copyManager.copyIn(
                       "COPY users FROM STDIN;", new StringReader("5\t5\t5\n6\t6\t6\n7\t7\t7\n")));
-      assertTrue(exception.getMessage().contains("Commit size: 20 has exceeded the limit: 10"));
+      assertTrue(
+          exception.getMessage(),
+          exception.getMessage().contains("Commit size: 20 has exceeded the limit: 10"));
     } finally {
       System.getProperties().remove("copy_in_commit_limit");
     }
@@ -351,20 +424,27 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
 
     byte[] payload = "5\t5\t5\n".getBytes(StandardCharsets.UTF_8);
     try (Connection connection = DriverManager.getConnection(createUrl())) {
-      CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
-      // Start a copy operation and then try to execute a query during the copy.
-      CopyIn copyOperation = copyManager.copyIn("COPY users FROM STDIN;");
-      copyOperation.writeToCopy(payload, 0, payload.length);
-      copyOperation.flushCopy();
-
-      // Use reflection to get hold of the underlying stream, so we can send an invalid message.
-      java.lang.reflect.Field queryExecutorField =
-          CopyOperationImpl.class.getDeclaredField("queryExecutor");
-      queryExecutorField.setAccessible(true);
-      QueryExecutorImpl queryExecutor = (QueryExecutorImpl) queryExecutorField.get(copyOperation);
+      PgConnection pgConnection = connection.unwrap(PgConnection.class);
+      QueryExecutorImpl queryExecutor = (QueryExecutorImpl) pgConnection.getQueryExecutor();
+      // Use reflection to get hold of the underlying stream, so we can send any message that we
+      // want.
       java.lang.reflect.Field pgStreamField = QueryExecutorBase.class.getDeclaredField("pgStream");
       pgStreamField.setAccessible(true);
       PGStream stream = (PGStream) pgStreamField.get(queryExecutor);
+
+      String sql = "copy users from stdin;";
+      int length = sql.length() + 5;
+      stream.sendChar('Q');
+      stream.sendInteger4(length);
+      stream.send(sql.getBytes(StandardCharsets.UTF_8));
+      stream.sendChar(0);
+      stream.flush();
+
+      // Send a CopyData message and then a 'Q', which is not allowed.
+      stream.sendChar('d');
+      stream.sendInteger4(payload.length + 4);
+      stream.send(payload);
+
       stream.sendChar('Q');
       // Length = 4 + 8 + 1 = 13
       // (msg length 4 bytes, 8 bytes for SELECT 1, 1 byte for \0)
@@ -373,18 +453,261 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
       stream.sendChar(0);
       stream.flush();
 
-      // PGAdapter drops the connection if an invalid message is received during a COPY. This is a
-      // safety measure as there is no other error handling in the COPY protocol, and the server
-      // could otherwise have been completely flushed with garbage if it continued to receive
-      // messages after receiving an invalid message.
-      SQLException exception = assertThrows(SQLException.class, copyOperation::endCopy);
-      assertEquals("Database connection failed when ending copy", exception.getMessage());
+      boolean receivedErrorMessage = false;
+      StringBuilder errorMessage = new StringBuilder();
+      while (!receivedErrorMessage) {
+        int command = stream.receiveChar();
+        if (command == 'E') {
+          receivedErrorMessage = true;
+          // Read and ignore the length.
+          stream.receiveInteger4();
+          while (stream.receiveChar() > 0) {
+            errorMessage.append(stream.receiveString()).append('\n');
+          }
+        } else {
+          // Just skip everything in the message.
+          length = stream.receiveInteger4();
+          stream.skip(length - 4);
+        }
+      }
+      assertEquals(
+          "ERROR\n"
+              + "XX000\n"
+              + "Expected CopyData ('d'), CopyDone ('c') or CopyFail ('f') messages, got: 'Q'\n",
+          errorMessage.toString());
+
+      stream.sendChar('c');
+      stream.sendInteger4(4);
+      stream.sendChar('x');
+      stream.sendInteger4(4);
+      stream.flush();
     }
 
     List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
     assertTrue(commitRequests.isEmpty());
   }
 
+  private static boolean isPsqlAvailable() {
+    ProcessBuilder builder = new ProcessBuilder();
+    String[] psqlCommand = new String[] {"psql", "--version"};
+    builder.command(psqlCommand);
+    try {
+      Process process = builder.start();
+      int res = process.waitFor();
+
+      return res == 0;
+    } catch (Exception ignored) {
+      return false;
+    }
+  }
+
+  @Test
+  public void testCopyInBatchPsql() throws Exception {
+    assumeTrue("This test requires psql to be installed", isPsqlAvailable());
+    setupCopyInformationSchemaResults();
+
+    String host = useDomainSocket ? "/tmp" : "localhost";
+    ProcessBuilder builder = new ProcessBuilder();
+    String[] psqlCommand =
+        new String[] {"psql", "-h", host, "-p", String.valueOf(pgServer.getLocalPort())};
+    builder.command(psqlCommand);
+    Process process = builder.start();
+    String errors;
+    String output;
+
+    try (OutputStreamWriter writer = new OutputStreamWriter(process.getOutputStream());
+        BufferedReader reader =
+            new BufferedReader(new InputStreamReader(process.getInputStream()));
+        BufferedReader errorReader =
+            new BufferedReader(new InputStreamReader(process.getErrorStream()))) {
+      writer.write(
+          "SELECT 1\\;copy users from stdin\\;copy users from stdin\\;SELECT 2;\n"
+              + "1\t2\t3\n"
+              + "\\.\n"
+              + "4\t5\t6\n"
+              + "\\.\n"
+              + "\n"
+              + "\\q\n");
+      writer.flush();
+      errors = errorReader.lines().collect(Collectors.joining("\n"));
+      output = reader.lines().collect(Collectors.joining("\n"));
+    }
+
+    assertEquals("", errors);
+    assertEquals(" C \n---\n 2\n(1 row)\n", output);
+    int res = process.waitFor();
+    assertEquals(0, res);
+
+    // The batch is executed as one implicit transaction.
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+    CommitRequest commitRequest = commitRequests.get(0);
+    assertEquals(1, commitRequest.getMutationsCount());
+    // We buffer 2 mutations, but the proto builder in the Java client library combines these two
+    // mutations into 1 mutation with 2 rows.
+    assertEquals(2, commitRequest.getMutations(0).getInsert().getValuesCount());
+    assertEquals(OperationCase.INSERT, commitRequest.getMutations(0).getOperationCase());
+    assertEquals(3, commitRequest.getMutations(0).getInsert().getColumnsCount());
+    assertEquals(
+        "1", commitRequest.getMutations(0).getInsert().getValues(0).getValues(0).getStringValue());
+    assertEquals(
+        "2", commitRequest.getMutations(0).getInsert().getValues(0).getValues(1).getStringValue());
+    assertEquals(
+        "3", commitRequest.getMutations(0).getInsert().getValues(0).getValues(2).getStringValue());
+    assertEquals(OperationCase.INSERT, commitRequest.getMutations(0).getOperationCase());
+    assertEquals(
+        "4", commitRequest.getMutations(0).getInsert().getValues(1).getValues(0).getStringValue());
+    assertEquals(
+        "5", commitRequest.getMutations(0).getInsert().getValues(1).getValues(1).getStringValue());
+    assertEquals(
+        "6", commitRequest.getMutations(0).getInsert().getValues(1).getValues(2).getStringValue());
+  }
+
+  @Test
+  public void testCopyInBatch() throws Exception {
+    setupCopyInformationSchemaResults();
+
+    byte[] row1 = "5\t6\t7\n".getBytes(StandardCharsets.UTF_8);
+    byte[] row2 = "6\t7\t8\n".getBytes(StandardCharsets.UTF_8);
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      PgConnection pgConnection = connection.unwrap(PgConnection.class);
+      QueryExecutorImpl queryExecutor = (QueryExecutorImpl) pgConnection.getQueryExecutor();
+      // Use reflection to get hold of the underlying stream, so we can send any message that we
+      // want.
+      java.lang.reflect.Field pgStreamField = QueryExecutorBase.class.getDeclaredField("pgStream");
+      pgStreamField.setAccessible(true);
+      PGStream stream = (PGStream) pgStreamField.get(queryExecutor);
+
+      String sql = "SELECT 1;copy users from stdin;copy users from stdin;SELECT 2;";
+      int length = sql.length() + 5;
+      stream.sendChar('Q');
+      stream.sendInteger4(length);
+      stream.send(sql.getBytes(StandardCharsets.UTF_8));
+      stream.sendChar(0);
+      stream.flush();
+
+      // Send CopyData + CopyDone for the first copy operation.
+      stream.sendChar('d');
+      stream.sendInteger4(row1.length + 4);
+      stream.send(row1);
+      stream.sendChar('c');
+      stream.sendInteger4(4);
+      stream.flush();
+
+      // Send CopyData + CopyDone for the second copy operation.
+      stream.sendChar('d');
+      stream.sendInteger4(row2.length + 4);
+      stream.send(row2);
+      stream.sendChar('c');
+      stream.sendInteger4(4);
+      stream.flush();
+
+      boolean receivedReadyForQuery = false;
+      while (!receivedReadyForQuery) {
+        int command = stream.receiveChar();
+        if (command == 'Z') {
+          receivedReadyForQuery = true;
+        } else {
+          // Just skip everything in the message.
+          length = stream.receiveInteger4();
+          stream.skip(length - 4);
+        }
+      }
+    }
+
+    assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+    CommitRequest commitRequest = mockSpanner.getRequestsOfType(CommitRequest.class).get(0);
+    assertEquals(1, commitRequest.getMutationsCount());
+    assertEquals(2, commitRequest.getMutations(0).getInsert().getValuesCount());
+    assertEquals(
+        "5", commitRequest.getMutations(0).getInsert().getValues(0).getValues(0).getStringValue());
+    assertEquals(
+        "6", commitRequest.getMutations(0).getInsert().getValues(0).getValues(1).getStringValue());
+    assertEquals(
+        "7", commitRequest.getMutations(0).getInsert().getValues(0).getValues(2).getStringValue());
+    assertEquals(
+        "6", commitRequest.getMutations(0).getInsert().getValues(1).getValues(0).getStringValue());
+    assertEquals(
+        "7", commitRequest.getMutations(0).getInsert().getValues(1).getValues(1).getStringValue());
+    assertEquals(
+        "8", commitRequest.getMutations(0).getInsert().getValues(1).getValues(2).getStringValue());
+  }
+
+  @Test
+  public void testCopyInBatchWithCopyFail() throws Exception {
+    setupCopyInformationSchemaResults();
+
+    byte[] row1 = "5\t6\t7\n".getBytes(StandardCharsets.UTF_8);
+    byte[] row2 = "8\t9\t10\n".getBytes(StandardCharsets.UTF_8);
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      PgConnection pgConnection = connection.unwrap(PgConnection.class);
+      QueryExecutorImpl queryExecutor = (QueryExecutorImpl) pgConnection.getQueryExecutor();
+      // Use reflection to get hold of the underlying stream, so we can send any message that we
+      // want.
+      java.lang.reflect.Field pgStreamField = QueryExecutorBase.class.getDeclaredField("pgStream");
+      pgStreamField.setAccessible(true);
+      PGStream stream = (PGStream) pgStreamField.get(queryExecutor);
+
+      String sql = "SELECT 1;copy users from stdin;copy users from stdin;SELECT 2;";
+      int length = sql.getBytes(StandardCharsets.UTF_8).length + 5;
+      stream.sendChar('Q');
+      stream.sendInteger4(length);
+      stream.send(sql.getBytes(StandardCharsets.UTF_8));
+      stream.sendChar(0);
+      stream.flush();
+
+      // Send CopyData + CopyDone for the first copy operation.
+      stream.sendChar('d');
+      stream.sendInteger4(row1.length + 4);
+      stream.send(row1);
+      stream.sendChar('c');
+      stream.sendInteger4(4);
+      stream.flush();
+
+      // Send CopyData + CopyFail for the second copy operation.
+      stream.sendChar('d');
+      stream.sendInteger4(row2.length + 4);
+      stream.send(row2);
+      String error = "Changed my mind";
+      stream.sendChar('f');
+      stream.sendInteger4(4 + error.getBytes(StandardCharsets.UTF_8).length + 1);
+      stream.send(error.getBytes(StandardCharsets.UTF_8));
+      stream.sendChar(0);
+      stream.flush();
+
+      boolean receivedReadyForQuery = false;
+      boolean receivedErrorMessage = false;
+      StringBuilder errorMessage = new StringBuilder();
+      while (!receivedReadyForQuery) {
+        int command = stream.receiveChar();
+        if (command == 'Z') {
+          receivedReadyForQuery = true;
+        } else if (command == 'E') {
+          receivedErrorMessage = true;
+          // Read and ignore the length.
+          stream.receiveInteger4();
+          while (stream.receiveChar() > 0) {
+            errorMessage.append(stream.receiveString()).append('\n');
+          }
+        } else {
+          // Just skip everything in the message.
+          length = stream.receiveInteger4();
+          stream.skip(length - 4);
+        }
+      }
+      assertTrue(receivedErrorMessage);
+      assertEquals(
+          "ERROR\n"
+              + "P0001\n"
+              + "CANCELLED: Changed my mind\n"
+              + "ERROR\n"
+              + "P0001\n"
+              + "INVALID_ARGUMENT: current transaction is aborted, commands ignored until end of transaction block\n",
+          errorMessage.toString());
+    }
+    assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+  }
+
   private void setupCopyInformationSchemaResults() {
     setupCopyInformationSchemaResults(true);
   }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
index 9c3b6ce3..a132b0db 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
@@ -476,12 +476,11 @@ public class ITJdbcTest implements IntegrationTest {
                   copyManager.copyIn(
                       "copy all_types from stdin;",
                       new FileInputStream("./src/test/resources/all_types_data.txt")));
-      // The JDBC driver CopyManager takes the COPY protocol quite literally, and as the COPY
-      // protocol does not include any error handling, the JDBC driver will just send all data to
-      // the server and ignore any error messages the server might send during the copy operation.
-      // PGAdapter therefore drops the connection if it continues to receive CopyData messages after
-      // it sent back an error message.
-      assertTrue(exception.getMessage().contains("Database connection failed"));
+      assertEquals(
+          "ERROR: FAILED_PRECONDITION: Record count: 2001 has exceeded the limit: 2000.\n\n"
+              + "The number of mutations per record is equal to the number of columns in the record plus the number of indexed columns in the record. The maximum number of mutations in one transaction is 20000.\n\n"
+              + "Execute `SET AUTOCOMMIT_DML_MODE='PARTITIONED_NON_ATOMIC'` before executing a large COPY operation to instruct PGAdapter to automatically break large transactions into multiple smaller. This will make the COPY operation non-atomic.\n\n",
+          exception.getMessage());
     }
 
     // Verify that the table is still empty.
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/InvalidMessagesTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/InvalidMessagesTest.java
index 636b68d8..3b16ada0 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/InvalidMessagesTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/InvalidMessagesTest.java
@@ -115,11 +115,14 @@ public class InvalidMessagesTest extends AbstractMockServerTest {
         outputStream.flush();
 
         // Read until the end of the stream. The stream should be closed by the backend.
-        int bytesRead = 0;
-        while (inputStream.read() > -1 && bytesRead < 1 << 16) {
-          bytesRead++;
+        try {
+          int bytesRead = 0;
+          while (inputStream.read() > -1 && bytesRead < 1 << 16) {
+            bytesRead++;
+          }
+          assertEquals(-1, inputStream.read());
+        } catch (IOException ignore) {
         }
-        assertEquals(-1, inputStream.read());
       }
     }
   }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
index d4e7fe94..9fe7b413 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
@@ -15,15 +15,21 @@
 package com.google.cloud.spanner.pgadapter.statements;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.AbstractStatementParser;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
@@ -32,15 +38,133 @@ import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.DdlTransactio
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResult;
 import com.google.cloud.spanner.pgadapter.statements.local.ListDatabasesStatement;
 import com.google.cloud.spanner.pgadapter.statements.local.LocalStatement;
+import com.google.cloud.spanner.pgadapter.utils.CopyDataReceiver;
+import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import com.google.common.collect.ImmutableList;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
+import java.util.concurrent.RejectedExecutionException;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 @RunWith(JUnit4.class)
 public class BackendConnectionTest {
+  private final AbstractStatementParser PARSER =
+      AbstractStatementParser.getInstance(Dialect.POSTGRESQL);
+
+  @Test
+  public void testCopyPropagatesNonSpannerException() {
+    String sql = "COPY users FROM STDIN";
+    Statement statement = Statement.of(sql);
+    ParsedStatement parsedStatement = PARSER.parse(statement);
+    Connection spannerConnection = mock(Connection.class);
+    CopyDataReceiver receiver = mock(CopyDataReceiver.class);
+    MutationWriter writer = mock(MutationWriter.class);
+    ExecutorService executor = mock(ExecutorService.class);
+    RejectedExecutionException rejectedExecutionException = new RejectedExecutionException();
+    doThrow(rejectedExecutionException).when(executor).execute(any(Runnable.class));
+
+    BackendConnection backendConnection =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    Future<StatementResult> result =
+        backendConnection.executeCopy(parsedStatement, statement, receiver, writer, executor);
+    backendConnection.flush();
+
+    assertTrue(result.isDone());
+    ExecutionException executionException = assertThrows(ExecutionException.class, result::get);
+    assertSame(rejectedExecutionException, executionException.getCause());
+  }
+
+  @Test
+  public void testHasDmlOrCopyStatementsAfter() {
+    CopyDataReceiver receiver = mock(CopyDataReceiver.class);
+    MutationWriter writer = mock(MutationWriter.class);
+    ExecutorService executor = mock(ExecutorService.class);
+    Connection spannerConnection = mock(Connection.class);
+    String updateSql = "insert into foo values (1)";
+    Statement updateStatement = Statement.of(updateSql);
+    ParsedStatement parsedUpdateStatement = PARSER.parse(updateStatement);
+    String selectSql = "select * from foo";
+    Statement selectStatement = Statement.of(selectSql);
+    ParsedStatement parsedSelectStatement = PARSER.parse(selectStatement);
+    String copySql = "copy foo from stdin";
+    Statement copyStatement = Statement.of(copySql);
+    ParsedStatement parsedCopyStatement = PARSER.parse(copyStatement);
+    String clientSideSql = "set spanner.statement_tag='foo'";
+    Statement clientSideStatement = Statement.of(clientSideSql);
+    ParsedStatement parsedClientSideStatement = PARSER.parse(clientSideStatement);
+    String unknownSql = "merge foo";
+    Statement unknownStatement = Statement.of(unknownSql);
+    ParsedStatement parsedUnknownStatement = PARSER.parse(unknownStatement);
+
+    BackendConnection onlyDmlStatements =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    onlyDmlStatements.execute(parsedUpdateStatement, updateStatement);
+    onlyDmlStatements.execute(parsedUpdateStatement, updateStatement);
+    assertTrue(onlyDmlStatements.hasDmlOrCopyStatementsAfter(0));
+    assertTrue(onlyDmlStatements.hasDmlOrCopyStatementsAfter(1));
+
+    BackendConnection onlyCopyStatements =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    onlyCopyStatements.executeCopy(parsedCopyStatement, copyStatement, receiver, writer, executor);
+    onlyCopyStatements.executeCopy(parsedCopyStatement, copyStatement, receiver, writer, executor);
+    assertTrue(onlyCopyStatements.hasDmlOrCopyStatementsAfter(0));
+    assertTrue(onlyCopyStatements.hasDmlOrCopyStatementsAfter(1));
+
+    BackendConnection dmlAndCopyStatements =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    dmlAndCopyStatements.execute(parsedUpdateStatement, updateStatement);
+    dmlAndCopyStatements.executeCopy(
+        parsedCopyStatement, copyStatement, receiver, writer, executor);
+    assertTrue(dmlAndCopyStatements.hasDmlOrCopyStatementsAfter(0));
+    assertTrue(dmlAndCopyStatements.hasDmlOrCopyStatementsAfter(1));
+
+    BackendConnection onlySelectStatements =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    onlySelectStatements.execute(parsedSelectStatement, selectStatement);
+    onlySelectStatements.execute(parsedSelectStatement, selectStatement);
+    assertFalse(onlySelectStatements.hasDmlOrCopyStatementsAfter(0));
+    assertFalse(onlySelectStatements.hasDmlOrCopyStatementsAfter(1));
+
+    BackendConnection onlyClientSideStatements =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    onlyClientSideStatements.execute(parsedClientSideStatement, clientSideStatement);
+    onlyClientSideStatements.execute(parsedClientSideStatement, clientSideStatement);
+    assertFalse(onlyClientSideStatements.hasDmlOrCopyStatementsAfter(0));
+    assertFalse(onlyClientSideStatements.hasDmlOrCopyStatementsAfter(1));
+
+    BackendConnection onlyUnknownStatements =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    onlyUnknownStatements.execute(parsedUnknownStatement, unknownStatement);
+    onlyUnknownStatements.execute(parsedUnknownStatement, unknownStatement);
+    assertFalse(onlyUnknownStatements.hasDmlOrCopyStatementsAfter(0));
+    assertFalse(onlyUnknownStatements.hasDmlOrCopyStatementsAfter(1));
+
+    BackendConnection dmlAndSelectStatements =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    dmlAndSelectStatements.execute(parsedUpdateStatement, updateStatement);
+    dmlAndSelectStatements.execute(parsedSelectStatement, selectStatement);
+    assertTrue(dmlAndSelectStatements.hasDmlOrCopyStatementsAfter(0));
+    assertFalse(dmlAndSelectStatements.hasDmlOrCopyStatementsAfter(1));
+
+    BackendConnection copyAndSelectStatements =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    copyAndSelectStatements.executeCopy(
+        parsedCopyStatement, copyStatement, receiver, writer, executor);
+    copyAndSelectStatements.execute(parsedSelectStatement, selectStatement);
+    assertTrue(copyAndSelectStatements.hasDmlOrCopyStatementsAfter(0));
+    assertFalse(copyAndSelectStatements.hasDmlOrCopyStatementsAfter(1));
+
+    BackendConnection copyAndUnknownStatements =
+        new BackendConnection(spannerConnection, DdlTransactionMode.Batch, ImmutableList.of());
+    copyAndUnknownStatements.executeCopy(
+        parsedCopyStatement, copyStatement, receiver, writer, executor);
+    copyAndUnknownStatements.execute(parsedUnknownStatement, unknownStatement);
+    assertTrue(copyAndUnknownStatements.hasDmlOrCopyStatementsAfter(0));
+    assertFalse(copyAndUnknownStatements.hasDmlOrCopyStatementsAfter(1));
+  }
 
   @Test
   public void testExecuteLocalStatement() throws ExecutionException, InterruptedException {
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java
index 0344746d..f44c88bb 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java
@@ -26,8 +26,10 @@ import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.ReadContext;
 import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerException;
@@ -47,11 +49,15 @@ import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
+import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Bytes;
 import java.io.ByteArrayInputStream;
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -90,9 +96,15 @@ public class StatementTest {
 
   @AfterClass
   public static void cleanup() {
+    deleteLogFile();
+  }
+
+  private static void deleteLogFile() {
     // TODO: Make error log file configurable and turn off writing to a file during tests.
-    File outputFile = new File("output.txt");
-    outputFile.delete();
+    try {
+      Files.deleteIfExists(new File("output.txt").toPath());
+    } catch (IOException ignore) {
+    }
   }
 
   @Test
@@ -411,8 +423,12 @@ public class StatementTest {
     String sql = "COPY keyvalue FROM STDIN;";
     CopyStatement statement =
         new CopyStatement(
-            connectionHandler, mock(OptionsMetadata.class), parse(sql), Statement.of(sql));
-    statement.execute();
+            connectionHandler, mock(OptionsMetadata.class), "", parse(sql), Statement.of(sql));
+
+    BackendConnection backendConnection =
+        new BackendConnection(connection, DdlTransactionMode.Batch, ImmutableList.of());
+    statement.executeAsync(backendConnection);
+    backendConnection.flush();
 
     byte[] payload = "2 3\n".getBytes();
     MutationWriter mutationWriter = statement.getMutationWriter();
@@ -470,21 +486,95 @@ public class StatementTest {
     assertThrows(IllegalStateException.class, intermediateStatement::getStatementResult);
   }
 
+  @Test
+  public void testCopyBatchSizeLimit() throws Exception {
+    setupQueryInformationSchemaResults();
+    BackendConnection backendConnection =
+        new BackendConnection(connection, DdlTransactionMode.Batch, ImmutableList.of());
+
+    byte[] payload = Files.readAllBytes(Paths.get("./src/test/resources/batch-size-test.txt"));
+
+    String sql = "COPY keyvalue FROM STDIN;";
+    CopyStatement copyStatement =
+        new CopyStatement(
+            connectionHandler, mock(OptionsMetadata.class), "", parse(sql), Statement.of(sql));
+
+    assertFalse(copyStatement.isExecuted());
+    copyStatement.executeAsync(backendConnection);
+    assertTrue(copyStatement.isExecuted());
+
+    backendConnection.flush();
+    MutationWriter mw = copyStatement.getMutationWriter();
+    //    // Inject a mock DatabaseClient for now.
+    //    // TODO: Fix this once we can use the Connection API.
+    //    Field databaseClientField = MutationWriter.class.getDeclaredField("databaseClient");
+    //    databaseClientField.setAccessible(true);
+    //    databaseClientField.set(mw, mock(DatabaseClient.class));
+    mw.addCopyData(payload);
+    mw.close();
+
+    assertEquals("TEXT", copyStatement.getFormatType());
+    assertEquals('\t', copyStatement.getDelimiterChar());
+    assertFalse(copyStatement.hasException());
+    assertEquals(12L, copyStatement.getUpdateCount());
+    assertEquals(12L, mw.getRowCount());
+
+    copyStatement.close();
+  }
+
+  @Test
+  public void testCopyDataRowLengthMismatchLimit() throws Exception {
+    setupQueryInformationSchemaResults();
+    BackendConnection backendConnection =
+        new BackendConnection(connection, DdlTransactionMode.Batch, ImmutableList.of());
+
+    byte[] payload = "1\t'one'\n2".getBytes();
+
+    String sql = "COPY keyvalue FROM STDIN;";
+    CopyStatement copyStatement =
+        new CopyStatement(
+            connectionHandler, mock(OptionsMetadata.class), "", parse(sql), Statement.of(sql));
+
+    assertFalse(copyStatement.isExecuted());
+    copyStatement.executeAsync(backendConnection);
+    assertTrue(copyStatement.isExecuted());
+
+    backendConnection.flush();
+    MutationWriter mw = copyStatement.getMutationWriter();
+    mw.addCopyData(payload);
+    mw.close();
+
+    SpannerException thrown = assertThrows(SpannerException.class, copyStatement::getUpdateCount);
+    assertEquals(ErrorCode.INVALID_ARGUMENT, thrown.getErrorCode());
+    assertEquals(
+        "INVALID_ARGUMENT: Invalid COPY data: Row length mismatched. Expected 2 columns, but only found 1",
+        thrown.getMessage());
+
+    copyStatement.close();
+
+    deleteLogFile();
+  }
+
   private void setupQueryInformationSchemaResults() {
+    DatabaseClient databaseClient = mock(DatabaseClient.class);
+    ReadContext singleUseReadContext = mock(ReadContext.class);
+    when(databaseClient.singleUse()).thenReturn(singleUseReadContext);
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
+    when(connection.getDatabaseClient()).thenReturn(databaseClient);
     ResultSet spannerType = mock(ResultSet.class);
     when(spannerType.getString("column_name")).thenReturn("key", "value");
     when(spannerType.getString("data_type")).thenReturn("bigint", "character varying");
     when(spannerType.next()).thenReturn(true, true, false);
-    when(connection.executeQuery(
+    when(singleUseReadContext.executeQuery(
             ArgumentMatchers.argThat(
                 statement ->
                     statement != null && statement.getSql().startsWith("SELECT column_name"))))
         .thenReturn(spannerType);
 
-    ResultSet countResult = Mockito.mock(ResultSet.class);
+    ResultSet countResult = mock(ResultSet.class);
     when(countResult.getLong(0)).thenReturn(2L);
     when(countResult.next()).thenReturn(true, false);
-    when(connection.executeQuery(
+    when(singleUseReadContext.executeQuery(
             ArgumentMatchers.argThat(
                 statement ->
                     statement != null && statement.getSql().startsWith("SELECT COUNT(*)"))))
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiverTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiverTest.java
new file mode 100644
index 00000000..a0658916
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiverTest.java
@@ -0,0 +1,72 @@
+// Copyright 2022 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.utils;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.SpannerException;
+import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
+import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
+import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
+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 CopyDataReceiverTest {
+
+  @Test
+  public void testCopyStatementWithException() {
+    SpannerException exception =
+        SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT, "Invalid copy statement");
+    CopyStatement statement = mock(CopyStatement.class);
+    when(statement.hasException()).thenReturn(true);
+    when(statement.getException()).thenReturn(exception);
+    ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+
+    CopyDataReceiver receiver = new CopyDataReceiver(statement, connectionHandler);
+    SpannerException spannerException = assertThrows(SpannerException.class, receiver::handleCopy);
+
+    assertSame(exception, spannerException);
+  }
+
+  @Test
+  public void testCopyStatementWithCopyFailed() {
+    DataOutputStream outputStream = new DataOutputStream(new ByteArrayOutputStream());
+    CopyStatement statement = mock(CopyStatement.class);
+    ConnectionMetadata connectionMetadata = mock(ConnectionMetadata.class);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
+    ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
+    when(connectionHandler.getStatus())
+        .thenReturn(ConnectionStatus.AUTHENTICATED, ConnectionStatus.COPY_FAILED);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+
+    CopyDataReceiver receiver = new CopyDataReceiver(statement, connectionHandler);
+    SpannerException spannerException = assertThrows(SpannerException.class, receiver::handleCopy);
+
+    assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode());
+    assertEquals("INTERNAL: Copy failed with unknown reason", spannerException.getMessage());
+  }
+}
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java
index 0842a7bf..4edfe88d 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/MutationWriterTest.java
@@ -34,6 +34,7 @@ import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.Value;
 import com.google.cloud.spanner.connection.Connection;
+import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter.CopyTransactionMode;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -128,11 +129,12 @@ public class MutationWriterTest {
             false);
 
     mutationWriter.addCopyData("1\t\"One\"\n2\t\"Two\"\n".getBytes(StandardCharsets.UTF_8));
+    mutationWriter.commit();
     mutationWriter.close();
 
-    long updateCount = mutationWriter.call();
+    StatementResult updateCount = mutationWriter.call();
 
-    assertEquals(2L, updateCount);
+    assertEquals(2L, updateCount.getUpdateCount().longValue());
     List<Mutation> expectedMutations =
         ImmutableList.of(
             Mutation.newInsertBuilder("numbers").set("number").to(1L).set("name").to("One").build(),
@@ -208,11 +210,12 @@ public class MutationWriterTest {
       mutationWriter.addCopyData(
           "1\t\"One\"\n2\t\"Two\"\n3\t\"Three\"\n4\t\"Four\"\n5\t\"Five\"\n"
               .getBytes(StandardCharsets.UTF_8));
+      mutationWriter.commit();
       mutationWriter.close();
 
-      long updateCount = mutationWriter.call();
+      StatementResult updateCount = mutationWriter.call();
 
-      assertEquals(5L, updateCount);
+      assertEquals(5L, updateCount.getUpdateCount().longValue());
       verify(databaseClient, times(3)).write(anyIterable());
     } finally {
       System.getProperties().remove("copy_in_mutation_limit");
@@ -285,11 +288,12 @@ public class MutationWriterTest {
       mutationWriter.addCopyData(
           "1\t\"One\"\n2\t\"Two\"\n3\t\"Three\"\n4\t\"Four\"\n5\t\"Five\"\n"
               .getBytes(StandardCharsets.UTF_8));
+      mutationWriter.commit();
       mutationWriter.close();
 
-      long updateCount = mutationWriter.call();
+      StatementResult updateCount = mutationWriter.call();
 
-      assertEquals(5L, updateCount);
+      assertEquals(5L, updateCount.getUpdateCount().longValue());
       // We expect two batches, because:
       // 1. The commit limit is 80 bytes. That is for safety halved down to 40 bytes.
       // 2. Each record is 20 bytes.
@@ -329,6 +333,7 @@ public class MutationWriterTest {
     mutationWriter.addCopyData(
         "1\t\"One\"\n2\t\"Two\"\n3\t\"Three\"\n4\t\"Four\"\n5\t\"Five\"\n"
             .getBytes(StandardCharsets.UTF_8));
+    mutationWriter.commit();
     mutationWriter.close();
 
     ExecutorService executor = Executors.newFixedThreadPool(2);
@@ -344,9 +349,9 @@ public class MutationWriterTest {
           mutationWriter.close();
           return null;
         });
-    Future<Long> updateCount = executor.submit(mutationWriter);
+    Future<StatementResult> updateCount = executor.submit(mutationWriter);
 
-    assertEquals(5L, updateCount.get().longValue());
+    assertEquals(5L, updateCount.get().getUpdateCount().longValue());
     List<Mutation> expectedMutations =
         ImmutableList.of(
             Mutation.newInsertBuilder("numbers").set("number").to(1L).set("name").to("One").build(),
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
index 6bea11f1..8ef439ba 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
@@ -14,15 +14,15 @@
 
 package com.google.cloud.spanner.pgadapter.wireprotocol;
 
+import static com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage.MAX_INVALID_MESSAGE_COUNT;
 import static org.junit.Assert.assertArrayEquals;
 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.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -33,8 +33,10 @@ import static org.mockito.Mockito.when;
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.ReadContext;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerException;
+import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractStatementParser;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
@@ -52,7 +54,6 @@ import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
 import com.google.cloud.spanner.pgadapter.statements.ExtendedQueryProtocolHandler;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePortalStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
-import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage.ManuallyCreatedToken;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage.PreparedType;
@@ -935,7 +936,8 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    Exception testException = new Exception("test error");
+    SpannerException testException =
+        SpannerExceptionFactory.newSpannerException(ErrorCode.INVALID_ARGUMENT, "test error");
     when(intermediatePortalStatement.hasException()).thenReturn(true);
     when(intermediatePortalStatement.getException()).thenReturn(testException);
     when(connectionHandler.getPortal(anyString())).thenReturn(intermediatePortalStatement);
@@ -1112,10 +1114,8 @@ public class ProtocolTest {
     byte[] value = Bytes.concat(messageMetadata, length);
 
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
-    ByteArrayOutputStream result = new ByteArrayOutputStream();
-    DataOutputStream outputStream = new DataOutputStream(result);
+    DataOutputStream outputStream = mock(DataOutputStream.class);
 
-    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
@@ -1127,11 +1127,8 @@ public class ProtocolTest {
 
     message.send();
 
-    // ReadyResponse
-    DataInputStream outputResult = inputStreamFromOutputStream(result);
-    assertEquals('Z', outputResult.readByte());
-    assertEquals(5, outputResult.readInt());
-    assertEquals('I', outputResult.readByte());
+    verify(extendedQueryProtocolHandler).flush();
+    verify(extendedQueryProtocolHandler, never()).sync();
   }
 
   @Test
@@ -1143,13 +1140,10 @@ public class ProtocolTest {
     byte[] value = Bytes.concat(messageMetadata, length);
 
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
-    ByteArrayOutputStream result = new ByteArrayOutputStream();
-    DataOutputStream outputStream = new DataOutputStream(result);
+    DataOutputStream outputStream = mock(DataOutputStream.class);
 
-    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
-    when(connection.isInTransaction()).thenReturn(true);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
@@ -1159,11 +1153,8 @@ public class ProtocolTest {
 
     message.send();
 
-    // ReadyResponse
-    DataInputStream outputResult = inputStreamFromOutputStream(result);
-    assertEquals('Z', outputResult.readByte());
-    assertEquals(5, outputResult.readInt());
-    assertEquals('T', outputResult.readByte());
+    verify(extendedQueryProtocolHandler).flush();
+    verify(extendedQueryProtocolHandler, never()).sync();
   }
 
   @Test
@@ -1283,6 +1274,26 @@ public class ProtocolTest {
     verify(mw).addCopyData(payload);
   }
 
+  @Test
+  public void testCopyDataMessageWithNoCopyStatement() throws Exception {
+    byte[] messageMetadata = {'d'};
+    byte[] payload = "This is the payload".getBytes();
+    byte[] length = intToBytes(4 + payload.length);
+    byte[] value = Bytes.concat(messageMetadata, length, payload);
+
+    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
+
+    when(connectionHandler.getActiveStatement()).thenReturn(null);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
+
+    WireMessage message = ControlMessage.create(connectionHandler);
+    // This should be a no-op.
+    message.sendPayload();
+  }
+
   @Test
   public void testMultipleCopyDataMessages() throws Exception {
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
@@ -1303,16 +1314,10 @@ public class ProtocolTest {
     DataInputStream inputStream2 = new DataInputStream(new ByteArrayInputStream(value2));
     DataInputStream inputStream3 = new DataInputStream(new ByteArrayInputStream(value3));
 
-    ResultSet spannerType = mock(ResultSet.class);
-    when(spannerType.getString("column_name")).thenReturn("key", "value");
-    when(spannerType.getString("data_type")).thenReturn("bigint", "character varying");
-    when(spannerType.next()).thenReturn(true, true, false);
-    when(connection.executeQuery(any(Statement.class))).thenReturn(spannerType);
-
     String sql = "COPY keyvalue FROM STDIN;";
     CopyStatement copyStatement =
-        new CopyStatement(connectionHandler, options, parse(sql), Statement.of(sql));
-    copyStatement.execute();
+        new CopyStatement(connectionHandler, options, "", parse(sql), Statement.of(sql));
+    copyStatement.executeAsync(mock(BackendConnection.class));
 
     when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
@@ -1355,24 +1360,22 @@ public class ProtocolTest {
     DataOutputStream outputStream = new DataOutputStream(result);
 
     CopyStatement copyStatement = mock(CopyStatement.class);
+    MutationWriter mutationWriter = mock(MutationWriter.class);
+    when(copyStatement.getMutationWriter()).thenReturn(mutationWriter);
     when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
-    MutationWriter mb = mock(MutationWriter.class);
-    when(copyStatement.getMutationWriter()).thenReturn(mb);
     WireMessage message = ControlMessage.create(connectionHandler);
 
     assertEquals(CopyDoneMessage.class, message.getClass());
     CopyDoneMessage messageSpy = (CopyDoneMessage) spy(message);
-    doNothing()
-        .when(messageSpy)
-        .sendSpannerResult(any(IntermediateStatement.class), any(QueryMode.class), anyLong());
 
     messageSpy.send();
-    verify(messageSpy).sendSpannerResult(copyStatement, QueryMode.SIMPLE, 0L);
+    verify(messageSpy).sendPayload();
+    verify(mutationWriter).commit();
   }
 
   @Test
@@ -1415,8 +1418,8 @@ public class ProtocolTest {
     String sql = "COPY keyvalue FROM STDIN;";
     CopyStatement copyStatement =
         new CopyStatement(
-            connectionHandler, mock(OptionsMetadata.class), parse(sql), Statement.of(sql));
-    copyStatement.execute();
+            connectionHandler, mock(OptionsMetadata.class), "", parse(sql), Statement.of(sql));
+    copyStatement.executeAsync(mock(BackendConnection.class));
 
     MutationWriter mw = copyStatement.getMutationWriter();
     mw.addCopyData(payload);
@@ -1428,113 +1431,6 @@ public class ProtocolTest {
     verify(resultSet, never()).close();
   }
 
-  @Test
-  public void testCopyBatchSizeLimit() throws Exception {
-    setupQueryInformationSchemaResults();
-    when(connection.getDatabaseClient()).thenReturn(mock(DatabaseClient.class));
-
-    byte[] payload = Files.readAllBytes(Paths.get("./src/test/resources/batch-size-test.txt"));
-
-    String sql = "COPY keyvalue FROM STDIN;";
-    CopyStatement copyStatement =
-        new CopyStatement(
-            connectionHandler, mock(OptionsMetadata.class), parse(sql), Statement.of(sql));
-
-    assertFalse(copyStatement.isExecuted());
-    copyStatement.execute();
-    assertTrue(copyStatement.isExecuted());
-
-    MutationWriter mw = copyStatement.getMutationWriter();
-    mw.addCopyData(payload);
-    mw.close();
-
-    assertEquals("TEXT", copyStatement.getFormatType());
-    assertEquals('\t', copyStatement.getDelimiterChar());
-    assertFalse(copyStatement.hasException());
-    assertEquals(12L, copyStatement.getUpdateCount());
-    assertEquals(12L, mw.getRowCount());
-
-    copyStatement.close();
-  }
-
-  @Test
-  public void testCopyDataRowLengthMismatchLimit() throws Exception {
-    setupQueryInformationSchemaResults();
-
-    byte[] payload = "1\t'one'\n2".getBytes();
-
-    String sql = "COPY keyvalue FROM STDIN;";
-    CopyStatement copyStatement =
-        new CopyStatement(
-            connectionHandler, mock(OptionsMetadata.class), parse(sql), Statement.of(sql));
-
-    assertFalse(copyStatement.isExecuted());
-    copyStatement.execute();
-    assertTrue(copyStatement.isExecuted());
-
-    MutationWriter mw = copyStatement.getMutationWriter();
-    mw.addCopyData(payload);
-    mw.close();
-
-    SpannerException thrown = assertThrows(SpannerException.class, copyStatement::getUpdateCount);
-    assertEquals(ErrorCode.INVALID_ARGUMENT, thrown.getErrorCode());
-    assertEquals(
-        "INVALID_ARGUMENT: Invalid COPY data: Row length mismatched. Expected 2 columns, but only found 1",
-        thrown.getMessage());
-
-    copyStatement.close();
-  }
-
-  @Test
-  public void testCopyResumeErrorOutputFile() throws Exception {
-    setupQueryInformationSchemaResults();
-
-    byte[] payload = Files.readAllBytes(Paths.get("./src/test/resources/test-copy-output.txt"));
-
-    String sql = "COPY keyvalue FROM STDIN;";
-    CopyStatement copyStatement =
-        new CopyStatement(
-            connectionHandler, mock(OptionsMetadata.class), parse(sql), Statement.of(sql));
-    assertFalse(copyStatement.isExecuted());
-    copyStatement.execute();
-    assertTrue(copyStatement.isExecuted());
-
-    MutationWriter mw = copyStatement.getMutationWriter();
-    mw.addCopyData(payload);
-
-    SpannerException thrown = assertThrows(SpannerException.class, copyStatement::getUpdateCount);
-    assertEquals(ErrorCode.INVALID_ARGUMENT, thrown.getErrorCode());
-    assertEquals(
-        "INVALID_ARGUMENT: Invalid input syntax for type INT64:\"'5'\"", thrown.getMessage());
-
-    copyStatement.close();
-  }
-
-  @Test
-  public void testCopyResumeErrorStartOutputFile() throws Exception {
-    setupQueryInformationSchemaResults();
-
-    byte[] payload =
-        Files.readAllBytes(Paths.get("./src/test/resources/test-copy-start-output.txt"));
-
-    String sql = "COPY keyvalue FROM STDIN;";
-    CopyStatement copyStatement =
-        new CopyStatement(connectionHandler, options, parse(sql), Statement.of(sql));
-    assertFalse(copyStatement.isExecuted());
-    copyStatement.execute();
-    assertTrue(copyStatement.isExecuted());
-
-    MutationWriter mw = copyStatement.getMutationWriter();
-    mw.addCopyData(payload);
-
-    SpannerException thrown = assertThrows(SpannerException.class, copyStatement::getUpdateCount);
-    assertEquals(ErrorCode.INVALID_ARGUMENT, thrown.getErrorCode());
-    assertEquals(
-        "INVALID_ARGUMENT: Invalid input syntax for type INT64:\"'1'\"", thrown.getMessage());
-
-    copyStatement.close();
-  }
-
   @Test
   public void testFunctionCallMessageThrowsException() throws Exception {
     byte[] messageMetadata = {'F'};
@@ -1801,13 +1697,184 @@ public class ProtocolTest {
     assertThrows(IllegalStateException.class, describeMessage::getPortalMetadata);
   }
 
+  @Test
+  public void testSkipMessage() throws Exception {
+    byte[] messageMetadata = {0, 0, 0, 45};
+    String payload = "INSERT INTO users (name) VALUES ('test')\0";
+    byte[] value = Bytes.concat(messageMetadata, payload.getBytes());
+
+    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
+    ByteArrayOutputStream result = new ByteArrayOutputStream();
+    DataOutputStream outputStream = new DataOutputStream(result);
+
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+
+    SkipMessage message = new SkipMessage(connectionHandler);
+    message.send();
+
+    // Verify that nothing was written to the output.
+    assertEquals(0, result.size());
+    assertEquals("Skip", message.getMessageName());
+    assertEquals("", message.getIdentifier());
+    assertEquals("Length: 45", message.getPayloadString());
+  }
+
+  @Test
+  public void testFlushSkippedInCopyMode() throws Exception {
+    byte[] messageMetadata = {FlushMessage.IDENTIFIER, 0, 0, 0, 4};
+
+    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(messageMetadata));
+    ByteArrayOutputStream result = new ByteArrayOutputStream();
+    DataOutputStream outputStream = new DataOutputStream(result);
+
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
+
+    ControlMessage message = ControlMessage.create(connectionHandler);
+
+    assertEquals(SkipMessage.class, message.getClass());
+    // Verify that nothing was written to the output.
+    assertEquals(0, result.size());
+  }
+
+  @Test
+  public void testSyncSkippedInCopyMode() throws Exception {
+    byte[] messageMetadata = {SyncMessage.IDENTIFIER, 0, 0, 0, 4};
+
+    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(messageMetadata));
+    ByteArrayOutputStream result = new ByteArrayOutputStream();
+    DataOutputStream outputStream = new DataOutputStream(result);
+
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
+
+    ControlMessage message = ControlMessage.create(connectionHandler);
+
+    assertEquals(SkipMessage.class, message.getClass());
+    // Verify that nothing was written to the output.
+    assertEquals(0, result.size());
+  }
+
+  @Test
+  public void testCopyDataSkippedInNormalMode() throws Exception {
+    byte[] messageMetadata = {CopyDataMessage.IDENTIFIER, 0, 0, 0, 4};
+    String payload = "1\t'One'\n2\t'Two'\n";
+    byte[] value = Bytes.concat(messageMetadata, payload.getBytes());
+
+    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
+    ByteArrayOutputStream result = new ByteArrayOutputStream();
+    DataOutputStream outputStream = new DataOutputStream(result);
+
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.AUTHENTICATED);
+
+    ControlMessage message = ControlMessage.create(connectionHandler);
+
+    assertEquals(SkipMessage.class, message.getClass());
+    // Verify that nothing was written to the output.
+    assertEquals(0, result.size());
+  }
+
+  @Test
+  public void testCopyDoneSkippedInNormalMode() throws Exception {
+    byte[] messageMetadata = {CopyDoneMessage.IDENTIFIER, 0, 0, 0, 4};
+
+    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(messageMetadata));
+    ByteArrayOutputStream result = new ByteArrayOutputStream();
+    DataOutputStream outputStream = new DataOutputStream(result);
+
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.AUTHENTICATED);
+
+    ControlMessage message = ControlMessage.create(connectionHandler);
+
+    assertEquals(SkipMessage.class, message.getClass());
+    // Verify that nothing was written to the output.
+    assertEquals(0, result.size());
+  }
+
+  @Test
+  public void testCopyFailSkippedInNormalMode() throws Exception {
+    byte[] messageMetadata = {CopyFailMessage.IDENTIFIER, 0, 0, 0, 4};
+
+    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(messageMetadata));
+    ByteArrayOutputStream result = new ByteArrayOutputStream();
+    DataOutputStream outputStream = new DataOutputStream(result);
+
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.AUTHENTICATED);
+
+    ControlMessage message = ControlMessage.create(connectionHandler);
+
+    assertEquals(SkipMessage.class, message.getClass());
+    // Verify that nothing was written to the output.
+    assertEquals(0, result.size());
+  }
+
+  @Test
+  public void testRepeatedCopyDataInNormalMode_TerminatesConnectionAndReturnsError()
+      throws Exception {
+    String payload = "1\t'One'\n2\t'Two'\n";
+    byte[] messageMetadata = {
+      CopyDataMessage.IDENTIFIER,
+      0,
+      0,
+      0,
+      (byte) (4 + payload.getBytes(StandardCharsets.UTF_8).length)
+    };
+    byte[] value = new byte[0];
+    for (int i = 0; i <= MAX_INVALID_MESSAGE_COUNT; i++) {
+      value = Bytes.concat(value, messageMetadata, payload.getBytes());
+    }
+
+    DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
+    ByteArrayOutputStream result = new ByteArrayOutputStream();
+    DataOutputStream outputStream = new DataOutputStream(result);
+
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
+    when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
+    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.AUTHENTICATED);
+    doCallRealMethod().when(connectionHandler).increaseInvalidMessageCount();
+    when(connectionHandler.getInvalidMessageCount()).thenCallRealMethod();
+
+    for (int i = 0; i < MAX_INVALID_MESSAGE_COUNT; i++) {
+      ControlMessage message = ControlMessage.create(connectionHandler);
+      assertEquals(SkipMessage.class, message.getClass());
+      // Verify that nothing was written to the output.
+      assertEquals(0, result.size());
+      verify(connectionHandler, never()).setStatus(ConnectionStatus.TERMINATED);
+    }
+
+    ControlMessage.create(connectionHandler);
+    verify(connectionHandler).setStatus(ConnectionStatus.TERMINATED);
+    byte[] resultBytes = result.toByteArray();
+    assertEquals('E', resultBytes[0]);
+  }
+
   private void setupQueryInformationSchemaResults() {
+    DatabaseClient databaseClient = mock(DatabaseClient.class);
+    ReadContext singleUseReadContext = mock(ReadContext.class);
+    when(databaseClient.singleUse()).thenReturn(singleUseReadContext);
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
+    when(connection.getDatabaseClient()).thenReturn(databaseClient);
     ResultSet spannerType = mock(ResultSet.class);
     when(spannerType.getString("column_name")).thenReturn("key", "value");
     when(spannerType.getString("data_type")).thenReturn("bigint", "character varying");
     when(spannerType.next()).thenReturn(true, true, false);
-    when(connection.executeQuery(
+    when(singleUseReadContext.executeQuery(
             ArgumentMatchers.argThat(
                 statement ->
                     statement != null && statement.getSql().startsWith("SELECT column_name"))))
@@ -1816,7 +1883,7 @@ public class ProtocolTest {
     ResultSet countResult = mock(ResultSet.class);
     when(countResult.getLong(0)).thenReturn(2L);
     when(countResult.next()).thenReturn(true, false);
-    when(connection.executeQuery(
+    when(singleUseReadContext.executeQuery(
             ArgumentMatchers.argThat(
                 statement ->
                     statement != null && statement.getSql().startsWith("SELECT COUNT(*)"))))