Merge pull request #303 from databricks/jprakash-db/PECO-1758

[ PECO - 1758 ] - [OSS JDBC] Make batch APIs work in JDBC

Author: jprakash-db

src/main/java/com/databricks/jdbc/core/DatabricksPreparedStatement.java

@@ -21,7 +21,8 @@
 
 public class DatabricksPreparedStatement extends DatabricksStatement implements PreparedStatement {
   private final String sql;
-  private final DatabricksParameterMetaData databricksParameterMetaData;
+  private DatabricksParameterMetaData databricksParameterMetaData;
+  List<DatabricksParameterMetaData> databricksBatchParameterMetaData;
   private final boolean supportManyParameters;
 
   private final int CHUNK_SIZE = 8192;
@@ -32,6 +33,7 @@ public DatabricksPreparedStatement(DatabricksConnection connection, String sql)
     this.supportManyParameters =
         connection.getSession().getConnectionContext().supportManyParameters();
     this.databricksParameterMetaData = new DatabricksParameterMetaData();
+    this.databricksBatchParameterMetaData = new ArrayList<>();
   }
 
   private void checkLength(int targetLength, int sourceLength) throws SQLException {
@@ -56,6 +58,15 @@ private void checkLength(long targetLength, long sourceLength) throws SQLExcepti
     }
   }
 
+  private void checkIfBatchOperation() throws DatabricksSQLException {
+    if (!this.databricksBatchParameterMetaData.isEmpty()) {
+      String errorMessage =
+          "Batch must either be executed with executeBatch() or cleared with clearBatch()";
+      LoggingUtil.log(LogLevel.ERROR, errorMessage);
+      throw new DatabricksSQLException(errorMessage);
+    }
+  }
+
   private byte[] readByteStream(InputStream x, int length) throws SQLException {
     if (x == null) {
       String errorMessage = "InputStream cannot be null";
@@ -77,16 +88,38 @@ private byte[] readByteStream(InputStream x, int length) throws SQLException {
   @Override
   public ResultSet executeQuery() throws SQLException {
     LoggingUtil.log(LogLevel.DEBUG, "public ResultSet executeQuery()");
+    checkIfBatchOperation();
     return interpolateIfRequiredAndExecute(StatementType.QUERY);
   }
 
   @Override
   public int executeUpdate() throws SQLException {
     LoggingUtil.log(LogLevel.DEBUG, "public int executeUpdate()");
+    checkIfBatchOperation();
     interpolateIfRequiredAndExecute(StatementType.UPDATE);
     return (int) resultSet.getUpdateCount();
   }
 
+  @Override
+  public int[] executeBatch() {
+    LoggingUtil.log(LogLevel.DEBUG, "public int executeBatch()");
+    int[] updateCount = new int[databricksBatchParameterMetaData.size()];
+
+    for (int i = 0; i < databricksBatchParameterMetaData.size(); i++) {
+      DatabricksParameterMetaData databricksParameterMetaData =
+          databricksBatchParameterMetaData.get(i);
+      try {
+        executeInternal(
+            sql, databricksParameterMetaData.getParameterBindings(), StatementType.UPDATE, false);
+        updateCount[i] = (int) resultSet.getUpdateCount();
+      } catch (SQLException e) {
+        LoggingUtil.log(LogLevel.ERROR, e.getMessage());
+        updateCount[i] = -1;
+      }
+    }
+    return updateCount;
+  }
+
   @Override
   public void setNull(int parameterIndex, int sqlType) throws SQLException {
     LoggingUtil.log(LogLevel.DEBUG, "public void setNull(int parameterIndex, int sqlType)");
@@ -264,15 +297,24 @@ private void setObject(int parameterIndex, Object x, String databricksType) {
   public boolean execute() throws SQLException {
     LoggingUtil.log(LogLevel.DEBUG, "public boolean execute()");
     checkIfClosed();
+    checkIfBatchOperation();
     interpolateIfRequiredAndExecute(StatementType.SQL);
     return shouldReturnResultSet(sql);
   }
 
   @Override
-  public void addBatch() throws SQLException {
+  public void addBatch() {
     LoggingUtil.log(LogLevel.DEBUG, "public void addBatch()");
-    throw new UnsupportedOperationException(
-        "Not implemented in DatabricksPreparedStatement - addBatch()");
+    this.databricksBatchParameterMetaData.add(databricksParameterMetaData);
+    this.databricksParameterMetaData = new DatabricksParameterMetaData();
+  }
+
+  @Override
+  public void clearBatch() throws DatabricksSQLException {
+    LoggingUtil.log(LogLevel.DEBUG, "public void clearBatch()");
+    checkIfClosed();
+    this.databricksParameterMetaData = new DatabricksParameterMetaData();
+    this.databricksBatchParameterMetaData = new ArrayList<>();
   }
 
   @Override

src/main/java/com/databricks/jdbc/core/DatabricksStatement.java

@@ -429,7 +429,10 @@ public void handleResultSetClose(IDatabricksResultSet resultSet) throws SQLExcep
   }
 
   DatabricksResultSet executeInternal(
-      String sql, Map<Integer, ImmutableSqlParameter> params, StatementType statementType)
+      String sql,
+      Map<Integer, ImmutableSqlParameter> params,
+      StatementType statementType,
+      boolean closeStatement)
       throws SQLException {
     String stackTraceMessage =
         format(
@@ -444,7 +447,9 @@ DatabricksResultSet executeInternal(
               ? futureResultSet.get() // Wait indefinitely when timeout is 0
               : futureResultSet.get(timeoutInSeconds, TimeUnit.SECONDS);
     } catch (TimeoutException e) {
-      this.close(); // Close the statement
+      if (closeStatement) {
+        this.close(); // Close the statement
+      }
       futureResultSet.cancel(true); // Cancel execution run
       throw new DatabricksTimeoutException(
           "Statement execution timed-out. " + stackTraceMessage, e);
@@ -466,6 +471,12 @@ DatabricksResultSet executeInternal(
     return resultSet;
   }
 
+  DatabricksResultSet executeInternal(
+      String sql, Map<Integer, ImmutableSqlParameter> params, StatementType statementType)
+      throws SQLException {
+    return executeInternal(sql, params, statementType, true);
+  }
+
   // Todo : Add timeout tests in the subsequent PR
   CompletableFuture<DatabricksResultSet> getFutureResult(
       String sql, Map<Integer, ImmutableSqlParameter> params, StatementType statementType) {

src/test/java/com/databricks/jdbc/core/DatabricksPreparedStatementTest.java

@@ -33,6 +33,8 @@ public class DatabricksPreparedStatementTest {
   private static final String WAREHOUSE_ID = "erg6767gg";
   private static final String STATEMENT =
       "SELECT * FROM orders WHERE user_id = ? AND shard = ? AND region_code = ? AND namespace = ?";
+  private static final String BATCH_STATEMENT =
+      "INSERT INTO orders (user_id, shard, region_code, namespace) VALUES (?, ?, ?, ?)";
   private static final String JDBC_URL =
       "jdbc:databricks://adb-565757575.18.azuredatabricks.net:4423/default;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/erg6767gg;";
   private static final String JDBC_URL_WITH_MANY_PARAMETERS =
@@ -147,6 +149,43 @@ public void testExecuteUpdateStatement() throws Exception {
     assertTrue(statement.isClosed());
   }
 
+  @Test
+  public void testExecuteBatchStatement() throws Exception {
+    IDatabricksConnectionContext connectionContext =
+        DatabricksConnectionContext.parse(JDBC_URL, new Properties());
+    DatabricksConnection connection = new DatabricksConnection(connectionContext, client);
+    DatabricksPreparedStatement statement =
+        new DatabricksPreparedStatement(connection, BATCH_STATEMENT);
+
+    // Setting to execute a batch of 4 statements
+    for (int i = 1; i <= 4; i++) {
+      statement.setLong(1, (long) 100);
+      statement.setShort(2, (short) 10);
+      statement.setByte(3, (byte) 15);
+      statement.setString(4, "value");
+      statement.addBatch();
+    }
+
+    when(client.executeStatement(
+            eq(BATCH_STATEMENT),
+            eq(new Warehouse(WAREHOUSE_ID)),
+            any(HashMap.class),
+            eq(StatementType.UPDATE),
+            any(IDatabricksSession.class),
+            eq(statement)))
+        .thenReturn(resultSet);
+
+    when(resultSet.getUpdateCount()).thenReturn(1L);
+
+    int[] expectedCountsResult = {1, 1, 1, 1};
+    int[] updateCounts = statement.executeBatch();
+
+    assertArrayEquals(expectedCountsResult, updateCounts);
+    assertFalse(statement.isClosed());
+    statement.close();
+    assertTrue(statement.isClosed());
+  }
+
   public static ImmutableSqlParameter getSqlParam(
       int parameterIndex, Object x, String databricksType) {
     return ImmutableSqlParameter.builder()
@@ -412,7 +451,6 @@ void testUnsupportedMethods() throws DatabricksSQLException {
     assertThrows(
         UnsupportedOperationException.class, () -> preparedStatement.setTime(1, null, null));
     assertThrows(UnsupportedOperationException.class, () -> preparedStatement.setBytes(1, null));
-    assertThrows(UnsupportedOperationException.class, () -> preparedStatement.addBatch());
     assertThrows(
         SQLFeatureNotSupportedException.class, () -> preparedStatement.setObject(1, null, null));
     assertThrows(

src/test/java/com/databricks/jdbc/local/DriverTester.java

@@ -250,4 +250,50 @@ void testAllPurposeClusters_errorHandling() throws Exception {
     con.close();
     System.out.println("Connection closed successfully......");
   }
+
+  @Test
+  void testSimbaBatchFunction() throws Exception {
+
+    String jdbcUrl =
+        "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/dd43ee29fedd958d;";
+    Connection con = DriverManager.getConnection(jdbcUrl, "[email protected]", "xx");
+    System.out.println("Connection established......");
+
+    //
+    // Batch Statement Testing
+    //
+    String sqlStatement =
+        "INSERT INTO ___________________first.`jprakash-test`.diamonds (carat, cut, color, clarity) VALUES (?, ?, ?, ?)";
+    PreparedStatement pstmt = con.prepareStatement(sqlStatement);
+    for (int i = 1; i <= 3; i++) {
+      pstmt.setFloat(1, 0.23f);
+      pstmt.setString(2, "OK");
+      pstmt.setString(3, "E");
+      pstmt.setString(4, "SI2");
+      pstmt.addBatch();
+    }
+
+    pstmt.setString(1, "Shaama");
+    pstmt.setString(2, "Bad");
+    pstmt.setString(3, "F");
+    pstmt.setString(4, "SI6");
+    pstmt.addBatch();
+
+    for (int i = 1; i <= 3; i++) {
+      pstmt.setFloat(1, 0.23f);
+      pstmt.setString(2, "Bad");
+      pstmt.setString(3, "F");
+      pstmt.setString(4, "SI6");
+      pstmt.addBatch();
+    }
+
+    // Execute the batch
+    int[] updateCounts = pstmt.executeBatch();
+
+    // Process the update counts
+    for (int count : updateCounts) {
+      System.out.println("Update count: " + count);
+    }
+    con.close();
+  }
 }