Merge pull request #193 from databricks/abortstatement

[PECO-1683] Add cancel Statement support

Author: gopalldb

src/main/java/com/databricks/jdbc/client/DatabricksClient.java

@@ -59,10 +59,16 @@ DatabricksResultSet executeStatement(
    * Closes a statement in Databricks server
    *
    * @param statementId statement which should be closed
-   * @return response for statement execution
    */
   void closeStatement(String statementId) throws DatabricksSQLException;
 
+  /**
+   * Cancels a statement in Databricks server
+   *
+   * @param statementId statement which should be aborted
+   */
+  void cancelStatement(String statementId) throws DatabricksSQLException;
+
   /**
    * Fetches the chunk details for given chunk index and statement-Id.
    *

src/main/java/com/databricks/jdbc/client/impl/sdk/DatabricksSdkClient.java

@@ -205,6 +205,14 @@ public void closeStatement(String statementId) {
     workspaceClient.apiClient().DELETE(path, request, Void.class, getHeaders());
   }
 
+  @Override
+  public void cancelStatement(String statementId) {
+    LOGGER.debug("public void cancelStatement(String statementId = {})", statementId);
+    CancelStatementRequest request = new CancelStatementRequest().setStatementId(statementId);
+    String path = String.format(CANCEL_STATEMENT_PATH_WITH_ID, request.getStatementId());
+    workspaceClient.apiClient().POST(path, request, Void.class, getHeaders());
+  }
+
   @Override
   public Collection<ExternalLink> getResultChunks(String statementId, long chunkIndex) {
     LOGGER.debug(

src/main/java/com/databricks/jdbc/client/impl/sdk/PathConstants.java

@@ -7,5 +7,6 @@ public class PathConstants {
   public static final String STATEMENT_PATH = BASE_PATH + "statements/";
   public static final String DELETE_SESSION_PATH_WITH_ID = STATEMENT_PATH + "sessions/%s";
   public static final String STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "%s";
+  public static final String CANCEL_STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "%s/cancel";
   public static final String RESULT_CHUNK_PATH = STATEMENT_PATH_WITH_ID + "/result/chunks/%s";
 }

src/main/java/com/databricks/jdbc/client/impl/thrift/DatabricksThriftServiceClient.java

@@ -119,6 +119,15 @@ public void closeStatement(String statementId) throws DatabricksSQLException {
         "closeStatement for all purpose cluster not implemented");
   }
 
+  @Override
+  public void cancelStatement(String statementId) throws DatabricksSQLException {
+    LOGGER.debug(
+        "public void cancelStatement(String statementId = {}) for all purpose cluster",
+        statementId);
+    throw new DatabricksSQLFeatureNotImplementedException(
+        "abortStatement for all purpose cluster not implemented");
+  }
+
   @Override
   public Collection<ExternalLink> getResultChunks(String statementId, long chunkIndex)
       throws DatabricksSQLException {

src/main/java/com/databricks/jdbc/client/sqlexec/CancelStatementRequest.java

@@ -0,0 +1,40 @@
+package com.databricks.jdbc.client.sqlexec;
+
+import com.databricks.sdk.support.ToStringer;
+import java.util.Objects;
+
+public class CancelStatementRequest {
+  private String statementId;
+
+  public CancelStatementRequest() {}
+
+  public CancelStatementRequest setStatementId(String statementId) {
+    this.statementId = statementId;
+    return this;
+  }
+
+  public String getStatementId() {
+    return this.statementId;
+  }
+
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    } else if (o != null && this.getClass() == o.getClass()) {
+      CancelStatementRequest that = (CancelStatementRequest) o;
+      return Objects.equals(this.statementId, that.statementId);
+    } else {
+      return false;
+    }
+  }
+
+  public int hashCode() {
+    return Objects.hash(new Object[] {this.statementId});
+  }
+
+  public String toString() {
+    return (new ToStringer(CancelStatementRequest.class))
+        .add("statementId", this.statementId)
+        .toString();
+  }
+}

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

@@ -132,8 +132,14 @@ public void setQueryTimeout(int seconds) throws SQLException {
   @Override
   public void cancel() throws SQLException {
     LOGGER.debug("public void cancel()");
-    throw new DatabricksSQLFeatureNotSupportedException(
-        "Not implemented in DatabricksStatement - cancel()");
+    checkIfClosed();
+
+    if (statementId != null) {
+      this.connection.getSession().getDatabricksClient().cancelStatement(statementId);
+    } else {
+      WarningUtil.addWarning(
+          warnings, "The statement you are trying to cancel does not have an ID yet.");
+    }
   }
 
   @Override

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

@@ -167,6 +167,31 @@ public void testExecuteStatement() throws Exception {
     assertEquals(STATEMENT_ID, statement.getStatementId());
   }
 
+  @Test
+  public void testCloseStatement() throws DatabricksSQLException {
+    String path = String.format(STATEMENT_PATH_WITH_ID, STATEMENT_ID);
+    IDatabricksConnectionContext connectionContext =
+        DatabricksConnectionContext.parse(JDBC_URL, new Properties());
+    DatabricksSdkClient databricksSdkClient =
+        new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient);
+    CloseStatementRequest request = new CloseStatementRequest().setStatementId(STATEMENT_ID);
+    databricksSdkClient.closeStatement(STATEMENT_ID);
+
+    verify(apiClient).DELETE(eq(path), eq(request), eq(Void.class), eq(headers));
+  }
+
+  @Test
+  public void testCancelStatement() throws DatabricksSQLException {
+    String path = String.format(CANCEL_STATEMENT_PATH_WITH_ID, STATEMENT_ID);
+    IDatabricksConnectionContext connectionContext =
+        DatabricksConnectionContext.parse(JDBC_URL, new Properties());
+    DatabricksSdkClient databricksSdkClient =
+        new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient);
+    CancelStatementRequest request = new CancelStatementRequest().setStatementId(STATEMENT_ID);
+    databricksSdkClient.cancelStatement(STATEMENT_ID);
+    verify(apiClient).POST(eq(path), eq(request), eq(Void.class), eq(headers));
+  }
+
   private StatementParameterListItem getParam(String type, String value, int ordinal) {
     return new PositionalStatementParameterListItem()
         .setOrdinal(ordinal)

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

@@ -81,6 +81,12 @@ public void testExecuteStatement() throws Exception {
 
     when(resultSet.hasUpdateCount()).thenReturn(true);
     assertFalse(statement.execute(STATEMENT, Statement.NO_GENERATED_KEYS));
+
+    assertFalse(statement.isClosed());
+    statement.cancel();
+
+    statement.close();
+    assertThrows(DatabricksSQLException.class, () -> statement.cancel());
   }
 
   @Test
@@ -199,7 +205,6 @@ public void testFeatureNotSupported() throws SQLException {
         () -> statement.isWrapperFor(java.sql.Connection.class));
     assertThrows(
         DatabricksSQLFeatureNotSupportedException.class, () -> statement.setCursorName("name"));
-    assertThrows(DatabricksSQLFeatureNotSupportedException.class, statement::cancel);
     assertThrows(DatabricksSQLFeatureNotSupportedException.class, statement::getMaxFieldSize);
     assertThrows(DatabricksSQLFeatureNotSupportedException.class, statement::getMoreResults);
     assertThrows(