diff --git a/src/integrationTest/java/integration/tests/StatementTest.java b/src/integrationTest/java/integration/tests/StatementTest.java index 42a441cba..f8a0579c6 100644 --- a/src/integrationTest/java/integration/tests/StatementTest.java +++ b/src/integrationTest/java/integration/tests/StatementTest.java @@ -30,9 +30,11 @@ import java.util.stream.IntStream; import static java.lang.String.format; +import static java.sql.Statement.SUCCESS_NO_INFO; import static java.util.stream.Collectors.joining; import static java.util.stream.Collectors.toList; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -150,6 +152,17 @@ private List selectIntValues(Connection connection, int limit) throws S return result; } + @Test + void shouldExecuteBatch() throws SQLException { + int size = 10; + try (Connection connection = createConnection(); Statement insert = connection.createStatement()) { + for (int i = 0; i < size; i++) { + insert.addBatch(format("INSERT INTO statement_test(id) values (%d)", i)); + } + assertArrayEquals(IntStream.generate(() -> SUCCESS_NO_INFO).limit(size).toArray(), insert.executeBatch()); + } + } + @Test void shouldThrowExceptionWhenTryingToExecuteQueryThatWouldReturnMultipleResultSets() throws SQLException { try (Connection connection = createConnection(); Statement statement = connection.createStatement()) { diff --git a/src/main/java/com/firebolt/jdbc/statement/FireboltStatement.java b/src/main/java/com/firebolt/jdbc/statement/FireboltStatement.java index d883433d8..38faf2fef 100644 --- a/src/main/java/com/firebolt/jdbc/statement/FireboltStatement.java +++ b/src/main/java/com/firebolt/jdbc/statement/FireboltStatement.java @@ -19,9 +19,11 @@ import java.sql.SQLException; import java.sql.SQLWarning; import java.sql.Statement; +import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.HashSet; +import java.util.LinkedList; import java.util.List; import java.util.Optional; import java.util.Set; @@ -45,6 +47,7 @@ public class FireboltStatement extends JdbcBase implements Statement { private StatementResultWrapper firstUnclosedStatementResult; private int queryTimeout = 0; // zero means that there is no limit private String runningStatementLabel; + private final List batchStatements = new LinkedList<>(); public FireboltStatement(FireboltStatementService statementService, FireboltProperties sessionProperties, FireboltConnection connection) { @@ -417,24 +420,26 @@ public int getResultSetType() { } @Override - @NotImplemented public void addBatch(String sql) throws SQLException { - // Batch are not supported by the driver - throw new FireboltUnsupportedOperationException(); + batchStatements.add(sql); } @Override - @NotImplemented public void clearBatch() throws SQLException { - // Batch are not supported by the driver - throw new FireboltUnsupportedOperationException(); + batchStatements.clear(); } @Override - @NotImplemented public int[] executeBatch() throws SQLException { - // Batch are not supported by the driver - throw new FireboltUnsupportedOperationException(); + List result = new ArrayList<>(); + for (String sql : batchStatements) { + for (StatementInfoWrapper query : StatementUtil.parseToStatementInfoWrappers(sql)) { + @SuppressWarnings("java:S6912") // Use "addBatch" and "executeBatch" to execute multiple SQL statements in a single call - this is the implementation of executeBatch + Optional rs = execute(List.of(query)); + result.add(rs.map(x -> 0).orElse(SUCCESS_NO_INFO)); + } + } + return result.stream().mapToInt(Integer::intValue).toArray(); } @Override diff --git a/src/test/java/com/firebolt/jdbc/statement/FireboltStatementTest.java b/src/test/java/com/firebolt/jdbc/statement/FireboltStatementTest.java index 0c5eaa63f..93cd0f986 100644 --- a/src/test/java/com/firebolt/jdbc/statement/FireboltStatementTest.java +++ b/src/test/java/com/firebolt/jdbc/statement/FireboltStatementTest.java @@ -38,6 +38,7 @@ import static java.lang.String.format; import static java.sql.Statement.CLOSE_CURRENT_RESULT; +import static java.sql.Statement.SUCCESS_NO_INFO; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; @@ -71,9 +72,6 @@ class FireboltStatementTest { private static Stream unsupported() { return Stream.of( Arguments.of("setCursorName", (Executable) () -> statement.setCursorName("my_cursor")), - Arguments.of("addBatch", (Executable) () -> statement.addBatch("insert into people (id, name) values (?, ?))")), - Arguments.of("clearBatch", (Executable) () -> statement.clearBatch()), - Arguments.of("executeBatch", (Executable) () -> statement.executeBatch()), Arguments.of("getGeneratedKeys", (Executable) () -> statement.getGeneratedKeys()), Arguments.of("executeUpdate(auto generated keys)", (Executable) () -> statement.executeUpdate("insert", Statement.RETURN_GENERATED_KEYS)), Arguments.of("executeUpdate(column indexes)", (Executable) () -> statement.executeUpdate("insert", new int[0])), @@ -438,4 +436,46 @@ void shouldLimitStringByMaxFieldSize(String inputText, int maxFieldSize, String assertFalse(rs.next()); } + + @Test + void shouldExecuteEmptyBatch() throws SQLException { + FireboltConnection connection = mock(FireboltConnection.class); + FireboltStatement fireboltStatement = new FireboltStatement(fireboltStatementService, fireboltProperties, connection); + assertArrayEquals(new int[0], fireboltStatement.executeBatch()); + } + + @Test + void shouldExecuteBatch() throws SQLException { + FireboltConnection connection = mock(FireboltConnection.class); + FireboltStatement fireboltStatement = new FireboltStatement(fireboltStatementService, fireboltProperties, connection); + when(fireboltStatementService.execute(any(), any(), anyBoolean(), any())).thenReturn( + Optional.of(mock(FireboltResultSet.class)), Optional.of(mock(FireboltResultSet.class)), Optional.empty(), Optional.empty(), Optional.of(mock(FireboltResultSet.class)) + ); + + fireboltStatement.addBatch("SELECT 1; SELECT 2;"); + fireboltStatement.addBatch("INSERT INTO PEOPLE (id, name) VALUES (1, 'Adam')"); + fireboltStatement.addBatch("INSERT INTO PEOPLE (id, name) VALUES (1, 'Eve')"); + fireboltStatement.addBatch("SELECT 3"); + + int[] actual = fireboltStatement.executeBatch(); + assertArrayEquals(new int[] {0, 0, SUCCESS_NO_INFO, SUCCESS_NO_INFO, 0}, actual); + } + + @Test + void shouldClearBatch() throws SQLException { + FireboltConnection connection = mock(FireboltConnection.class); + FireboltStatement fireboltStatement = new FireboltStatement(fireboltStatementService, fireboltProperties, connection); + when(fireboltStatementService.execute(any(), any(), anyBoolean(), any())).thenReturn( + Optional.empty(), Optional.empty() + ); + + fireboltStatement.addBatch("SELECT 1; SELECT 2;"); + fireboltStatement.clearBatch(); + assertArrayEquals(new int[0], fireboltStatement.executeBatch()); + + fireboltStatement.addBatch("INSERT INTO PEOPLE (id, name) VALUES (1, 'Adam')"); + fireboltStatement.addBatch("INSERT INTO PEOPLE (id, name) VALUES (1, 'Eve')"); + + assertArrayEquals(new int[] {SUCCESS_NO_INFO, SUCCESS_NO_INFO}, fireboltStatement.executeBatch()); + } } \ No newline at end of file