FIR-31967: batch related functionality in Statement (#397)

firebolt-db · Apr 24, 2024 · 4365d50 · 4365d50
1 parent 645d209
commit 4365d50
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 12 deletions.
diff --git 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<Integer> 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
@@ -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<String> 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<Integer> 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<ResultSet> 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
@@ -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<Arguments> 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());
+    }
 }