[CALCITE-6781] The isUpdateCapable method of calcite.avatica will incorrectly traverse the returned result value

Author: caicancai

core/src/main/java/org/apache/calcite/avatica/AvaticaConnection.java

@@ -614,18 +614,19 @@ private void isUpdateCapable(final AvaticaStatement statement)
       return;
     }
     if (signature.statementType.canUpdate() && statement.updateCount == -1) {
-      statement.openResultSet.next();
-      Object obj = statement.openResultSet.getObject(ROWCOUNT_COLUMN_NAME);
-      if (obj instanceof Number) {
-        statement.updateCount = ((Number) obj).intValue();
-      } else if (obj instanceof List) {
-        @SuppressWarnings("unchecked")
-        final List<Number> numbers = (List<Number>) obj;
-        statement.updateCount = numbers.get(0).intValue();
-      } else {
-        throw HELPER.createException("Not a valid return result.");
+      if (statement.openResultSet.next()) {
+        Object obj = statement.openResultSet.getObject(ROWCOUNT_COLUMN_NAME);
+        if (obj instanceof Number) {
+          statement.updateCount = ((Number) obj).intValue();
+        } else if (obj instanceof List) {
+          @SuppressWarnings("unchecked")
+          final List<Number> numbers = (List<Number>) obj;
+          statement.updateCount = numbers.get(0).intValue();
+        } else {
+          throw HELPER.createException("Not a valid return result.");
+        }
+        statement.openResultSet = null;
       }
-      statement.openResultSet = null;
     }
   }
 

core/src/test/java/org/apache/calcite/avatica/AvaticaConnectionTest.java

@@ -20,7 +20,10 @@
 import org.junit.Test;
 import org.mockito.Mockito;
 
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.sql.SQLException;
+import java.util.Collections;
 import java.util.Properties;
 
 /**
@@ -71,6 +74,49 @@ public void testNumExecuteRetries() {
     Assert.assertEquals(10, connection.getNumStatementRetries(props));
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6781">[CALCITE-6781]
+   * The isUpdateCapable method of calcite.avatica will incorrectly traverse
+   * the returned result value</a>.
+   */
+  @Test
+  public void testIsUpdateCapableSkipsRowCountWhenResultSetHasNoRows() throws Exception {
+    AvaticaConnection connection = Mockito.mock(
+        AvaticaConnection.class, Mockito.CALLS_REAL_METHODS);
+    AvaticaStatement statement = Mockito.mock(AvaticaStatement.class);
+    AvaticaResultSet resultSet = Mockito.mock(AvaticaResultSet.class);
+
+    Meta.Signature signature = new Meta.Signature(Collections.<ColumnMetaData>emptyList(), null,
+        Collections.<AvaticaParameter>emptyList(), Collections.<String, Object>emptyMap(), null,
+        Meta.StatementType.INSERT);
+
+    Mockito.when(statement.getSignature()).thenReturn(signature);
+    Mockito.when(resultSet.next()).thenReturn(false);
+    statement.updateCount = -1;
+    statement.openResultSet = resultSet;
+
+    invokeIsUpdateCapable(connection, statement);
+
+    Assert.assertEquals(-1, statement.updateCount);
+    Assert.assertSame(resultSet, statement.openResultSet);
+    Mockito.verify(resultSet, Mockito.never()).getObject(AvaticaConnection.ROWCOUNT_COLUMN_NAME);
+  }
+
+  private static void invokeIsUpdateCapable(
+      AvaticaConnection connection, AvaticaStatement statement) throws Exception {
+    Method method = AvaticaConnection.class
+        .getDeclaredMethod("isUpdateCapable", AvaticaStatement.class);
+    method.setAccessible(true);
+    try {
+      method.invoke(connection, statement);
+    } catch (InvocationTargetException e) {
+      if (e.getCause() instanceof SQLException) {
+        throw (SQLException) e.getCause();
+      }
+      throw e;
+    }
+  }
+
 }
 
 // End AvaticaConnectionTest.java