Repository: GoogleCloudPlatform/pgadapter
Pull request #122 (2022-04-22 06:52:25)

feat: support fetching rows in multiple steps
Adds support for letting the frontend send multiple Execute messages with maxRows
set to a non-zero value. This prevents the entire result set from being transferred
from the backend to the frontend in one go, and can reduce memory usage.


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 0a792eaa..98d73398 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -25,6 +25,7 @@ import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse.Severity;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
+import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse.Status;
 import com.google.cloud.spanner.pgadapter.wireoutput.TerminateResponse;
 import com.google.cloud.spanner.pgadapter.wireprotocol.BootstrapMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
@@ -140,7 +141,7 @@ public class ConnectionHandler extends Thread {
 
       try {
         this.connectionMetadata = new ConnectionMetadata(input, output);
-        this.message = BootstrapMessage.create(this);
+        this.message = this.server.recordMessage(BootstrapMessage.create(this));
         this.message.send();
         while (this.status == ConnectionStatus.UNAUTHENTICATED) {
           try {
@@ -380,7 +381,7 @@ public class ConnectionHandler extends Thread {
   }
 
   public void setMessageState(WireMessage message) {
-    this.message = message;
+    this.message = this.server.recordMessage(message);
   }
 
   public ConnectionMetadata getConnectionMetadata() {
@@ -401,11 +402,22 @@ public class ConnectionHandler extends Thread {
 
   /** Status of a {@link ConnectionHandler} */
   public enum ConnectionStatus {
-    UNAUTHENTICATED,
-    IDLE,
-    COPY_IN,
-    TERMINATED,
-    TRANSACTION_ABORTED
+    UNAUTHENTICATED(Status.IDLE),
+    IDLE(Status.IDLE),
+    TRANSACTION(Status.TRANSACTION),
+    COPY_IN(Status.IDLE),
+    TERMINATED(Status.IDLE),
+    TRANSACTION_ABORTED(Status.FAILED);
+
+    private final ReadyResponse.Status readyResponseStatus;
+
+    ConnectionStatus(ReadyResponse.Status readyResponseStatus) {
+      this.readyResponseStatus = readyResponseStatus;
+    }
+
+    public ReadyResponse.Status getReadyResponseStatus() {
+      return this.readyResponseStatus;
+    }
   }
 
   /**
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java b/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
index 676aa0b0..93ee8737 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
@@ -25,6 +25,7 @@ import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.TextFormat;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse.Severity;
+import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
 import java.io.BufferedOutputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
@@ -36,6 +37,8 @@ import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -53,6 +56,13 @@ public class ProxyServer extends AbstractApiService {
   private ServerSocket serverSocket;
   private int localPort;
 
+  /** The server will keep track of all messages it receives if it started in DEBUG mode. */
+  private static final int MAX_DEBUG_MESSAGES = 100_000;
+
+  private final boolean debugMode;
+  private final ConcurrentLinkedQueue<WireMessage> debugMessages = new ConcurrentLinkedQueue<>();
+  private final AtomicInteger debugMessageCount = new AtomicInteger();
+
   /**
    * Instantiates the ProxyServer from CLI-gathered metadata.
    *
@@ -62,6 +72,7 @@ public class ProxyServer extends AbstractApiService {
     this.options = optionsMetadata;
     this.localPort = optionsMetadata.getProxyPort();
     this.properties = new Properties();
+    this.debugMode = optionsMetadata.isDebugMode();
     addConnectionProperties();
   }
 
@@ -77,6 +88,7 @@ public class ProxyServer extends AbstractApiService {
     this.options = optionsMetadata;
     this.localPort = optionsMetadata.getProxyPort();
     this.properties = properties;
+    this.debugMode = optionsMetadata.isDebugMode();
     addConnectionProperties();
   }
 
@@ -270,6 +282,29 @@ public class ProxyServer extends AbstractApiService {
     return String.format("ProxyServer[port: %d]", getLocalPort());
   }
 
+  ConcurrentLinkedQueue<WireMessage> getDebugMessages() {
+    return debugMessages;
+  }
+
+  void clearDebugMessages() {
+    synchronized (debugMessages) {
+      debugMessages.clear();
+      debugMessageCount.set(0);
+    }
+  }
+
+  WireMessage recordMessage(WireMessage message) {
+    if (debugMode) {
+      if (debugMessageCount.get() >= MAX_DEBUG_MESSAGES) {
+        throw new IllegalStateException(
+            "Received too many debug messages. Did you turn on DEBUG mode by accident?");
+      }
+      debugMessages.add(message);
+      debugMessageCount.incrementAndGet();
+    }
+    return message;
+  }
+
   /**
    * The PostgreSQL wire protocol can send data in both binary and text format. When using text
    * format, the {@link Server} will normally send output back to the client using a format
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
index 9778022b..df025cea 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
@@ -58,6 +58,7 @@ public class OptionsMetadata {
   private static final String OPTION_SPANNER_ENDPOINT = "e";
   private static final String OPTION_JDBC_PROPERTIES = "r";
   private static final String OPTION_SERVER_VERSION = "v";
+  private static final String OPTION_DEBUG_MODE = "debug";
 
   private final CommandMetadataParser commandMetadataParser;
   private final String connectionURL;
@@ -71,6 +72,7 @@ public class OptionsMetadata {
   private final JSONObject commandMetadataJSON;
   private final Map<String, String> propertyMap;
   private final String serverVersion;
+  private final boolean debugMode;
 
   public OptionsMetadata(String[] args) {
     CommandLine commandLine = buildOptions(args);
@@ -88,6 +90,7 @@ public class OptionsMetadata {
     this.propertyMap = parseProperties(commandLine.getOptionValue(OPTION_JDBC_PROPERTIES, ""));
     this.disableLocalhostCheck = commandLine.hasOption(OPTION_DISABLE_LOCALHOST_CHECK);
     this.serverVersion = commandLine.getOptionValue(OPTION_SERVER_VERSION, DEFAULT_SERVER_VERSION);
+    this.debugMode = commandLine.hasOption(OPTION_DEBUG_MODE);
   }
 
   public OptionsMetadata(
@@ -111,6 +114,7 @@ public class OptionsMetadata {
     this.propertyMap = new HashMap<>();
     this.disableLocalhostCheck = false;
     this.serverVersion = DEFAULT_SERVER_VERSION;
+    this.debugMode = false;
   }
 
   private Map<String, String> parseProperties(String propertyOptions) {
@@ -348,6 +352,12 @@ public class OptionsMetadata {
         "This option specifies what server_version PG Adapter should claim to be. If not specified "
             + " it will default to version "
             + DEFAULT_SERVER_VERSION);
+    options.addOption(
+        OPTION_DEBUG_MODE,
+        "debug-mode",
+        false,
+        "-- ONLY USE FOR DEBUGGING -- This option only intended for debugging. It will "
+            + "instruct the server to keep track of all messages it receives.");
 
     CommandLineParser parser = new DefaultParser();
     HelpFormatter help = new HelpFormatter();
@@ -368,6 +378,10 @@ public class OptionsMetadata {
     return this.binaryFormat;
   }
 
+  public boolean isDebugMode() {
+    return this.debugMode;
+  }
+
   public JSONObject getCommandMetadataJSON() {
     return this.commandMetadataJSON;
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/parsers/ArrayParser.java b/src/main/java/com/google/cloud/spanner/pgadapter/parsers/ArrayParser.java
index 726d8e7a..12d4469c 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/parsers/ArrayParser.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/parsers/ArrayParser.java
@@ -69,9 +69,10 @@ class ArrayParser extends Parser<List<?>> {
       case INT64:
         return value.getInt64Array();
       case STRING:
-        return value.getStringArray();
       case PG_NUMERIC:
-        return value.getNumericArray();
+        // Get numeric arrays as a string array instead of as an array of BigDecimal, as numeric
+        // arrays could contain 'NaN' values, which are not supported by BigDecimal.
+        return value.getStringArray();
       case TIMESTAMP:
         return value.getTimestampArray();
       case FLOAT64:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/parsers/NumericParser.java b/src/main/java/com/google/cloud/spanner/pgadapter/parsers/NumericParser.java
index e25de469..f2efe4f1 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/parsers/NumericParser.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/parsers/NumericParser.java
@@ -22,6 +22,8 @@ import java.nio.charset.StandardCharsets;
 import org.postgresql.util.ByteConverter;
 
 /** Translate from wire protocol to {@link Number}. */
+// TODO(230559811): Change this to use String as type, and only convert to BigDecimal for binary
+// transfer.
 class NumericParser extends Parser<Number> {
   NumericParser(ResultSet item, int position) {
     // This should be either a BigDecimal value or a Double.NaN.
@@ -34,7 +36,16 @@ class NumericParser extends Parser<Number> {
   }
 
   NumericParser(Object item) {
-    this.item = (Number) item;
+    if (item instanceof Number) {
+      this.item = (Number) item;
+    } else if (item instanceof String) {
+      String stringValue = (String) item;
+      if (stringValue.equalsIgnoreCase("NaN")) {
+        this.item = Double.NaN;
+      } else {
+        this.item = new BigDecimal(stringValue);
+      }
+    }
   }
 
   NumericParser(byte[] item, FormatCode formatCode) {
@@ -45,7 +56,7 @@ class NumericParser extends Parser<Number> {
           if (stringValue.equalsIgnoreCase("NaN")) {
             this.item = Double.NaN;
           } else {
-            this.item = new BigDecimal(new String(item));
+            this.item = new BigDecimal(stringValue);
           }
           break;
         case BINARY:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java
index 2ef651f2..59b87ce8 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/CopyStatement.java
@@ -26,7 +26,7 @@ import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.AutocommitDmlMode;
-import com.google.cloud.spanner.connection.Connection;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser;
 import com.google.cloud.spanner.pgadapter.parsers.copy.TokenMgrError;
@@ -63,8 +63,10 @@ public class CopyStatement extends IntermediateStatement {
   private final ExecutorService executor = Executors.newSingleThreadExecutor();
 
   public CopyStatement(
-      OptionsMetadata options, ParsedStatement parsedStatement, Connection connection) {
-    super(options, parsedStatement, connection);
+      ConnectionHandler connectionHandler,
+      OptionsMetadata options,
+      ParsedStatement parsedStatement) {
+    super(connectionHandler, options, parsedStatement);
   }
 
   public Exception getException() {
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePortalStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePortalStatement.java
index 7b53a2bb..e5dd694e 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePortalStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePortalStatement.java
@@ -19,7 +19,7 @@ import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
-import com.google.cloud.spanner.connection.Connection;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.DescribePortalMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
@@ -41,8 +41,10 @@ public class IntermediatePortalStatement extends IntermediatePreparedStatement {
   protected List<Short> resultFormatCodes;
 
   public IntermediatePortalStatement(
-      OptionsMetadata options, ParsedStatement parsedStatement, Connection connection) {
-    super(options, parsedStatement, connection);
+      ConnectionHandler connectionHandler,
+      OptionsMetadata options,
+      ParsedStatement parsedStatement) {
+    super(connectionHandler, options, parsedStatement);
     this.statement = Statement.of(parsedStatement.getSqlWithoutComments());
     this.parameterFormatCodes = new ArrayList<>();
     this.resultFormatCodes = new ArrayList<>();
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java
index b4aeab4e..3375fdcf 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java
@@ -20,8 +20,10 @@ import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
-import com.google.cloud.spanner.connection.Connection;
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeStatementMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
@@ -42,8 +44,10 @@ public class IntermediatePreparedStatement extends IntermediateStatement {
   protected Statement statement;
 
   public IntermediatePreparedStatement(
-      OptionsMetadata options, ParsedStatement parsedStatement, Connection connection) {
-    super(options, parsedStatement, connection);
+      ConnectionHandler connectionHandler,
+      OptionsMetadata options,
+      ParsedStatement parsedStatement) {
+    super(connectionHandler, options, parsedStatement);
     this.parameterDataTypes = null;
   }
 
@@ -78,8 +82,18 @@ public class IntermediatePreparedStatement extends IntermediateStatement {
     // don't need to do that once more.
     if (getStatementResult(0) == null) {
       try {
-        StatementResult result = connection.execute(this.statement);
-        this.updateResultCount(0, result);
+        if (!connection.isInTransaction()
+            // TODO(230579451): Refactor to use ClientSideStatement information.
+            && this.parsedStatement.getType().equals(StatementType.CLIENT_SIDE)
+            && (this.commands.get(0).equals("ROLLBACK") || this.commands.get(0).equals("COMMIT"))) {
+          // TODO(230579929): Return warning that no transaction if connection status == IDLE.
+          connectionHandler.setStatus(ConnectionStatus.IDLE);
+        } else {
+          StatementResult result = connection.execute(this.statement);
+          this.updateResultCount(0, result);
+          connectionHandler.setStatus(
+              connection.isInTransaction() ? ConnectionStatus.TRANSACTION : ConnectionStatus.IDLE);
+        }
       } catch (SpannerException e) {
         handleExecutionException(0, e);
       }
@@ -98,7 +112,7 @@ public class IntermediatePreparedStatement extends IntermediateStatement {
   public IntermediatePortalStatement bind(
       byte[][] parameters, List<Short> parameterFormatCodes, List<Short> resultFormatCodes) {
     IntermediatePortalStatement portal =
-        new IntermediatePortalStatement(this.options, this.parsedStatement, this.connection);
+        new IntermediatePortalStatement(this.connectionHandler, this.options, this.parsedStatement);
     portal.setParameterFormatCodes(parameterFormatCodes);
     portal.setResultFormatCodes(resultFormatCodes);
     Statement.Builder builder = Statement.newBuilder(this.parsedStatement.getSqlWithoutComments());
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/ReadyResponse.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/ReadyResponse.java
index 39d59e1b..a3900b29 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/ReadyResponse.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/ReadyResponse.java
@@ -14,6 +14,7 @@
 
 package com.google.cloud.spanner.pgadapter.wireoutput;
 
+import com.google.api.client.util.Preconditions;
 import com.google.api.core.InternalApi;
 import java.io.DataOutputStream;
 import java.io.IOException;
@@ -29,7 +30,7 @@ public class ReadyResponse extends WireOutput {
 
   public ReadyResponse(DataOutputStream output, Status status) {
     super(output, 5);
-    this.status = status;
+    this.status = Preconditions.checkNotNull(status);
   }
 
   @Override
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
index 91b229ec..8000f145 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
@@ -26,7 +26,6 @@ import com.google.cloud.spanner.pgadapter.wireoutput.DataRowResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse.State;
 import com.google.cloud.spanner.pgadapter.wireoutput.PortalSuspendedResponse;
-import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
 import java.io.DataInputStream;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -138,8 +137,11 @@ public abstract class ControlMessage extends WireMessage {
    * @throws Exception if there is some issue in the sending of the error messages.
    */
   protected void handleError(Exception e) throws Exception {
-    new ErrorResponse(this.outputStream, e, State.InternalError).send(false);
-    new ReadyResponse(this.outputStream, ReadyResponse.Status.IDLE).send(false);
+    new ErrorResponse(this.outputStream, e, State.RaiseException).send(false);
+    if (connection.getSpannerConnection().isInTransaction()) {
+      connection.getSpannerConnection().rollbackAsync();
+      connection.setStatus(ConnectionStatus.TRANSACTION_ABORTED);
+    }
     this.outputStream.flush();
   }
 
@@ -167,7 +169,7 @@ public abstract class ControlMessage extends WireMessage {
         if (state.hasMoreRows()) {
           new PortalSuspendedResponse(this.outputStream).send(false);
         } else {
-          statement.getStatementResult(resultIndex).close();
+          statement.close(resultIndex);
           new CommandCompleteResponse(this.outputStream, "SELECT " + state.getNumberOfRowsSent())
               .send(false);
         }
@@ -194,7 +196,7 @@ public abstract class ControlMessage extends WireMessage {
    * @param mode Specific Query Mode required for this specific message for Postgres
    * @param maxRows Maximum number of rows requested
    * @return An adapted representation with specific metadata which PG wire requires.
-   * @throws Exception
+   * @throws com.google.cloud.spanner.SpannerException if traversing the {@link ResultSet} fails.
    */
   public SendResultSetState sendResultSet(
       int resultIndex, IntermediateStatement describedResult, QueryMode mode, long maxRows)
@@ -202,7 +204,6 @@ public abstract class ControlMessage extends WireMessage {
     long rows = 0;
     boolean hasData = describedResult.isHasMoreData(resultIndex);
     ResultSet resultSet = describedResult.getStatementResult(resultIndex);
-    // TODO optimize loop
     while (hasData) {
       new DataRowResponse(
               this.outputStream,
@@ -212,11 +213,7 @@ public abstract class ControlMessage extends WireMessage {
               mode)
           .send(false);
       rows++;
-      try {
-        hasData = resultSet.next();
-      } catch (Exception e) {
-        System.err.println("Something went wrong with getting next!");
-      }
+      hasData = resultSet.next();
       if (rows == maxRows) {
         break;
       }
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 cbbd3139..ce5319e2 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
@@ -75,6 +75,10 @@ public class DescribeMessage extends ControlMessage {
     return this.name;
   }
 
+  public PreparedType getType() {
+    return type;
+  }
+
   @Override
   protected int getHeaderLength() {
     return 5;
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ExecuteMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ExecuteMessage.java
index 42064e98..66b2154c 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ExecuteMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ExecuteMessage.java
@@ -81,8 +81,12 @@ public class ExecuteMessage extends ControlMessage {
     if (this.statement.hasException(0)) {
       this.handleError(this.statement.getException(0));
     } else {
-      this.sendSpannerResult(0, this.statement, QueryMode.EXTENDED, this.maxRows);
-      this.outputStream.flush();
+      try {
+        this.sendSpannerResult(0, this.statement, QueryMode.EXTENDED, this.maxRows);
+        this.outputStream.flush();
+      } catch (Exception e) {
+        handleError(e);
+      }
     }
     this.connection.cleanUp(this.statement);
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
index f59584c1..007b1513 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
@@ -49,9 +49,7 @@ public class ParseMessage extends ControlMessage {
     }
     this.statement =
         new IntermediatePreparedStatement(
-            connection.getServer().getOptions(),
-            parsedStatement,
-            connection.getSpannerConnection());
+            connection, connection.getServer().getOptions(), parsedStatement);
     this.statement.setParameterDataTypes(this.parameterDataTypes);
   }
 
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java
index 9db6a217..7b5fa312 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/QueryMessage.java
@@ -52,10 +52,7 @@ public class QueryMessage extends ControlMessage {
     this.isCopy = StatementParser.isCommand(COPY, parsedStatement.getSqlWithoutComments());
     if (isCopy) {
       this.statement =
-          new CopyStatement(
-              connection.getServer().getOptions(),
-              parsedStatement,
-              this.connection.getSpannerConnection());
+          new CopyStatement(connection, connection.getServer().getOptions(), parsedStatement);
     } else if (!connection.getServer().getOptions().requiresMatcher()) {
       this.statement =
           new IntermediateStatement(
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 697e67a9..e0f3bee3 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
@@ -16,6 +16,7 @@ package com.google.cloud.spanner.pgadapter.wireprotocol;
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
 import com.google.cloud.spanner.pgadapter.wireoutput.MD5AuthenticationRequest;
 import java.io.IOException;
 import java.text.MessageFormat;
@@ -48,6 +49,7 @@ public class StartupMessage extends BootstrapMessage {
           this.connection.getConnectionId(),
           this.connection.getSecret(),
           this.connection.getServer().getOptions());
+      this.connection.setStatus(ConnectionStatus.IDLE);
     } else {
       new MD5AuthenticationRequest(this.outputStream, 0).send();
     }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SyncMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SyncMessage.java
index 1759804e..b59833af 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SyncMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SyncMessage.java
@@ -17,7 +17,6 @@ package com.google.cloud.spanner.pgadapter.wireprotocol;
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
-import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse.Status;
 import java.text.MessageFormat;
 
 /**
@@ -35,10 +34,7 @@ public class SyncMessage extends ControlMessage {
 
   @Override
   protected void sendPayload() throws Exception {
-    boolean inTransaction = connection.getSpannerConnection().isInTransaction();
-    new ReadyResponse(
-            this.outputStream, inTransaction ? Status.TRANSACTION : ReadyResponse.Status.IDLE)
-        .send();
+    new ReadyResponse(this.outputStream, connection.getStatus().getReadyResponseStatus()).send();
   }
 
   @Override


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 e3a25ebc..ded14f9f 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
@@ -61,7 +61,7 @@ public class IntermediateStatement {
   protected final Connection connection;
   protected long[] updateCounts;
   protected final ImmutableList<ParsedStatement> statements;
-  private ConnectionHandler connectionHandler;
+  protected final ConnectionHandler connectionHandler;
   private ExecutionStatus executionStatus;
 
   private static final char STATEMENT_DELIMITER = ';';
@@ -71,15 +71,17 @@ public class IntermediateStatement {
       OptionsMetadata options,
       ParsedStatement parsedStatement,
       ConnectionHandler connectionHandler) {
-    this(options, parsedStatement, connectionHandler.getSpannerConnection());
-    this.connectionHandler = connectionHandler;
+    this(connectionHandler, options, parsedStatement);
   }
 
   protected IntermediateStatement(
-      OptionsMetadata options, ParsedStatement parsedStatement, Connection connection) {
+      ConnectionHandler connectionHandler,
+      OptionsMetadata options,
+      ParsedStatement parsedStatement) {
+    this.connectionHandler = connectionHandler;
     this.options = options;
     this.parsedStatement = replaceKnownUnsupportedQueries(parsedStatement);
-    this.connection = connection;
+    this.connection = connectionHandler.getSpannerConnection();
     this.statements = parseStatements(parsedStatement);
     this.commands = StatementParser.parseCommands(this.statements);
     this.commandTags = new ArrayList<>(this.commands);
@@ -203,8 +205,9 @@ public class IntermediateStatement {
   }
 
   public void close(int index) throws Exception {
-    if (this.getStatementResult(index) != null) {
-      this.getStatementResult(index).close();
+    if (this.statementResults != null && this.statementResults[index] != null) {
+      this.statementResults[index].close();
+      this.statementResults[index] = null;
     }
   }
 
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
index 92dde79d..0d9a4eda 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
@@ -27,6 +27,7 @@ import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.admin.database.v1.MockDatabaseAdminImpl;
 import com.google.cloud.spanner.connection.SpannerPool;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
 import com.google.common.collect.ImmutableList;
 import com.google.longrunning.Operation;
 import com.google.protobuf.Any;
@@ -57,7 +58,10 @@ import java.net.InetSocketAddress;
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
 import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.BeforeClass;
@@ -86,6 +90,8 @@ public abstract class AbstractMockServerTest {
 
   protected static final Statement SELECT1 = Statement.of("SELECT 1");
   protected static final Statement SELECT2 = Statement.of("SELECT 2");
+  protected static final Statement SELECT_FIVE_ROWS =
+      Statement.of("SELECT * FROM TableWithFiveRows");
   protected static final Statement INVALID_SELECT = Statement.of("SELECT foo");
   private static final ResultSetMetadata SELECT1_METADATA =
       ResultSetMetadata.newBuilder()
@@ -114,6 +120,27 @@ public abstract class AbstractMockServerTest {
                   .build())
           .setMetadata(SELECT1_METADATA)
           .build();
+  private static final com.google.spanner.v1.ResultSet SELECT_FIVE_ROWS_RESULTSET =
+      com.google.spanner.v1.ResultSet.newBuilder()
+          .addAllRows(
+              ImmutableList.of(
+                  ListValue.newBuilder()
+                      .addValues(Value.newBuilder().setStringValue("1").build())
+                      .build(),
+                  ListValue.newBuilder()
+                      .addValues(Value.newBuilder().setStringValue("2").build())
+                      .build(),
+                  ListValue.newBuilder()
+                      .addValues(Value.newBuilder().setStringValue("3").build())
+                      .build(),
+                  ListValue.newBuilder()
+                      .addValues(Value.newBuilder().setStringValue("4").build())
+                      .build(),
+                  ListValue.newBuilder()
+                      .addValues(Value.newBuilder().setStringValue("5").build())
+                      .build()))
+          .setMetadata(SELECT1_METADATA)
+          .build();
   protected static final Statement UPDATE_STATEMENT =
       Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2");
   protected static final int UPDATE_COUNT = 2;
@@ -215,6 +242,13 @@ public abstract class AbstractMockServerTest {
   private static Server spannerServer;
   protected static ProxyServer pgServer;
 
+  protected <T extends WireMessage> List<T> getWireMessagesOfType(Class<T> type) {
+    return pgServer.getDebugMessages().stream()
+        .filter(msg -> msg.getClass().equals(type))
+        .map(msg -> (T) msg)
+        .collect(Collectors.toList());
+  }
+
   @BeforeClass
   public static void startMockSpannerAndPgAdapterServers() throws Exception {
     doStartMockSpannerAndPgAdapterServers(Collections.emptyList());
@@ -226,6 +260,8 @@ public abstract class AbstractMockServerTest {
     mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions.
     mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET));
     mockSpanner.putStatementResult(StatementResult.query(SELECT2, SELECT2_RESULTSET));
+    mockSpanner.putStatementResult(
+        StatementResult.query(SELECT_FIVE_ROWS, SELECT_FIVE_ROWS_RESULTSET));
     mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT));
     mockSpanner.putStatementResult(StatementResult.update(INSERT_STATEMENT, INSERT_COUNT));
     mockSpanner.putStatementResult(
@@ -276,6 +312,7 @@ public abstract class AbstractMockServerTest {
                 "-d",
                 "d",
                 "-jdbc",
+                "-debug",
                 "-c",
                 "", // empty credentials file, as we are using a plain text connection.
                 "-s",
@@ -310,7 +347,7 @@ public abstract class AbstractMockServerTest {
       }
     }
     spannerServer.shutdown();
-    spannerServer.awaitTermination();
+    spannerServer.awaitTermination(10L, TimeUnit.SECONDS);
   }
 
   /**
@@ -339,6 +376,9 @@ public abstract class AbstractMockServerTest {
   public void clearRequests() {
     mockSpanner.clearRequests();
     mockDatabaseAdmin.reset();
+    if (pgServer != null) {
+      pgServer.clearDebugMessages();
+    }
   }
 
   protected void addDdlResponseToSpannerAdmin() {
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
index dff92b11..64811f3a 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
@@ -17,6 +17,7 @@ package com.google.cloud.spanner.pgadapter;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
@@ -553,4 +554,111 @@ public class ITJdbcTest implements IntegrationTest {
       }
     }
   }
+
+  @Test
+  public void testFetchSize() throws SQLException {
+    writeExtraTestRows();
+    try (Connection connection = DriverManager.getConnection(getConnectionUrl())) {
+      // Fetch size is only respected in a transaction.
+      connection.setAutoCommit(false);
+      try (PreparedStatement statement =
+          connection.prepareStatement("select * from all_types order by col_bigint")) {
+        // Fetch two rows at a time from the PG server.
+        statement.setFetchSize(2);
+        try (ResultSet resultSet = statement.executeQuery()) {
+          int rowCount = 0;
+          while (resultSet.next()) {
+            assertEquals(++rowCount, resultSet.getLong(1));
+            for (int col = 2; col <= resultSet.getMetaData().getColumnCount(); col++) {
+              if (resultSet.getLong(1) == 3L) {
+                assertNull(resultSet.getObject(col));
+              } else {
+                assertNotNull(resultSet.getObject(col));
+              }
+            }
+          }
+          assertEquals(5, rowCount);
+        }
+      }
+      connection.commit();
+    }
+  }
+
+  private void writeExtraTestRows() {
+    testEnv.write(
+        database.getId().getDatabase(),
+        Arrays.asList(
+            Mutation.newInsertBuilder("all_types")
+                .set("col_bigint")
+                .to(2L)
+                .set("col_bool")
+                .to(false)
+                .set("col_bytea")
+                .to(ByteArray.copyFrom("foo"))
+                .set("col_float8")
+                .to(-3.14d)
+                .set("col_int")
+                .to(Integer.MAX_VALUE)
+                .set("col_numeric")
+                .to(new BigDecimal("-3.14"))
+                .set("col_timestamptz")
+                .to(Timestamp.parseTimestamp("2022-04-22T19:27:30+02:00"))
+                .set("col_varchar")
+                .to("bar")
+                .build(),
+            Mutation.newInsertBuilder("all_types")
+                .set("col_bigint")
+                .to(3L)
+                .set("col_bool")
+                .to((Boolean) null)
+                .set("col_bytea")
+                .to((ByteArray) null)
+                .set("col_float8")
+                .to((Double) null)
+                .set("col_int")
+                .to((Long) null)
+                .set("col_numeric")
+                .to((BigDecimal) null)
+                .set("col_timestamptz")
+                .to((Timestamp) null)
+                .set("col_varchar")
+                .to((String) null)
+                .build(),
+            Mutation.newInsertBuilder("all_types")
+                .set("col_bigint")
+                .to(4L)
+                .set("col_bool")
+                .to(true)
+                .set("col_bytea")
+                .to(ByteArray.copyFrom(""))
+                .set("col_float8")
+                .to(0d)
+                .set("col_int")
+                .to(0)
+                .set("col_numeric")
+                .to(BigDecimal.ZERO)
+                .set("col_timestamptz")
+                .to(Timestamp.parseTimestamp("0001-01-01T00:00:00Z"))
+                .set("col_varchar")
+                .to("")
+                .build(),
+            Mutation.newInsertBuilder("all_types")
+                .set("col_bigint")
+                .to(5L)
+                .set("col_bool")
+                .to(true)
+                .set("col_bytea")
+                .to(ByteArray.copyFrom(""))
+                .set("col_float8")
+                .to(0d)
+                .set("col_int")
+                .to(0)
+                .set("col_numeric")
+                .to(BigDecimal.ZERO)
+                .set("col_timestamptz")
+                .to(Timestamp.parseTimestamp("0001-01-01T00:00:00Z"))
+                .set("col_varchar")
+                .to("")
+                .build()));
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
index 3c05d8e5..cd428506 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -18,13 +18,20 @@ import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.ByteArray;
 import com.google.cloud.Date;
 import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
 import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.RandomResultSetGenerator;
+import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage.PreparedType;
+import com.google.cloud.spanner.pgadapter.wireprotocol.DescribeMessage;
+import com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage;
+import com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage;
 import com.google.protobuf.Value;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
@@ -33,6 +40,7 @@ import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import com.google.spanner.v1.Type;
 import com.google.spanner.v1.TypeAnnotationCode;
 import com.google.spanner.v1.TypeCode;
+import io.grpc.Status;
 import java.math.BigDecimal;
 import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
@@ -51,6 +59,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 import org.junit.BeforeClass;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -60,11 +69,21 @@ import org.postgresql.jdbc.PgStatement;
 
 @RunWith(JUnit4.class)
 public class JdbcMockServerTest extends AbstractMockServerTest {
+  private static final int RANDOM_RESULTS_ROW_COUNT = 10;
+  private static final Statement SELECT_RANDOM = Statement.of("select * from random_table");
 
   @BeforeClass
   public static void loadPgJdbcDriver() throws Exception {
     // Make sure the PG JDBC driver is loaded.
     Class.forName("org.postgresql.Driver");
+
+    addRandomResultResults();
+  }
+
+  private static void addRandomResultResults() {
+    // TODO(230579459): Add dialect argument once RandomResultSetGenerator supports PostgreSQL.
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(RANDOM_RESULTS_ROW_COUNT);
+    mockSpanner.putStatementResult(StatementResult.query(SELECT_RANDOM, generator.generate()));
   }
 
   /**
@@ -538,4 +557,171 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
       }
     }
   }
+
+  @Test
+  public void testCursorSuccess() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      connection.setAutoCommit(false);
+      try (PreparedStatement statement = connection.prepareStatement(SELECT_FIVE_ROWS.getSql())) {
+        // Fetch two rows at a time from the PG server.
+        statement.setFetchSize(2);
+        try (ResultSet resultSet = statement.executeQuery()) {
+          int index = 0;
+          while (resultSet.next()) {
+            assertEquals(++index, resultSet.getInt(1));
+          }
+          assertEquals(5, index);
+        }
+      }
+      connection.commit();
+    }
+    // The ExecuteSql request should only be sent once to Cloud Spanner.
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    ExecuteSqlRequest executeRequest =
+        mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
+    assertEquals(QueryMode.NORMAL, executeRequest.getQueryMode());
+    assertEquals(SELECT_FIVE_ROWS.getSql(), executeRequest.getSql());
+
+    // PGAdapter should receive 5 Execute messages:
+    // 1. BEGIN
+    // 2. Execute - fetch rows 1, 2
+    // 3. Execute - fetch rows 3, 4
+    // 4. Execute - fetch rows 5
+    // 5. COMMIT
+    if (pgServer != null) {
+      List<DescribeMessage> describeMessages = getWireMessagesOfType(DescribeMessage.class);
+      assertEquals(1, describeMessages.size());
+      DescribeMessage describeMessage = describeMessages.get(0);
+      assertEquals(PreparedType.Portal, describeMessage.getType());
+
+      List<ExecuteMessage> executeMessages = getWireMessagesOfType(ExecuteMessage.class);
+      assertEquals(5, executeMessages.size());
+      assertEquals("", executeMessages.get(0).getName());
+      for (ExecuteMessage executeMessage : executeMessages.subList(1, executeMessages.size() - 1)) {
+        assertEquals(describeMessage.getName(), executeMessage.getName());
+        assertEquals(2, executeMessage.getMaxRows());
+      }
+      assertEquals("", executeMessages.get(executeMessages.size() - 1).getName());
+
+      List<ParseMessage> parseMessages = getWireMessagesOfType(ParseMessage.class);
+      assertEquals(3, parseMessages.size());
+      assertEquals("BEGIN", parseMessages.get(0).getStatement().getSql());
+      assertEquals(SELECT_FIVE_ROWS.getSql(), parseMessages.get(1).getStatement().getSql());
+      assertEquals("COMMIT", parseMessages.get(2).getStatement().getSql());
+    }
+  }
+
+  @Test
+  public void testCursorFailsHalfway() throws SQLException {
+    mockSpanner.setExecuteStreamingSqlExecutionTime(
+        SimulatedExecutionTime.ofStreamException(Status.DATA_LOSS.asRuntimeException(), 2));
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      connection.setAutoCommit(false);
+      try (PreparedStatement statement = connection.prepareStatement(SELECT_FIVE_ROWS.getSql())) {
+        // Fetch one row at a time from the PG server.
+        statement.setFetchSize(1);
+        try (ResultSet resultSet = statement.executeQuery()) {
+          // The first row should succeed.
+          assertTrue(resultSet.next());
+          // The second row should fail.
+          assertThrows(SQLException.class, resultSet::next);
+        }
+      }
+      connection.rollback();
+    }
+    // The ExecuteSql request should only be sent once to Cloud Spanner.
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    ExecuteSqlRequest executeRequest =
+        mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
+    assertEquals(QueryMode.NORMAL, executeRequest.getQueryMode());
+    assertEquals(SELECT_FIVE_ROWS.getSql(), executeRequest.getSql());
+
+    // PGAdapter should receive 4 Execute messages:
+    // 1. BEGIN
+    // 2. Execute - fetch row 1
+    // 3. Execute - fetch row 2 -- This fails with a DATA_LOSS error
+    // The JDBC driver does not send a ROLLBACK
+    if (pgServer != null) {
+      List<DescribeMessage> describeMessages = getWireMessagesOfType(DescribeMessage.class);
+      assertEquals(1, describeMessages.size());
+      DescribeMessage describeMessage = describeMessages.get(0);
+      assertEquals(PreparedType.Portal, describeMessage.getType());
+
+      List<ExecuteMessage> executeMessages = getWireMessagesOfType(ExecuteMessage.class);
+      assertEquals(4, executeMessages.size());
+      assertEquals("", executeMessages.get(0).getName());
+      for (ExecuteMessage executeMessage : executeMessages.subList(1, executeMessages.size() - 1)) {
+        assertEquals(describeMessage.getName(), executeMessage.getName());
+        assertEquals(1, executeMessage.getMaxRows());
+      }
+      assertEquals("", executeMessages.get(executeMessages.size() - 1).getName());
+
+      List<ParseMessage> parseMessages = getWireMessagesOfType(ParseMessage.class);
+      assertEquals(3, parseMessages.size());
+      assertEquals("BEGIN", parseMessages.get(0).getStatement().getSql());
+      assertEquals(SELECT_FIVE_ROWS.getSql(), parseMessages.get(1).getStatement().getSql());
+      assertEquals("ROLLBACK", parseMessages.get(2).getStatement().getSql());
+    }
+  }
+
+  @Test
+  @Ignore("RandomResultSetGenerator does not yet support PostgreSQL")
+  public void testRandomResults() throws SQLException {
+    final int fetchSize = 3;
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      connection.setAutoCommit(false);
+      try (PreparedStatement statement = connection.prepareStatement(SELECT_RANDOM.getSql())) {
+        statement.setFetchSize(fetchSize);
+        try (ResultSet resultSet = statement.executeQuery()) {
+          int rowCount = 0;
+          while (resultSet.next()) {
+            for (int col = 0; col < resultSet.getMetaData().getColumnCount(); col++) {
+              resultSet.getObject(col + 1);
+            }
+            rowCount++;
+          }
+          assertEquals(RANDOM_RESULTS_ROW_COUNT, rowCount);
+        }
+      }
+      connection.commit();
+    }
+    // The ExecuteSql request should only be sent once to Cloud Spanner.
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    ExecuteSqlRequest executeRequest =
+        mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
+    assertEquals(QueryMode.NORMAL, executeRequest.getQueryMode());
+    assertEquals(SELECT_RANDOM.getSql(), executeRequest.getSql());
+
+    // PGAdapter should receive 5 Execute messages:
+    // 1. BEGIN
+    // 2. Execute - fetch rows 1, 2
+    // 3. Execute - fetch rows 3, 4
+    // 4. Execute - fetch rows 5
+    // 5. COMMIT
+    if (pgServer != null) {
+      List<DescribeMessage> describeMessages = getWireMessagesOfType(DescribeMessage.class);
+      assertEquals(1, describeMessages.size());
+      DescribeMessage describeMessage = describeMessages.get(0);
+      assertEquals(PreparedType.Portal, describeMessage.getType());
+
+      List<ExecuteMessage> executeMessages = getWireMessagesOfType(ExecuteMessage.class);
+      int expectedExecuteMessageCount =
+          RANDOM_RESULTS_ROW_COUNT / fetchSize
+              + ((RANDOM_RESULTS_ROW_COUNT % fetchSize) > 0 ? 1 : 0)
+              + 2;
+      assertEquals(expectedExecuteMessageCount, executeMessages.size());
+      assertEquals("", executeMessages.get(0).getName());
+      for (ExecuteMessage executeMessage : executeMessages.subList(1, executeMessages.size() - 1)) {
+        assertEquals(describeMessage.getName(), executeMessage.getName());
+        assertEquals(fetchSize, executeMessage.getMaxRows());
+      }
+      assertEquals("", executeMessages.get(executeMessages.size() - 1).getName());
+
+      List<ParseMessage> parseMessages = getWireMessagesOfType(ParseMessage.class);
+      assertEquals(3, parseMessages.size());
+      assertEquals("BEGIN", parseMessages.get(0).getStatement().getSql());
+      assertEquals(SELECT_RANDOM.getSql(), parseMessages.get(1).getStatement().getSql());
+      assertEquals("COMMIT", parseMessages.get(2).getStatement().getSql());
+    }
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
index 34dd88e3..eb0f73b2 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
@@ -1013,7 +1013,7 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
+    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.IDLE);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
@@ -1042,8 +1042,7 @@ public class ProtocolTest {
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
-    when(connection.isInTransaction()).thenReturn(true);
+    when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.TRANSACTION);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
@@ -1230,6 +1229,7 @@ public class ProtocolTest {
 
   @Test
   public void testMultipleCopyDataMessages() throws Exception {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     when(connectionHandler.getStatus()).thenReturn(ConnectionStatus.COPY_IN);
 
     byte[] messageMetadata = {'d'};
@@ -1254,7 +1254,7 @@ public class ProtocolTest {
     when(connection.executeQuery(any(Statement.class))).thenReturn(spannerType);
 
     CopyStatement copyStatement =
-        new CopyStatement(options, parse("COPY keyvalue FROM STDIN;"), connection);
+        new CopyStatement(connectionHandler, options, parse("COPY keyvalue FROM STDIN;"));
     copyStatement.execute();
 
     when(connectionHandler.getActiveStatement()).thenReturn(copyStatement);
@@ -1358,7 +1358,7 @@ public class ProtocolTest {
 
     CopyStatement copyStatement =
         new CopyStatement(
-            mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"), connection);
+            connectionHandler, mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"));
     copyStatement.execute();
 
     MutationWriter mw = copyStatement.getMutationWriter();
@@ -1379,7 +1379,7 @@ public class ProtocolTest {
 
     CopyStatement copyStatement =
         new CopyStatement(
-            mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"), connection);
+            connectionHandler, mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"));
 
     assertFalse(copyStatement.isExecuted());
     copyStatement.execute();
@@ -1411,7 +1411,7 @@ public class ProtocolTest {
 
     CopyStatement copyStatement =
         new CopyStatement(
-            mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"), connection);
+            connectionHandler, mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"));
 
     assertFalse(copyStatement.isExecuted());
     copyStatement.execute();
@@ -1443,7 +1443,7 @@ public class ProtocolTest {
 
     CopyStatement copyStatement =
         new CopyStatement(
-            mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"), connection);
+            connectionHandler, mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"));
     assertFalse(copyStatement.isExecuted());
     copyStatement.execute();
     assertTrue(copyStatement.isExecuted());
@@ -1484,7 +1484,7 @@ public class ProtocolTest {
     }
 
     CopyStatement copyStatement =
-        new CopyStatement(options, parse("COPY keyvalue FROM STDIN;"), connection);
+        new CopyStatement(connectionHandler, options, parse("COPY keyvalue FROM STDIN;"));
     assertFalse(copyStatement.isExecuted());
     copyStatement.execute();
     assertTrue(copyStatement.isExecuted());
@@ -1756,6 +1756,7 @@ public class ProtocolTest {
   }
 
   private void setupQueryInformationSchemaResults() {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     ResultSet spannerType = mock(ResultSet.class);
     when(spannerType.getString("column_name")).thenReturn("key", "value");
     when(spannerType.getString("data_type")).thenReturn("bigint", "character varying");
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java
index d64b833b..8247ecec 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java
@@ -247,6 +247,7 @@ public class StatementTest {
 
   @Test
   public void testPreparedStatement() {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     String sqlStatement = "SELECT * FROM users WHERE age > $2 AND age < $3 AND name = $1";
     int[] parameterDataTypes = new int[] {Oid.VARCHAR, Oid.INT8, Oid.INT4};
 
@@ -264,7 +265,7 @@ public class StatementTest {
     when(connection.execute(statement)).thenReturn(statementResult);
 
     IntermediatePreparedStatement intermediateStatement =
-        new IntermediatePreparedStatement(options, parse(sqlStatement), connection);
+        new IntermediatePreparedStatement(connectionHandler, options, parse(sqlStatement));
     intermediateStatement.setParameterDataTypes(parameterDataTypes);
 
     assertEquals(sqlStatement, intermediateStatement.getSql());
@@ -288,7 +289,7 @@ public class StatementTest {
     int[] parameterDataTypes = new int[] {Oid.JSON};
 
     IntermediatePreparedStatement intermediateStatement =
-        new IntermediatePreparedStatement(options, parse(sqlStatement), connection);
+        new IntermediatePreparedStatement(connectionHandler, options, parse(sqlStatement));
     intermediateStatement.setParameterDataTypes(parameterDataTypes);
 
     byte[][] parameters = {"{}".getBytes()};
@@ -300,23 +301,25 @@ public class StatementTest {
 
   @Test
   public void testPreparedStatementDescribeDoesNotThrowException() {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     String sqlStatement = "SELECT * FROM users WHERE name = $1 AND age > $2 AND age < $3";
     when(connection.analyzeQuery(Statement.of(sqlStatement), QueryAnalyzeMode.PLAN))
         .thenReturn(resultSet);
 
     IntermediatePreparedStatement intermediateStatement =
-        new IntermediatePreparedStatement(options, parse(sqlStatement), connection);
+        new IntermediatePreparedStatement(connectionHandler, options, parse(sqlStatement));
 
     intermediateStatement.describe();
   }
 
   @Test
   public void testPortalStatement() {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     String sqlStatement = "SELECT * FROM users WHERE age > $1 AND age < $2 AND name = $3";
     when(connection.executeQuery(Statement.of(sqlStatement))).thenReturn(resultSet);
 
     IntermediatePortalStatement intermediateStatement =
-        new IntermediatePortalStatement(options, parse(sqlStatement), connection);
+        new IntermediatePortalStatement(connectionHandler, options, parse(sqlStatement));
 
     intermediateStatement.describe();
 
@@ -352,10 +355,11 @@ public class StatementTest {
 
   @Test
   public void testPortalStatementDescribePropagatesFailure() {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     String sqlStatement = "SELECT * FROM users WHERE age > $1 AND age < $2 AND name = $3";
 
     IntermediatePortalStatement intermediateStatement =
-        new IntermediatePortalStatement(options, parse(sqlStatement), connection);
+        new IntermediatePortalStatement(connectionHandler, options, parse(sqlStatement));
 
     when(connection.executeQuery(Statement.of(sqlStatement)))
         .thenThrow(
@@ -456,11 +460,12 @@ public class StatementTest {
 
   @Test
   public void testCopyInvalidBuildMutation() throws Exception {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     setupQueryInformationSchemaResults();
 
     CopyStatement statement =
         new CopyStatement(
-            mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"), connection);
+            connectionHandler, mock(OptionsMetadata.class), parse("COPY keyvalue FROM STDIN;"));
     statement.execute();
 
     byte[] payload = "2 3\n".getBytes();
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatementTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatementTest.java
index a88eaa6f..72599d5c 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatementTest.java
@@ -67,7 +67,6 @@ public class IntermediateStatementTest {
     StatementResult result = mock(StatementResult.class);
     when(result.getResultType()).thenReturn(ResultType.RESULT_SET);
     when(result.getResultSet()).thenReturn(resultSet);
-    when(result.getUpdateCount()).thenThrow(new IllegalStateException());
 
     statement.updateResultCount(0, result);
 
@@ -84,7 +83,6 @@ public class IntermediateStatementTest {
             mock(OptionsMetadata.class), parse("update bar set foo=1"), connectionHandler);
     StatementResult result = mock(StatementResult.class);
     when(result.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
-    when(result.getResultSet()).thenThrow(new IllegalStateException());
     when(result.getUpdateCount()).thenReturn(100L);
 
     statement.updateResultCount(0, result);
@@ -104,8 +102,6 @@ public class IntermediateStatementTest {
             connectionHandler);
     StatementResult result = mock(StatementResult.class);
     when(result.getResultType()).thenReturn(ResultType.NO_RESULT);
-    when(result.getResultSet()).thenThrow(new IllegalStateException());
-    when(result.getUpdateCount()).thenThrow(new IllegalStateException());
 
     statement.updateResultCount(0, result);