Repository: GoogleCloudPlatform/pgadapter
Pull request #345 (2022-09-05 13:39:00)

fix: use main thread for CopyDataReceiver
Use the main thread for the `CopyDataReceiver` to ensure that all `CopyData` messages are sent to this receiver, and do not end up in the buffer for the main connection handler. The latter could happen if the client would start sending `CopyData` messages before the server had sent a `CopyInResponse`. This again could lead to the first `CopyData` message getting lost.


Patch:
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 a2c284de..63160374 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -357,7 +357,7 @@ public class ConnectionHandler extends Thread {
         exception,
         () ->
             String.format("Exception on connection handler with ID %s: %s", getName(), exception));
-    DataOutputStream output = getConnectionMetadata().peekOutputStream();
+    DataOutputStream output = getConnectionMetadata().getOutputStream();
     if (this.status == ConnectionStatus.TERMINATED) {
       new ErrorResponse(output, exception).send();
       new TerminateResponse(output).send();
@@ -402,11 +402,12 @@ public class ConnectionHandler extends Thread {
     this.portalsMap.put(portalName, portal);
   }
 
-  public void closePortal(String portalName) {
+  public void closePortal(String portalName) throws Exception {
     if (!hasPortal(portalName)) {
-      throw new IllegalStateException("Unregistered statement: " + portalName);
+      throw new IllegalStateException("Unregistered portal: " + portalName);
     }
-    this.portalsMap.remove(portalName);
+    IntermediatePortalStatement portal = this.portalsMap.get(portalName);
+    portal.close();
   }
 
   public boolean hasPortal(String portalName) {
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/ConnectionMetadata.java b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/ConnectionMetadata.java
index 7269b59e..7ebd88d1 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/ConnectionMetadata.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/ConnectionMetadata.java
@@ -15,7 +15,6 @@
 package com.google.cloud.spanner.pgadapter.metadata;
 
 import com.google.api.core.InternalApi;
-import com.google.cloud.Tuple;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Stopwatch;
 import java.io.BufferedInputStream;
@@ -25,65 +24,44 @@ import java.io.DataOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.util.Stack;
 import java.util.concurrent.TimeUnit;
 
 @InternalApi
 public class ConnectionMetadata implements AutoCloseable {
   private static final int SOCKET_BUFFER_SIZE = 1 << 16;
 
-  private final InputStream rawInputStream;
-  private final OutputStream rawOutputStream;
-  private final Stack<DataInputStream> inputStream = new Stack<>();
-  private final Stack<DataOutputStream> outputStream = new Stack<>();
+  private final DataInputStream inputStream;
+  private final DataOutputStream outputStream;
 
   /**
    * Creates a {@link DataInputStream} and a {@link DataOutputStream} from the given raw streams and
    * pushes these as the current streams to use for communication for a connection.
    */
   public ConnectionMetadata(InputStream rawInputStream, OutputStream rawOutputStream) {
-    this.rawInputStream = Preconditions.checkNotNull(rawInputStream);
-    this.rawOutputStream = Preconditions.checkNotNull(rawOutputStream);
-    pushNewStreams();
+    this.inputStream =
+        new DataInputStream(
+            new BufferedInputStream(
+                Preconditions.checkNotNull(rawInputStream), SOCKET_BUFFER_SIZE));
+    this.outputStream =
+        new DataOutputStream(
+            new BufferedOutputStream(
+                Preconditions.checkNotNull(rawOutputStream), SOCKET_BUFFER_SIZE));
   }
 
   @Override
   public void close() throws Exception {
-    this.rawInputStream.close();
-    this.rawOutputStream.close();
-  }
-
-  /**
-   * Creates a new buffered {@link DataInputStream} and {@link DataOutputStream} tuple to use for
-   * the connection. This is used for the COPY sub-protocol to prevent mixing the buffers used for
-   * the normal protocol with the COPY protocol, as that would cause multithreaded access to those
-   * buffers.
-   */
-  public Tuple<DataInputStream, DataOutputStream> pushNewStreams() {
-    return Tuple.of(
-        this.inputStream.push(
-            new DataInputStream(new BufferedInputStream(this.rawInputStream, SOCKET_BUFFER_SIZE))),
-        this.outputStream.push(
-            new DataOutputStream(
-                new BufferedOutputStream(this.rawOutputStream, SOCKET_BUFFER_SIZE))));
-  }
-
-  /**
-   * Pops the current {@link DataInputStream} and {@link DataOutputStream} from the connection. This
-   * is done when the COPY sub-protocol has finished.
-   */
-  public Tuple<DataInputStream, DataOutputStream> popStreams() {
-    return Tuple.of(this.inputStream.pop(), this.outputStream.pop());
+    this.inputStream.close();
+    this.outputStream.close();
   }
 
   /** Returns the current {@link DataInputStream} for the connection. */
-  public DataInputStream peekInputStream() {
-    return inputStream.peek();
+  public DataInputStream getInputStream() {
+    return inputStream;
   }
 
   /** Returns the current {@link DataOutputStream} for the connection. */
-  public DataOutputStream peekOutputStream() {
-    return outputStream.peek();
+  public DataOutputStream getOutputStream() {
+    return outputStream;
   }
 
   /**
@@ -92,16 +70,15 @@ public class ConnectionMetadata implements AutoCloseable {
    * become available in the buffer.
    */
   public char peekNextByte(long maxWaitMillis) throws IOException {
-    DataInputStream dataInputStream = inputStream.peek();
     Stopwatch stopwatch = Stopwatch.createStarted();
-    while (dataInputStream.available() == 0
+    while (inputStream.available() == 0
         && stopwatch.elapsed(TimeUnit.MILLISECONDS) < maxWaitMillis) {
       Thread.yield();
     }
-    if (dataInputStream.available() > 0) {
-      dataInputStream.mark(1);
-      char result = (char) dataInputStream.readUnsignedByte();
-      dataInputStream.reset();
+    if (inputStream.available() > 0) {
+      inputStream.mark(1);
+      char result = (char) inputStream.readUnsignedByte();
+      inputStream.reset();
 
       return result;
     }
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 3e65b554..bd910e59 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
@@ -308,6 +308,9 @@ public class BackendConnection {
     }
   }
 
+  private static final ListeningExecutorService DATA_RECEIVER_EXECUTOR =
+      MoreExecutors.newDirectExecutorService();
+
   /**
    * 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
@@ -343,8 +346,9 @@ public class BackendConnection {
         // 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);
+        ListenableFuture<Void> copyDataReceiverFuture =
+            DATA_RECEIVER_EXECUTOR.submit(copyDataReceiver);
 
         // Make sure both the front-end CopyDataReceiver and the backend MutationWriter processes
         // have finished before we proceed.
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 37e42f97..56267cc8 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
@@ -130,7 +130,7 @@ public class ExtendedQueryProtocolHandler {
         }
       }
     } finally {
-      connectionHandler.getConnectionMetadata().peekOutputStream().flush();
+      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 e9b800c2..5578d788 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
@@ -61,52 +61,44 @@ public class CopyDataReceiver implements Callable<Void> {
     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 {
-        // Push a new set of input/output streams for the connection. This ensures that the COPY
-        // process does not use the buffers of the normal protocol.
-        this.connectionHandler.getConnectionMetadata().pushNewStreams();
-        this.connectionHandler.addActiveStatement(copyStatement);
-        new CopyInResponse(
-                this.connectionHandler.getConnectionMetadata().peekOutputStream(),
-                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();
-          }
-          // 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
-          // respond with an ErrorResponse. That is why we do not check for COPY_FAILED here, and do
-          // not return an ErrorResponse.
-          if (connectionHandler.getStatus() == ConnectionStatus.COPY_DONE) {
-            new CommandCompleteResponse(
-                    this.connectionHandler.getConnectionMetadata().peekOutputStream(),
-                    "COPY " + copyStatement.getUpdateCount(ResultNotReadyBehavior.BLOCK))
-                .send();
-          }
-          // Throw an exception if the COPY failed. This ensures that the BackendConnection receives
-          // an error and marks the current (implicit) transaction as aborted.
-          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");
-            }
+        this.connectionHandler.setStatus(ConnectionStatus.COPY_IN);
+        // Loop here until COPY_IN mode has finished.
+        while (this.connectionHandler.getStatus() == ConnectionStatus.COPY_IN) {
+          this.connectionHandler.handleMessages();
+        }
+        // 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
+        // respond with an ErrorResponse. That is why we do not check for COPY_FAILED here, and do
+        // not return an ErrorResponse.
+        if (connectionHandler.getStatus() == ConnectionStatus.COPY_DONE) {
+          new CommandCompleteResponse(
+                  this.connectionHandler.getConnectionMetadata().getOutputStream(),
+                  "COPY " + copyStatement.getUpdateCount(ResultNotReadyBehavior.BLOCK))
+              .send();
+        }
+        // Throw an exception if the COPY failed. This ensures that the BackendConnection receives
+        // an error and marks the current (implicit) transaction as aborted.
+        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);
         }
       } finally {
-        // Pop the COPY input/output streams to return to the normal protocol and streams.
-        this.connectionHandler.getConnectionMetadata().popStreams();
+        this.connectionHandler.removeActiveStatement(copyStatement);
+        this.copyStatement.close();
+        this.connectionHandler.setStatus(initialConnectionStatus);
       }
     }
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/BootstrapMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/BootstrapMessage.java
index 32c9c1bc..8a5b20b7 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/BootstrapMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/BootstrapMessage.java
@@ -50,11 +50,11 @@ public abstract class BootstrapMessage extends WireMessage {
    * @throws Exception If construction or reading fails.
    */
   public static BootstrapMessage create(ConnectionHandler connection) throws Exception {
-    int length = connection.getConnectionMetadata().peekInputStream().readInt();
+    int length = connection.getConnectionMetadata().getInputStream().readInt();
     if (length > MAX_BOOTSTRAP_MESSAGE_LENGTH) {
       throw new IllegalArgumentException("Invalid bootstrap message length: " + length);
     }
-    int protocol = connection.getConnectionMetadata().peekInputStream().readInt();
+    int protocol = connection.getConnectionMetadata().getInputStream().readInt();
     switch (protocol) {
       case SSLMessage.IDENTIFIER:
         return new SSLMessage(connection);
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 08a2aea0..fe160787 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
@@ -85,7 +85,7 @@ public abstract class ControlMessage extends WireMessage {
   private final ManuallyCreatedToken manuallyCreatedToken;
 
   public ControlMessage(ConnectionHandler connection) throws IOException {
-    super(connection, connection.getConnectionMetadata().peekInputStream().readInt());
+    super(connection, connection.getConnectionMetadata().getInputStream().readInt());
     this.manuallyCreatedToken = null;
   }
 
@@ -108,7 +108,7 @@ public abstract class ControlMessage extends WireMessage {
    */
   public static ControlMessage create(ConnectionHandler connection) throws Exception {
     boolean validMessage = true;
-    char nextMsg = (char) connection.getConnectionMetadata().peekInputStream().readUnsignedByte();
+    char nextMsg = (char) connection.getConnectionMetadata().getInputStream().readUnsignedByte();
     try {
       if (connection.getStatus() == ConnectionStatus.COPY_IN) {
         switch (nextMsg) {
@@ -176,7 +176,7 @@ public abstract class ControlMessage extends WireMessage {
         connection.increaseInvalidMessageCount();
         if (connection.getInvalidMessageCount() > MAX_INVALID_MESSAGE_COUNT) {
           new ErrorResponse(
-                  connection.getConnectionMetadata().peekOutputStream(),
+                  connection.getConnectionMetadata().getOutputStream(),
                   PGException.newBuilder()
                       .setMessage(
                           String.format(
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
index 6f92855a..7449bccc 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SkipMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SkipMessage.java
@@ -44,7 +44,7 @@ public class SkipMessage extends ControlMessage {
     int skipLength = this.length - 4;
     int skipped = 0;
     // Read and skip bytes until we have reached the total message length.
-    DataInputStream input = connection.getConnectionMetadata().peekInputStream();
+    DataInputStream input = connection.getConnectionMetadata().getInputStream();
     // If the stream is potentially invalid we'll use a more pessimistic skipping algorithm than if
     // we just received a message that we want to skip. We do this because the Unix domain socket
     // implementation might block on skip(n) if the stream has no bytes more to read.
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/StartupMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/StartupMessage.java
index b4366989..c4a7235a 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/StartupMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/StartupMessage.java
@@ -89,7 +89,7 @@ public class StartupMessage extends BootstrapMessage {
           .initSessionSetting(parameter.getKey(), parameter.getValue());
     }
     sendStartupMessage(
-        connection.getConnectionMetadata().peekOutputStream(),
+        connection.getConnectionMetadata().getOutputStream(),
         connection.getConnectionId(),
         connection.getSecret(),
         connection.getExtendedQueryProtocolHandler().getBackendConnection().getSessionState());
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/WireMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/WireMessage.java
index 97205970..a8dbd0f8 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/WireMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/WireMessage.java
@@ -40,8 +40,8 @@ public abstract class WireMessage {
   public WireMessage(ConnectionHandler connection, int length) {
     Preconditions.checkArgument(length >= 4);
     this.connection = connection;
-    this.inputStream = connection.getConnectionMetadata().peekInputStream();
-    this.outputStream = connection.getConnectionMetadata().peekOutputStream();
+    this.inputStream = connection.getConnectionMetadata().getInputStream();
+    this.outputStream = connection.getConnectionMetadata().getOutputStream();
     this.length = length;
   }
 


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 a7e7a603..a9549af0 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
@@ -108,7 +108,7 @@ public class IntermediateStatement {
     this.connection = connectionHandler.getSpannerConnection();
     this.command = parseCommand(this.parsedStatement.getSqlWithoutComments());
     this.commandTag = this.command;
-    this.outputStream = connectionHandler.getConnectionMetadata().peekOutputStream();
+    this.outputStream = connectionHandler.getConnectionMetadata().getOutputStream();
   }
 
   /**
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 8e51af89..23335f01 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
@@ -15,6 +15,8 @@
 package com.google.cloud.spanner.pgadapter;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -24,6 +26,7 @@ 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.statements.IntermediatePortalStatement;
 import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
 import java.io.ByteArrayOutputStream;
 import java.io.DataOutputStream;
@@ -84,6 +87,96 @@ public class ConnectionHandlerTest {
     verify(socket).close();
   }
 
+  @Test
+  public void testTerminateClosesAllPortals() throws Exception {
+    ProxyServer server = mock(ProxyServer.class);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+    IntermediatePortalStatement portal1 = mock(IntermediatePortalStatement.class);
+    IntermediatePortalStatement portal2 = mock(IntermediatePortalStatement.class);
+
+    ConnectionHandler connection = new ConnectionHandler(server, socket);
+    connection.registerPortal("portal1", portal1);
+    connection.registerPortal("portal2", portal2);
+
+    connection.terminate();
+
+    verify(portal1).close();
+    verify(portal2).close();
+  }
+
+  @Test
+  public void testTerminateIgnoresPortalCloseError() throws Exception {
+    ProxyServer server = mock(ProxyServer.class);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+    IntermediatePortalStatement portal1 = mock(IntermediatePortalStatement.class);
+    IntermediatePortalStatement portal2 = mock(IntermediatePortalStatement.class);
+    doThrow(new Exception("test")).when(portal2).close();
+
+    ConnectionHandler connection = new ConnectionHandler(server, socket);
+    connection.registerPortal("portal1", portal1);
+    connection.registerPortal("portal2", portal2);
+
+    connection.terminate();
+
+    verify(portal1).close();
+    verify(portal2).close();
+  }
+
+  @Test
+  public void testGetPortal() {
+    ProxyServer server = mock(ProxyServer.class);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+    IntermediatePortalStatement portal1 = mock(IntermediatePortalStatement.class);
+
+    ConnectionHandler connection = new ConnectionHandler(server, socket);
+    connection.registerPortal("portal1", portal1);
+
+    assertSame(portal1, connection.getPortal("portal1"));
+  }
+
+  @Test
+  public void testGetUnknownPortal() {
+    ProxyServer server = mock(ProxyServer.class);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+
+    ConnectionHandler connection = new ConnectionHandler(server, socket);
+    assertThrows(IllegalStateException.class, () -> connection.getPortal("unknown portal"));
+  }
+
+  @Test
+  public void testClosePortal() throws Exception {
+    ProxyServer server = mock(ProxyServer.class);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+    IntermediatePortalStatement portal1 = mock(IntermediatePortalStatement.class);
+
+    ConnectionHandler connection = new ConnectionHandler(server, socket);
+    connection.registerPortal("portal1", portal1);
+
+    connection.closePortal("portal1");
+    verify(portal1).close();
+  }
+
+  @Test
+  public void testCloseUnknownPortal() {
+    ProxyServer server = mock(ProxyServer.class);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+
+    ConnectionHandler connection = new ConnectionHandler(server, socket);
+    assertThrows(IllegalStateException.class, () -> connection.closePortal("unknown portal"));
+  }
+
   @Test
   public void testHandleMessages_NonFatalException() throws Exception {
     ProxyServer server = mock(ProxyServer.class);
@@ -92,7 +185,7 @@ public class ConnectionHandlerTest {
     when(socket.getInetAddress()).thenReturn(address);
     DataOutputStream dataOutputStream = new DataOutputStream(new ByteArrayOutputStream());
     ConnectionMetadata connectionMetadata = mock(ConnectionMetadata.class);
-    when(connectionMetadata.peekOutputStream()).thenReturn(dataOutputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(dataOutputStream);
     WireMessage message = mock(WireMessage.class);
     when(server.recordMessage(message)).thenReturn(message);
     doThrow(
@@ -122,7 +215,7 @@ public class ConnectionHandlerTest {
     when(socket.getInetAddress()).thenReturn(address);
     DataOutputStream dataOutputStream = new DataOutputStream(new ByteArrayOutputStream());
     ConnectionMetadata connectionMetadata = mock(ConnectionMetadata.class);
-    when(connectionMetadata.peekOutputStream()).thenReturn(dataOutputStream);
+    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();
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 c8f87a81..4e80eeb6 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java
@@ -65,8 +65,8 @@ public final class ControlMessageTest {
         new DataInputStream(
             new ByteArrayInputStream(new byte[] {(byte) QUERY_IDENTIFIER, 0, 0, 0, 5, 0}));
 
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(intermediateStatement.getStatementType()).thenReturn(StatementType.UPDATE);
     when(intermediateStatement.getCommandTag()).thenReturn("INSERT");
@@ -111,8 +111,8 @@ public final class ControlMessageTest {
         new DataInputStream(
             new ByteArrayInputStream(new byte[] {(byte) QUERY_IDENTIFIER, 0, 0, 0, 5, 0}));
 
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     ExecuteMessage executeMessage =
         new ExecuteMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN);
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxMockServerTest.java
index 3e32dd4f..795df045 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxMockServerTest.java
@@ -987,8 +987,6 @@ public class PgxMockServerTest extends AbstractMockServerTest {
     }
   }
 
-  @Ignore(
-      "pgx copy implementation seems buggy (CopyDone message can be sent before all data has been sent)")
   @Test
   public void testCopyIn() {
     CopyInMockServerTest.setupCopyInformationSchemaResults(mockSpanner, true);
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandlerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandlerTest.java
index 9f56cee0..68a3d7ef 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandlerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandlerTest.java
@@ -65,7 +65,7 @@ public class ExtendedQueryProtocolHandlerTest {
   @Test
   public void testFlush() throws Exception {
     ConnectionMetadata connectionMetadata = mock(ConnectionMetadata.class);
-    when(connectionMetadata.peekOutputStream()).thenReturn(mock(DataOutputStream.class));
+    when(connectionMetadata.getOutputStream()).thenReturn(mock(DataOutputStream.class));
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     ParseMessage parseMessage = mock(ParseMessage.class);
     BindMessage bindMessage = mock(BindMessage.class);
@@ -96,7 +96,7 @@ public class ExtendedQueryProtocolHandlerTest {
   @Test
   public void testSync() throws Exception {
     ConnectionMetadata connectionMetadata = mock(ConnectionMetadata.class);
-    when(connectionMetadata.peekOutputStream()).thenReturn(mock(DataOutputStream.class));
+    when(connectionMetadata.getOutputStream()).thenReturn(mock(DataOutputStream.class));
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     ParseMessage parseMessage = mock(ParseMessage.class);
     BindMessage bindMessage = mock(BindMessage.class);
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 78c6be11..7806580d 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
@@ -418,8 +418,8 @@ public class StatementTest {
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getServer()).thenReturn(server);
     when(server.getOptions()).thenReturn(options);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     assertEquals(QueryMessage.class, message.getClass());
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
index 51787418..a0658916 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiverTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/CopyDataReceiverTest.java
@@ -57,7 +57,7 @@ public class CopyDataReceiverTest {
     DataOutputStream outputStream = new DataOutputStream(new ByteArrayOutputStream());
     CopyStatement statement = mock(CopyStatement.class);
     ConnectionMetadata connectionMetadata = mock(ConnectionMetadata.class);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     ConnectionHandler connectionHandler = mock(ConnectionHandler.class);
     when(connectionHandler.getStatus())
         .thenReturn(ConnectionStatus.AUTHENTICATED, ConnectionStatus.COPY_FAILED);
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 8210f3aa..eab676ee 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
@@ -156,8 +156,8 @@ public class ProtocolTest {
     when(backendConnection.getConnectionState()).thenReturn(ConnectionState.IDLE);
     when(connectionHandler.getStatement("")).thenReturn(intermediatePortalStatement);
     when(connectionHandler.getPortal("")).thenReturn(intermediatePortalStatement);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     assertEquals(QueryMessage.class, message.getClass());
@@ -182,8 +182,8 @@ public class ProtocolTest {
     when(connectionHandler.getServer()).thenReturn(server);
     when(server.getOptions()).thenReturn(options);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     assertEquals(QueryMessage.class, message.getClass());
@@ -200,8 +200,8 @@ public class ProtocolTest {
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     assertThrows(IOException.class, () -> ControlMessage.create(connectionHandler));
   }
@@ -244,8 +244,8 @@ public class ProtocolTest {
     when(server.getOptions()).thenReturn(options);
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -308,8 +308,8 @@ public class ProtocolTest {
     when(server.getOptions()).thenReturn(options);
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -373,8 +373,8 @@ public class ProtocolTest {
     when(server.getOptions()).thenReturn(options);
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -423,8 +423,8 @@ public class ProtocolTest {
     when(server.getOptions()).thenReturn(options);
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -481,8 +481,8 @@ public class ProtocolTest {
     when(server.getOptions()).thenReturn(options);
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -527,8 +527,8 @@ public class ProtocolTest {
     when(server.getOptions()).thenReturn(options);
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -574,8 +574,8 @@ public class ProtocolTest {
     when(server.getOptions()).thenReturn(options);
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -634,8 +634,8 @@ public class ProtocolTest {
     DataOutputStream outputStream = new DataOutputStream(result);
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -723,8 +723,8 @@ public class ProtocolTest {
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     assertEquals(BindMessage.class, message.getClass());
@@ -793,8 +793,8 @@ public class ProtocolTest {
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     assertEquals(BindMessage.class, message.getClass());
@@ -823,8 +823,8 @@ public class ProtocolTest {
     when(connectionHandler.getPortal(anyString())).thenReturn(intermediatePortalStatement);
     when(intermediatePortalStatement.getSql()).thenReturn("select * from foo");
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -861,8 +861,8 @@ public class ProtocolTest {
 
     when(connectionHandler.getStatement(anyString())).thenReturn(intermediatePreparedStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -899,8 +899,8 @@ public class ProtocolTest {
 
     when(connectionHandler.getPortal(anyString())).thenReturn(intermediatePortalStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
     when(extendedQueryProtocolHandler.getBackendConnection()).thenReturn(backendConnection);
@@ -948,8 +948,8 @@ public class ProtocolTest {
     when(intermediatePortalStatement.getException()).thenReturn(testException);
     when(connectionHandler.getPortal(anyString())).thenReturn(intermediatePortalStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
     when(extendedQueryProtocolHandler.getBackendConnection()).thenReturn(backendConnection);
@@ -988,8 +988,8 @@ public class ProtocolTest {
 
     when(connectionHandler.getPortal(anyString())).thenReturn(intermediatePortalStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     assertEquals(CloseMessage.class, message.getClass());
@@ -1026,8 +1026,8 @@ public class ProtocolTest {
 
     when(connectionHandler.getStatement(anyString())).thenReturn(intermediatePortalStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     assertEquals(CloseMessage.class, message.getClass());
@@ -1059,8 +1059,8 @@ public class ProtocolTest {
 
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.AUTHENTICATED);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
     when(extendedQueryProtocolHandler.getBackendConnection()).thenReturn(backendConnection);
@@ -1092,8 +1092,8 @@ public class ProtocolTest {
 
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.AUTHENTICATED);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
     when(extendedQueryProtocolHandler.getBackendConnection()).thenReturn(backendConnection);
@@ -1123,8 +1123,8 @@ public class ProtocolTest {
     DataOutputStream outputStream = mock(DataOutputStream.class);
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
 
@@ -1151,8 +1151,8 @@ public class ProtocolTest {
     when(connectionHandler.getExtendedQueryProtocolHandler())
         .thenReturn(extendedQueryProtocolHandler);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     assertEquals(FlushMessage.class, message.getClass());
@@ -1187,8 +1187,8 @@ public class ProtocolTest {
     when(backendConnection.getConnectionState()).thenReturn(ConnectionState.TRANSACTION);
     OptionsMetadata options = mock(OptionsMetadata.class);
     when(server.getOptions()).thenReturn(options);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     ExtendedQueryProtocolHandler extendedQueryProtocolHandler =
         new ExtendedQueryProtocolHandler(connectionHandler, backendConnection);
@@ -1232,8 +1232,8 @@ public class ProtocolTest {
     DataOutputStream outputStream = new DataOutputStream(result);
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     assertEquals(TerminateMessage.class, message.getClass());
@@ -1246,7 +1246,7 @@ public class ProtocolTest {
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(messageMetadata));
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
 
     assertThrows(IllegalStateException.class, () -> ControlMessage.create(connectionHandler));
   }
@@ -1264,8 +1264,8 @@ public class ProtocolTest {
     when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     MutationWriter mw = mock(MutationWriter.class);
     when(copyStatement.getMutationWriter()).thenReturn(mw);
@@ -1292,8 +1292,8 @@ public class ProtocolTest {
     when(connectionHandler.getActiveStatement()).thenReturn(null);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
     // This should be a no-op.
@@ -1328,10 +1328,10 @@ public class ProtocolTest {
 
     when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     {
-      when(connectionMetadata.peekInputStream()).thenReturn(inputStream1);
+      when(connectionMetadata.getInputStream()).thenReturn(inputStream1);
       WireMessage message = ControlMessage.create(connectionHandler);
       assertEquals(CopyDataMessage.class, message.getClass());
       assertArrayEquals(payload1, ((CopyDataMessage) message).getPayload());
@@ -1339,7 +1339,7 @@ public class ProtocolTest {
       copyDataMessage.send();
     }
     {
-      when(connectionMetadata.peekInputStream()).thenReturn(inputStream2);
+      when(connectionMetadata.getInputStream()).thenReturn(inputStream2);
       WireMessage message = ControlMessage.create(connectionHandler);
       assertEquals(CopyDataMessage.class, message.getClass());
       assertArrayEquals(payload2, ((CopyDataMessage) message).getPayload());
@@ -1347,7 +1347,7 @@ public class ProtocolTest {
       copyDataMessage.send();
     }
     {
-      when(connectionMetadata.peekInputStream()).thenReturn(inputStream3);
+      when(connectionMetadata.getInputStream()).thenReturn(inputStream3);
       WireMessage message = ControlMessage.create(connectionHandler);
       assertEquals(CopyDataMessage.class, message.getClass());
       assertArrayEquals(payload3, ((CopyDataMessage) message).getPayload());
@@ -1372,8 +1372,8 @@ public class ProtocolTest {
     when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
 
@@ -1404,8 +1404,8 @@ public class ProtocolTest {
     when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
 
@@ -1482,8 +1482,8 @@ public class ProtocolTest {
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = ControlMessage.create(connectionHandler);
 
@@ -1535,8 +1535,8 @@ public class ProtocolTest {
     when(server.getOptions()).thenReturn(options);
     when(options.getServerVersion()).thenReturn("13.4");
     when(options.shouldAuthenticate()).thenReturn(false);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = BootstrapMessage.create(connectionHandler);
     assertEquals(StartupMessage.class, message.getClass());
@@ -1629,8 +1629,8 @@ public class ProtocolTest {
     DataOutputStream outputStream = new DataOutputStream(result);
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = BootstrapMessage.create(connectionHandler);
     assertEquals(CancelMessage.class, message.getClass());
@@ -1656,8 +1656,8 @@ public class ProtocolTest {
     DataOutputStream outputStream = new DataOutputStream(result);
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = BootstrapMessage.create(connectionHandler);
     assertEquals(SSLMessage.class, message.getClass());
@@ -1682,8 +1682,8 @@ public class ProtocolTest {
     DataOutputStream outputStream = new DataOutputStream(result);
 
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
     WireMessage message = BootstrapMessage.create(connectionHandler);
     assertEquals(SSLMessage.class, message.getClass());
@@ -1723,8 +1723,8 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
 
     SkipMessage message = SkipMessage.createForValidStream(connectionHandler);
@@ -1745,8 +1745,8 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
 
@@ -1765,8 +1765,8 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
 
@@ -1787,8 +1787,8 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.AUTHENTICATED);
 
@@ -1807,8 +1807,8 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.AUTHENTICATED);
 
@@ -1827,8 +1827,8 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    when(connectionMetadata.getInputStream()).thenReturn(inputStream);
+    when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.AUTHENTICATED);
 
@@ -1859,8 +1859,8 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    when(connectionMetadata.peekInputStream()).thenReturn(inputStream);
-    when(connectionMetadata.peekOutputStream()).thenReturn(outputStream);
+    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();