diff --git a/core/src/main/java/org/apache/calcite/avatica/AvaticaConnection.java b/core/src/main/java/org/apache/calcite/avatica/AvaticaConnection.java index f740d85ff..bb067a6f7 100644 --- a/core/src/main/java/org/apache/calcite/avatica/AvaticaConnection.java +++ b/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 numbers = (List) 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 numbers = (List) obj; + statement.updateCount = numbers.get(0).intValue(); + } else { + throw HELPER.createException("Not a valid return result."); + } + statement.openResultSet = null; } - statement.openResultSet = null; } } diff --git a/core/src/test/java/org/apache/calcite/avatica/AvaticaConnectionTest.java b/core/src/test/java/org/apache/calcite/avatica/AvaticaConnectionTest.java index 9482eba2a..463711a10 100644 --- a/core/src/test/java/org/apache/calcite/avatica/AvaticaConnectionTest.java +++ b/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 + * [CALCITE-6781] + * The isUpdateCapable method of calcite.avatica will incorrectly traverse + * the returned result value. + */ + @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.emptyList(), null, + Collections.emptyList(), Collections.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