Repository: GoogleCloudPlatform/pgadapter
Pull request #356 (2022-09-12 07:31:55)

fix: cancel requests were ignored
Cancel requests for a (long-running) query or other statement was ignored by PGAdapter. This fix makes sure that both long-running queries, updates and COPY operations can be cancelled.

Fixes b/245128300

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#356](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (78682b5) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/245cd7cc2d422a1b5d0adf1dd3159cf82a0de6b4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (245cd7c) will **decrease** coverage by `0.26%`.
> The diff coverage is `34.04%`.

> :exclamation: Current head 78682b5 differs from pull request most recent head a788daf. Consider uploading reports for the commit a788daf to get more accurate results

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #356      +/-   ##
========================================================
- Coverage                 85.82%   85.55%   -0.27%     
  Complexity                 1749     1749              
========================================================
  Files                       119      119              
  Lines                      5827     5843      +16     
  Branches                    794      797       +3     
========================================================
- Hits                       5001     4999       -2     
- Misses                      599      615      +16     
- Partials                    227      229       +2     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `85.55% <34.04%> (-0.27%)` | :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/356?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...loud/spanner/pgadapter/statements/DdlExecutor.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvRGRsRXhlY3V0b3IuamF2YQ==) | `91.27% <0.00%> (ø)` | |
| [...apter/statements/ExtendedQueryProtocolHandler.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvRXh0ZW5kZWRRdWVyeVByb3RvY29sSGFuZGxlci5qYXZh) | `93.02% <0.00%> (-4.54%)` | :arrow_down: |
| [...er/pgadapter/statements/IntermediateStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvSW50ZXJtZWRpYXRlU3RhdGVtZW50LmphdmE=) | `93.33% <0.00%> (-3.22%)` | :arrow_down: |
| [.../cloud/spanner/pgadapter/utils/MutationWriter.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL011dGF0aW9uV3JpdGVyLmphdmE=) | `90.73% <0.00%> (-1.35%)` | :arrow_down: |
| [...spanner/pgadapter/wireprotocol/ControlMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9Db250cm9sTWVzc2FnZS5qYXZh) | `86.69% <0.00%> (-0.87%)` | :arrow_down: |
| [...panner/pgadapter/wireprotocol/DescribeMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9EZXNjcmliZU1lc3NhZ2UuamF2YQ==) | `85.71% <0.00%> (ø)` | |
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `90.48% <20.00%> (-0.43%)` | :arrow_down: |
| [...gle/cloud/spanner/pgadapter/ConnectionHandler.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL0Nvbm5lY3Rpb25IYW5kbGVyLmphdmE=) | `77.54% <42.10%> (-1.54%)` | :arrow_down: |
| [...loud/spanner/pgadapter/utils/CopyDataReceiver.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0NvcHlEYXRhUmVjZWl2ZXIuamF2YQ==) | `94.11% <50.00%> (-5.89%)` | :arrow_down: |
| [...ud/spanner/pgadapter/error/PGExceptionFactory.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL2Vycm9yL1BHRXhjZXB0aW9uRmFjdG9yeS5qYXZh) | `96.00% <66.66%> (-4.00%)` | :arrow_down: |
| ... and [4 more](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/356/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | |

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


Patch:
diff --git a/clirr-ignored-differences.xml b/clirr-ignored-differences.xml
index a5d8d47d..00c5b6cd 100644
--- a/clirr-ignored-differences.xml
+++ b/clirr-ignored-differences.xml
@@ -55,6 +55,29 @@
     <className>com/google/cloud/spanner/pgadapter/wireprotocol/AbstractQueryProtocolMessage</className>
     <method>java.lang.String getSql()</method>
   </difference>
+  <!-- ActiveStatement has been removed as it is not needed. Instead, we keep track of which connections are registered and just use the cancel feature built into the Connection API. -->
+  <difference>
+    <differenceType>7002</differenceType>
+    <className>com/google/cloud/spanner/pgadapter/ConnectionHandler</className>
+    <method>void addActiveStatement(com.google.cloud.spanner.pgadapter.statements.IntermediateStatement)</method>
+  </difference>
+  <difference>
+    <differenceType>7002</differenceType>
+    <className>com/google/cloud/spanner/pgadapter/ConnectionHandler</className>
+    <method>com.google.cloud.spanner.pgadapter.statements.IntermediateStatement getActiveStatement()</method>
+  </difference>
+  <difference>
+    <differenceType>7002</differenceType>
+    <className>com/google/cloud/spanner/pgadapter/ConnectionHandler</className>
+    <method>void removeActiveStatement(com.google.cloud.spanner.pgadapter.statements.IntermediateStatement)</method>
+  </difference>
+  <difference>
+    <differenceType>7006</differenceType>
+    <className>com/google/cloud/spanner/pgadapter/ConnectionHandler</className>
+    <method>void cancelActiveStatement(int, int)</method>
+    <from>void</from>
+    <to>boolean</to>
+  </difference>
   <difference>
     <differenceType>7005</differenceType>
     <className>com/google/cloud/spanner/connection/ConnectionOptionsHelper</className>
@@ -68,6 +91,7 @@
     <to>void connectToSpanner(java.lang.String, com.google.auth.Credentials)</to>
   </difference>
 
+  7006: com.google.cloud.spanner.pgadapter.ConnectionHandler: Return type of method 'public void cancelActiveStatement(int, int)' has been changed to boolean
   <!-- Add ignores for all sub packages, as these are considered internal. -->
   <difference>
     <differenceType>1001</differenceType>
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 fc14fbe4..6925c8ce 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -39,6 +39,7 @@ import com.google.cloud.spanner.pgadapter.error.Severity;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.SslMode;
+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;
@@ -50,6 +51,7 @@ import com.google.cloud.spanner.pgadapter.wireoutput.TerminateResponse;
 import com.google.cloud.spanner.pgadapter.wireprotocol.BootstrapMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.SSLMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.spanner.admin.database.v1.InstanceName;
 import com.google.spanner.v1.DatabaseName;
 import java.io.DataOutputStream;
@@ -91,11 +93,10 @@ public class ConnectionHandler extends Thread {
   private Socket socket;
   private final Map<String, IntermediatePreparedStatement> statementsMap = new HashMap<>();
   private final Map<String, IntermediatePortalStatement> portalsMap = new HashMap<>();
-  private static final Map<Integer, IntermediateStatement> activeStatementsMap =
+  private static final Map<Integer, ConnectionHandler> CONNECTION_HANDLERS =
       new ConcurrentHashMap<>();
-  private static final Map<Integer, Integer> connectionToSecretMapping = new ConcurrentHashMap<>();
   private volatile ConnectionStatus status = ConnectionStatus.UNAUTHENTICATED;
-  private int connectionId;
+  private final int connectionId;
   private final int secret;
   // Separate the following from the threat ID generator, since PG connection IDs are maximum
   //  32 bytes, and shouldn't be incremented on failed startups.
@@ -107,12 +108,21 @@ public class ConnectionHandler extends Thread {
   private DatabaseId databaseId;
   private WellKnownClient wellKnownClient;
   private ExtendedQueryProtocolHandler extendedQueryProtocolHandler;
+  private CopyStatement activeCopyStatement;
 
   ConnectionHandler(ProxyServer server, Socket socket) {
+    this(server, socket, null);
+  }
+
+  /** Constructor only for testing. */
+  @VisibleForTesting
+  ConnectionHandler(ProxyServer server, Socket socket, Connection spannerConnection) {
     super("ConnectionHandler-" + CONNECTION_HANDLER_ID_GENERATOR.incrementAndGet());
     this.server = server;
     this.socket = socket;
     this.secret = new SecureRandom().nextInt();
+    this.connectionId = incrementingConnectionId.incrementAndGet();
+    CONNECTION_HANDLERS.put(this.connectionId, this);
     setDaemon(true);
     logger.log(
         Level.INFO,
@@ -120,6 +130,7 @@ public class ConnectionHandler extends Thread {
             String.format(
                 "Connection handler with ID %s created for client %s",
                 getName(), socket.getInetAddress().getHostAddress()));
+    this.spannerConnection = spannerConnection;
   }
 
   void createSSLSocket() throws IOException {
@@ -498,60 +509,50 @@ public class ConnectionHandler extends Thread {
     return this.portalsMap.containsKey(portalName);
   }
 
-  /**
-   * Add a currently executing statement to a buffer. This is only used in case a statement in
-   * flight is cancelled. It must be saved separately, as a new connection is spawned to issue a
-   * cancellation (as per Postgres protocol standard). This means some sort of IPC is required,
-   * which in this case is a global Map.
-   *
-   * @param statement Currently executing statement to be saved.
-   */
-  public synchronized void addActiveStatement(IntermediateStatement statement) {
-    activeStatementsMap.put(this.connectionId, statement);
-  }
-
-  /**
-   * Remove a statement from the buffer if it is currently executing. For more information on this
-   * use-case, read addActiveStatement comment.
-   *
-   * @param statement The statement to be removed.
-   */
-  public synchronized void removeActiveStatement(IntermediateStatement statement) {
-    if (activeStatementsMap.get(this.connectionId) == statement) {
-      activeStatementsMap.remove(this.connectionId);
-    }
-  }
-
   /**
    * To be used by a cancellation command to cancel a currently running statement, as contained in a
    * specific connection identified by connectionId. Since cancellation is a flimsy contract at
    * best, it is not imperative that the cancellation run, but it should be attempted nonetheless.
    *
-   * @param connectionId The connection owhose statement must be cancelled.
-   * @param secret The secret value linked to this connection. If it does not match, we cannot
-   *     cancel.
+   * @param connectionId The connection whose statement must be cancelled.
+   * @param secret The secret value linked to the connection that is being cancelled. If it does not
+   *     match, we cannot cancel.
+   * @return true if the statement was cancelled.
    */
-  public synchronized void cancelActiveStatement(int connectionId, int secret) {
-    int expectedSecret = ConnectionHandler.connectionToSecretMapping.get(connectionId);
-    if (secret != expectedSecret) {
+  public boolean cancelActiveStatement(int connectionId, int secret) {
+    if (connectionId == this.connectionId) {
+      // You can't cancel your own statement.
+      return false;
+    }
+    ConnectionHandler connectionToCancel = CONNECTION_HANDLERS.get(connectionId);
+    if (connectionToCancel == null) {
+      logger.log(
+          Level.WARNING,
+          () ->
+              MessageFormat.format(
+                  "User attempted to cancel an unknown connection. Connection: {0}", connectionId));
+      return false;
+    }
+    if (secret != connectionToCancel.secret) {
       logger.log(
           Level.WARNING,
           () ->
               MessageFormat.format(
                   "User attempted to cancel a connection with the incorrect secret."
-                      + "Connection: {}, Secret: {}, Expected Secret: {}",
-                  connectionId,
-                  secret,
-                  expectedSecret));
+                      + "Connection: {0}, Secret: {1}, Expected Secret: {2}",
+                  connectionId, secret, connectionToCancel.secret));
       // Since the user does not accept a response, there is no need to except here: simply return.
-      return;
+      return false;
     }
-    if (activeStatementsMap.containsKey(connectionId)) {
-      IntermediateStatement statement = activeStatementsMap.remove(connectionId);
-      // We can mostly ignore the exception since cancel does not expect any result (positive or
-      // otherwise)
-      statement.getConnection().cancel();
+    // We can mostly ignore the exception since cancel does not expect any result (positive or
+    // otherwise)
+    try {
+      connectionToCancel.getSpannerConnection().cancel();
+      connectionToCancel.interrupt();
+      return true;
+    } catch (Throwable ignore) {
     }
+    return false;
   }
 
   public IntermediatePreparedStatement getStatement(String statementName) {
@@ -581,6 +582,18 @@ public class ConnectionHandler extends Thread {
     }
   }
 
+  public void setActiveCopyStatement(CopyStatement copyStatement) {
+    this.activeCopyStatement = copyStatement;
+  }
+
+  public CopyStatement getActiveCopyStatement() {
+    return this.activeCopyStatement;
+  }
+
+  public void clearActiveCopyStatement() {
+    this.activeCopyStatement = null;
+  }
+
   public boolean hasStatement(String statementName) {
     return this.statementsMap.containsKey(statementName);
   }
@@ -598,10 +611,6 @@ public class ConnectionHandler extends Thread {
   }
 
   public int getConnectionId() {
-    if (this.connectionId == 0) {
-      this.connectionId = ConnectionHandler.incrementingConnectionId.incrementAndGet();
-      ConnectionHandler.connectionToSecretMapping.put(this.connectionId, this.secret);
-    }
     return this.connectionId;
   }
 
@@ -646,10 +655,6 @@ public class ConnectionHandler extends Thread {
     return extendedQueryProtocolHandler;
   }
 
-  public synchronized IntermediateStatement getActiveStatement() {
-    return activeStatementsMap.get(this.connectionId);
-  }
-
   public synchronized ConnectionStatus getStatus() {
     return status;
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/error/PGExceptionFactory.java b/src/main/java/com/google/cloud/spanner/pgadapter/error/PGExceptionFactory.java
index ed9eb8b9..fabc9998 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/error/PGExceptionFactory.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/error/PGExceptionFactory.java
@@ -28,12 +28,25 @@ public class PGExceptionFactory {
    * SQLState#RaiseException}.
    */
   public static PGException newPGException(String message) {
+    return newPGException(message, SQLState.RaiseException);
+  }
+
+  /**
+   * Creates a basic {@link PGException} with {@link Severity#ERROR} and the specified {@link
+   * SQLState}.
+   */
+  public static PGException newPGException(String message, SQLState sqlState) {
     return PGException.newBuilder(message)
         .setSeverity(Severity.ERROR)
-        .setSQLState(SQLState.RaiseException)
+        .setSQLState(sqlState)
         .build();
   }
 
+  /** Creates a new exception that indicates that the current query was cancelled by the client. */
+  public static PGException newQueryCancelledException() {
+    return newPGException("Query cancelled", SQLState.QueryCanceled);
+  }
+
   /** Converts the given {@link SpannerException} to a {@link PGException}. */
   public static PGException toPGException(SpannerException spannerException) {
     return newPGException(extractMessage(spannerException));
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 bd910e59..90d7c331 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
@@ -44,6 +44,7 @@ import com.google.cloud.spanner.connection.ConnectionOptionsHelper;
 import com.google.cloud.spanner.connection.ResultSetHelper;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
+import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.DdlTransactionMode;
 import com.google.cloud.spanner.pgadapter.session.SessionState;
@@ -228,7 +229,11 @@ public class BackendConnection {
             throw exception;
           }
         }
-        result.setException(spannerException);
+        if (spannerException.getErrorCode() == ErrorCode.CANCELLED || Thread.interrupted()) {
+          result.setException(PGExceptionFactory.newQueryCancelledException());
+        } else {
+          result.setException(spannerException);
+        }
         throw spannerException;
       } catch (Exception exception) {
         result.setException(exception);
@@ -360,8 +365,8 @@ public class BackendConnection {
         result.setException(executionException.getCause());
         throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
       } catch (InterruptedException interruptedException) {
-        result.setException(SpannerExceptionFactory.propagateInterrupt(interruptedException));
-        throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
+        result.setException(PGExceptionFactory.newQueryCancelledException());
+        throw PGExceptionFactory.newQueryCancelledException();
       } catch (Exception exception) {
         result.setException(exception);
         throw exception;
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/DdlExecutor.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/DdlExecutor.java
index bec201df..ef947b46 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/DdlExecutor.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/DdlExecutor.java
@@ -30,6 +30,7 @@ import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStateme
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.ConnectionOptionsHelper;
 import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.statements.SimpleParser.TableOrIndexName;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
@@ -88,7 +89,7 @@ class DdlExecutor {
         }
         return Boolean.TRUE;
       } catch (InterruptedException interruptedException) {
-        throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
+        throw PGExceptionFactory.newQueryCancelledException();
       }
     }
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java
index a2bc681d..ae61125d 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java
@@ -15,6 +15,7 @@
 package com.google.cloud.spanner.pgadapter.statements;
 
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.wireprotocol.AbstractQueryProtocolMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.SyncMessage;
 import com.google.common.annotations.VisibleForTesting;
@@ -129,6 +130,9 @@ public class ExtendedQueryProtocolHandler {
           break;
         }
       }
+      if (Thread.interrupted()) {
+        throw PGExceptionFactory.newQueryCancelledException();
+      }
     } finally {
       connectionHandler.getConnectionMetadata().getOutputStream().flush();
       messages.clear();
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
index 5578d788..cba2c259 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiver.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiver.java
@@ -19,6 +19,7 @@ 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.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement.ResultNotReadyBehavior;
 import com.google.cloud.spanner.pgadapter.wireoutput.CommandCompleteResponse;
@@ -61,7 +62,7 @@ public class CopyDataReceiver implements Callable<Void> {
     if (copyStatement.hasException()) {
       throw copyStatement.getException();
     } else {
-      this.connectionHandler.addActiveStatement(copyStatement);
+      this.connectionHandler.setActiveCopyStatement(copyStatement);
       new CopyInResponse(
               this.connectionHandler.getConnectionMetadata().getOutputStream(),
               copyStatement.getTableColumns().size(),
@@ -73,6 +74,9 @@ public class CopyDataReceiver implements Callable<Void> {
         // Loop here until COPY_IN mode has finished.
         while (this.connectionHandler.getStatus() == ConnectionStatus.COPY_IN) {
           this.connectionHandler.handleMessages();
+          if (Thread.interrupted()) {
+            throw PGExceptionFactory.newQueryCancelledException();
+          }
         }
         // Return CommandComplete if the COPY succeeded. This should not be cached until a flush.
         // Note that if an error occurred during the COPY, the message handler will automatically
@@ -96,7 +100,7 @@ public class CopyDataReceiver implements Callable<Void> {
           }
         }
       } finally {
-        this.connectionHandler.removeActiveStatement(copyStatement);
+        this.connectionHandler.clearActiveCopyStatement();
         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 76d4aa19..1c792fce 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
@@ -31,6 +31,7 @@ import com.google.cloud.spanner.Type;
 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.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser.CopyOptions;
 import com.google.cloud.spanner.pgadapter.session.CopySettings;
 import com.google.cloud.spanner.pgadapter.session.SessionState;
@@ -45,6 +46,7 @@ import com.google.common.util.concurrent.MoreExecutors;
 import io.grpc.Context;
 import java.io.Closeable;
 import java.io.IOException;
+import java.io.InterruptedIOException;
 import java.io.PipedOutputStream;
 import java.util.ArrayList;
 import java.util.Iterator;
@@ -152,6 +154,12 @@ public class MutationWriter implements Callable<StatementResult>, Closeable {
     }
     try {
       this.payload.write(payload);
+    } catch (InterruptedIOException interruptedIOException) {
+      // The IO operation was interrupted. This indicates that the user wants to cancel the COPY
+      // operation. Re-instate the interrupted flag on the current thread and throw an exception to
+      // indicate that the operation should be cancelled.
+      Thread.currentThread().interrupt();
+      throw PGExceptionFactory.newQueryCancelledException();
     } catch (IOException e) {
       // 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.
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 83205690..354c3110 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
@@ -487,6 +487,9 @@ public abstract class ControlMessage extends WireMessage {
       prefixSent.get();
       long rows = 0L;
       while (hasData) {
+        if (Thread.interrupted()) {
+          throw PGExceptionFactory.newQueryCancelledException();
+        }
         WireOutput wireOutput = describedResult.createDataRowResponse(resultSet, mode);
         synchronized (describedResult) {
           wireOutput.send(false);
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 25613a83..d10b73e9 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
@@ -41,7 +41,7 @@ public class CopyDataMessage extends ControlMessage {
     int dataLength = this.length - 4;
     this.payload = new byte[dataLength];
     this.inputStream.readFully(this.payload);
-    this.statement = (CopyStatement) connection.getActiveStatement();
+    this.statement = connection.getActiveCopyStatement();
   }
 
   @Override
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 bc9a09ef..f02f9b5f 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
@@ -33,7 +33,7 @@ public class CopyDoneMessage extends ControlMessage {
 
   public CopyDoneMessage(ConnectionHandler connection) throws Exception {
     super(connection);
-    this.statement = (CopyStatement) connection.getActiveStatement();
+    this.statement = connection.getActiveCopyStatement();
   }
 
   @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 0c55b5d4..b2dbfe86 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
@@ -20,7 +20,6 @@ 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 java.text.MessageFormat;
 
@@ -40,12 +39,7 @@ public class CopyFailMessage extends ControlMessage {
   public CopyFailMessage(ConnectionHandler connection) throws Exception {
     super(connection);
     this.errorMessage = this.readAll();
-    IntermediateStatement activeStatement = connection.getActiveStatement();
-    if (activeStatement instanceof CopyStatement) {
-      this.statement = (CopyStatement) activeStatement;
-    } else {
-      this.statement = null;
-    }
+    this.statement = connection.getActiveCopyStatement();
   }
 
   @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 81eb90a2..f6942a80 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
@@ -18,6 +18,7 @@ import com.google.api.core.InternalApi;
 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.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.metadata.DescribePortalMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeStatementMetadata;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
@@ -174,7 +175,7 @@ public class DescribeMessage extends AbstractQueryProtocolMessage {
     } catch (ExecutionException executionException) {
       throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
     } catch (InterruptedException interruptedException) {
-      throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
+      throw PGExceptionFactory.newQueryCancelledException();
     }
   }
 


Test patch:
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 c9cc62dc..effdeee5 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
@@ -31,6 +31,7 @@ import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
+import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
@@ -205,6 +206,10 @@ public class IntermediateStatement {
     return this.connection;
   }
 
+  public ConnectionHandler getConnectionHandler() {
+    return this.connectionHandler;
+  }
+
   public String getStatement() {
     return this.parsedStatement.getSqlWithoutComments();
   }
@@ -220,7 +225,10 @@ public class IntermediateStatement {
       } catch (ExecutionException executionException) {
         setException(SpannerExceptionFactory.asSpannerException(executionException.getCause()));
       } catch (InterruptedException interruptedException) {
-        setException(SpannerExceptionFactory.propagateInterrupt(interruptedException));
+        // TODO(b/246193644): Switch to PGException
+        setException(
+            SpannerExceptionFactory.asSpannerException(
+                PGExceptionFactory.newQueryCancelledException()));
       } finally {
         this.futureStatementResult = null;
       }
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 a7364ee5..7035976b 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
@@ -26,6 +26,7 @@ import static org.mockito.Mockito.when;
 
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
 import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.error.SQLState;
@@ -294,12 +295,50 @@ public class ConnectionHandlerTest {
   }
 
   @Test
-  public void testRestartConnectionWithSsl_CreatesSslSocket() {
+  public void testCancelActiveStatement() {
     ProxyServer server = mock(ProxyServer.class);
     Socket socket = mock(Socket.class);
     InetAddress address = mock(InetAddress.class);
     when(socket.getInetAddress()).thenReturn(address);
+    when(address.getHostAddress()).thenReturn("address1");
+    Connection spannerConnection = mock(Connection.class);
 
+    ConnectionHandler connectionHandler =
+        new ConnectionHandler(server, socket, mock(Connection.class));
+    ConnectionHandler connectionHandlerToCancel =
+        new ConnectionHandler(server, socket, spannerConnection);
+
+    // Cancelling yourself is not allowed.
+    assertFalse(
+        connectionHandler.cancelActiveStatement(
+            connectionHandler.getConnectionId(), connectionHandler.getSecret()));
+    // Cancelling a random non-existing connection should not work.
+    assertFalse(connectionHandler.cancelActiveStatement(100, 100));
+    // Cancelling another connecting using the wrong secret is not allowed.
+    assertFalse(
+        connectionHandler.cancelActiveStatement(
+            connectionHandlerToCancel.getConnectionId(),
+            connectionHandlerToCancel.getSecret() - 1));
+
+    assertTrue(
+        connectionHandler.cancelActiveStatement(
+            connectionHandlerToCancel.getConnectionId(), connectionHandlerToCancel.getSecret()));
+
+    // The method should just return false if an error occurs.
+    doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.INTERNAL, "test error"))
+        .when(spannerConnection)
+        .cancel();
+    assertFalse(
+        connectionHandler.cancelActiveStatement(
+            connectionHandlerToCancel.getConnectionId(), connectionHandlerToCancel.getSecret()));
+  }
+
+  @Test
+  public void testRestartConnectionWithSsl_CreatesSslSocket() {
+    ProxyServer server = mock(ProxyServer.class);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
     AtomicBoolean calledCreateSSLSocket = new AtomicBoolean();
     ConnectionHandler connection =
         new ConnectionHandler(server, socket) {
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 041f7c40..fee0f46d 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
@@ -31,7 +31,10 @@ import static org.mockito.Mockito.when;
 
 import com.google.cloud.spanner.DatabaseId;
 import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.ErrorCode;
 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;
@@ -39,6 +42,8 @@ import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
+import com.google.cloud.spanner.pgadapter.error.PGException;
+import com.google.cloud.spanner.pgadapter.error.SQLState;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResult;
@@ -355,6 +360,31 @@ public class BackendConnectionTest {
     assertSame(executionException.getCause(), error);
   }
 
+  @Test
+  public void testCancelledException() {
+    Connection connection = mock(Connection.class);
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    Statement statement = Statement.of("select foo from bar");
+    SpannerException error =
+        SpannerExceptionFactory.newSpannerException(ErrorCode.CANCELLED, "query cancelled");
+    when(connection.execute(statement)).thenThrow(error);
+
+    BackendConnection backendConnection =
+        new BackendConnection(
+            DatabaseId.of("p", "i", "d"),
+            connection,
+            mock(OptionsMetadata.class),
+            EMPTY_LOCAL_STATEMENTS);
+    Future<StatementResult> resultFuture = backendConnection.execute(parsedStatement, statement);
+    backendConnection.flush();
+
+    ExecutionException executionException =
+        assertThrows(ExecutionException.class, resultFuture::get);
+    assertEquals(PGException.class, executionException.getCause().getClass());
+    PGException pgException = (PGException) executionException.getCause();
+    assertEquals(SQLState.QueryCanceled, pgException.getSQLState());
+  }
+
   @Test
   public void testDdlExceptionInBatch() {
     String sql1 = "create table foo (id bigint primary key)";
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 a1e130ed..85470a6a 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
@@ -422,7 +422,7 @@ public class StatementTest {
     backendConnection.flush();
 
     assertTrue(intermediateStatement.hasException());
-    SpannerException exception = (SpannerException) intermediateStatement.getException();
+    SpannerException exception = intermediateStatement.getException();
     assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode());
   }
 
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 7526574d..5bfcfcc9 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
@@ -1265,7 +1265,7 @@ public class ProtocolTest {
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
 
     CopyStatement copyStatement = mock(CopyStatement.class);
-    when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
+    when(connectionHandler.getActiveCopyStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
@@ -1293,7 +1293,7 @@ public class ProtocolTest {
 
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
 
-    when(connectionHandler.getActiveStatement()).thenReturn(null);
+    when(connectionHandler.getActiveCopyStatement()).thenReturn(null);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
@@ -1330,7 +1330,7 @@ public class ProtocolTest {
         new CopyStatement(connectionHandler, options, "", parse(sql), Statement.of(sql));
     copyStatement.executeAsync(mock(BackendConnection.class));
 
-    when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
+    when(connectionHandler.getActiveCopyStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
@@ -1373,7 +1373,7 @@ public class ProtocolTest {
     CopyStatement copyStatement = mock(CopyStatement.class);
     MutationWriter mutationWriter = mock(MutationWriter.class);
     when(copyStatement.getMutationWriter()).thenReturn(mutationWriter);
-    when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
+    when(connectionHandler.getActiveCopyStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
@@ -1405,7 +1405,7 @@ public class ProtocolTest {
     CopyStatement copyStatement = mock(CopyStatement.class);
     MutationWriter mutationWriter = mock(MutationWriter.class);
     when(copyStatement.getMutationWriter()).thenReturn(mutationWriter);
-    when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
+    when(connectionHandler.getActiveCopyStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);