Repository: GoogleCloudPlatform/pgadapter
Pull request #108 (2022-04-08 23:48:53)

feat: all statement batching
Support batching of all statement types. Response messages are sent for each statement in the batch. 
Cover most of the cases (not only the short-term solution to python driver) with the exception:
1. The remaining statements in a dml/ddl batch are skipped and throw exceptions if the dml/ddl batch fails in the middle. The preceding statements in dml/ddl batch may not take effect even though they succeed. The whole batch resumes after the dml/ddl batch.


libingye816: @olavloite Do you have an idea why clirr and e2e-psql-xxx are failing? 
olavloite: The e2e psql tests are failing because some of those tests check whether they get an expected error when executing a mixed batch. That is no longer the case, as this PR introduces support for that. The psql tests work with text files that contain the expected output. If you take a look at the original draft PR, you will see that I changed the expected output of some of those (see for example dml-ddl-batching.txt: https://github.com/GoogleCloudPlatform/pgadapter/pull/30/files)

Clirr is failing because this PR introduces breaking changes in the public API of PGAdapter. I think we should disable clirr for this repository. It does not really make very much sense, as this is not a library that is imported as a dependency into other applications.

libingye816: > The e2e psql tests are failing because some of those tests check whether they get an expected error when executing a mixed batch. That is no longer the case, as this PR introduces support for that. The psql tests work with text files that contain the expected output. If you take a look at the original draft PR, you will see that I changed the expected output of some of those (see for example dml-ddl-batching.txt: https://github.com/GoogleCloudPlatform/pgadapter/pull/30/files)

I have already updated those expected output files. All tests were passed at last Friday, so I am surprised that they fail again. The psql receives error while expecting some outputs. Looks like the batch support doesn't take effect. 
libingye816: I fixed the logic of statement split to ignore escaped quotes. It should work now. @olavloite 
libingye816: Hi @olavloite Can you give a second round of quick review?

Patch:
diff --git a/.ci/evaluate-with-psql.sh b/.ci/evaluate-with-psql.sh
index 18439ba9..e44e6001 100644
--- a/.ci/evaluate-with-psql.sh
+++ b/.ci/evaluate-with-psql.sh
@@ -106,9 +106,9 @@ echo "------Test \"-c option WITH DML batching\"------"
 diff -i -w -s .ci/e2e-result/with-dml-batching.txt .ci/e2e-expected/with-dml-batching.txt
 RETURN_CODE=$((${RETURN_CODE}||$?))
 
-echo "------Test \"-c option invalid begin/commit batching\"------"
-/usr/lib/postgresql/"${PSQL_VERSION}"/bin/psql -h localhost -p 4242 -d "${GOOGLE_CLOUD_DATABASE_WITH_VERSION}" -c "$(cat .ci/e2e-batching/invalid-commit-batch.txt)" &> .ci/e2e-result/invalid-commit-batching.txt
-diff -i -w -s .ci/e2e-result/invalid-commit-batching.txt .ci/e2e-expected/invalid-commit-batching.txt
+echo "------Test \"-c option begin/commit batching\"------"
+/usr/lib/postgresql/"${PSQL_VERSION}"/bin/psql -h localhost -p 4242 -d "${GOOGLE_CLOUD_DATABASE_WITH_VERSION}" -c "$(cat .ci/e2e-batching/begin-commit-batch.txt)" &> .ci/e2e-result/begin-commit-batching.txt
+diff -i -w -s .ci/e2e-result/begin-commit-batching.txt .ci/e2e-expected/begin-commit-batching.txt
 RETURN_CODE=$((${RETURN_CODE}||$?))
 
 echo "------Test \"COPY FROM STDIN\"------"
diff --git a/.ci/run-with-credentials.sh b/.ci/run-with-credentials.sh
index 2d18eca9..c54df692 100755
--- a/.ci/run-with-credentials.sh
+++ b/.ci/run-with-credentials.sh
@@ -30,19 +30,19 @@ set +e
 
 case ${JOB_TYPE} in
 units)
-  mvn verify -Dclirr.skip=true -DskipITs=true -Ptest-all
+  mvn verify -B -Dclirr.skip=true -DskipITs=true -Ptest-all
   ;;
 lint)
-  mvn com.coveo:fmt-maven-plugin:check
+  mvn -B com.coveo:fmt-maven-plugin:check
   ;;
 clirr)
-  mvn clirr:check
+  mvn -B clirr:check
   ;;
 integration)
-  mvn verify -Dclirr.skip=true -DskipITs=false -DPG_ADAPTER_HOST="https://${GOOGLE_CLOUD_ENDPOINT}" -DPG_ADAPTER_INSTANCE="${GOOGLE_CLOUD_INSTANCE}" -DPG_ADAPTER_DATABASE="${GOOGLE_CLOUD_DATABASE}"
+  mvn verify -B -Dclirr.skip=true -DskipITs=false -DPG_ADAPTER_HOST="https://${GOOGLE_CLOUD_ENDPOINT}" -DPG_ADAPTER_INSTANCE="${GOOGLE_CLOUD_INSTANCE}" -DPG_ADAPTER_DATABASE="${GOOGLE_CLOUD_DATABASE}"
   ;;
 uber-jar-build)
-  mvn package -Pshade -DskipTests
+  mvn package -Pshade -DskipTests -B -Dorg.slf4j.simpleLogger.log.org.apache.maven.cli.transfer.Slf4jMavenTransferListener=warn
   ;;
 uber-jar-release)
   PGADAPTER_VERSION="$(mvn org.apache.maven.plugins:maven-help-plugin:3.2.0:evaluate -Dexpression=project.version -q -DforceStdout)"
@@ -73,6 +73,8 @@ e2e-psql)
   done
 
 #  start PGAdapter
+  echo "Starting PGAdapter"
+  ls -lh target
   UBER_JAR="pgadapter.jar"
   (java -jar target/"${UBER_JAR}" -p "${GOOGLE_CLOUD_PROJECT}" -i "${GOOGLE_CLOUD_INSTANCE}" -d "${GOOGLE_CLOUD_DATABASE_WITH_VERSION}" -e "${GOOGLE_CLOUD_ENDPOINT}" -s 4242 -q > /dev/null 2>&1) &
   BACK_PID=$!
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 220c70ab..55f9b13a 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -55,15 +55,6 @@ jobs:
         java-version: 11
     - run: java -version
     - run: .ci/run-with-credentials.sh lint
-  clirr:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v2
-    - uses: actions/setup-java@v1
-      with:
-        java-version: 8
-    - run: java -version
-    - run: .ci/run-with-credentials.sh clirr
   e2e-psql-v11:
     runs-on: ubuntu-latest
     steps:
@@ -160,7 +151,7 @@ jobs:
     runs-on: ubuntu-latest
     # Only releases on the postgresql-dialect branch
     if: github.head_ref == 'postgresql-dialect'
-    needs: [ units, lint, clirr, integration, e2e-psql-v11, e2e-psql-v12, e2e-psql-v13 ]
+    needs: [ units, lint, integration, e2e-psql-v11, e2e-psql-v12, e2e-psql-v13 ]
     steps:
     - uses: actions/checkout@v2
     - uses: actions/setup-java@v1
@@ -180,7 +171,7 @@ jobs:
     runs-on: ubuntu-latest
     # Only releases on the postgresql-dialect branch
     if: github.head_ref == 'postgresql-dialect'
-    needs: [ units, lint, clirr, integration, e2e-psql-v11, e2e-psql-v12, e2e-psql-v13 ]
+    needs: [ units, lint, integration, e2e-psql-v11, e2e-psql-v12, e2e-psql-v13 ]
     steps:
     - uses: actions/checkout@v2
     - uses: actions/setup-java@v1
diff --git a/.gitignore b/.gitignore
index 6f02bed5..9b80f76d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,7 @@ target/
 *.iml
 *.class
 *.lst
+output.txt
 
 src/test/golang/**/*.h
 src/test/golang/**/*.so
diff --git a/pom.xml b/pom.xml
index 991bf778..5336f6cd 100644
--- a/pom.xml
+++ b/pom.xml
@@ -67,7 +67,7 @@
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-spanner</artifactId>
-      <version>6.21.0</version>
+      <version>6.23.1</version>
     </dependency>
     <dependency>
       <groupId>org.postgresql</groupId>
@@ -93,14 +93,14 @@
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-spanner</artifactId>
-      <version>6.21.0</version>
+      <version>6.23.1</version>
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
     <dependency>
       <groupId>com.google.api</groupId>
       <artifactId>gax-grpc</artifactId>
-      <version>2.12.2</version>
+      <version>2.13.0</version>
       <classifier>testlib</classifier>
       <scope>test</scope>
     </dependency>
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 46bde684..bbbf1fe3 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -235,8 +235,10 @@ public class ConnectionHandler extends Thread {
 
   /** Closes portals and statements if the result of an execute was the end of a transaction. */
   public void cleanUp(IntermediateStatement statement) throws Exception {
-    if (!statement.isHasMoreData() && statement.isBound()) {
-      statement.close();
+    for (int index = 0; index < statement.getStatementCount(); index++) {
+      if (!statement.isHasMoreData(index) && statement.isBound()) {
+        statement.close(index);
+      }
     }
     // TODO when we have transaction data from jdbcConnection, close all portals if done
   }
@@ -401,7 +403,8 @@ public class ConnectionHandler extends Thread {
     UNAUTHENTICATED,
     IDLE,
     COPY_IN,
-    TERMINATED
+    TERMINATED,
+    TRANSACTION_ABORTED
   }
 
   /**
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 f73f99fd..efa860d3 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
@@ -23,9 +23,9 @@ import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.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.connection.StatementResult.ResultType;
 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;
@@ -66,10 +66,9 @@ public class CopyStatement extends IntermediateStatement {
     super(options, parsedStatement, connection);
   }
 
-  @Override
   public Exception getException() {
     // Do not clear exceptions on a CopyStatement.
-    return this.exception;
+    return this.exceptions[0];
   }
 
   @Override
@@ -81,7 +80,6 @@ public class CopyStatement extends IntermediateStatement {
     super.close();
   }
 
-  @Override
   public Long getUpdateCount() {
     try {
       return updateCount.get();
@@ -93,8 +91,8 @@ public class CopyStatement extends IntermediateStatement {
   }
 
   @Override
-  public ResultType getResultType() {
-    return ResultType.UPDATE_COUNT;
+  public StatementType getStatementType(int index) {
+    return StatementType.UPDATE;
   }
 
   /** @return Mapping of table column names to column type. */
@@ -288,10 +286,14 @@ public class CopyStatement extends IntermediateStatement {
     return 0;
   }
 
-  @Override
   public void handleExecutionException(SpannerException e) {
+    this.handleExecutionException(0, e);
+  }
+
+  @Override
+  public void handleExecutionException(int index, SpannerException e) {
     executor.shutdownNow();
-    super.handleExecutionException(e);
+    super.handleExecutionException(index, e);
   }
 
   @Override
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 680ac061..6bf180b0 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
@@ -14,6 +14,7 @@
 
 package com.google.cloud.spanner.pgadapter.statements;
 
+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;
@@ -86,8 +87,8 @@ public class IntermediatePortalStatement extends IntermediatePreparedStatement {
       // DescribePortal message without a following Execute message is extremely rare, as that would
       // only happen if the client is ill-behaved, or if the client crashes between the
       // DescribePortal and Execute.
-      this.statementResult = connection.executeQuery(this.statement);
-      this.hasMoreData = this.statementResult.next();
+      ResultSet statementResult = connection.executeQuery(this.statement);
+      setStatementResult(0, statementResult);
       return new DescribePortalMetadata(statementResult);
     } catch (SpannerException e) {
       logger.log(Level.SEVERE, e, e::getMessage);
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 968d9905..eab578a4 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
@@ -73,12 +73,12 @@ public class IntermediatePreparedStatement extends IntermediateStatement {
     this.executed = true;
     // If the portal has already been described, the statement has already been executed, and we
     // don't need to do that once more.
-    if (this.statementResult == null) {
+    if (getStatementResult(0) == null) {
       try {
         StatementResult result = connection.execute(this.statement);
-        this.updateResultCount(result);
+        this.updateResultCount(0, result);
       } catch (SpannerException e) {
-        handleExecutionException(e);
+        handleExecutionException(0, e);
       }
     }
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/MatcherStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/MatcherStatement.java
index 75600019..150b197f 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/MatcherStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/MatcherStatement.java
@@ -16,11 +16,9 @@ package com.google.cloud.spanner.pgadapter.statements;
 
 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.commands.Command;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
-import org.json.simple.JSONObject;
 
 /**
  * Meant to be utilized when running as a proxy for any interactive terminal tool either PSQL or
@@ -34,13 +32,7 @@ public class MatcherStatement extends IntermediateStatement {
       OptionsMetadata options,
       ParsedStatement parsedStatement,
       ConnectionHandler connectionHandler) {
-    super(
-        options,
-        translateSQL(
-            parsedStatement,
-            connectionHandler.getSpannerConnection(),
-            connectionHandler.getServer().getOptions().getCommandMetadataJSON()),
-        connectionHandler.getSpannerConnection());
+    super(options, translateSQL(parsedStatement, connectionHandler), connectionHandler);
   }
 
   @Override
@@ -57,10 +49,12 @@ public class MatcherStatement extends IntermediateStatement {
    *     gives out the original Statement.
    */
   private static ParsedStatement translateSQL(
-      ParsedStatement parsedStatement, Connection connection, JSONObject commandMetadataJSON) {
+      ParsedStatement parsedStatement, ConnectionHandler connectionHandler) {
     for (Command currentCommand :
         Command.getCommands(
-            parsedStatement.getSqlWithoutComments(), connection, commandMetadataJSON)) {
+            parsedStatement.getSqlWithoutComments(),
+            connectionHandler.getSpannerConnection(),
+            connectionHandler.getServer().getOptions().getCommandMetadataJSON())) {
       if (currentCommand.is()) {
         return PARSER.parse(Statement.of(currentCommand.translate()));
       }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/StatementParser.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/StatementParser.java
index f42c9731..f167e9be 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/StatementParser.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/StatementParser.java
@@ -14,7 +14,9 @@
 
 package com.google.cloud.spanner.pgadapter.utils;
 
+import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 
 public class StatementParser {
 
@@ -64,4 +66,14 @@ public class StatementParser {
     return Character.isSpaceChar(query.charAt(command.length()))
         && query.substring(0, command.length()).equalsIgnoreCase(command);
   }
+
+  public static ImmutableList<String> parseCommands(ImmutableList<ParsedStatement> statements) {
+    Preconditions.checkNotNull(statements);
+
+    ImmutableList.Builder<String> builder = ImmutableList.builder();
+    for (ParsedStatement stat : statements) {
+      builder.add(parseCommand(stat.getSqlWithoutComments()));
+    }
+    return builder.build();
+  }
 }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/DataRowResponse.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/DataRowResponse.java
index 9575ea04..afaf9acc 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/DataRowResponse.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/DataRowResponse.java
@@ -31,21 +31,20 @@ public class DataRowResponse extends WireOutput {
   private static final int COLUMN_NUMBER_LENGTH = 2;
   private static final int COLUMN_SIZE_LENGTH = 4;
 
-  private final IntermediateStatement statement;
   private final QueryMode mode;
   private final ResultSet resultSet;
   private final List<byte[]> columns;
 
   public DataRowResponse(
       DataOutputStream output,
+      int resultIndex,
       IntermediateStatement statement,
       OptionsMetadata options,
       QueryMode mode)
       throws Exception {
     super(output, 0);
-    this.statement = statement;
     this.mode = mode;
-    this.resultSet = statement.getStatementResult();
+    this.resultSet = statement.getStatementResult(resultIndex);
     this.columns = new ArrayList<>(this.resultSet.getColumnCount());
     this.length = HEADER_LENGTH + COLUMN_NUMBER_LENGTH;
     // TODO profile this with immense rows/column count to see if it's an aread of optimization.
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 cd6e3161..8979e46e 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
@@ -150,34 +150,37 @@ public abstract class ControlMessage extends WireMessage {
    *
    * <p>NOTE: This method does not flush the output stream.
    */
-  public boolean sendSpannerResult(IntermediateStatement statement, QueryMode mode, long maxRows)
+  public boolean sendSpannerResult(
+      int resultIndex, IntermediateStatement statement, QueryMode mode, long maxRows)
       throws Exception {
-    String command = statement.getCommand();
-    switch (statement.getResultType()) {
-      case NO_RESULT:
+    String command = statement.getCommandTag(resultIndex);
+    switch (statement.getStatementType(resultIndex)) {
+      case DDL:
+      case CLIENT_SIDE:
         new CommandCompleteResponse(this.outputStream, command).send(false);
         return false;
-      case RESULT_SET:
-        SendResultSetState state = sendResultSet(statement, mode, maxRows);
-        statement.setHasMoreData(state.hasMoreRows());
+      case QUERY:
+        SendResultSetState state = sendResultSet(resultIndex, statement, mode, maxRows);
+        statement.setHasMoreData(resultIndex, state.hasMoreRows());
         if (state.hasMoreRows()) {
           new PortalSuspendedResponse(this.outputStream).send(false);
         } else {
-          statement.getStatementResult().close();
+          statement.getStatementResult(resultIndex).close();
           new CommandCompleteResponse(this.outputStream, "SELECT " + state.getNumberOfRowsSent())
               .send(false);
         }
         return state.hasMoreRows();
-      case UPDATE_COUNT:
+      case UPDATE:
         // For an INSERT command, the tag is INSERT oid rows, where rows is the number of rows
         // inserted. oid used to be the object ID of the inserted row if rows was 1 and the target
         // table had OIDs, but OIDs system columns are not supported anymore; therefore oid is
         // always 0.
-        command += ("INSERT".equals(command) ? " 0 " : " ") + statement.getUpdateCount();
+        command += ("INSERT".equals(command) ? " 0 " : " ") + statement.getUpdateCount(resultIndex);
         new CommandCompleteResponse(this.outputStream, command).send(false);
         return false;
       default:
-        throw new IllegalStateException("Unknown result type: " + statement.getResultType());
+        throw new IllegalStateException(
+            "Unknown statement type: " + statement.getStatement(resultIndex));
     }
   }
 
@@ -192,14 +195,19 @@ public abstract class ControlMessage extends WireMessage {
    * @throws Exception
    */
   public SendResultSetState sendResultSet(
-      IntermediateStatement describedResult, QueryMode mode, long maxRows) throws Exception {
+      int resultIndex, IntermediateStatement describedResult, QueryMode mode, long maxRows)
+      throws Exception {
     long rows = 0;
-    boolean hasData = describedResult.isHasMoreData();
-    ResultSet resultSet = describedResult.getStatementResult();
+    boolean hasData = describedResult.isHasMoreData(resultIndex);
+    ResultSet resultSet = describedResult.getStatementResult(resultIndex);
     // TODO optimize loop
     while (hasData) {
       new DataRowResponse(
-              this.outputStream, describedResult, this.connection.getServer().getOptions(), mode)
+              this.outputStream,
+              resultIndex,
+              describedResult,
+              this.connection.getServer().getOptions(),
+              mode)
           .send(false);
       rows++;
       try {
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDataMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDataMessage.java
index 67950633..56067fc2 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDataMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDataMessage.java
@@ -50,7 +50,7 @@ public class CopyDataMessage extends ControlMessage {
     }
     // If backend error occurred during copy-in mode, drop any subsequent CopyData messages.
     MutationWriter mutationWriter = this.statement.getMutationWriter();
-    if (!statement.hasException()) {
+    if (!statement.hasException(0)) {
       try {
         mutationWriter.addCopyData(this.payload);
       } catch (SpannerException exception) {
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDoneMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDoneMessage.java
index 17b14679..f1ac7884 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDoneMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyDoneMessage.java
@@ -43,11 +43,11 @@ public class CopyDoneMessage extends ControlMessage {
     if (this.statement != null) {
       MutationWriter mutationWriter = this.statement.getMutationWriter();
       statement.close();
-      if (!statement.hasException()) {
+      if (!statement.hasException(0)) {
         try {
           long rowCount = this.statement.getUpdateCount();
-          statement.addUpdateCount(rowCount); // Increase the row count of number of rows copied.
-          this.sendSpannerResult(this.statement, QueryMode.SIMPLE, 0L);
+          statement.setUpdateCount(0, rowCount); // Set the row count of number of rows copied.
+          this.sendSpannerResult(0, this.statement, QueryMode.SIMPLE, 0L);
           this.outputStream.flush();
         } catch (Exception e) {
           // Spanner returned an error when trying to commit the batch of mutations.
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyFailMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyFailMessage.java
index dab3404a..597323fd 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyFailMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/CopyFailMessage.java
@@ -50,7 +50,7 @@ public class CopyFailMessage extends ControlMessage {
       mutationWriter.rollback();
       mutationWriter.closeErrorFile();
       statement.close();
-      if (!statement.hasException()) {
+      if (!statement.hasException(0)) {
         new ErrorResponse(this.outputStream, new Exception(this.errorMessage), State.IOError)
             .send();
       }
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 60322df8..238d7afa 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
@@ -84,15 +84,16 @@ public class DescribeMessage extends ControlMessage {
    * @throws Exception if sending the message back to the client causes an error.
    */
   public void handleDescribePortal() throws Exception {
-    if (this.statement.hasException()) {
-      this.handleError(this.statement.getException());
+    if (this.statement.hasException(0)) {
+      this.handleError(this.statement.getException(0));
     } else {
-      switch (this.statement.getResultType()) {
-        case UPDATE_COUNT:
-        case NO_RESULT:
+      switch (this.statement.getStatementType(0)) {
+        case UPDATE:
+        case DDL:
+        case CLIENT_SIDE:
           new NoDataResponse(this.outputStream).send();
           break;
-        case RESULT_SET:
+        case QUERY:
           new RowDescriptionResponse(
                   this.outputStream,
                   this.statement,
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 83e3a01e..6a5988c1 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
@@ -76,10 +76,10 @@ public class ExecuteMessage extends ControlMessage {
    * @throws Exception if sending the message back to the client causes an error.
    */
   private void handleExecute() throws Exception {
-    if (this.statement.hasException()) {
-      this.handleError(this.statement.getException());
+    if (this.statement.hasException(0)) {
+      this.handleError(this.statement.getException(0));
     } else {
-      this.sendSpannerResult(this.statement, QueryMode.EXTENDED, this.maxRows);
+      this.sendSpannerResult(0, this.statement, QueryMode.EXTENDED, this.maxRows);
       this.outputStream.flush();
     }
     this.connection.cleanUp(this.statement);
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 7dca2ac8..afbf1b83 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
@@ -26,6 +26,8 @@ import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
 import com.google.cloud.spanner.pgadapter.statements.MatcherStatement;
 import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.cloud.spanner.pgadapter.wireoutput.CopyInResponse;
+import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse;
+import com.google.cloud.spanner.pgadapter.wireoutput.ErrorResponse.State;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse.Status;
 import com.google.cloud.spanner.pgadapter.wireoutput.RowDescriptionResponse;
@@ -54,9 +56,7 @@ public class QueryMessage extends ControlMessage {
     } else if (!connection.getServer().getOptions().requiresMatcher()) {
       this.statement =
           new IntermediateStatement(
-              connection.getServer().getOptions(),
-              parsedStatement,
-              this.connection.getSpannerConnection());
+              connection.getServer().getOptions(), parsedStatement, this.connection);
     } else {
       this.statement =
           new MatcherStatement(
@@ -69,7 +69,7 @@ public class QueryMessage extends ControlMessage {
   protected void sendPayload() throws Exception {
     this.statement.execute();
     this.handleQuery();
-    if (!this.statement.getCommand().equalsIgnoreCase(COPY)) {
+    if (!this.statement.getCommand(0).equalsIgnoreCase(COPY)) {
       this.connection.removeActiveStatement(this.statement);
     }
   }
@@ -99,13 +99,16 @@ public class QueryMessage extends ControlMessage {
    * error, handle error, otherwise sends the result and if contains result set, send row
    * description)
    *
-   * @throws Exception
+   * @throws Exception if handling the query fails
    */
   public void handleQuery() throws Exception {
-    if (this.statement.hasException()) {
-      this.handleError(this.statement.getException());
-    } else {
-      if (this.statement.getCommand().equalsIgnoreCase(COPY)) {
+    // Skip unexecuted statements, as no response needs be returned
+    for (int index = 0; index < statement.getStatements().size(); index++) {
+      if (this.statement.hasException(index)) {
+        new ErrorResponse(
+                this.outputStream, this.statement.getException(index), State.InternalError)
+            .send();
+      } else if (this.statement.getCommand(index).equalsIgnoreCase(COPY)) {
         CopyStatement copyStatement = (CopyStatement) this.statement;
         new CopyInResponse(
                 this.outputStream,
@@ -116,21 +119,31 @@ public class QueryMessage extends ControlMessage {
 
         // Return early as we do not respond with CommandComplete after a COPY command.
         return;
-      } else if (this.statement.containsResultSet()) {
-        new RowDescriptionResponse(
-                this.outputStream,
-                this.statement,
-                this.statement.getStatementResult(),
-                this.connection.getServer().getOptions(),
-                QueryMode.SIMPLE)
-            .send(false);
+      } else {
+        if (this.statement.containsResultSet(index)) {
+          new RowDescriptionResponse(
+                  this.outputStream,
+                  this.statement,
+                  this.statement.getStatementResult(index),
+                  this.connection.getServer().getOptions(),
+                  QueryMode.SIMPLE)
+              .send(false);
+        }
+        this.sendSpannerResult(index, this.statement, QueryMode.SIMPLE, 0L);
+      }
+    }
+    boolean inTransaction = connection.getSpannerConnection().isInTransaction();
+    Status transactionStatus = Status.IDLE;
+    if (inTransaction) {
+      if (connection.getStatus() == ConnectionStatus.TRANSACTION_ABORTED) {
+        transactionStatus = Status.FAILED;
+        // Actively rollback the aborted transaction but still block clients
+        connection.getSpannerConnection().rollback();
+      } else {
+        transactionStatus = Status.TRANSACTION;
       }
-      this.sendSpannerResult(this.statement, QueryMode.SIMPLE, 0L);
-      boolean inTransaction = connection.getSpannerConnection().isInTransaction();
-      new ReadyResponse(
-              this.outputStream, inTransaction ? Status.TRANSACTION : ReadyResponse.Status.IDLE)
-          .send();
     }
+    new ReadyResponse(this.outputStream, transactionStatus).send();
     this.connection.cleanUp(this.statement);
   }
 }


Test patch:
diff --git a/.ci/e2e-batching/begin-commit-batch.txt b/.ci/e2e-batching/begin-commit-batch.txt
new file mode 100644
index 00000000..ed2532f8
--- /dev/null
+++ b/.ci/e2e-batching/begin-commit-batch.txt
@@ -0,0 +1,1 @@
+BEGIN TRANSACTION; INSERT INTO users (id, age, name) VALUES (100, 100, 'test'); COMMIT;
\ No newline at end of file
diff --git a/.ci/e2e-batching/invalid-commit-batch.txt b/.ci/e2e-batching/invalid-commit-batch.txt
deleted file mode 100644
index 337eaff4..00000000
--- a/.ci/e2e-batching/invalid-commit-batch.txt
+++ /dev/null
@@ -1,1 +0,0 @@
-BEGIN TRANSACTION; INSERT INTO users (id, age, name) VALUES (99, 99, 'test'); COMMIT;
\ No newline at end of file
diff --git a/.ci/e2e-expected/begin-commit-batching.txt b/.ci/e2e-expected/begin-commit-batching.txt
new file mode 100644
index 00000000..2a4398d8
--- /dev/null
+++ b/.ci/e2e-expected/begin-commit-batching.txt
@@ -0,0 +1,1 @@
+COMMIT
diff --git a/.ci/e2e-expected/copy-from-stdin.txt b/.ci/e2e-expected/copy-from-stdin.txt
index 2ce9e991..d179c686 100644
--- a/.ci/e2e-expected/copy-from-stdin.txt
+++ b/.ci/e2e-expected/copy-from-stdin.txt
@@ -1,18 +1,20 @@
- id | age | name
-----+-----+------
-  1 |   1 | John
-  2 |  20 | Joe 
-  3 |  23 | Jack
-  7 |   7 | 7
-  8 |   8 | 8
-  9 |   9 | 9
- 10 |  10 | 10
- 11 |  11 | 11
- 12 |  12 | 12
- 13 |  13 | 13
- 14 |  14 | 14
- 15 |  15 | 15
- 16 |  16 | 16
- 17 |  17 | 17
-(14 rows)
+  id | age | name
+-----+-----+--------
+   1 |   1 | John
+   2 |  20 | Joe
+   3 |  23 | Jack
+   7 |   7 | 7
+   8 |   8 | 8
+   9 |   9 | 9
+  10 |  10 | 10
+  11 |  11 | 11
+  12 |  12 | 12
+  13 |  13 | 13
+  14 |  14 | 14
+  15 |  15 | 15
+  16 |  16 | 16
+  17 |  17 | 17
+  99 |  99 | person
+ 100 | 100 | test
+(16 rows)
 
diff --git a/.ci/e2e-expected/ddl-dml-batching.txt b/.ci/e2e-expected/ddl-dml-batching.txt
index f5f0ce84..08beba40 100644
--- a/.ci/e2e-expected/ddl-dml-batching.txt
+++ b/.ci/e2e-expected/ddl-dml-batching.txt
@@ -1,1 +1,1 @@
-ERROR: FAILED_PRECONDITION: Executing updates is not allowed for DDL batches. "CREATE TABLE testtable (id integer PRIMARY KEY, data integer); INSERT INTO testtable (id, data) VALUES (1, 2); INSERT INTO testtable (id, data) VALUES (3, 4)"
+INSERT 0 1
\ No newline at end of file
diff --git a/.ci/e2e-expected/dml-batching.txt b/.ci/e2e-expected/dml-batching.txt
index b0e8ed90..c4c01584 100644
--- a/.ci/e2e-expected/dml-batching.txt
+++ b/.ci/e2e-expected/dml-batching.txt
@@ -1,1 +1,1 @@
-INSERT 0 5
+INSERT 0 1
diff --git a/.ci/e2e-expected/invalid-commit-batching.txt b/.ci/e2e-expected/invalid-commit-batching.txt
deleted file mode 100644
index d4e73782..00000000
--- a/.ci/e2e-expected/invalid-commit-batching.txt
+++ /dev/null
@@ -1,1 +0,0 @@
-ERROR: FAILED_PRECONDITION: This connection has an active batch and cannot begin a transaction "BEGIN TRANSACTION; INSERT INTO users (id, age, name) VALUES (99, 99, 'test'); COMMIT"
diff --git a/.ci/e2e-expected/select-ddl-batching.txt b/.ci/e2e-expected/select-ddl-batching.txt
index 0cebc3da..4a6f7e32 100644
--- a/.ci/e2e-expected/select-ddl-batching.txt
+++ b/.ci/e2e-expected/select-ddl-batching.txt
@@ -1,1 +1,4 @@
-ERROR: FAILED_PRECONDITION: Executing queries is not allowed for DDL batches. "CREATE TABLE test (id integer PRIMARY KEY); SELECT * FROM test; SELECT * FROM test"
+ id
+----
+(0 rows)
+
diff --git a/.ci/e2e-expected/select-dml-batching.txt b/.ci/e2e-expected/select-dml-batching.txt
index 6ad7b053..18a673e3 100644
--- a/.ci/e2e-expected/select-dml-batching.txt
+++ b/.ci/e2e-expected/select-dml-batching.txt
@@ -1,1 +1,13 @@
-ERROR: INVALID_ARGUMENT: Statement type is not supported for batching "SELECT * FROM users; INSERT INTO users (id, age, name) VALUES (99, 99, 'person'); SELECT * FROM users"
+ id | age |  name
+----+-----+--------
+  1 |   1 | John
+  2 |  20 | Joe
+  3 |  23 | Jack
+  7 |   7 | 7
+  8 |   8 | 8
+  9 |   9 | 9
+ 10 |  10 | 10
+ 11 |  11 | 11
+ 99 |  99 | person
+(9 rows)
+
diff --git a/.ci/e2e-expected/with-dml-batching.txt b/.ci/e2e-expected/with-dml-batching.txt
index 63a4bd1d..709e9b81 100644
--- a/.ci/e2e-expected/with-dml-batching.txt
+++ b/.ci/e2e-expected/with-dml-batching.txt
@@ -8,5 +8,6 @@
   9
  10
  11
-(8 rows)
+ 99
+(9 rows)
 
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
index 07bbd3a5..2f1141c9 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
@@ -17,20 +17,26 @@ package com.google.cloud.spanner.pgadapter.statements;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.SpannerBatchUpdateException;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractStatementParser;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.PostgreSQLStatementParser;
 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.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * Data type to store simple SQL statement with designated metadata. Allows manipulation of
@@ -42,45 +48,42 @@ public class IntermediateStatement {
       (PostgreSQLStatementParser) AbstractStatementParser.getInstance(Dialect.POSTGRESQL);
 
   protected final OptionsMetadata options;
-  private final ResultType resultType;
-  protected ResultSet statementResult;
-  protected boolean hasMoreData;
-  protected SpannerException exception;
+  private final ImmutableList<StatementType> statementTypes;
+  protected ResultSet[] statementResults;
+  protected final boolean[] hasMoreData;
+  protected final SpannerException[] exceptions;
   protected final ParsedStatement parsedStatement;
-  protected final String command;
-  protected boolean executed;
+  protected final ImmutableList<String> commands;
+  protected final List<String> commandTags;
+  protected boolean executed = false;
   protected final Connection connection;
-  protected Long updateCount;
-  protected final ImmutableList<String> statements;
+  protected long[] updateCounts;
+  protected final ImmutableList<ParsedStatement> statements;
+  private ConnectionHandler connectionHandler;
+  private ExecutionStatus executionStatus;
 
   private static final char STATEMENT_DELIMITER = ';';
   private static final char SINGLE_QUOTE = '\'';
 
   public IntermediateStatement(
-      OptionsMetadata options, ParsedStatement parsedStatement, Connection connection) {
-    this(
-        options,
-        parsedStatement,
-        connection,
-        parseStatements(parsedStatement.getSqlWithoutComments()));
+      OptionsMetadata options,
+      ParsedStatement parsedStatement,
+      ConnectionHandler connectionHandler) {
+    this(options, parsedStatement, connectionHandler.getSpannerConnection());
+    this.connectionHandler = connectionHandler;
   }
 
   protected IntermediateStatement(
-      OptionsMetadata options,
-      ParsedStatement parsedStatement,
-      Connection connection,
-      ImmutableList<String> statements) {
+      OptionsMetadata options, ParsedStatement parsedStatement, Connection connection) {
     this.options = options;
     this.parsedStatement = replaceKnownUnsupportedQueries(parsedStatement);
-    this.statements = statements;
-    this.command = StatementParser.parseCommand(this.parsedStatement.getSqlWithoutComments());
     this.connection = connection;
-    // Note: This determines the result type based on the first statement in the SQL statement. That
-    // means that it assumes that if this is a batch of statements, all the statements in the batch
-    // will have the same type of result (that is; they are all DML statements, all DDL statements,
-    // all queries, etc.). That is a safe assumption for now, as PgAdapter currently only supports
-    // all-DML and all-DDL batches.
-    this.resultType = determineResultType(this.parsedStatement);
+    this.statements = parseStatements(parsedStatement);
+    this.commands = StatementParser.parseCommands(this.statements);
+    this.commandTags = new ArrayList<>(this.commands);
+    this.statementTypes = determineStatementTypes(this.statements);
+    this.hasMoreData = new boolean[this.statements.size()];
+    this.exceptions = new SpannerException[this.statements.size()];
   }
 
   protected ParsedStatement replaceKnownUnsupportedQueries(ParsedStatement parsedStatement) {
@@ -95,24 +98,35 @@ public class IntermediateStatement {
     return parsedStatement;
   }
 
+  /** @return The number of SQL statements in this {@link IntermediateStatement} */
+  public int getStatementCount() {
+    return statements.size();
+  }
+
   /**
-   * Determines the result type based on the given sql string. The sql string must already been
-   * stripped of any comments that might precede the actual sql string.
+   * Determines the statement types based on the given statement strings. The statement string must
+   * already been stripped of any comments that might precede the actual statement string.
    *
-   * @param parsedStatement The parsed statement to determine the type of result for
-   * @return The {@link ResultType} that the given sql string will produce
+   * @param statements The parsed statements to determine the type for
+   * @return The list of {@link StatementType} that the given statement strings will produce
    */
-  protected static ResultType determineResultType(ParsedStatement parsedStatement) {
-    if (parsedStatement.isUpdate()) {
-      return ResultType.UPDATE_COUNT;
-    } else if (parsedStatement.isQuery()) {
-      return ResultType.RESULT_SET;
-    } else {
-      return ResultType.NO_RESULT;
+  protected static ImmutableList<StatementType> determineStatementTypes(
+      List<ParsedStatement> statements) {
+    ImmutableList.Builder<StatementType> builder = ImmutableList.builder();
+    for (ParsedStatement stmt : statements) {
+      if (stmt.isUpdate()) {
+        builder.add(StatementType.UPDATE);
+      } else if (stmt.isQuery()) {
+        builder.add(StatementType.QUERY);
+      } else if (stmt.isDdl()) {
+        builder.add(StatementType.DDL);
+      } else {
+        builder.add(StatementType.CLIENT_SIDE);
+      }
     }
+    return builder.build();
   }
 
-  // Split statements by ';' delimiter, but ignore anything that is nested with '' or "".
   private static ImmutableList<String> splitStatements(String sql) {
     // First check trivial cases with only one statement.
     int firstIndexOfDelimiter = sql.indexOf(STATEMENT_DELIMITER);
@@ -125,13 +139,16 @@ public class IntermediateStatement {
 
     ImmutableList.Builder<String> builder = ImmutableList.builder();
     // TODO: Fix this parsing, as it does not take all types of quotes into consideration.
-    boolean quoteEscape = false;
+    boolean insideQuotes = false;
+    boolean currentCharacterIsEscaped = false;
     int index = 0;
     for (int i = 0; i < sql.length(); ++i) {
-      if (sql.charAt(i) == SINGLE_QUOTE) {
-        quoteEscape = !quoteEscape;
+      // ignore escaped quotes
+      if (sql.charAt(i) == SINGLE_QUOTE && !currentCharacterIsEscaped) {
+        insideQuotes = !insideQuotes;
       }
-      if (sql.charAt(i) == STATEMENT_DELIMITER && !quoteEscape) {
+      // skip semicolon inside quotes
+      if (sql.charAt(i) == STATEMENT_DELIMITER && !insideQuotes) {
         String stmt = sql.substring(index, i).trim();
         // Statements with only ';' character are empty and dropped.
         if (stmt.length() > 0) {
@@ -139,6 +156,11 @@ public class IntermediateStatement {
         }
         index = i + 1;
       }
+      if (currentCharacterIsEscaped) {
+        currentCharacterIsEscaped = false;
+      } else if (sql.charAt(i) == '\\') {
+        currentCharacterIsEscaped = true;
+      }
     }
 
     if (index < sql.length()) {
@@ -147,9 +169,14 @@ public class IntermediateStatement {
     return builder.build();
   }
 
-  protected static ImmutableList<String> parseStatements(String sql) {
+  protected static ImmutableList<ParsedStatement> parseStatements(ParsedStatement stmt) {
+    String sql = stmt.getSqlWithoutComments();
     Preconditions.checkNotNull(sql);
-    return splitStatements(sql);
+    ImmutableList.Builder<ParsedStatement> builder = ImmutableList.builder();
+    for (String statement : splitStatements(sql)) {
+      builder.add(PARSER.parse(Statement.of(statement)));
+    }
+    return builder.build();
   }
 
   /**
@@ -168,14 +195,20 @@ public class IntermediateStatement {
    * @throws Exception if closing fails server-side.
    */
   public void close() throws Exception {
-    if (this.getStatementResult() != null) {
-      this.getStatementResult().close();
+    for (int index = 0; index < this.getStatementCount(); index++) {
+      close(index);
+    }
+  }
+
+  public void close(int index) throws Exception {
+    if (this.getStatementResult(index) != null) {
+      this.getStatementResult(index).close();
     }
   }
 
   /** @return True if this is a select statement, false otherwise. */
-  public boolean containsResultSet() {
-    return this.resultType == ResultType.RESULT_SET;
+  public boolean containsResultSet(int index) {
+    return this.statementTypes.get(index) == StatementType.QUERY;
   }
 
   /** @return True if this statement was executed, False otherwise. */
@@ -183,55 +216,81 @@ public class IntermediateStatement {
     return executed;
   }
 
-  /** @return The number of items that were modified by this execution. */
-  public Long getUpdateCount() {
-    return this.updateCount;
+  /**
+   * @return The number of items that were modified by this execution for DML. 0 for DDL and -1 for
+   *     QUERY.
+   */
+  public long getUpdateCount(int index) {
+    if (this.updateCounts == null) {
+      return 0;
+    }
+    return this.updateCounts[index];
   }
 
-  public void addUpdateCount(long count) {
-    if (this.updateCount == null) {
-      this.updateCount = 0L;
+  public void setUpdateCount(int index, long updateCount) {
+    if (this.updateCounts == null) {
+      this.updateCounts = new long[getStatementCount()];
     }
-    this.updateCount += count;
+    this.updateCounts[index] = updateCount;
+    this.hasMoreData[index] = false;
   }
 
   /** @return True if at some point in execution, and exception was thrown. */
-  public boolean hasException() {
-    return this.exception != null;
+  public boolean hasException(int index) {
+    return this.exceptions[index] != null;
   }
 
   /** @return True if only a subset of the available data has been returned. */
-  public boolean isHasMoreData() {
-    return this.hasMoreData;
+  public boolean isHasMoreData(int index) {
+    return this.hasMoreData[index];
   }
 
-  public void setHasMoreData(boolean hasMoreData) {
-    this.hasMoreData = hasMoreData;
+  public void setHasMoreData(int index, boolean hasMoreData) {
+    this.hasMoreData[index] = hasMoreData;
   }
 
   public Connection getConnection() {
     return this.connection;
   }
 
-  public List<String> getStatements() {
+  public List<ParsedStatement> getStatements() {
     return this.statements;
   }
 
-  public ResultSet getStatementResult() {
-    return this.statementResult;
+  public String getStatement(int index) {
+    return this.statements.get(index).getSqlWithoutComments();
+  }
+
+  public ResultSet getStatementResult(int index) {
+    if (this.statementResults == null) {
+      return null;
+    }
+    return this.statementResults[index];
+  }
+
+  protected void setStatementResult(int index, ResultSet resultSet) {
+    if (this.statementResults == null) {
+      this.statementResults = new ResultSet[getStatementCount()];
+    }
+    this.statementResults[index] = resultSet;
+    this.hasMoreData[index] = resultSet.next();
+    if (this.updateCounts == null) {
+      this.updateCounts = new long[getStatementCount()];
+    }
+    this.updateCounts[index] = -1;
   }
 
-  public ResultType getResultType() {
-    return this.resultType;
+  public StatementType getStatementType(int index) {
+    return this.statementTypes.get(index);
   }
 
   public String getSql() {
     return this.parsedStatement.getSqlWithoutComments();
   }
 
-  public Exception getException() {
-    Exception e = this.exception;
-    this.exception = null;
+  public Exception getException(int index) {
+    Exception e = this.exceptions[index];
+    this.exceptions[index] = null;
     return e;
   }
 
@@ -240,20 +299,16 @@ public class IntermediateStatement {
    * execution (including results and update counts). An array of updateCounts is needed in the case
    * of updateBatchResultCount.
    */
-  protected void updateResultCount(StatementResult result) {
+  protected void updateResultCount(int index, StatementResult result) {
     switch (result.getResultType()) {
       case RESULT_SET:
-        this.statementResult = result.getResultSet();
-        this.hasMoreData = this.statementResult.next();
+        setStatementResult(index, result.getResultSet());
         break;
       case UPDATE_COUNT:
-        this.updateCount = result.getUpdateCount();
-        this.hasMoreData = false;
-        this.statementResult = null;
+        setUpdateCount(index, result.getUpdateCount());
         break;
       case NO_RESULT:
-        this.hasMoreData = false;
-        this.statementResult = null;
+        this.hasMoreData[index] = false;
         break;
       default:
         throw SpannerExceptionFactory.newSpannerException(
@@ -261,13 +316,10 @@ public class IntermediateStatement {
     }
   }
 
-  protected void updateBatchResultCount(long[] updateCounts) {
-    this.updateCount = 0L;
-    for (int i = 0; i < updateCounts.length; ++i) {
-      this.updateCount += updateCounts[i];
+  protected void updateBatchResultCount(int fromIndex, long[] updateCounts) {
+    for (int index = fromIndex; index < fromIndex + updateCounts.length; index++) {
+      setUpdateCount(index, updateCounts[index - fromIndex]);
     }
-    this.hasMoreData = false;
-    this.statementResult = null;
   }
 
   /**
@@ -275,50 +327,164 @@ public class IntermediateStatement {
    *
    * @param e The exception to store.
    */
-  protected void handleExecutionException(SpannerException e) {
-    this.exception = e;
-    this.hasMoreData = false;
-    this.statementResult = null;
+  protected void handleExecutionException(int index, SpannerException e) {
+    this.exceptions[index] = e;
+    this.hasMoreData[index] = false;
+  }
+
+  private void handleExecutionExceptionAndTransactionStatus(int index, SpannerException e) {
+    if (executionStatus == ExecutionStatus.EXPLICIT_TRANSACTION) {
+      connectionHandler.setStatus(ConnectionStatus.TRANSACTION_ABORTED);
+    }
+    handleExecutionException(index, e);
   }
 
   /** Execute the SQL statement, storing metadata. */
   public void execute() {
-    this.executed = true;
-    try {
-      if (statements.size() > 1) {
-        // TODO: Clean this up a little once the statement parsing in the client library has been
-        // released: https://github.com/googleapis/java-spanner/pull/1690
-        // Also, the restriction that mixed batches are not allowed will be removed in a future PR.
-        if (resultType == ResultType.UPDATE_COUNT) {
-          connection.startBatchDml();
-        } else if (resultType == ResultType.NO_RESULT) {
-          connection.startBatchDdl();
-        } else {
-          throw SpannerExceptionFactory.newSpannerException(
-              ErrorCode.INVALID_ARGUMENT, "Statement type is not supported for batching");
-        }
-        for (String stmt : statements) {
-          connection.execute(Statement.of(stmt));
-        }
-        long[] updateCounts = connection.runBatch();
-        updateBatchResultCount(updateCounts);
+    executed = true;
+    if (connection.isInTransaction()) {
+      executionStatus = ExecutionStatus.EXPLICIT_TRANSACTION;
+    } else {
+      executionStatus = ExecutionStatus.AUTOCOMMIT;
+    }
+    int index = 0;
+    while (index < getStatementCount()) {
+      if (connectionHandler.getStatus() == ConnectionStatus.TRANSACTION_ABORTED
+          && !"ROLLBACK".equals(getCommand(index))
+          && !"COMMIT".equals(getCommand(index))) {
+        handleExecutionException(
+            index,
+            SpannerExceptionFactory.newSpannerException(
+                ErrorCode.INVALID_ARGUMENT,
+                "ERROR: current transaction is aborted, commands ignored until end of transaction block"));
+        index++;
+        continue;
+      }
+      if (getStatementType(index) == StatementType.DDL
+          && executionStatus == ExecutionStatus.EXPLICIT_TRANSACTION) {
+        handleExecutionExceptionAndTransactionStatus(
+            index,
+            SpannerExceptionFactory.newSpannerException(
+                ErrorCode.INVALID_ARGUMENT,
+                "ERROR: DDL statements are only allowed outside explicit transactions."));
+        index++;
+        continue;
+      }
+      boolean canUseBatch = false;
+      if (index < (getStatementCount() - 1)) {
+        StatementType statementType = getStatementType(index);
+        StatementType nextStatementType = getStatementType(index + 1);
+        canUseBatch = canBeBatchedTogether(statementType, nextStatementType);
+      }
+      if (canUseBatch) {
+        index += executeStatementsInBatch(index);
+      } else {
+        executeSingleStatement(index);
+        index++;
+      }
+    }
+  }
+
+  private void executeSingleStatement(int index) {
+    // Before executing the statement, handle specific statements that change the transaction status
+    String command = getCommand(index);
+    String statement = getStatement(index);
+    if ("BEGIN".equals(command)) {
+      if (executionStatus == ExecutionStatus.EXPLICIT_TRANSACTION) {
+        // Executing a BEGIN statement when an explicit transaction is already active is a no-op
+        return;
       } else {
-        StatementResult result =
-            connection.execute(Statement.of(this.parsedStatement.getSqlWithoutComments()));
-        updateResultCount(result);
+        executionStatus = ExecutionStatus.EXPLICIT_TRANSACTION;
       }
+    }
+    if ("COMMIT".equals(command) || "ROLLBACK".equals(command)) {
+      if (connectionHandler.getStatus() == ConnectionStatus.TRANSACTION_ABORTED) {
+        connectionHandler.setStatus(ConnectionStatus.IDLE);
+        // COMMIT rollbacks aborted transaction
+        statement = "ROLLBACK";
+        commandTags.set(index, "ROLLBACK");
+      }
+      if (executionStatus == ExecutionStatus.AUTOCOMMIT) {
+        // Executing a COMMIT/ROLLBACK statement in autocommit mode is a no-op
+        return;
+      } else {
+        executionStatus = ExecutionStatus.AUTOCOMMIT;
+      }
+    }
+    try {
+      StatementResult result = this.connection.execute(Statement.of(statement));
+      updateResultCount(index, result);
     } catch (SpannerException e) {
-      if (statements.size() > 1) {
-        SpannerException exception =
-            SpannerExceptionFactory.newSpannerException(
-                e.getErrorCode(),
-                e.getMessage() + " \"" + this.parsedStatement.getSqlWithoutComments() + "\"",
-                e);
-        handleExecutionException(exception);
+      handleExecutionExceptionAndTransactionStatus(index, e);
+    }
+  }
+
+  /**
+   * Executes the statements from fromIndex in a DML/DDL batch. The batch will consist of all
+   * statements from fromIndex till the first statement that is of a different type than the
+   * statement at fromIndex. That is; If the first statement is a DML statement, the batch will
+   * contain all statements that follow until it encounters a statement that is not a DML statement.
+   * The same also applies to DDL statements. Query statements and other statements can not be
+   * batched.
+   *
+   * @param fromIndex The index of the statements array where the batch should start
+   * @return The number of statements included in the batch.
+   */
+  private int executeStatementsInBatch(int fromIndex) {
+    Preconditions.checkArgument(fromIndex < getStatementCount() - 1);
+    Preconditions.checkArgument(
+        canBeBatchedTogether(getStatementType(fromIndex), getStatementType(fromIndex + 1)));
+    StatementType batchType = getStatementType(fromIndex);
+    switch (batchType) {
+      case UPDATE:
+        connection.startBatchDml();
+        break;
+      case DDL:
+        connection.startBatchDdl();
+        break;
+      default:
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT, "Statement type is not supported for batching");
+    }
+    int index = fromIndex;
+    while (index < getStatementCount()) {
+      StatementType statementType = getStatementType(index);
+      if (canBeBatchedTogether(batchType, statementType)) {
+        connection.execute(Statement.of(getStatement(index)));
+        index++;
       } else {
-        handleExecutionException(e);
+        // End the batch here, as the statement type on this index can not be batched together with
+        // the other statements in the batch.
+        break;
       }
     }
+    try {
+      long[] counts = connection.runBatch();
+      updateBatchResultCount(fromIndex, counts);
+    } catch (SpannerBatchUpdateException e) {
+      long[] counts = e.getUpdateCounts();
+      updateBatchResultCount(fromIndex, counts);
+      handleExecutionExceptionAndTransactionStatus(fromIndex + counts.length, e);
+      // Skip remaining statements in DML/DDL batch
+      // TODO: May want to start off from the failure in order to be compatible with the
+      // behavior of PG.
+      for (int i = fromIndex + counts.length + 1; i < index; i++) {
+        handleExecutionException(
+            i,
+            SpannerExceptionFactory.newSpannerException(
+                ErrorCode.INTERNAL,
+                "ERROR: Statement is not executed due to the preceding error in DML/DDL batch."));
+      }
+    }
+    return index - fromIndex;
+  }
+
+  private boolean canBeBatchedTogether(StatementType statementType1, StatementType statementType2) {
+    if (Objects.equals(statementType1, StatementType.QUERY)
+        || Objects.equals(statementType1, StatementType.CLIENT_SIDE)) {
+      return false;
+    }
+    return Objects.equals(statementType1, statementType2);
   }
 
   /**
@@ -339,8 +505,13 @@ public class IntermediateStatement {
   }
 
   /** @return the extracted command (first word) from the SQL statement. */
-  public String getCommand() {
-    return this.command;
+  public String getCommand(int index) {
+    return this.commands.get(index);
+  }
+
+  /** @return the extracted command (first word) from the really executed SQL statement. */
+  public String getCommandTag(int index) {
+    return this.commandTags.get(index);
   }
 
   /* Used for testing purposes */
@@ -348,9 +519,8 @@ public class IntermediateStatement {
     return (statements.size() > 1);
   }
 
-  public enum ResultType {
-    UPDATE_COUNT,
-    RESULT_SET,
-    NO_RESULT
+  enum ExecutionStatus {
+    EXPLICIT_TRANSACTION,
+    AUTOCOMMIT
   }
 }
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 18d8a098..92dde79d 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
@@ -17,32 +17,24 @@ package com.google.cloud.spanner.pgadapter;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
-import com.google.api.gax.core.NoCredentialsProvider;
-import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
-import com.google.api.gax.longrunning.OperationSnapshot;
-import com.google.api.gax.longrunning.OperationTimedPollAlgorithm;
-import com.google.api.gax.retrying.RetrySettings;
-import com.google.api.gax.rpc.UnaryCallSettings;
-import com.google.cloud.NoCredentials;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
-import com.google.cloud.spanner.MockDatabaseAdminServiceImpl;
-import com.google.cloud.spanner.MockOperationsServiceImpl;
 import com.google.cloud.spanner.MockSpannerServiceImpl;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
-import com.google.cloud.spanner.admin.database.v1.DatabaseAdminClient;
-import com.google.cloud.spanner.admin.database.v1.DatabaseAdminSettings;
+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.common.collect.ImmutableList;
+import com.google.longrunning.Operation;
+import com.google.protobuf.Any;
+import com.google.protobuf.Empty;
 import com.google.protobuf.ListValue;
 import com.google.protobuf.NullValue;
 import com.google.protobuf.Value;
-import com.google.spanner.admin.database.v1.CreateDatabaseRequest;
-import com.google.spanner.admin.database.v1.InstanceName;
+import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
 import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.SpannerGrpc;
 import com.google.spanner.v1.StructType;
@@ -52,22 +44,23 @@ import com.google.spanner.v1.TypeAnnotationCode;
 import com.google.spanner.v1.TypeCode;
 import io.grpc.Context;
 import io.grpc.Contexts;
-import io.grpc.ManagedChannelBuilder;
 import io.grpc.Metadata;
 import io.grpc.Server;
 import io.grpc.ServerCall;
 import io.grpc.ServerCall.Listener;
 import io.grpc.ServerCallHandler;
 import io.grpc.ServerInterceptor;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
 import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
 import java.net.InetSocketAddress;
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
+import java.util.Collections;
 import java.util.logging.Logger;
 import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.BeforeClass;
-import org.threeten.bp.Duration;
 
 /**
  * Abstract base class for tests that verify that PgAdapter is receiving wire protocol requests
@@ -93,6 +86,7 @@ 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 INVALID_SELECT = Statement.of("SELECT foo");
   private static final ResultSetMetadata SELECT1_METADATA =
       ResultSetMetadata.newBuilder()
           .setRowType(
@@ -125,6 +119,8 @@ public abstract class AbstractMockServerTest {
   protected static final int UPDATE_COUNT = 2;
   protected static final Statement INSERT_STATEMENT = Statement.of("INSERT INTO FOO VALUES (1)");
   protected static final int INSERT_COUNT = 1;
+  protected static final Statement INVALID_DML = Statement.of("INSERT INTO FOO VALUES ('abc')");
+  protected static final Statement INVALID_DDL = Statement.of("CREATE TABLE FOO (id int64)");
 
   protected static final ResultSetMetadata ALL_TYPES_METADATA =
       ResultSetMetadata.newBuilder()
@@ -211,14 +207,21 @@ public abstract class AbstractMockServerTest {
                   .build())
           .build();
 
+  protected static final StatusRuntimeException EXCEPTION =
+      Status.INVALID_ARGUMENT.withDescription("Statement is invalid.").asRuntimeException();
+
   protected static MockSpannerServiceImpl mockSpanner;
-  protected static MockOperationsServiceImpl mockOperationsService;
-  protected static MockDatabaseAdminServiceImpl mockDatabaseAdmin;
+  protected static MockDatabaseAdminImpl mockDatabaseAdmin;
   private static Server spannerServer;
   protected static ProxyServer pgServer;
 
   @BeforeClass
   public static void startMockSpannerAndPgAdapterServers() throws Exception {
+    doStartMockSpannerAndPgAdapterServers(Collections.emptyList());
+  }
+
+  protected static void doStartMockSpannerAndPgAdapterServers(
+      Iterable<String> extraPGAdapterOptions) throws Exception {
     mockSpanner = new MockSpannerServiceImpl();
     mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions.
     mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET));
@@ -227,16 +230,17 @@ public abstract class AbstractMockServerTest {
     mockSpanner.putStatementResult(StatementResult.update(INSERT_STATEMENT, INSERT_COUNT));
     mockSpanner.putStatementResult(
         MockSpannerServiceImpl.StatementResult.detectDialectResult(Dialect.POSTGRESQL));
+    mockSpanner.putStatementResult(StatementResult.exception(INVALID_SELECT, EXCEPTION));
+    mockSpanner.putStatementResult(StatementResult.exception(INVALID_DML, EXCEPTION));
+    mockSpanner.putStatementResult(StatementResult.exception(INVALID_DDL, EXCEPTION));
 
-    mockOperationsService = new MockOperationsServiceImpl();
-    mockDatabaseAdmin = new MockDatabaseAdminServiceImpl(mockOperationsService);
+    mockDatabaseAdmin = new MockDatabaseAdminImpl();
 
     InetSocketAddress address = new InetSocketAddress("localhost", 0);
     spannerServer =
         NettyServerBuilder.forAddress(address)
             .addService(mockSpanner)
             .addService(mockDatabaseAdmin)
-            .addService(mockOperationsService)
             .intercept(
                 new ServerInterceptor() {
                   @Override
@@ -262,10 +266,6 @@ public abstract class AbstractMockServerTest {
             .build()
             .start();
 
-    // Create the test database on the mock server. This should be replaced by a simple feature in
-    // the mock server to just add a database instead of having to simulate the creation of it.
-    createDatabase();
-
     ImmutableList.Builder<String> argsListBuilder =
         ImmutableList.<String>builder()
             .add(
@@ -284,64 +284,12 @@ public abstract class AbstractMockServerTest {
                 String.format("localhost:%d", spannerServer.getPort()),
                 "-r",
                 "usePlainText=true;");
+    argsListBuilder.addAll(extraPGAdapterOptions);
     String[] args = argsListBuilder.build().toArray(new String[0]);
     pgServer = new ProxyServer(new OptionsMetadata(args));
     pgServer.startServer();
   }
 
-  private static void createDatabase() throws Exception {
-    // TODO: Replace this entire method with a feature in the test framework to just manually add a
-    // database instead of having to create it.
-    DatabaseAdminSettings.Builder builder =
-        DatabaseAdminSettings.newBuilder()
-            .setCredentialsProvider(NoCredentialsProvider.create())
-            .setTransportChannelProvider(
-                InstantiatingGrpcChannelProvider.newBuilder()
-                    .setEndpoint(String.format("localhost:%d", spannerServer.getPort()))
-                    .setCredentials(NoCredentials.getInstance())
-                    .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
-                    .build());
-
-    RetrySettings longRunningInitialRetrySettings =
-        RetrySettings.newBuilder()
-            .setInitialRpcTimeout(Duration.ofMillis(600L))
-            .setMaxRpcTimeout(Duration.ofMillis(6000L))
-            .setInitialRetryDelay(Duration.ofMillis(20L))
-            .setMaxRetryDelay(Duration.ofMillis(45L))
-            .setRetryDelayMultiplier(1.5)
-            .setRpcTimeoutMultiplier(1.5)
-            .setTotalTimeout(Duration.ofMinutes(48L))
-            .build();
-    builder
-        .createDatabaseOperationSettings()
-        .setInitialCallSettings(
-            UnaryCallSettings
-                .<CreateDatabaseRequest, OperationSnapshot>newUnaryCallSettingsBuilder()
-                .setRetrySettings(longRunningInitialRetrySettings)
-                .build());
-    builder
-        .createDatabaseOperationSettings()
-        .setPollingAlgorithm(
-            OperationTimedPollAlgorithm.create(
-                RetrySettings.newBuilder()
-                    .setInitialRpcTimeout(Duration.ofMillis(20L))
-                    .setInitialRetryDelay(Duration.ofMillis(10L))
-                    .setMaxRetryDelay(Duration.ofMillis(150L))
-                    .setMaxRpcTimeout(Duration.ofMillis(150L))
-                    .setMaxAttempts(10)
-                    .setTotalTimeout(Duration.ofMillis(5000L))
-                    .setRetryDelayMultiplier(1.3)
-                    .setRpcTimeoutMultiplier(1.3)
-                    .build()));
-
-    DatabaseAdminClient client = DatabaseAdminClient.create(builder.build());
-    client
-        .createDatabaseAsync(
-            InstanceName.newBuilder().setProject("p").setInstance("i").build(), "CREATE DATABASE d")
-        .get();
-    client.close();
-  }
-
   @AfterClass
   public static void stopMockSpannerAndPgAdapterServers() throws Exception {
     pgServer.stopServer();
@@ -390,5 +338,20 @@ public abstract class AbstractMockServerTest {
   @Before
   public void clearRequests() {
     mockSpanner.clearRequests();
+    mockDatabaseAdmin.reset();
+  }
+
+  protected void addDdlResponseToSpannerAdmin() {
+    mockDatabaseAdmin.addResponse(
+        Operation.newBuilder()
+            .setDone(true)
+            .setResponse(Any.pack(Empty.getDefaultInstance()))
+            .setMetadata(Any.pack(UpdateDatabaseDdlMetadata.getDefaultInstance()))
+            .build());
+  }
+
+  protected void addDdlExceptionToSpannerAdmin() {
+    mockDatabaseAdmin.addException(
+        Status.INVALID_ARGUMENT.withDescription("Statement is invalid.").asRuntimeException());
   }
 }
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 a75aa03e..ac11304c 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ControlMessageTest.java
@@ -17,13 +17,13 @@ package com.google.cloud.spanner.pgadapter;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.when;
 
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.TextFormat;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
-import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement.ResultType;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
@@ -65,9 +65,9 @@ public final class ControlMessageTest {
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
-    when(intermediateStatement.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
-    when(intermediateStatement.getCommand()).thenReturn("INSERT");
-    when(intermediateStatement.getUpdateCount()).thenReturn(1L);
+    when(intermediateStatement.getStatementType(0)).thenReturn(StatementType.UPDATE);
+    when(intermediateStatement.getCommandTag(0)).thenReturn("INSERT");
+    when(intermediateStatement.getUpdateCount(0)).thenReturn(1L);
     when(connectionHandler.getSpannerConnection()).thenReturn(connection);
 
     JSONParser parser = new JSONParser();
@@ -86,7 +86,7 @@ public final class ControlMessageTest {
     when(connectionHandler.getServer()).thenReturn(server);
 
     ControlMessage controlMessage = ControlMessage.create(connectionHandler);
-    controlMessage.sendSpannerResult(intermediateStatement, QueryMode.SIMPLE, 0L);
+    controlMessage.sendSpannerResult(0, intermediateStatement, QueryMode.SIMPLE, 0L);
 
     DataInputStream outputReader =
         new DataInputStream(new ByteArrayInputStream(buffer.toByteArray()));
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
index 86e12929..38dea69a 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/CopyInMockServerTest.java
@@ -194,10 +194,11 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
   public void testCopyInExceedsCommitSizeLimit_BatchesInNonAtomicMode()
       throws SQLException, IOException {
     setupCopyInformationSchemaResults();
-
     try (Connection connection = DriverManager.getConnection(createUrl())) {
       System.setProperty("copy_in_commit_limit", "10");
-      connection.createStatement().execute("set autocommit_dml_mode='partitioned_non_atomic'");
+      connection
+          .createStatement()
+          .execute("set spanner.autocommit_dml_mode='partitioned_non_atomic'");
       CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
       copyManager.copyIn("COPY users FROM STDIN;", new StringReader("5\t5\t5\n6\t6\t6\n7\t7\t7\n"));
     } finally {
@@ -222,7 +223,9 @@ public class CopyInMockServerTest extends AbstractMockServerTest {
 
     try (Connection connection = DriverManager.getConnection(createUrl())) {
       System.setProperty("copy_in_commit_limit", "10");
-      connection.createStatement().execute("set autocommit_dml_mode='partitioned_non_atomic'");
+      connection
+          .createStatement()
+          .execute("set spanner.autocommit_dml_mode='partitioned_non_atomic'");
       CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
       SQLException exception =
           assertThrows(
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/EmulatedPsqlMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/EmulatedPsqlMockServerTest.java
new file mode 100644
index 00000000..a68bbfa1
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/EmulatedPsqlMockServerTest.java
@@ -0,0 +1,76 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.cloud.spanner.pgadapter;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.Statement;
+import com.google.common.collect.ImmutableList;
+import com.google.spanner.v1.ExecuteBatchDmlRequest;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class EmulatedPsqlMockServerTest extends AbstractMockServerTest {
+  private static final String INSERT1 = "insert into foo values (1)";
+  private static final String INSERT2 = "insert into foo values (2)";
+
+  @BeforeClass
+  public static void loadPgJdbcDriver() throws Exception {
+    // Make sure the PG JDBC driver is loaded.
+    Class.forName("org.postgresql.Driver");
+  }
+
+  @BeforeClass
+  public static void startMockSpannerAndPgAdapterServers() throws Exception {
+    doStartMockSpannerAndPgAdapterServers(ImmutableList.of("-q"));
+
+    mockSpanner.putStatementResults(
+        StatementResult.update(Statement.of(INSERT1), 1L),
+        StatementResult.update(Statement.of(INSERT2), 1L));
+  }
+
+  /**
+   * Creates a JDBC connection string that instructs the PG JDBC driver to use the default simple
+   * mode for queries and DML statements. This makes the JDBC driver behave in (much) the same way
+   * as psql.
+   */
+  private String createUrl() {
+    return String.format(
+        "jdbc:postgresql://localhost:%d/?preferQueryMode=simple", pgServer.getLocalPort());
+  }
+
+  @Test
+  public void testTwoInserts() throws SQLException {
+    String sql = "insert into foo values (1); insert into foo values (2);";
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      connection.createStatement().execute(String.format("%s; %s", INSERT1, INSERT2));
+    }
+
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+    ExecuteBatchDmlRequest request =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0);
+    assertEquals(2, request.getStatementsCount());
+    assertEquals(INSERT1, request.getStatements(0).getSql());
+    assertEquals(INSERT2, request.getStatements(1).getSql());
+  }
+}
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 b79ca51c..057924b3 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
@@ -391,7 +391,7 @@ public class ITJdbcTest implements IntegrationTest {
    * This must be run with the system property -Dcopy_in_insert_or_update=true set, or otherwise it
    * will eventually fail on a unique key constraint violation, as the primary key value is a random
    * number. <code>
-   *   psql -h localhost -d knut-test-db -c "copy (select (random()*1000000000)::bigint, random()<0.5, md5(random()::text || clock_timestamp()::text)::bytea, random()*123456789, (random()*999999)::int, (random()*999999)::numeric, now()-random()*interval '50 year', md5(random()::text || clock_timestamp()::text)::varchar from generate_series(1, 1000000) s(i)) to stdout" | psql -h localhost -p 5433 -d test -c "set autocommit_dml_mode='partitioned_non_atomic'" -c "copy all_types from stdin;"
+   *   psql -h localhost -d knut-test-db -c "copy (select (random()*1000000000)::bigint, random()<0.5, md5(random()::text || clock_timestamp()::text)::bytea, random()*123456789, (random()*999999)::int, (random()*999999)::numeric, now()-random()*interval '50 year', md5(random()::text || clock_timestamp()::text)::varchar from generate_series(1, 1000000) s(i)) to stdout" | psql -h localhost -p 5433 -d test -c "set spanner.autocommit_dml_mode='partitioned_non_atomic'" -c "copy all_types from stdin;"
    * </code>
    */
   @Test
@@ -461,7 +461,9 @@ public class ITJdbcTest implements IntegrationTest {
     testEnv.write(databaseId, Collections.singleton(Mutation.delete("all_types", KeySet.all())));
 
     try (Connection connection = DriverManager.getConnection(getConnectionUrl())) {
-      connection.createStatement().execute("set autocommit_dml_mode='partitioned_non_atomic'");
+      connection
+          .createStatement()
+          .execute("set spanner.autocommit_dml_mode='partitioned_non_atomic'");
 
       PGConnection pgConnection = connection.unwrap(PGConnection.class);
       CopyManager copyManager = pgConnection.getCopyAPI();
@@ -485,4 +487,70 @@ public class ITJdbcTest implements IntegrationTest {
       connection.createStatement().execute("delete from all_types");
     }
   }
+
+  @Test
+  public void testTwoDmlStatementsInSimpleMode() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(getConnectionUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        // Statement#execute(String) returns false if the result is an update count or no result.
+        assertFalse(
+            statement.execute(
+                "INSERT INTO numbers VALUES (2, 'Two'); UPDATE numbers SET name=name || ' - Updated';"));
+
+        // Note that we have sent two DML statements to the database in one string. These should be
+        // treated as separate statements, and there should therefore be two results coming back
+        // from the server. That is; The first update count should be 1 (the INSERT), and the second
+        // should be 2 (the UPDATE).
+        assertEquals(1, statement.getUpdateCount());
+
+        // The following is a prime example of how not to design an API, but this is how JDBC works.
+        // getMoreResults() returns true if the next result is a ResultSet. However, if the next
+        // result is an update count, it returns false, and we have to check getUpdateCount() to
+        // verify whether there were any more results.
+        assertFalse(statement.getMoreResults());
+        assertEquals(2, statement.getUpdateCount());
+
+        // There are no more results. This is indicated by getMoreResults returning false AND
+        // getUpdateCount returning -1.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+
+        // Read back the data to verify.
+        try (ResultSet resultSet =
+            statement.executeQuery("SELECT name FROM numbers ORDER BY num")) {
+          assertTrue(resultSet.next());
+          assertEquals("One - Updated", resultSet.getString("name"));
+          assertTrue(resultSet.next());
+          assertEquals("Two - Updated", resultSet.getString("name"));
+          assertFalse(resultSet.next());
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testDmlAndQueryInBatchInSimpleMode() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(getConnectionUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        assertFalse(
+            statement.execute(
+                "INSERT INTO numbers VALUES (2, 'Two'); SELECT name FROM numbers ORDER BY num;"));
+        assertEquals(1, statement.getUpdateCount());
+
+        assertTrue(statement.getMoreResults());
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals("One", resultSet.getString("name"));
+          assertTrue(resultSet.next());
+          assertEquals("Two", resultSet.getString("name"));
+          assertFalse(resultSet.next());
+        }
+
+        // There are no more results. This is indicated by getMoreResults returning false AND
+        // getUpdateCount returning -1.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+  }
 }
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 2c76044a..486caf0d 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -386,6 +386,7 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
   @Test
   public void testDdl() throws SQLException {
     String sql = "CREATE TABLE foo (id bigint primary key)";
+    addDdlResponseToSpannerAdmin();
 
     try (Connection connection = DriverManager.getConnection(createUrl())) {
       try (java.sql.Statement statement = connection.createStatement()) {
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java
index f89accdf..170b81d8 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java
@@ -14,6 +14,8 @@
 
 package com.google.cloud.spanner.pgadapter;
 
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
@@ -22,8 +24,11 @@ import static org.junit.Assert.assertTrue;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
+import com.google.spanner.v1.CommitRequest;
+import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
+import com.google.spanner.v1.RollbackRequest;
 import java.math.BigDecimal;
 import java.sql.Connection;
 import java.sql.DriverManager;
@@ -172,83 +177,391 @@ public class JdbcSimpleModeMockServerTest extends AbstractMockServerTest {
   }
 
   @Test
-  public void testMixedBatch() throws SQLException {
-    // TODO: Change this test case once we support mixed batches.
+  public void testTwoDmlStatements() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        // Statement#execute(String) returns false if the result is an update count or no result.
+        assertFalse(
+            statement.execute(String.format("%s; %s;", INSERT_STATEMENT, UPDATE_STATEMENT)));
+
+        // Note that we have sent two DML statements to the database in one string. These should be
+        // treated as separate statements, and there should therefore be two results coming back
+        // from the server. That is; The first update count should be 1 (the INSERT), and the second
+        // should be 2 (the UPDATE).
+        assertEquals(1, statement.getUpdateCount());
+
+        // The following is a prime example of how not to design an API, but this is how JDBC works.
+        // getMoreResults() returns true if the next result is a ResultSet. However, if the next
+        // result is an update count, it returns false, and we have to check getUpdateCount() to
+        // verify whether there were any more results.
+        assertFalse(statement.getMoreResults());
+        assertEquals(2, statement.getUpdateCount());
+
+        // There are no more results. This is indicated by getMoreResults returning false AND
+        // getUpdateCount returning -1.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+
+    // Verify that the DML statements were batched together by PgAdapter.
+    List<ExecuteBatchDmlRequest> requests =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
+    assertEquals(1, requests.size());
+    ExecuteBatchDmlRequest request = requests.get(0);
+    assertEquals(2, request.getStatementsCount());
+    assertEquals(INSERT_STATEMENT.getSql(), request.getStatements(0).getSql());
+    assertEquals(UPDATE_STATEMENT.getSql(), request.getStatements(1).getSql());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+  }
+
+  @Test
+  public void testBatchInActiveTransaction() throws SQLException {
+    String sql = String.format("%s; %s; %s;", INSERT_STATEMENT, "COMMIT", SELECT2);
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        // Start an explicit transaction before executing batch
+        assertFalse(statement.execute("BEGIN; SELECT 1"));
+        assertEquals(0, statement.getUpdateCount());
+
+        assertTrue(statement.getMoreResults());
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(1L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+
+        // Execute Batch
+        assertFalse(statement.execute(sql));
+        assertEquals(1, statement.getUpdateCount());
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(0, statement.getUpdateCount());
+
+        assertTrue(statement.getMoreResults());
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(2L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(3, requests.size());
+    assertEquals(SELECT1.getSql(), requests.get(0).getSql());
+    assertTrue(requests.get(0).getTransaction().hasBegin());
+    assertEquals(INSERT_STATEMENT.getSql(), requests.get(1).getSql());
+    assertFalse(requests.get(1).getTransaction().hasBegin());
+    assertTrue(requests.get(1).getTransaction().hasId());
+    assertEquals(SELECT2.getSql(), requests.get(2).getSql());
+    assertTrue(requests.get(2).getTransaction().hasSingleUse());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+  }
+
+  @Test
+  public void testErrorHandlingInTwoDmlStatements() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        SQLException exception =
+            assertThrows(
+                SQLException.class,
+                () -> statement.execute(String.format("%s; %s;", INSERT_STATEMENT, INVALID_DML)));
+        assertEquals("ERROR: INVALID_ARGUMENT: Statement is invalid.", exception.getMessage());
+      }
+    }
+
+    // Verify that the DML statements were batched together by PgAdapter.
+    List<ExecuteBatchDmlRequest> requests =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
+    assertEquals(1, requests.size());
+    ExecuteBatchDmlRequest request = requests.get(0);
+    assertEquals(2, request.getStatementsCount());
+    assertEquals(INSERT_STATEMENT.getSql(), request.getStatements(0).getSql());
+    assertEquals(INVALID_DML.getSql(), request.getStatements(1).getSql());
+
+    // Verify that the implicit transaction is rolled back due to the exception
+    List<RollbackRequest> rollbackRequests = mockSpanner.getRequestsOfType(RollbackRequest.class);
+    assertEquals(1, rollbackRequests.size());
+  }
+
+  @Test
+  public void testErrorHandlingInExplicitTransactionWithCommit() throws SQLException {
     String sql =
-        "CREATE TABLE testtable (id integer PRIMARY KEY, data integer); INSERT INTO testtable (id, data) VALUES (1, 2); INSERT INTO testtable (id, data) VALUES (3, 4);";
+        String.format(
+            "%s; %s; %s; %s; %s;",
+            INSERT_STATEMENT, "BEGIN", UPDATE_STATEMENT, INVALID_DML, "COMMIT");
     try (Connection connection = DriverManager.getConnection(createUrl())) {
-      try (Statement statement = connection.createStatement()) {
+      try (java.sql.Statement statement = connection.createStatement()) {
         SQLException exception = assertThrows(SQLException.class, () -> statement.execute(sql));
-        assertEquals(
-            String.format(
-                "ERROR: FAILED_PRECONDITION: Executing updates is not allowed for DDL batches. \"%s\"",
-                sql.substring(0, sql.length() - 1)),
-            exception.getMessage());
+        assertThat(
+            exception.getMessage(), containsString("INVALID_ARGUMENT: Statement is invalid."));
       }
     }
+
+    // Verify that the DML statements were batched together by PgAdapter.
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(1, requests.size());
+    assertEquals(INSERT_STATEMENT.getSql(), requests.get(0).getSql());
+
+    List<ExecuteBatchDmlRequest> batchDmlRequests =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
+    assertEquals(1, requests.size());
+    ExecuteBatchDmlRequest request = batchDmlRequests.get(0);
+    assertEquals(2, request.getStatementsCount());
+    assertEquals(UPDATE_STATEMENT.getSql(), request.getStatements(0).getSql());
+    assertEquals(INVALID_DML.getSql(), request.getStatements(1).getSql());
+
+    List<RollbackRequest> rollbackRequests = mockSpanner.getRequestsOfType(RollbackRequest.class);
+    assertEquals(1, rollbackRequests.size());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
   }
 
   @Test
-  public void testSelectInDdlBatch() throws SQLException {
-    // TODO: Change this test case once we support mixed batches.
+  public void testErrorHandlingInExplicitTransactionWithoutCommit() throws SQLException {
     String sql =
-        "CREATE TABLE test (id integer PRIMARY KEY); SELECT * FROM test; SELECT * FROM test;";
+        String.format("%s; %s; %s; %s;", INSERT_STATEMENT, "BEGIN", UPDATE_STATEMENT, INVALID_DML);
     try (Connection connection = DriverManager.getConnection(createUrl())) {
-      try (Statement statement = connection.createStatement()) {
+      try (java.sql.Statement statement = connection.createStatement()) {
         SQLException exception = assertThrows(SQLException.class, () -> statement.execute(sql));
-        assertEquals(
-            String.format(
-                "ERROR: FAILED_PRECONDITION: Executing queries is not allowed for DDL batches. \"%s\"",
-                sql.substring(0, sql.length() - 1)),
-            exception.getMessage());
+        assertThat(
+            exception.getMessage(), containsString("INVALID_ARGUMENT: Statement is invalid."));
       }
     }
+
+    // Verify that the DML statements were batched together by PgAdapter.
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(1, requests.size());
+    assertEquals(INSERT_STATEMENT.getSql(), requests.get(0).getSql());
+
+    List<ExecuteBatchDmlRequest> batchDmlRequests =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
+    assertEquals(1, requests.size());
+    ExecuteBatchDmlRequest request = batchDmlRequests.get(0);
+    assertEquals(2, request.getStatementsCount());
+    assertEquals(UPDATE_STATEMENT.getSql(), request.getStatements(0).getSql());
+    assertEquals(INVALID_DML.getSql(), request.getStatements(1).getSql());
+
+    // The aborted transaction should be rolled back by pgadapter
+    List<RollbackRequest> rollbackRequests = mockSpanner.getRequestsOfType(RollbackRequest.class);
+    assertEquals(1, rollbackRequests.size());
+
+    // BEGIN statement commits the implicit transaction
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
   }
 
   @Test
-  public void testTransactionStatementsInBatch() throws SQLException {
-    // TODO: Change this test case once we support mixed batches.
+  public void testErrorHandlingOfDmlBatchAfterCommit() throws SQLException {
     String sql =
-        "BEGIN TRANSACTION; INSERT INTO users (id, age, name) VALUES (99, 99, 'test'); COMMIT;";
+        String.format(
+            "%s; %s; %s; %s; %s;",
+            INSERT_STATEMENT, "COMMIT", UPDATE_STATEMENT, INVALID_DML, SELECT1);
     try (Connection connection = DriverManager.getConnection(createUrl())) {
-      try (Statement statement = connection.createStatement()) {
+      try (java.sql.Statement statement = connection.createStatement()) {
         SQLException exception = assertThrows(SQLException.class, () -> statement.execute(sql));
-        assertEquals(
-            String.format(
-                "ERROR: FAILED_PRECONDITION: This connection has an active batch and cannot begin a transaction \"%s\"",
-                sql.substring(0, sql.length() - 1)),
-            exception.getMessage());
+        assertThat(
+            exception.getMessage(), containsString("INVALID_ARGUMENT: Statement is invalid."));
       }
     }
+
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(2, requests.size());
+    assertEquals(INSERT_STATEMENT.getSql(), requests.get(0).getSql());
+    assertEquals(SELECT1.getSql(), requests.get(1).getSql());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+
+    List<ExecuteBatchDmlRequest> batchDmlRequests =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
+    assertEquals(1, batchDmlRequests.size());
+    ExecuteBatchDmlRequest request = batchDmlRequests.get(0);
+    assertEquals(2, request.getStatementsCount());
+    assertEquals(UPDATE_STATEMENT.getSql(), request.getStatements(0).getSql());
+    assertEquals(INVALID_DML.getSql(), request.getStatements(1).getSql());
   }
 
   @Test
-  public void testSelectAtStartOfBatch() throws SQLException {
-    // TODO: Change this test case once we support mixed batches.
+  public void testErrorHandlingInAutocommitMode() throws SQLException {
     String sql =
-        "SELECT * FROM users; INSERT INTO users (id, age, name) VALUES (99, 99, 'person'); SELECT * FROM users;";
+        String.format(
+            "%s; %s; %s; %s; %s;", INSERT_STATEMENT, "COMMIT", SELECT1, INVALID_SELECT, SELECT2);
     try (Connection connection = DriverManager.getConnection(createUrl())) {
-      try (Statement statement = connection.createStatement()) {
+      try (java.sql.Statement statement = connection.createStatement()) {
         SQLException exception = assertThrows(SQLException.class, () -> statement.execute(sql));
-        assertEquals(
-            String.format(
-                "ERROR: INVALID_ARGUMENT: Statement type is not supported for batching \"%s\"",
-                sql.substring(0, sql.length() - 1)),
-            exception.getMessage());
+        assertThat(
+            exception.getMessage(), containsString("INVALID_ARGUMENT: Statement is invalid."));
+      }
+    }
+
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(4, requests.size());
+    assertEquals(INSERT_STATEMENT.getSql(), requests.get(0).getSql());
+    assertEquals(SELECT1.getSql(), requests.get(1).getSql());
+    assertTrue(requests.get(1).getTransaction().hasSingleUse());
+    assertEquals(INVALID_SELECT.getSql(), requests.get(2).getSql());
+    assertTrue(requests.get(2).getTransaction().hasSingleUse());
+    assertEquals(SELECT2.getSql(), requests.get(3).getSql());
+    assertTrue(requests.get(3).getTransaction().hasSingleUse());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+
+    List<RollbackRequest> rollbackRequests = mockSpanner.getRequestsOfType(RollbackRequest.class);
+    assertEquals(0, rollbackRequests.size());
+  }
+
+  @Test
+  public void testTwoQueries() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        // Statement#execute(String) returns true if the result is a result set.
+        assertTrue(statement.execute("SELECT 1; SELECT 2;"));
+
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(1L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        // getMoreResults() returns true if the next result is a ResultSet.
+        assertTrue(statement.getMoreResults());
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(2L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        // getMoreResults() should now return false. We should also check getUpdateCount() as that
+        // method should return -1 to indicate that there is also no update count available.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
       }
     }
+
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(2, requests.size());
+    assertEquals(SELECT1.getSql(), requests.get(0).getSql());
+    assertTrue(requests.get(0).getTransaction().hasSingleUse());
+    assertEquals(SELECT2.getSql(), requests.get(1).getSql());
+    assertTrue(requests.get(1).getTransaction().hasSingleUse());
   }
 
   @Test
-  public void testDdl() throws SQLException {
-    String sql = "CREATE TABLE foo (id bigint primary key)";
+  public void testMixedBatch() throws SQLException {
+    String sql =
+        "CREATE TABLE foo (id bigint primary key); INSERT INTO FOO VALUES (1); UPDATE FOO SET BAR=1 WHERE BAZ=2;";
+    addDdlResponseToSpannerAdmin();
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (Statement statement = connection.createStatement()) {
+        // Statement#execute(String) returns false if the result was either an update count or there
+        // was no result. Statement#getUpdateCount() returns 0 if there was no result.
+        assertFalse(statement.execute(sql));
+        assertEquals(0, statement.getUpdateCount());
+
+        // getMoreResults() returns false as the next result is an update count.
+        assertFalse(statement.getMoreResults());
+        assertEquals(1, statement.getUpdateCount());
+
+        // getMoreResults() returns false as the next result is an update count.
+        assertFalse(statement.getMoreResults());
+        assertEquals(2, statement.getUpdateCount());
+
+        // getMoreResults() should now return false. We should also check getUpdateCount() as that
+        // method should return -1 to indicate that there is also no update count available.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
 
+    List<UpdateDatabaseDdlRequest> updateDatabaseDdlRequests =
+        mockDatabaseAdmin.getRequests().stream()
+            .filter(request -> request instanceof UpdateDatabaseDdlRequest)
+            .map(UpdateDatabaseDdlRequest.class::cast)
+            .collect(Collectors.toList());
+    assertEquals(1, updateDatabaseDdlRequests.size());
+    assertEquals(1, updateDatabaseDdlRequests.get(0).getStatementsCount());
+    assertEquals(
+        "CREATE TABLE foo (id bigint primary key)",
+        updateDatabaseDdlRequests.get(0).getStatements(0));
+
+    // Verify that the DML statements were batched together by PgAdapter.
+    List<ExecuteBatchDmlRequest> requests =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
+    assertEquals(1, requests.size());
+    ExecuteBatchDmlRequest request = requests.get(0);
+    assertEquals(2, request.getStatementsCount());
+    assertEquals(INSERT_STATEMENT.getSql(), request.getStatements(0).getSql());
+    assertEquals(UPDATE_STATEMENT.getSql(), request.getStatements(1).getSql());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+  }
+
+  @Test
+  public void testErrorHandlingOfDdl() throws SQLException {
+    addDdlExceptionToSpannerAdmin();
+    String sql = String.format("%s; %s; %s", INSERT_STATEMENT, INVALID_DDL, UPDATE_STATEMENT);
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (Statement statement = connection.createStatement()) {
+        SQLException exception = assertThrows(SQLException.class, () -> statement.execute(sql));
+        assertThat(
+            exception.getMessage(), containsString("INVALID_ARGUMENT: Statement is invalid."));
+      }
+    }
+
+    // Verify that the execution does not halt and the last statement is executed.
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(2, requests.size());
+    assertEquals(INSERT_STATEMENT.getSql(), requests.get(0).getSql());
+    assertEquals(UPDATE_STATEMENT.getSql(), requests.get(1).getSql());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(2, commitRequests.size());
+  }
+
+  @Test
+  public void testSelectInDdlBatch() throws SQLException {
+    String sql = "CREATE TABLE foo (id bigint primary key); SELECT 1; SELECT 2;";
+    addDdlResponseToSpannerAdmin();
     try (Connection connection = DriverManager.getConnection(createUrl())) {
       try (Statement statement = connection.createStatement()) {
         // Statement#execute(String) returns false if the result was either an update count or there
         // was no result. Statement#getUpdateCount() returns 0 if there was no result.
         assertFalse(statement.execute(sql));
         assertEquals(0, statement.getUpdateCount());
+
+        assertTrue(statement.getMoreResults());
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(1L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        assertTrue(statement.getMoreResults());
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(2L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        // getMoreResults() should now return false. We should also check getUpdateCount() as that
+        // method should return -1 to indicate that there is also no update count available.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
       }
     }
 
@@ -259,6 +572,335 @@ public class JdbcSimpleModeMockServerTest extends AbstractMockServerTest {
             .collect(Collectors.toList());
     assertEquals(1, updateDatabaseDdlRequests.size());
     assertEquals(1, updateDatabaseDdlRequests.get(0).getStatementsCount());
-    assertEquals(sql, updateDatabaseDdlRequests.get(0).getStatements(0));
+    assertEquals(
+        "CREATE TABLE foo (id bigint primary key)",
+        updateDatabaseDdlRequests.get(0).getStatements(0));
+
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(2, requests.size());
+    assertEquals(SELECT1.getSql(), requests.get(0).getSql());
+    // Check that DDL statements close the implicit transaction
+    assertTrue(requests.get(0).getTransaction().hasSingleUse());
+    assertEquals(SELECT2.getSql(), requests.get(1).getSql());
+    // Check that DDL statements close the implicit transaction
+    assertTrue(requests.get(1).getTransaction().hasSingleUse());
+  }
+
+  @Test
+  public void testBeginAndDml() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        // batch a BEGIN [TRANSACTION] statement together with an update statement.
+        assertFalse(statement.execute(String.format("BEGIN; %s;", UPDATE_STATEMENT)));
+        assertEquals(0, statement.getUpdateCount());
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(2, statement.getUpdateCount());
+
+        // getMoreResults() should now return false. We should also check getUpdateCount() as that
+        // method should return -1 to indicate that there is also no update count available.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+
+        // Execute an insert statement as part of the transaction.
+        assertFalse(statement.execute(INSERT_STATEMENT.getSql()));
+        assertEquals(1, statement.getUpdateCount());
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+
+        statement.execute("COMMIT");
+      }
+    }
+
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(2, requests.size());
+    // The first statement in the transaction will include the BeginTransaction option.
+    assertEquals(UPDATE_STATEMENT.getSql(), requests.get(0).getSql());
+    assertTrue(requests.get(0).getTransaction().hasBegin());
+    assertEquals(INSERT_STATEMENT.getSql(), requests.get(1).getSql());
+    assertTrue(requests.get(1).getTransaction().hasId());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+  }
+
+  @Test
+  public void testTransactionStatementsInBatch() throws SQLException {
+    String sql =
+        "BEGIN TRANSACTION; INSERT INTO FOO VALUES (1); UPDATE FOO SET BAR=1 WHERE BAZ=2; COMMIT;";
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (Statement statement = connection.createStatement()) {
+        // Send a complete transaction as a single batch.
+        assertFalse(statement.execute(sql));
+        // The BEGIN statement should not return anything.
+        assertEquals(0, statement.getUpdateCount());
+        // INSERT
+        assertFalse(statement.getMoreResults());
+        assertEquals(1, statement.getUpdateCount());
+        // UPDATE
+        assertFalse(statement.getMoreResults());
+        assertEquals(2, statement.getUpdateCount());
+        // COMMIT
+        assertFalse(statement.getMoreResults());
+        assertEquals(0, statement.getUpdateCount());
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+
+    // As the two DML statements are adjacent in the batch, they can be combined into a BatchDML
+    // request. This is the first statement in the transaction and will include the BeginTransaction
+    // option.
+    List<ExecuteBatchDmlRequest> batchDmlRequests =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
+    assertEquals(1, batchDmlRequests.size());
+    ExecuteBatchDmlRequest batchDmlRequest = batchDmlRequests.get(0);
+    assertTrue(batchDmlRequest.getTransaction().hasBegin());
+    assertEquals(INSERT_STATEMENT.getSql(), batchDmlRequest.getStatements(0).getSql());
+    assertEquals(UPDATE_STATEMENT.getSql(), batchDmlRequest.getStatements(1).getSql());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+  }
+
+  @Test
+  public void testBeginInExplicitTransaction() throws SQLException {
+    String sql =
+        "BEGIN TRANSACTION; INSERT INTO FOO VALUES (1); UPDATE FOO SET BAR=1 WHERE BAZ=2; COMMIT;";
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        // Start an explicit transaction before executing batch
+        assertFalse(statement.execute("BEGIN; SELECT 1"));
+        assertEquals(0, statement.getUpdateCount());
+
+        assertTrue(statement.getMoreResults());
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(1L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+
+        // The BEGIN statement should be no-op and not cause exception.
+        assertFalse(statement.execute(sql));
+        // The BEGIN statement should not return anything.
+        assertEquals(0, statement.getUpdateCount());
+        // INSERT
+        assertFalse(statement.getMoreResults());
+        assertEquals(1, statement.getUpdateCount());
+        // UPDATE
+        assertFalse(statement.getMoreResults());
+        assertEquals(2, statement.getUpdateCount());
+        // COMMIT
+        assertFalse(statement.getMoreResults());
+        assertEquals(0, statement.getUpdateCount());
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+
+    List<ExecuteBatchDmlRequest> batchDmlRequests =
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class);
+    assertEquals(1, batchDmlRequests.size());
+    ExecuteBatchDmlRequest batchDmlRequest = batchDmlRequests.get(0);
+    // Verify that the BEGIN statement is no-op
+    assertFalse(batchDmlRequest.getTransaction().hasBegin());
+    assertTrue(batchDmlRequest.getTransaction().hasId());
+    assertEquals(INSERT_STATEMENT.getSql(), batchDmlRequest.getStatements(0).getSql());
+    assertEquals(UPDATE_STATEMENT.getSql(), batchDmlRequest.getStatements(1).getSql());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(1, commitRequests.size());
+  }
+
+  @Test
+  public void testSelectAtStartOfBatch() throws SQLException {
+    String sql = "SELECT 1; INSERT INTO FOO VALUES (1); SELECT 2;";
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (Statement statement = connection.createStatement()) {
+        assertTrue(statement.execute(sql));
+
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(1L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        // getMoreResults() returns false as the next result is an update count.
+        assertFalse(statement.getMoreResults());
+        assertEquals(1, statement.getUpdateCount());
+
+        assertTrue(statement.getMoreResults());
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(2L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        // getMoreResults() should now return false. We should also check getUpdateCount() as that
+        // method should return -1 to indicate that there is also no update count available.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+
+      List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+      assertEquals(3, requests.size());
+      assertEquals(SELECT1.getSql(), requests.get(0).getSql());
+      assertTrue(requests.get(0).getTransaction().hasSingleUse());
+      assertEquals(INSERT_STATEMENT.getSql(), requests.get(1).getSql());
+      assertTrue(requests.get(1).getTransaction().hasBegin());
+      assertEquals(SELECT2.getSql(), requests.get(2).getSql());
+      assertTrue(requests.get(2).getTransaction().hasSingleUse());
+
+      List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+      assertEquals(1, commitRequests.size());
+    }
+  }
+
+  @Test
+  public void testTwoDdlStatements() throws SQLException {
+    addDdlResponseToSpannerAdmin();
+    addDdlResponseToSpannerAdmin();
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        // execute returns false if the result is an update count or no result.
+        assertFalse(
+            statement.execute(
+                String.format(
+                    "%s; %s;",
+                    "CREATE TABLE FOO (id bigint primary key)",
+                    "CREATE TABLE BAR (id bigint primary key)")));
+        assertEquals(0, statement.getUpdateCount());
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(0, statement.getUpdateCount());
+
+        // getMoreResults() should now return false. We should also check getUpdateCount() as that
+        // method should return -1 to indicate that there is also no update count available.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+
+    // Verify that the two DDL statements are sent as one batch.
+    List<UpdateDatabaseDdlRequest> requests =
+        mockDatabaseAdmin.getRequests().stream()
+            .filter(r -> r instanceof UpdateDatabaseDdlRequest)
+            .map(r -> (UpdateDatabaseDdlRequest) r)
+            .collect(Collectors.toList());
+    assertEquals(1, requests.size());
+    UpdateDatabaseDdlRequest request = requests.get(0);
+    assertEquals(2, request.getStatementsCount());
+    assertEquals("CREATE TABLE FOO (id bigint primary key)", request.getStatements(0));
+    assertEquals("CREATE TABLE BAR (id bigint primary key)", request.getStatements(1));
+  }
+
+  @Test
+  public void testDdlInExplicitTransaction() throws SQLException {
+    addDdlResponseToSpannerAdmin();
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        // Start an explicit transaction before executing batch
+        assertFalse(statement.execute("BEGIN; SELECT 1"));
+        assertEquals(0, statement.getUpdateCount());
+
+        assertTrue(statement.getMoreResults());
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(1L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+
+        // Execute batch
+        SQLException exception =
+            assertThrows(
+                SQLException.class,
+                () -> statement.execute("SELECT 2; CREATE TABLE BAR (id bigint primary key);"));
+        assertThat(
+            exception.getMessage(),
+            containsString("DDL statements are only allowed outside explicit transactions."));
+      }
+    }
+
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(2, requests.size());
+    assertEquals(SELECT1.getSql(), requests.get(0).getSql());
+    assertTrue(requests.get(0).getTransaction().hasBegin());
+    // Verify that the explicit transaction is active while batch starts
+    assertEquals(SELECT2.getSql(), requests.get(1).getSql());
+    assertFalse(requests.get(1).getTransaction().hasBegin());
+    assertTrue(requests.get(1).getTransaction().hasId());
+  }
+
+  @Test
+  public void testJdbcBatch() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.addBatch(UPDATE_STATEMENT.getSql());
+        int[] updateCounts = statement.executeBatch();
+
+        assertEquals(2, updateCounts.length);
+        assertEquals(1, updateCounts[0]);
+        assertEquals(2, updateCounts[1]);
+      }
+    }
+
+    // The PostgreSQL JDBC driver will send the DML statements as separated statements to PG when
+    // executing a batch using simple mode. This means that Spanner will receive two separate DML
+    // requests.
+    assertTrue(mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).isEmpty());
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+
+    assertEquals(2, requests.size());
+    assertEquals(INSERT_STATEMENT.getSql(), requests.get(0).getSql());
+    assertEquals(UPDATE_STATEMENT.getSql(), requests.get(1).getSql());
+  }
+
+  @Test
+  public void testAbortedTransactionRollbackInAutocommit() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (java.sql.Statement statement = connection.createStatement()) {
+        assertFalse(statement.execute("BEGIN"));
+        assertEquals(0, statement.getUpdateCount());
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+
+        assertFalse(statement.execute(UPDATE_STATEMENT.getSql()));
+        assertEquals(2, statement.getUpdateCount());
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+
+        SQLException exception =
+            assertThrows(SQLException.class, () -> statement.execute(INVALID_DML.getSql()));
+        assertThat(
+            exception.getMessage(), containsString("INVALID_ARGUMENT: Statement is invalid."));
+
+        assertFalse(statement.execute("COMMIT"));
+        assertEquals(0, statement.getUpdateCount());
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(2, requests.size());
+    assertEquals(UPDATE_STATEMENT.getSql(), requests.get(0).getSql());
+    assertTrue(requests.get(0).getTransaction().hasBegin());
+    assertEquals(INVALID_DML.getSql(), requests.get(1).getSql());
+    assertTrue(requests.get(1).getTransaction().hasId());
+
+    List<RollbackRequest> rollbackRequests = mockSpanner.getRequestsOfType(RollbackRequest.class);
+    assertEquals(1, rollbackRequests.size());
   }
 }
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 bd1a610c..dcbca61b 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
@@ -20,6 +20,7 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.doNothing;
@@ -914,13 +915,14 @@ public class ProtocolTest {
 
     doReturn(false)
         .when(messageSpy)
-        .sendSpannerResult(any(IntermediatePortalStatement.class), any(QueryMode.class), anyLong());
+        .sendSpannerResult(
+            anyInt(), any(IntermediatePortalStatement.class), any(QueryMode.class), anyLong());
 
     messageSpy.send();
 
     verify(intermediatePortalStatement).execute();
     verify(messageSpy)
-        .sendSpannerResult(intermediatePortalStatement, QueryMode.EXTENDED, totalRows);
+        .sendSpannerResult(0, intermediatePortalStatement, QueryMode.EXTENDED, totalRows);
     verify(connectionHandler).cleanUp(intermediatePortalStatement);
   }
 
@@ -1310,10 +1312,11 @@ public class ProtocolTest {
     CopyDoneMessage messageSpy = (CopyDoneMessage) spy(message);
     doReturn(false)
         .when(messageSpy)
-        .sendSpannerResult(any(IntermediateStatement.class), any(QueryMode.class), anyLong());
+        .sendSpannerResult(
+            anyInt(), any(IntermediateStatement.class), any(QueryMode.class), anyLong());
 
     messageSpy.send();
-    verify(messageSpy).sendSpannerResult(copyStatement, QueryMode.SIMPLE, 0L);
+    verify(messageSpy).sendSpannerResult(0, copyStatement, QueryMode.SIMPLE, 0L);
   }
 
   @Test
@@ -1393,7 +1396,7 @@ public class ProtocolTest {
 
     assertEquals("TEXT", copyStatement.getFormatType());
     assertEquals('\t', copyStatement.getDelimiterChar());
-    assertFalse(copyStatement.hasException());
+    assertFalse(copyStatement.hasException(0));
     assertEquals(12L, copyStatement.getUpdateCount().longValue());
     assertEquals(12L, mw.getRowCount());
 
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 e4a1f703..792281d7 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java
@@ -33,6 +33,7 @@ import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.AbstractStatementParser;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
@@ -41,7 +42,6 @@ import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePortalStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
-import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement.ResultType;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage;
@@ -100,21 +100,22 @@ public class StatementTest {
     when(statementResult.getResultSet()).thenReturn(resultSet);
     when(resultSet.next()).thenReturn(true);
 
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     IntermediateStatement intermediateStatement =
-        new IntermediateStatement(options, parse("SELECT * FROM users"), connection);
+        new IntermediateStatement(options, parse("SELECT * FROM users"), connectionHandler);
 
     assertFalse(intermediateStatement.isExecuted());
-    assertEquals("SELECT", intermediateStatement.getCommand());
+    assertEquals("SELECT", intermediateStatement.getCommand(0));
 
     intermediateStatement.execute();
 
-    verify(connection, Mockito.times(1)).execute(Statement.of("SELECT * FROM users"));
-    assertTrue(intermediateStatement.containsResultSet());
+    verify(connection).execute(Statement.of("SELECT * FROM users"));
+    assertTrue(intermediateStatement.containsResultSet(0));
     assertTrue(intermediateStatement.isExecuted());
-    assertEquals(ResultType.RESULT_SET, intermediateStatement.getResultType());
-    assertEquals(resultSet, intermediateStatement.getStatementResult());
-    assertTrue(intermediateStatement.isHasMoreData());
-    assertFalse(intermediateStatement.hasException());
+    assertEquals(StatementType.QUERY, intermediateStatement.getStatementType(0));
+    assertEquals(resultSet, intermediateStatement.getStatementResult(0));
+    assertTrue(intermediateStatement.isHasMoreData(0));
+    assertFalse(intermediateStatement.hasException(0));
     assertEquals(0, intermediateStatement.getResultFormatCode(0));
 
     intermediateStatement.close();
@@ -129,24 +130,24 @@ public class StatementTest {
     when(connection.execute(Statement.of("UPDATE users SET name = someName WHERE id = 10")))
         .thenReturn(statementResult);
 
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     IntermediateStatement intermediateStatement =
         new IntermediateStatement(
-            options, parse("UPDATE users SET name = someName WHERE id = 10"), connection);
+            options, parse("UPDATE users SET name = someName WHERE id = 10"), connectionHandler);
 
     assertFalse(intermediateStatement.isExecuted());
-    assertEquals("UPDATE", intermediateStatement.getCommand());
+    assertEquals("UPDATE", intermediateStatement.getCommand(0));
 
     intermediateStatement.execute();
 
-    verify(connection, Mockito.times(1))
-        .execute(Statement.of("UPDATE users SET name = someName WHERE id = 10"));
-    assertFalse(intermediateStatement.containsResultSet());
-    assertEquals(1L, intermediateStatement.getUpdateCount().longValue());
+    verify(connection).execute(Statement.of("UPDATE users SET name = someName WHERE id = 10"));
+    assertFalse(intermediateStatement.containsResultSet(0));
+    assertEquals(1L, intermediateStatement.getUpdateCount(0));
     assertTrue(intermediateStatement.isExecuted());
-    assertEquals(ResultType.UPDATE_COUNT, intermediateStatement.getResultType());
-    assertNull(intermediateStatement.getStatementResult());
-    assertFalse(intermediateStatement.isHasMoreData());
-    assertFalse(intermediateStatement.hasException());
+    assertEquals(StatementType.UPDATE, intermediateStatement.getStatementType(0));
+    assertNull(intermediateStatement.getStatementResult(0));
+    assertFalse(intermediateStatement.isHasMoreData(0));
+    assertFalse(intermediateStatement.hasException(0));
     assertEquals(0, intermediateStatement.getResultFormatCode(0));
 
     intermediateStatement.close();
@@ -161,23 +162,25 @@ public class StatementTest {
     when(connection.execute(Statement.of("UPDATE users SET name = someName WHERE id = -1")))
         .thenReturn(statementResult);
 
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     IntermediateStatement intermediateStatement =
         new IntermediateStatement(
-            options, parse("UPDATE users SET name = someName WHERE id = -1"), connection);
+            options, parse("UPDATE users SET name = someName WHERE id = -1"), connectionHandler);
 
     assertFalse(intermediateStatement.isExecuted());
-    assertEquals("UPDATE", intermediateStatement.getCommand());
+    assertEquals("UPDATE", intermediateStatement.getCommand(0));
 
     intermediateStatement.execute();
 
-    verify(connection).execute(Statement.of("UPDATE users SET name = someName WHERE id = -1"));
-    assertFalse(intermediateStatement.containsResultSet());
-    assertEquals(0L, intermediateStatement.getUpdateCount().longValue());
+    Mockito.verify(connection)
+        .execute(Statement.of("UPDATE users SET name = someName WHERE id = -1"));
+    assertFalse(intermediateStatement.containsResultSet(0));
+    assertEquals(0L, intermediateStatement.getUpdateCount(0));
     assertTrue(intermediateStatement.isExecuted());
-    assertEquals(ResultType.UPDATE_COUNT, intermediateStatement.getResultType());
-    assertNull(intermediateStatement.getStatementResult());
-    assertFalse(intermediateStatement.isHasMoreData());
-    assertFalse(intermediateStatement.hasException());
+    assertEquals(StatementType.UPDATE, intermediateStatement.getStatementType(0));
+    assertNull(intermediateStatement.getStatementResult(0));
+    assertFalse(intermediateStatement.isHasMoreData(0));
+    assertFalse(intermediateStatement.hasException(0));
     assertEquals(0, intermediateStatement.getResultFormatCode(0));
 
     intermediateStatement.close();
@@ -191,34 +194,37 @@ public class StatementTest {
     when(connection.execute(Statement.of("CREATE TABLE users (name varchar(100) primary key)")))
         .thenReturn(statementResult);
 
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     IntermediateStatement intermediateStatement =
         new IntermediateStatement(
-            options, parse("CREATE TABLE users (name varchar(100) primary key)"), connection);
+            options,
+            parse("CREATE TABLE users (name varchar(100) primary key)"),
+            connectionHandler);
 
     assertFalse(intermediateStatement.isExecuted());
-    assertEquals("CREATE", intermediateStatement.getCommand());
+    assertEquals("CREATE", intermediateStatement.getCommand(0));
 
     intermediateStatement.execute();
 
     verify(connection).execute(Statement.of("CREATE TABLE users (name varchar(100) primary key)"));
-    assertFalse(intermediateStatement.containsResultSet());
-    assertNull(intermediateStatement.getUpdateCount());
+    assertFalse(intermediateStatement.containsResultSet(0));
+    assertEquals(0, intermediateStatement.getUpdateCount(0));
     assertTrue(intermediateStatement.isExecuted());
-    assertEquals(ResultType.NO_RESULT, intermediateStatement.getResultType());
-    assertNull(intermediateStatement.getStatementResult());
-    assertFalse(intermediateStatement.isHasMoreData());
-    assertFalse(intermediateStatement.hasException());
+    assertEquals(StatementType.DDL, intermediateStatement.getStatementType(0));
+    assertNull(intermediateStatement.getStatementResult(0));
+    assertFalse(intermediateStatement.isHasMoreData(0));
+    assertFalse(intermediateStatement.hasException(0));
     assertEquals(0, intermediateStatement.getResultFormatCode(0));
 
     intermediateStatement.close();
 
-    verify(resultSet, never()).close();
+    Mockito.verify(resultSet, never()).close();
   }
 
   @Test
   public void testDescribeBasicStatementThrowsException() {
     IntermediateStatement intermediateStatement =
-        new IntermediateStatement(options, parse("SELECT * FROM users"), connection);
+        new IntermediateStatement(options, parse("SELECT * FROM users"), connectionHandler);
 
     assertThrows(IllegalStateException.class, intermediateStatement::describe);
   }
@@ -230,13 +236,14 @@ public class StatementTest {
 
     when(connection.execute(Statement.of("SELECT * FROM users"))).thenThrow(thrownException);
 
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     IntermediateStatement intermediateStatement =
-        new IntermediateStatement(options, parse("SELECT * FROM users"), connection);
+        new IntermediateStatement(options, parse("SELECT * FROM users"), connectionHandler);
 
     intermediateStatement.execute();
 
-    assertTrue(intermediateStatement.hasException());
-    assertEquals(thrownException, intermediateStatement.getException());
+    assertTrue(intermediateStatement.hasException(0));
+    assertEquals(thrownException, intermediateStatement.getException(0));
   }
 
   @Test
@@ -271,7 +278,7 @@ public class StatementTest {
     verify(connection).execute(statement);
 
     assertEquals(sqlStatement, intermediatePortalStatement.getSql());
-    assertEquals("SELECT", intermediatePortalStatement.getCommand());
+    assertEquals("SELECT", intermediatePortalStatement.getCommand(0));
     assertFalse(intermediatePortalStatement.isExecuted());
     assertTrue(intermediateStatement.isBound());
   }
@@ -365,14 +372,13 @@ public class StatementTest {
     String sql =
         "INSERT INTO users (id) VALUES (1); INSERT INTO users (id) VALUES (2);INSERT INTO users (id) VALUES (3);";
     IntermediateStatement intermediateStatement =
-        new IntermediateStatement(options, parse(sql), connection);
+        new IntermediateStatement(options, parse(sql), connectionHandler);
 
     assertTrue(intermediateStatement.isBatchedQuery());
-    List<String> result = intermediateStatement.getStatements();
-    assertEquals(3, result.size(), 3);
-    assertEquals("INSERT INTO users (id) VALUES (1)", result.get(0));
-    assertEquals("INSERT INTO users (id) VALUES (2)", result.get(1));
-    assertEquals("INSERT INTO users (id) VALUES (3)", result.get(2));
+    assertEquals(3, intermediateStatement.getStatements().size(), 3);
+    assertEquals("INSERT INTO users (id) VALUES (1)", intermediateStatement.getStatement(0));
+    assertEquals("INSERT INTO users (id) VALUES (2)", intermediateStatement.getStatement(1));
+    assertEquals("INSERT INTO users (id) VALUES (3)", intermediateStatement.getStatement(2));
   }
 
   @Test
@@ -380,30 +386,28 @@ public class StatementTest {
     String sql =
         "BEGIN TRANSACTION; INSERT INTO users (id) VALUES (1); INSERT INTO users (id) VALUES (2); INSERT INTO users (id) VALUES (3); COMMIT;";
     IntermediateStatement intermediateStatement =
-        new IntermediateStatement(options, parse(sql), connection);
+        new IntermediateStatement(options, parse(sql), connectionHandler);
 
     assertTrue(intermediateStatement.isBatchedQuery());
-    List<String> result = intermediateStatement.getStatements();
-
-    assertEquals(5, result.size(), 5);
-    assertEquals("BEGIN TRANSACTION", result.get(0));
-    assertEquals("INSERT INTO users (id) VALUES (1)", result.get(1));
-    assertEquals("INSERT INTO users (id) VALUES (2)", result.get(2));
-    assertEquals("INSERT INTO users (id) VALUES (3)", result.get(3));
-    assertEquals("COMMIT", result.get(4));
+
+    assertEquals(5, intermediateStatement.getStatements().size(), 5);
+    assertEquals("BEGIN TRANSACTION", intermediateStatement.getStatement(0));
+    assertEquals("INSERT INTO users (id) VALUES (1)", intermediateStatement.getStatement(1));
+    assertEquals("INSERT INTO users (id) VALUES (2)", intermediateStatement.getStatement(2));
+    assertEquals("INSERT INTO users (id) VALUES (3)", intermediateStatement.getStatement(3));
+    assertEquals("COMMIT", intermediateStatement.getStatement(4));
   }
 
   @Test
   public void testBatchStatementsWithEmptyStatements() {
     String sql = "INSERT INTO users (id) VALUES (1); ;;; INSERT INTO users (id) VALUES (2);";
     IntermediateStatement intermediateStatement =
-        new IntermediateStatement(options, parse(sql), connection);
+        new IntermediateStatement(options, parse(sql), connectionHandler);
 
     assertTrue(intermediateStatement.isBatchedQuery());
-    List<String> result = intermediateStatement.getStatements();
-    assertEquals(2, result.size());
-    assertEquals("INSERT INTO users (id) VALUES (1)", result.get(0));
-    assertEquals("INSERT INTO users (id) VALUES (2)", result.get(1));
+    assertEquals(2, intermediateStatement.getStatements().size());
+    assertEquals("INSERT INTO users (id) VALUES (1)", intermediateStatement.getStatement(0));
+    assertEquals("INSERT INTO users (id) VALUES (2)", intermediateStatement.getStatement(1));
   }
 
   @Test
@@ -411,13 +415,15 @@ public class StatementTest {
     String sql =
         "INSERT INTO users (name) VALUES (';;test;;'); INSERT INTO users (name1, name2) VALUES ('''''', ';'';');";
     IntermediateStatement intermediateStatement =
-        new IntermediateStatement(options, parse(sql), connection);
+        new IntermediateStatement(options, parse(sql), connectionHandler);
 
     assertTrue(intermediateStatement.isBatchedQuery());
-    List<String> result = intermediateStatement.getStatements();
-    assertEquals(2, result.size());
-    assertEquals("INSERT INTO users (name) VALUES (';;test;;')", result.get(0));
-    assertEquals("INSERT INTO users (name1, name2) VALUES ('''''', ';'';')", result.get(1));
+    assertEquals(2, intermediateStatement.getStatements().size());
+    assertEquals(
+        "INSERT INTO users (name) VALUES (';;test;;')", intermediateStatement.getStatement(0));
+    assertEquals(
+        "INSERT INTO users (name1, name2) VALUES ('''''', ';'';')",
+        intermediateStatement.getStatement(1));
   }
 
   @Test
@@ -441,10 +447,12 @@ public class StatementTest {
     IntermediateStatement intermediateStatement = ((QueryMessage) message).getStatement();
 
     assertTrue(intermediateStatement.isBatchedQuery());
-    List<String> result = intermediateStatement.getStatements();
-    assertEquals(2, result.size());
-    assertEquals("INSERT INTO users (name) VALUES (';;test;;')", result.get(0));
-    assertEquals("INSERT INTO users (name1, name2) VALUES ('''''', ';'';')", result.get(1));
+    assertEquals(2, intermediateStatement.getStatements().size());
+    assertEquals(
+        "INSERT INTO users (name) VALUES (';;test;;')", intermediateStatement.getStatement(0));
+    assertEquals(
+        "INSERT INTO users (name1, name2) VALUES ('''''', ';'';')",
+        intermediateStatement.getStatement(1));
   }
 
   @Test
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 8d2b952c..a88eaa6f 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
@@ -30,14 +30,22 @@ import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStateme
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
+import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 
 @RunWith(JUnit4.class)
 public class IntermediateStatementTest {
+  @Rule public MockitoRule rule = MockitoJUnit.rule();
+
+  @Mock private ConnectionHandler connectionHandler;
+
   private static final AbstractStatementParser PARSER =
       AbstractStatementParser.getInstance(Dialect.POSTGRESQL);
 
@@ -49,9 +57,11 @@ public class IntermediateStatementTest {
 
   @Test
   public void testUpdateResultCount_ResultSet() {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
+
     IntermediateStatement statement =
         new IntermediateStatement(
-            mock(OptionsMetadata.class), parse("select foo from bar"), connection);
+            mock(OptionsMetadata.class), parse("select foo from bar"), connectionHandler);
     ResultSet resultSet = mock(ResultSet.class);
     when(resultSet.next()).thenReturn(true, false);
     StatementResult result = mock(StatementResult.class);
@@ -59,46 +69,48 @@ public class IntermediateStatementTest {
     when(result.getResultSet()).thenReturn(resultSet);
     when(result.getUpdateCount()).thenThrow(new IllegalStateException());
 
-    statement.updateResultCount(result);
+    statement.updateResultCount(0, result);
 
-    assertTrue(statement.hasMoreData);
-    assertNull(statement.updateCount);
-    assertSame(resultSet, statement.statementResult);
+    assertTrue(statement.hasMoreData[0]);
+    assertEquals(-1, statement.getUpdateCount(0));
+    assertSame(resultSet, statement.getStatementResult(0));
   }
 
   @Test
   public void testUpdateResultCount_UpdateCount() {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     IntermediateStatement statement =
         new IntermediateStatement(
-            mock(OptionsMetadata.class), parse("update bar set foo=1"), connection);
+            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(result);
+    statement.updateResultCount(0, result);
 
-    assertFalse(statement.hasMoreData);
-    assertEquals(100L, statement.updateCount.longValue());
-    assertNull(statement.statementResult);
+    assertFalse(statement.hasMoreData[0]);
+    assertEquals(100L, statement.getUpdateCount(0));
+    assertNull(statement.getStatementResult(0));
   }
 
   @Test
   public void testUpdateResultCount_NoResult() {
+    when(connectionHandler.getSpannerConnection()).thenReturn(connection);
     IntermediateStatement statement =
         new IntermediateStatement(
             mock(OptionsMetadata.class),
             parse("create table bar (foo bigint primary key)"),
-            connection);
+            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(result);
+    statement.updateResultCount(0, result);
 
-    assertFalse(statement.hasMoreData);
-    assertNull(statement.updateCount);
-    assertNull(statement.statementResult);
+    assertFalse(statement.hasMoreData[0]);
+    assertEquals(0, statement.getUpdateCount(0));
+    assertNull(statement.getStatementResult(0));
   }
 }