diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java index 9c7112f1c3..21cc3e431f 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java @@ -62,14 +62,17 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p parameterSchema == null ? new ArrayList<>() : ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields()); - + StatementType statementType = + resultSetSchema == null || resultSetSchema.getFields().isEmpty() + ? StatementType.IS_DML + : StatementType.SELECT; return new Signature( columnMetaData, sql, parameters, Collections.emptyMap(), null, // unnecessary, as SQL requests use ArrowFlightJdbcCursor - StatementType.SELECT); + statementType); } @Override @@ -105,7 +108,8 @@ public ExecuteResult execute( preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator()) .bind(typedValues); - if (statementHandle.signature == null) { + if (statementHandle.signature == null + || statementHandle.signature.statementType == StatementType.IS_DML) { // Update query long updatedCount = preparedStatement.executeUpdate(); return new ExecuteResult( diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java index 774ad0081e..0369c3a162 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java @@ -20,6 +20,8 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; import java.sql.Connection; @@ -83,6 +85,19 @@ public void testSimpleQueryNoParameterBinding() throws SQLException { } } + @Test + public void testSimpleQueryNoParameterBindingWithExecute() throws SQLException { + final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; + try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) { + boolean isResultSet = preparedStatement.execute(); + assertTrue(isResultSet); + final ResultSet resultSet = preparedStatement.getResultSet(); + CoreMockedSqlProducers.assertLegacyRegularSqlResultSet(resultSet); + assertFalse(preparedStatement.getMoreResults()); + assertEquals(-1, preparedStatement.getUpdateCount()); + } + } + @Test public void testQueryWithParameterBinding() throws SQLException { final String query = "Fake query with parameters"; @@ -174,6 +189,20 @@ public void testUpdateQuery() throws SQLException { } } + @Test + public void testUpdateQueryWithExecute() throws SQLException { + String query = "Fake update with execute"; + PRODUCER.addUpdateQuery(query, /*updatedRows*/ 42); + try (final PreparedStatement stmt = connection.prepareStatement(query)) { + boolean isResultSet = stmt.execute(); + assertFalse(isResultSet); + int updated = stmt.getUpdateCount(); + assertEquals(42, updated); + assertFalse(stmt.getMoreResults()); + assertEquals(-1, stmt.getUpdateCount()); + } + } + @Test public void testUpdateQueryWithParameters() throws SQLException { String query = "Fake update with parameters";