fix: Statements with no results would return an error (#57)

Statements that return no results, such as DDL statements, would cause an error as PGAdapter tried to get the update count of these. Fixes #56

Author: olavloite

src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java

@@ -203,13 +203,23 @@ public Exception getException() {
  * of updateBatchResultCount.
  */
  protected void updateResultCount(StatementResult result) {
- if (result.getResultType() == StatementResult.ResultType.RESULT_SET) {
- this.statementResult = result.getResultSet();
- this.hasMoreData = this.statementResult.next();
- } else {
- this.updateCount = result.getUpdateCount();
- this.hasMoreData = false;
- this.statementResult = null;
+ switch (result.getResultType()) {
+ case RESULT_SET:
+ this.statementResult = result.getResultSet();
+ this.hasMoreData = this.statementResult.next();
+ break;
+ case UPDATE_COUNT:
+ this.updateCount = result.getUpdateCount();
+ this.hasMoreData = false;
+ this.statementResult = null;
+ break;
+ case NO_RESULT:
+ this.hasMoreData = false;
+ this.statementResult = null;
+ break;
+ default:
+ throw SpannerExceptionFactory.newSpannerException(
+ ErrorCode.INTERNAL, "Unknown or unsupported result type: " + result.getResultType());
  }
  }
 

src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java

@@ -26,6 +26,7 @@
 import com.google.cloud.spanner.Statement;
 import com.google.protobuf.ListValue;
 import com.google.protobuf.Value;
+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;
@@ -56,6 +57,7 @@
 import java.util.Base64;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -477,4 +479,27 @@ public void testTwoQueries() throws SQLException {
  }
  }
  }
+
+ @Test
+ public void testDdl() throws SQLException {
+ String sql = "CREATE TABLE foo (id bigint primary key)";
+
+ try (Connection connection = DriverManager.getConnection(createUrl())) {
+ try (java.sql.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());
+ }
+ }
+
+ 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(sql, updateDatabaseDdlRequests.get(0).getStatements(0));
+ }
 }

src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java

@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import java.math.BigDecimal;
@@ -33,7 +34,9 @@
 import java.time.LocalDateTime;
 import java.time.OffsetDateTime;
 import java.time.ZoneOffset;
+import java.util.List;
 import java.util.TimeZone;
+import java.util.stream.Collectors;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -213,4 +216,27 @@ public void testSelectAtStartOfBatch() throws SQLException {
  }
  }
  }
+
+ @Test
+ public void testDdl() throws SQLException {
+ String sql = "CREATE TABLE foo (id bigint primary key)";
+
+ 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());
+ }
+ }
+
+ 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(sql, updateDatabaseDdlRequests.get(0).getStatements(0));
+ }
 }

src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java

@@ -28,6 +28,7 @@
 import com.google.cloud.spanner.Statement;
 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.ConnectionStatus;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
@@ -155,6 +156,7 @@ public void testQueryMessage() throws Exception {
  String expectedSQL = "SELECT * FROM users";
 
  Mockito.when(connection.execute(Statement.of(expectedSQL))).thenReturn(statementResult);
+ when(statementResult.getResultType()).thenReturn(ResultType.RESULT_SET);
  when(statementResult.getResultSet()).thenReturn(resultSet);
  Mockito.when(connectionHandler.getServer()).thenReturn(server);
  Mockito.when(server.getOptions()).thenReturn(options);
@@ -189,6 +191,7 @@ public void testQueryUsesPSQLStatementWhenPSQLModeSelectedMessage() throws Excep
  String expectedSQL = "SELECT * FROM users";
 
  when(connection.execute(Statement.of(expectedSQL))).thenReturn(statementResult);
+ when(statementResult.getResultType()).thenReturn(ResultType.RESULT_SET);
  Mockito.when(statementResult.getResultSet()).thenReturn(resultSet);
  Mockito.when(connectionHandler.getServer()).thenReturn(server);
  Mockito.when(server.getOptions()).thenReturn(options);
@@ -1120,8 +1123,8 @@ public void testQueryMessageInTransaction() throws Exception {
 
  when(connection.isInTransaction()).thenReturn(true);
  when(connectionHandler.getSpannerConnection()).thenReturn(connection);
- ResultSet resultSet = mock(ResultSet.class);
- when(statementResult.getResultSet()).thenReturn(resultSet);
+ when(statementResult.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
+ when(statementResult.getUpdateCount()).thenReturn(1L);
  when(connection.execute(Statement.of(expectedSQL))).thenReturn(statementResult);
  when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
  when(connectionHandler.getServer()).thenReturn(server);
@@ -1143,11 +1146,11 @@ public void testQueryMessageInTransaction() throws Exception {
  assertEquals('\0', outputResult.readByte());
  assertEquals('\0', outputResult.readByte());
  assertEquals('\0', outputResult.readByte());
- // 15 = 4 + "INSERT".length() + " 0 0".length() + 1 (header + command length + null terminator)
+ // 15 = 4 + "INSERT".length() + " 0 1".length() + 1 (header + command length + null terminator)
  assertEquals(15, outputResult.readByte());
  byte[] command = new byte[10];
  assertEquals(10, outputResult.read(command, 0, 10));
- assertEquals("INSERT 0 0", new String(command));
+ assertEquals("INSERT 0 1", new String(command));
  assertEquals('\0', outputResult.readByte());
  // ReadyResponse in transaction ('T')
  assertEquals('Z', outputResult.readByte());

src/test/java/com/google/cloud/spanner/pgadapter/StatementTest.java

@@ -16,6 +16,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.when;
@@ -136,7 +137,7 @@ public void testBasicUpdateStatement() throws Exception {
  assertEquals(intermediateStatement.getUpdateCount().longValue(), 1L);
  assertTrue(intermediateStatement.isExecuted());
  assertEquals(intermediateStatement.getResultType(), ResultType.UPDATE_COUNT);
- Assert.assertNull(intermediateStatement.getStatementResult());
+ assertNull(intermediateStatement.getStatementResult());
  assertFalse(intermediateStatement.isHasMoreData());
  assertFalse(intermediateStatement.hasException());
  assertEquals(intermediateStatement.getResultFormatCode(0), 0);
@@ -167,7 +168,7 @@ public void testBasicZeroUpdateCountResultStatement() throws Exception {
  assertEquals(intermediateStatement.getUpdateCount().longValue(), 0L);
  assertTrue(intermediateStatement.isExecuted());
  assertEquals(intermediateStatement.getResultType(), ResultType.UPDATE_COUNT);
- Assert.assertNull(intermediateStatement.getStatementResult());
+ assertNull(intermediateStatement.getStatementResult());
  assertFalse(intermediateStatement.isHasMoreData());
  assertFalse(intermediateStatement.hasException());
  assertEquals(intermediateStatement.getResultFormatCode(0), 0);
@@ -180,7 +181,7 @@ public void testBasicZeroUpdateCountResultStatement() throws Exception {
  @Test
  public void testBasicNoResultStatement() throws Exception {
  when(statementResult.getResultType()).thenReturn(StatementResult.ResultType.NO_RESULT);
- when(statementResult.getUpdateCount()).thenReturn(0L);
+ when(statementResult.getUpdateCount()).thenThrow(new IllegalStateException());
  when(connection.execute(Statement.of("CREATE TABLE users (name varchar(100) primary key)")))
  .thenReturn(statementResult);
 
@@ -195,10 +196,10 @@ public void testBasicNoResultStatement() throws Exception {
  Mockito.verify(connection, Mockito.times(1))
  .execute(Statement.of("CREATE TABLE users (name varchar(100) primary key)"));
  assertFalse(intermediateStatement.containsResultSet());
- assertEquals(intermediateStatement.getUpdateCount().longValue(), 0L);
+ assertNull(intermediateStatement.getUpdateCount());
  assertTrue(intermediateStatement.isExecuted());
  assertEquals(intermediateStatement.getResultType(), ResultType.NO_RESULT);
- Assert.assertNull(intermediateStatement.getStatementResult());
+ assertNull(intermediateStatement.getStatementResult());
  assertFalse(intermediateStatement.isHasMoreData());
  assertFalse(intermediateStatement.hasException());
  assertEquals(intermediateStatement.getResultFormatCode(0), 0);

src/test/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatementTest.java

@@ -0,0 +1,85 @@
+// 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.statements;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.connection.Connection;
+import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.connection.StatementResult.ResultType;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+
+@RunWith(JUnit4.class)
+public class IntermediateStatementTest {
+ @Mock private Connection connection;
+
+ @Test
+ public void testUpdateResultCount_ResultSet() {
+ IntermediateStatement statement = new IntermediateStatement("select foo from bar", connection);
+ 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);
+ when(result.getUpdateCount()).thenThrow(new IllegalStateException());
+
+ statement.updateResultCount(result);
+
+ assertTrue(statement.hasMoreData);
+ assertNull(statement.updateCount);
+ assertSame(resultSet, statement.statementResult);
+ }
+
+ @Test
+ public void testUpdateResultCount_UpdateCount() {
+ IntermediateStatement statement = new IntermediateStatement("update bar set foo=1", connection);
+ 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);
+
+ assertFalse(statement.hasMoreData);
+ assertEquals(100L, statement.updateCount.longValue());
+ assertNull(statement.statementResult);
+ }
+
+ @Test
+ public void testUpdateResultCount_NoResult() {
+ IntermediateStatement statement =
+ new IntermediateStatement("create table bar (foo bigint primary key)", connection);
+ 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);
+
+ assertFalse(statement.hasMoreData);
+ assertNull(statement.updateCount);
+ assertNull(statement.statementResult);
+ }
+}