fix: command should determine result type (#29)

* fix: command should determine result type

The result type was determined by looking at the update count returned
by the JDBC statement that was executed on Cloud Spanner. This is
however wrong, as:
1. An update count of 0 does not mean 'no result'. It means that zero
   rows were updated/deleted.
2. The result type is determined by the type of statement. A DML
   statement will always return an update count. A DDL statement will
   always return no result. A query will always return a result set.

* fix: remove unused code + add missing doc

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

@@ -37,7 +37,7 @@ public class IntermediatePreparedStatement extends IntermediateStatement {
   protected List<Integer> parameterDataTypes;
 
   public IntermediatePreparedStatement(String sql, Connection connection) throws SQLException {
-    super();
+    super(sql);
     SQLMetadata parsedSQL = Converter.toJDBCParams(sql);
     this.parameterCount = parsedSQL.getParameterCount();
     this.sql = parsedSQL.getSqlString();
@@ -54,7 +54,7 @@ protected IntermediatePreparedStatement(
       int totalParameters,
       SetMultimap<Integer, Integer> parameterIndexToPositions,
       Connection connection) {
-    super();
+    super(sql);
     this.sql = sql;
     this.command = parseCommand(sql);
     this.parameterCount = totalParameters;

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

@@ -14,7 +14,8 @@
 
 package com.google.cloud.spanner.pgadapter.statements;
 
-import com.google.cloud.spanner.jdbc.JdbcConstants;
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.connection.AbstractStatementParser;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeMetadata;
 import com.google.common.base.Preconditions;
 import java.sql.Connection;
@@ -30,9 +31,11 @@
  * statements which does not belong directly to Postgres, Spanner, etc.
  */
 public class IntermediateStatement {
+  private static final AbstractStatementParser PARSER =
+      AbstractStatementParser.getInstance(Dialect.POSTGRESQL);
 
   protected Statement statement;
-  protected ResultType resultType;
+  private final ResultType resultType;
   protected ResultSet statementResult;
   protected boolean hasMoreData;
   protected Exception exception;
@@ -47,40 +50,37 @@ public class IntermediateStatement {
   private static final char SINGLE_QUOTE = '\'';
 
   public IntermediateStatement(String sql, Connection connection) throws SQLException {
-    this();
     this.sql = sql;
     this.statements = parseStatements(sql);
     this.command = parseCommand(sql);
     this.connection = connection;
     this.statement = connection.createStatement();
+    // 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(sql);
   }
 
-  protected IntermediateStatement() {
-    this.executed = false;
-    this.exception = null;
-    this.resultType = null;
-    this.hasMoreData = false;
-    this.statementResult = null;
-    this.updateCount = null;
+  protected IntermediateStatement(String sql) {
+    this.resultType = determineResultType(sql);
   }
 
   /**
-   * Extracts what type of result exists within the statement. In JDBC a statement update count is
-   * positive if it is an update statement, 0 if there is no result, or negative if there are
-   * results (i.e.: select statement)
+   * 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.
    *
-   * @param statement The resulting statement from an execution.
-   * @return The statement result type.
-   * @throws SQLException If getUpdateCount fails.
+   * @param sql The sql string to determine the type of result for
+   * @return The {@link ResultType} that the given sql string will produce
    */
-  private static ResultType extractResultType(Statement statement) throws SQLException {
-    switch (statement.getUpdateCount()) {
-      case JdbcConstants.STATEMENT_NO_RESULT:
-        return ResultType.NO_RESULT;
-      case JdbcConstants.STATEMENT_RESULT_SET:
-        return ResultType.RESULT_SET;
-      default:
-        return ResultType.UPDATE_COUNT;
+  protected static ResultType determineResultType(String sql) {
+    if (PARSER.isUpdateStatement(sql)) {
+      return ResultType.UPDATE_COUNT;
+    } else if (PARSER.isQuery(sql)) {
+      return ResultType.RESULT_SET;
+    } else {
+      return ResultType.NO_RESULT;
     }
   }
 
@@ -209,7 +209,6 @@ public Exception getException() {
    * @throws SQLException If an issue occurred in extracting result metadata.
    */
   protected void updateResultCount() throws SQLException {
-    this.resultType = IntermediateStatement.extractResultType(this.statement);
     if (this.containsResultSet()) {
       this.statementResult = this.statement.getResultSet();
       this.hasMoreData = this.statementResult.next();
@@ -221,7 +220,6 @@ protected void updateResultCount() throws SQLException {
   }
 
   protected void updateBatchResultCount(int[] updateCounts) throws SQLException {
-    this.resultType = IntermediateStatement.extractResultType(this.statement);
     this.updateCount = 0;
     for (int i = 0; i < updateCounts.length; ++i) {
       this.updateCount += updateCounts[i];
@@ -239,7 +237,6 @@ protected void handleExecutionException(SQLException e) {
     this.exception = e;
     this.hasMoreData = false;
     this.statementResult = null;
-    this.resultType = ResultType.NO_RESULT;
   }
 
   /** Execute the SQL statement, storing metadata. */

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

@@ -30,7 +30,7 @@ public class MatcherStatement extends IntermediateStatement {
   private JSONObject commandMetadataJSON;
 
   public MatcherStatement(String sql, ConnectionHandler connectionHandler) throws SQLException {
-    super();
+    super(sql);
     this.connection = connectionHandler.getJdbcConnection();
     this.statement = this.connection.createStatement();
     this.commandMetadataJSON = connectionHandler.getServer().getOptions().getCommandMetadataJSON();

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

@@ -29,7 +29,7 @@ public class PSQLStatement extends IntermediateStatement {
   private JSONObject commandMetadataJSON;
 
   public PSQLStatement(String sql, ConnectionHandler connectionHandler) throws SQLException {
-    super();
+    super(sql);
     this.connection = connectionHandler.getJdbcConnection();
     this.statement = this.connection.createStatement();
     this.commandMetadataJSON = connectionHandler.getServer().getOptions().getCommandMetadataJSON();

src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/DescribeMessage.java

@@ -87,13 +87,21 @@ public void handleDescribePortal() throws Exception {
     if (this.statement.hasException()) {
       this.handleError(this.statement.getException());
     } else {
-      new RowDescriptionResponse(
-              this.outputStream,
-              this.statement,
-              ((DescribePortalMetadata) this.statement.describe()).getMetadata(),
-              this.connection.getServer().getOptions(),
-              QueryMode.EXTENDED)
-          .send();
+      switch (this.statement.getResultType()) {
+        case UPDATE_COUNT:
+        case NO_RESULT:
+          new NoDataResponse(this.outputStream).send();
+          break;
+        case RESULT_SET:
+          new RowDescriptionResponse(
+                  this.outputStream,
+                  this.statement,
+                  ((DescribePortalMetadata) this.statement.describe()).getMetadata(),
+                  this.connection.getServer().getOptions(),
+                  QueryMode.EXTENDED)
+              .send();
+          break;
+      }
     }
   }
 

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

@@ -23,12 +23,12 @@
 import static org.mockito.Mockito.when;
 
 import com.google.cloud.spanner.jdbc.CloudSpannerJdbcConnection;
-import com.google.cloud.spanner.jdbc.JdbcConstants;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.DescribePortalMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeStatementMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.TextFormat;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePortalStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
@@ -61,6 +61,8 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -135,6 +137,7 @@ public void testQueryMessage() throws Exception {
     String expectedSQL = "SELECT * FROM users";
 
     Mockito.when(connection.createStatement()).thenReturn(statement);
+    Mockito.when(statement.getResultSet()).thenReturn(mock(ResultSet.class));
     Mockito.when(connectionHandler.getServer()).thenReturn(server);
     Mockito.when(server.getOptions()).thenReturn(options);
     Mockito.when(options.requiresMatcher()).thenReturn(false);
@@ -168,6 +171,7 @@ public void testQueryUsesPSQLStatementWhenPSQLModeSelectedMessage() throws Excep
     String expectedSQL = "SELECT * FROM users";
 
     Mockito.when(connection.createStatement()).thenReturn(statement);
+    Mockito.when(statement.getResultSet()).thenReturn(mock(ResultSet.class));
     Mockito.when(connectionHandler.getServer()).thenReturn(server);
     Mockito.when(server.getOptions()).thenReturn(options);
     Mockito.when(options.requiresMatcher()).thenReturn(true);
@@ -1076,28 +1080,29 @@ public void testFlushMessageInTransaction() throws Exception {
 
   @Test
   public void testQueryMessageInTransaction() throws Exception {
-    byte[] messageMetadata = {'Q', 0, 0, 0, 24};
-    String payload = "SELECT * FROM users\0";
+    byte[] messageMetadata = {'Q', 0, 0, 0, 45};
+    String payload = "INSERT INTO users (name) VALUES ('test')\0";
     byte[] value = Bytes.concat(messageMetadata, payload.getBytes());
 
     DataInputStream inputStream = new DataInputStream(new ByteArrayInputStream(value));
     ByteArrayOutputStream result = new ByteArrayOutputStream();
     DataOutputStream outputStream = new DataOutputStream(result);
 
-    String expectedSQL = "SELECT * FROM users";
+    String expectedSQL = "INSERT INTO users (name) VALUES ('test')";
 
     CloudSpannerJdbcConnection cloudSpannerConnection = mock(CloudSpannerJdbcConnection.class);
     when(cloudSpannerConnection.isInTransaction()).thenReturn(true);
     when(connectionHandler.getJdbcConnection()).thenReturn(connection);
     when(connection.unwrap(CloudSpannerJdbcConnection.class)).thenReturn(cloudSpannerConnection);
     when(connection.createStatement()).thenReturn(statement);
-    // TODO: Remove the following mock result, as this is wrong, but is currently needed because of
-    // a bug in the way that PgAdapter determines the result type of a statement. That is fixed in
-    // https://github.com/GoogleCloudPlatform/pgadapter/pull/23
-    when(statement.getUpdateCount()).thenReturn(JdbcConstants.STATEMENT_NO_RESULT);
+    ResultSet resultSet = mock(ResultSet.class);
+    when(statement.getResultSet()).thenReturn(resultSet);
+    when(resultSet.getMetaData()).thenReturn(mock(ResultSetMetaData.class));
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
     when(connectionHandler.getServer()).thenReturn(server);
-    when(server.getOptions()).thenReturn(mock(OptionsMetadata.class));
+    OptionsMetadata options = mock(OptionsMetadata.class);
+    when(server.getOptions()).thenReturn(options);
+    when(options.getTextFormat()).thenReturn(TextFormat.POSTGRESQL);
     when(connectionMetadata.getInputStream()).thenReturn(inputStream);
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
 
@@ -1113,11 +1118,11 @@ public void testQueryMessageInTransaction() throws Exception {
     assertEquals('\0', outputResult.readByte());
     assertEquals('\0', outputResult.readByte());
     assertEquals('\0', outputResult.readByte());
-    // 11 = 4 + "SELECT".length() + 1 (header + command length + null terminator)
-    assertEquals(11, outputResult.readByte());
-    byte[] command = new byte[6];
-    assertEquals(6, outputResult.read(command, 0, 6));
-    assertEquals("SELECT", new String(command));
+    // 15 = 4 + "INSERT".length() + " 0 0".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('\0', outputResult.readByte());
     // ReadyResponse in transaction ('T')
     assertEquals('Z', outputResult.readByte());