Repository: GoogleCloudPlatform/pgadapter
Pull request #591 (2023-01-13 20:31:39)

fix: potential session leak when using prepared statements in autocommit mode
Creating multiple prepared statements in autocommit mode could cause a session leak in PGAdapter, as the underlying result sets were not closed. This again would hold on to the session, as there was no transaction that was being closed after creating the prepared statement, which would otherwise have returned the session to the pool.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#591](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (49dfb3c) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/5905213b9b662b1f1bb8afc1ef0ca4a5cdaf7097?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (5905213) will **decrease** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #591      +/-   ##
========================================================
- Coverage                 88.60%   88.59%   -0.01%     
- Complexity                 2103     2105       +2     
========================================================
  Files                       123      123              
  Lines                      6959     6964       +5     
  Branches                    982      982              
========================================================
+ Hits                       6166     6170       +4     
  Misses                      562      562              
- Partials                    231      232       +1     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `88.59% <100.00%> (-0.01%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `93.41% <ø> (ø)` | |
| [...er/pgadapter/statements/IntermediateStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvSW50ZXJtZWRpYXRlU3RhdGVtZW50LmphdmE=) | `98.85% <ø> (-0.08%)` | :arrow_down: |
| [...oud/spanner/pgadapter/metadata/DescribeResult.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL21ldGFkYXRhL0Rlc2NyaWJlUmVzdWx0LmphdmE=) | `100.00% <100.00%> (ø)` | |
| [...dapter/statements/IntermediatePortalStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvSW50ZXJtZWRpYXRlUG9ydGFsU3RhdGVtZW50LmphdmE=) | `93.75% <100.00%> (+1.25%)` | :arrow_up: |
| [...pter/statements/IntermediatePreparedStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvSW50ZXJtZWRpYXRlUHJlcGFyZWRTdGF0ZW1lbnQuamF2YQ==) | `95.08% <100.00%> (+0.25%)` | :arrow_up: |
| [...r/pgadapter/wireoutput/RowDescriptionResponse.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVvdXRwdXQvUm93RGVzY3JpcHRpb25SZXNwb25zZS5qYXZh) | `92.59% <100.00%> (ø)` | |
| [...panner/pgadapter/wireprotocol/DescribeMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9EZXNjcmliZU1lc3NhZ2UuamF2YQ==) | `85.50% <100.00%> (ø)` | |
| [...ava/com/google/cloud/spanner/pgadapter/Server.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/591?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL1NlcnZlci5qYXZh) | `58.06% <0.00%> (-3.23%)` | :arrow_down: |

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


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/DescribeResult.java b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/DescribeResult.java
index 8e43e050..80bb8467 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/DescribeResult.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/DescribeResult.java
@@ -16,6 +16,7 @@ package com.google.cloud.spanner.pgadapter.metadata;
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.error.SQLState;
 import com.google.cloud.spanner.pgadapter.parsers.Parser;
@@ -25,12 +26,12 @@ import javax.annotation.Nullable;
 
 @InternalApi
 public class DescribeResult {
-  @Nullable private final ResultSet resultSet;
+  @Nullable private final Type columns;
   private final int[] parameters;
 
   public DescribeResult(int[] givenParameterTypes, @Nullable ResultSet resultMetadata) {
-    this.resultSet = resultMetadata;
     this.parameters = extractParameters(givenParameterTypes, resultMetadata);
+    this.columns = resultMetadata == null ? null : resultMetadata.getType();
   }
 
   static int[] extractParameters(int[] givenParameterTypes, @Nullable ResultSet resultSet) {
@@ -80,7 +81,7 @@ public class DescribeResult {
     return parameters;
   }
 
-  public @Nullable ResultSet getResultSet() {
-    return resultSet;
+  public @Nullable Type getColumns() {
+    return columns;
   }
 }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
index ca733e3c..d3ba9aa3 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
@@ -615,7 +615,8 @@ public class BackendConnection {
     return execute.result;
   }
 
-  public Future<StatementResult> analyze(ParsedStatement parsedStatement, Statement statement) {
+  public ListenableFuture<StatementResult> analyze(
+      ParsedStatement parsedStatement, Statement statement) {
     Execute execute = new Execute(parsedStatement, statement, Function.identity(), true);
     bufferedStatements.add(execute);
     return execute.result;
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 80e3edd6..0c3a3b16 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
@@ -17,8 +17,10 @@ package com.google.cloud.spanner.pgadapter.statements;
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.connection.StatementResult.ResultType;
 import com.google.cloud.spanner.pgadapter.parsers.Parser;
 import com.google.cloud.spanner.pgadapter.parsers.Parser.FormatCode;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
 import java.util.List;
 import java.util.concurrent.Future;
 
@@ -89,6 +91,18 @@ public class IntermediatePortalStatement extends IntermediatePreparedStatement {
     }
   }
 
+  public void setStatementResult(StatementResult statementResult) {
+    super.setStatementResult(statementResult);
+    if (statementResult != null) {
+      if (statementResult.getResultType() == ResultType.RESULT_SET) {
+        this.hasMoreData = statementResult.getResultSet().next();
+      } else if (statementResult instanceof NoResult
+          && ((NoResult) statementResult).hasCommandTag()) {
+        this.commandTag = ((NoResult) statementResult).getCommandTag();
+      }
+    }
+  }
+
   /** Binds this portal to a set of parameter values. */
   public Statement bind(Statement statement) {
     // Make sure the results from any Describe message are propagated to the prepared statement
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 34e03e41..60cec223 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
@@ -24,6 +24,8 @@ import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeResult;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
@@ -97,18 +99,24 @@ public class IntermediatePreparedStatement extends IntermediateStatement {
 
   @Override
   public Future<StatementResult> describeAsync(BackendConnection backendConnection) {
-    Future<StatementResult> statementResultFuture =
+    ListenableFuture<StatementResult> statementResultFuture =
         backendConnection.analyze(this.parsedStatement, this.statement);
     setFutureStatementResult(statementResultFuture);
     this.describeResult =
-        Futures.lazyTransform(
+        Futures.transform(
             statementResultFuture,
-            result ->
-                new DescribeResult(
-                    this.givenParameterDataTypes,
-                    result.getResultType() == ResultType.RESULT_SET
-                        ? result.getResultSet()
-                        : null));
+            result -> {
+              DescribeResult describeResult;
+              if (result.getResultType() == ResultType.RESULT_SET) {
+                describeResult =
+                    new DescribeResult(this.givenParameterDataTypes, result.getResultSet());
+                result.getResultSet().close();
+              } else {
+                describeResult = new DescribeResult(this.givenParameterDataTypes, null);
+              }
+              return describeResult;
+            },
+            MoreExecutors.directExecutor());
     this.described = true;
     return statementResultFuture;
   }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/RowDescriptionResponse.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/RowDescriptionResponse.java
index 480f3ff5..47623f73 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/RowDescriptionResponse.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireoutput/RowDescriptionResponse.java
@@ -15,7 +15,6 @@
 package com.google.cloud.spanner.pgadapter.wireoutput;
 
 import com.google.api.core.InternalApi;
-import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
 import com.google.cloud.spanner.pgadapter.ProxyServer.DataFormat;
@@ -43,7 +42,7 @@ public class RowDescriptionResponse extends WireOutput {
   private static final byte DEFAULT_FLAG = 0;
 
   private final IntermediateStatement statement;
-  private final ResultSet resultSet;
+  private final Type columns;
   private final QueryMode mode;
   private final OptionsMetadata options;
   private final int columnCount;
@@ -51,25 +50,25 @@ public class RowDescriptionResponse extends WireOutput {
   public RowDescriptionResponse(
       DataOutputStream output,
       IntermediateStatement statement,
-      ResultSet resultSet,
+      Type columns,
       OptionsMetadata options,
       QueryMode mode)
       throws Exception {
-    super(output, calculateLength(resultSet));
+    super(output, calculateLength(columns));
     this.statement = statement;
-    this.resultSet = resultSet;
+    this.columns = columns;
     this.options = options;
     this.mode = mode;
-    this.columnCount = resultSet.getColumnCount();
+    this.columnCount = columns.getStructFields().size();
   }
 
-  private static int calculateLength(ResultSet resultSet) throws Exception {
+  private static int calculateLength(Type columns) {
     int length = HEADER_LENGTH + FIELD_NUMBER_LENGTH;
     for (int column_index = 0; /* columns start at 0 */
-        column_index < resultSet.getColumnCount();
+        column_index < columns.getStructFields().size();
         column_index++) {
       length +=
-          resultSet.getType().getStructFields().get(column_index).getName().length()
+          columns.getStructFields().get(column_index).getName().length()
               + NULL_TERMINATOR_LENGTH
               + TABLE_OID_LENGTH
               + COLUMN_INDEX_LENGTH
@@ -85,18 +84,18 @@ public class RowDescriptionResponse extends WireOutput {
   protected void sendPayload() throws Exception {
     this.outputStream.writeShort(this.columnCount);
     DataFormat defaultFormat = DataFormat.getDataFormat(0, this.statement, this.mode, this.options);
-    for (int column_index = 0; /* columns start at 0 */
-        column_index < this.columnCount;
-        column_index++) {
+    for (int columnIndex = 0; /* columns start at 0 */
+        columnIndex < this.columnCount;
+        columnIndex++) {
       this.outputStream.write(
-          this.resultSet.getType().getStructFields().get(column_index).getName().getBytes(UTF8));
+          this.columns.getStructFields().get(columnIndex).getName().getBytes(UTF8));
       // If it can be identified as a column of a table, the object ID of the table.
       this.outputStream.writeByte(DEFAULT_FLAG);
       // If it can be identified as a column of a table, the attribute number of the column
       this.outputStream.writeInt(DEFAULT_FLAG);
       // The object ID of the field's data type.
       this.outputStream.writeShort(DEFAULT_FLAG);
-      int oidType = getOidType(column_index);
+      int oidType = getOidType(columnIndex);
       this.outputStream.writeInt(oidType);
       this.outputStream.writeShort(getOidTypeSize(oidType));
       // The type modifier. The meaning of the modifier is type-specific.
@@ -105,7 +104,7 @@ public class RowDescriptionResponse extends WireOutput {
       short format =
           this.statement == null
               ? defaultFormat.getCode()
-              : this.statement.getResultFormatCode(column_index);
+              : this.statement.getResultFormatCode(columnIndex);
       this.outputStream.writeShort(format);
     }
   }
@@ -129,9 +128,9 @@ public class RowDescriptionResponse extends WireOutput {
             });
   }
 
-  int getOidType(int column_index) {
+  int getOidType(int columnIndex) {
     try {
-      Type type = this.resultSet.getColumnType(column_index);
+      Type type = this.columns.getStructFields().get(columnIndex).getType();
       return Parser.toOid(type);
     } catch (Exception ignored) {
       return Oid.UNSPECIFIED;
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 4edc9ba5..b308113e 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
@@ -147,7 +147,7 @@ public class DescribeMessage extends AbstractQueryProtocolMessage {
             new RowDescriptionResponse(
                     this.outputStream,
                     this.statement,
-                    getPortalMetadata().getResultSet(),
+                    getPortalMetadata().getResultSet().getType(),
                     this.connection.getServer().getOptions(),
                     this.queryMode)
                 .send(false);
@@ -191,11 +191,11 @@ public class DescribeMessage extends AbstractQueryProtocolMessage {
       if (isExtendedProtocol()) {
         DescribeResult metadata = this.statement.describe();
         new ParameterDescriptionResponse(this.outputStream, metadata.getParameters()).send(false);
-        if (metadata.getResultSet() != null && metadata.getResultSet().getColumnCount() > 0) {
+        if (metadata.getColumns() != null && metadata.getColumns().getStructFields().size() > 0) {
           new RowDescriptionResponse(
                   this.outputStream,
                   this.statement,
-                  metadata.getResultSet(),
+                  metadata.getColumns(),
                   this.connection.getServer().getOptions(),
                   this.queryMode)
               .send(false);


Test patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
index c870414f..8be705cc 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
@@ -33,7 +33,6 @@ import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeResult;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
-import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
 import com.google.cloud.spanner.pgadapter.utils.Converter;
 import com.google.cloud.spanner.pgadapter.wireoutput.DataRowResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.WireOutput;
@@ -251,14 +250,6 @@ public class IntermediateStatement {
 
   public void setStatementResult(StatementResult statementResult) {
     this.statementResult = statementResult;
-    if (statementResult != null) {
-      if (statementResult.getResultType() == ResultType.RESULT_SET) {
-        this.hasMoreData = statementResult.getResultSet().next();
-      } else if (statementResult instanceof NoResult
-          && ((NoResult) statementResult).hasCommandTag()) {
-        this.commandTag = ((NoResult) statementResult).getCommandTag();
-      }
-    }
   }
 
   protected void setFutureStatementResult(Future<StatementResult> result) {
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 f8925dd2..820f1ce9 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -737,6 +737,30 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  @Test(timeout = 60_000)
+  public void testMultiplePreparedStatements() throws SQLException {
+    // Execute more statements than there are sessions in the pool to verify that repeatedly
+    // creating a prepared statement with the same name does not cause a session leak.
+    final int numStatements = 1000;
+    String sql = "SELECT 1";
+
+    for (boolean autocommit : new boolean[] {true, false}) {
+      try (Connection connection = DriverManager.getConnection(createUrl())) {
+        // Verify that there's no session leak both in autocommit and transactional mode.
+        connection.setAutoCommit(autocommit);
+        // Force the use of prepared statements.
+        connection.unwrap(PGConnection.class).setPrepareThreshold(-1);
+        for (int i = 0; i < numStatements; i++) {
+          try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+            assertTrue(resultSet.next());
+            assertEquals(1L, resultSet.getLong(1));
+            assertFalse(resultSet.next());
+          }
+        }
+      }
+    }
+  }
+
   @Test
   public void testMultipleQueriesInTransaction() throws SQLException {
     String sql = "SELECT 1";
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 3cf44c76..44bce6c5 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
@@ -19,7 +19,6 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertSame;
-import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -102,14 +101,13 @@ public class IntermediateStatementTest {
         new IntermediateStatement(
             mock(OptionsMetadata.class), parse(sql), Statement.of(sql), connectionHandler);
     ResultSet resultSet = mock(ResultSet.class);
-    when(resultSet.next()).thenReturn(true, false);
     StatementResult result = mock(StatementResult.class);
-    when(result.getResultType()).thenReturn(ResultType.RESULT_SET);
     when(result.getResultSet()).thenReturn(resultSet);
 
     statement.setStatementResult(result);
 
-    assertTrue(statement.hasMoreData);
+    // Only IntermediatePortalStatements can return data.
+    assertFalse(statement.hasMoreData);
     assertEquals(-1, statement.getUpdateCount());
     assertSame(resultSet, statement.getStatementResult().getResultSet());
   }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/wireoutput/RowDescriptionTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/wireoutput/RowDescriptionTest.java
index 023610dd..5f0a6f6d 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/wireoutput/RowDescriptionTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/wireoutput/RowDescriptionTest.java
@@ -18,9 +18,7 @@ import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
-import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.ResultSet;
-import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.Type.StructField;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
@@ -60,15 +58,16 @@ public final class RowDescriptionTest {
 
   @Test
   public void OidTest() throws Exception {
-    when(metadata.getColumnCount()).thenReturn(0);
-    when(metadata.getColumnType(0)).thenReturn(Type.int64());
-    when(metadata.getColumnType(1)).thenReturn(Type.pgNumeric());
-    when(metadata.getColumnType(2)).thenReturn(Type.float64());
-    when(metadata.getColumnType(3)).thenReturn(Type.string());
-    when(metadata.getColumnType(4)).thenReturn(Type.bytes());
-    when(metadata.getColumnType(5)).thenReturn(Type.bool());
-    when(metadata.getColumnType(6)).thenReturn(Type.date());
-    when(metadata.getColumnType(7)).thenReturn(Type.timestamp());
+    Type type =
+        Type.struct(
+            StructField.of("f1", Type.int64()),
+            StructField.of("f1", Type.pgNumeric()),
+            StructField.of("f1", Type.float64()),
+            StructField.of("f1", Type.string()),
+            StructField.of("f1", Type.bytes()),
+            StructField.of("f1", Type.bool()),
+            StructField.of("f1", Type.date()),
+            StructField.of("f1", Type.timestamp()));
 
     JSONParser parser = new JSONParser();
     JSONObject commandMetadata = (JSONObject) parser.parse(EMPTY_COMMAND_JSON);
@@ -84,7 +83,7 @@ public final class RowDescriptionTest {
             commandMetadata);
     QueryMode mode = QueryMode.SIMPLE;
     RowDescriptionResponse response =
-        new RowDescriptionResponse(output, statement, metadata, options, mode);
+        new RowDescriptionResponse(output, statement, type, options, mode);
 
     // Types.BIGINT
     assertEquals(Oid.INT8, response.getOidType(0));
@@ -123,10 +122,7 @@ public final class RowDescriptionTest {
   public void SendPayloadNullStatementTest() throws Exception {
     int COLUMN_COUNT = 1;
     String COLUMN_NAME = "default-column-name";
-    Type rowType = Type.struct(StructField.of(COLUMN_NAME, Type.string()));
-    when(metadata.getColumnCount()).thenReturn(COLUMN_COUNT);
-    when(metadata.getType()).thenReturn(rowType);
-    when(metadata.getColumnType(Mockito.anyInt())).thenReturn(Type.int64());
+    Type rowType = Type.struct(StructField.of(COLUMN_NAME, Type.int64()));
 
     JSONParser parser = new JSONParser();
     JSONObject commandMetadata = (JSONObject) parser.parse(EMPTY_COMMAND_JSON);
@@ -142,7 +138,7 @@ public final class RowDescriptionTest {
             commandMetadata);
     QueryMode mode = QueryMode.EXTENDED;
     RowDescriptionResponse response =
-        new RowDescriptionResponse(output, null, metadata, options, mode);
+        new RowDescriptionResponse(output, null, rowType, options, mode);
     response.sendPayload();
     DataInputStream outputReader =
         new DataInputStream(new ByteArrayInputStream(buffer.toByteArray()));
@@ -176,9 +172,6 @@ public final class RowDescriptionTest {
         Type.struct(
             StructField.of("numeric-text", Type.pgNumeric()),
             StructField.of("numeric-binary", Type.pgNumeric()));
-    when(metadata.getColumnCount()).thenReturn(rowType.getStructFields().size());
-    when(metadata.getType()).thenReturn(rowType);
-    when(metadata.getColumnType(Mockito.anyInt())).thenReturn(Type.pgNumeric());
     when(statement.getResultFormatCode(0)).thenReturn((short) 0);
     when(statement.getResultFormatCode(1)).thenReturn((short) 1);
     JSONParser parser = new JSONParser();
@@ -195,7 +188,7 @@ public final class RowDescriptionTest {
             commandMetadata);
     QueryMode mode = QueryMode.EXTENDED;
     RowDescriptionResponse response =
-        new RowDescriptionResponse(output, statement, metadata, options, mode);
+        new RowDescriptionResponse(output, statement, rowType, options, mode);
     response.sendPayload();
     DataInputStream outputReader =
         new DataInputStream(new ByteArrayInputStream(buffer.toByteArray()));
@@ -247,10 +240,7 @@ public final class RowDescriptionTest {
   public void SendPayloadStatementWithBinaryOptionTest() throws Exception {
     int COLUMN_COUNT = 1;
     String COLUMN_NAME = "default-column-name";
-    Type rowType = Type.struct(StructField.of(COLUMN_NAME, Type.string()));
-    when(metadata.getColumnCount()).thenReturn(COLUMN_COUNT);
-    when(metadata.getType()).thenReturn(rowType);
-    when(metadata.getColumnType(Mockito.anyInt())).thenReturn(Type.int64());
+    Type rowType = Type.struct(StructField.of(COLUMN_NAME, Type.int64()));
     when(statement.getResultFormatCode(Mockito.anyInt())).thenReturn((short) 1);
     JSONParser parser = new JSONParser();
     JSONObject commandMetadata = (JSONObject) parser.parse(EMPTY_COMMAND_JSON);
@@ -266,7 +256,7 @@ public final class RowDescriptionTest {
             commandMetadata);
     QueryMode mode = QueryMode.EXTENDED;
     RowDescriptionResponse response =
-        new RowDescriptionResponse(output, statement, metadata, options, mode);
+        new RowDescriptionResponse(output, statement, rowType, options, mode);
     response.sendPayload();
     DataInputStream outputReader =
         new DataInputStream(new ByteArrayInputStream(buffer.toByteArray()));
@@ -300,10 +290,7 @@ public final class RowDescriptionTest {
     String COLUMN_NAME = "default-column-name";
     Type rowType =
         Type.struct(
-            StructField.of(COLUMN_NAME, Type.string()), StructField.of(COLUMN_NAME, Type.string()));
-    when(metadata.getColumnCount()).thenReturn(COLUMN_COUNT);
-    when(metadata.getType()).thenReturn(rowType);
-    when(metadata.getColumnType(Mockito.anyInt())).thenReturn(Type.int64());
+            StructField.of(COLUMN_NAME, Type.int64()), StructField.of(COLUMN_NAME, Type.int64()));
     when(statement.getResultFormatCode(Mockito.anyInt()))
         .thenReturn((short) 0)
         .thenReturn((short) 0)
@@ -322,7 +309,7 @@ public final class RowDescriptionTest {
             commandMetadata);
     QueryMode mode = QueryMode.EXTENDED;
     RowDescriptionResponse response =
-        new RowDescriptionResponse(output, statement, metadata, options, mode);
+        new RowDescriptionResponse(output, statement, rowType, options, mode);
     response.sendPayload();
     DataInputStream outputReader =
         new DataInputStream(new ByteArrayInputStream(buffer.toByteArray()));
@@ -353,15 +340,13 @@ public final class RowDescriptionTest {
 
   @Test
   public void testUnknownTypeReturnsUnspecified() throws Exception {
-    ResultSet resultSet = mock(ResultSet.class);
-    when(resultSet.getColumnType(0))
-        .thenThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.INTERNAL, "unknown type"));
+    Type type = Type.struct();
 
     RowDescriptionResponse response =
         new RowDescriptionResponse(
             mock(DataOutputStream.class),
             mock(IntermediateStatement.class),
-            resultSet,
+            type,
             mock(OptionsMetadata.class),
             QueryMode.SIMPLE);