[PECOBLR-384] Add implementation for fetching statement status for async execution

NEXT_CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Added
 - Support for fetching tables and views across all catalogs using SHOW TABLES FROM/IN ALL CATALOGS in the SQL Exec API.
 - Support for Token Exchange in OAuth flows where in third party tokens are exchanged for InHouse tokens.
+- Added support for polling of statementStatus and sqlState for async SQL execution.
 
 ### Updated
 - 

src/main/java/com/databricks/jdbc/api/ExecutionState.java

@@ -0,0 +1,20 @@
+package com.databricks.jdbc.api;
+
+/**
+ * Represents the possible states of a SQL statement execution in Databricks. This enum is used to
+ * track the progress and status of SQL statements executed through the Databricks JDBC API.
+ */
+public enum ExecutionState {
+  // The statement is in a pending state and has not yet started executing.
+  PENDING,
+  // The statement is currently executing.
+  RUNNING,
+  // The statement has completed successfully.
+  SUCCEEDED,
+  // The statement has completed with an error.
+  FAILED,
+  // The statement has been closed and is no longer available.
+  CLOSED,
+  // The statement has been cancelled by the user.
+  ABORTED;
+}

src/main/java/com/databricks/jdbc/api/IDatabricksResultSet.java

@@ -43,9 +43,19 @@ public interface IDatabricksResultSet extends ResultSet {
    * to monitor the execution progress and state of the statement.
    *
    * @return The current {@link StatementStatus} of the statement
+   * @deprecated Use {@link #getExecutionStatus()} instead.
    */
+  @Deprecated
   StatementStatus getStatementStatus();
 
+  /**
+   * Retrieves the current status of the statement associated with this result set. This can be used
+   * to monitor the execution progress and state of the statement.
+   *
+   * @return The current {@link StatementStatus} of the statement
+   */
+  IExecutionStatus getExecutionStatus();
+
   /**
    * Retrieves the number of rows affected by the SQL statement. For SELECT statements or statements
    * that don't modify data, this will return 0.

src/main/java/com/databricks/jdbc/api/IExecutionStatus.java

@@ -0,0 +1,24 @@
+package com.databricks.jdbc.api;
+
+public interface IExecutionStatus {
+  /**
+   * Returns the error message if the statement execution failed.
+   *
+   * @return the error message, or null if there was no error
+   */
+  String getErrorMessage();
+
+  /**
+   * Returns the SQL state code if the statement execution failed.
+   *
+   * @return the SQL state code, or null if there was no error
+   */
+  String getSqlState();
+
+  /**
+   * Returns the current state of the statement execution.
+   *
+   * @return the current state of the statement execution
+   */
+  ExecutionState getExecutionState();
+}

src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java

@@ -6,6 +6,7 @@
 import static com.databricks.jdbc.common.util.DatabricksTypeUtil.STRUCT;
 
 import com.databricks.jdbc.api.IDatabricksResultSet;
+import com.databricks.jdbc.api.IExecutionStatus;
 import com.databricks.jdbc.api.impl.arrow.ArrowStreamResult;
 import com.databricks.jdbc.api.impl.converters.ConverterHelper;
 import com.databricks.jdbc.api.impl.converters.ObjectConverter;
@@ -56,7 +57,7 @@ enum ResultSetType {
 
   private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DatabricksResultSet.class);
   protected static final String AFFECTED_ROWS_COUNT = "num_affected_rows";
-  private final StatementStatus statementStatus;
+  private final ExecutionStatus executionStatus;
   private final StatementId statementId;
   private final IExecutionResult executionResult;
   private final DatabricksResultSetMetaData resultSetMetaData;
@@ -81,7 +82,7 @@ public DatabricksResultSet(
       IDatabricksSession session,
       IDatabricksStatementInternal parentStatement)
       throws DatabricksSQLException {
-    this.statementStatus = statementStatus;
+    this.executionStatus = new ExecutionStatus(statementStatus);
     this.statementId = statementId;
     if (resultData != null) {
       this.executionResult =
@@ -122,7 +123,7 @@ public DatabricksResultSet(
       IExecutionResult executionResult,
       DatabricksResultSetMetaData resultSetMetaData,
       boolean complexDatatypeSupport) {
-    this.statementStatus = statementStatus;
+    this.executionStatus = new ExecutionStatus(statementStatus);
     this.statementId = statementId;
     this.executionResult = executionResult;
     this.resultSetMetaData = resultSetMetaData;
@@ -143,7 +144,7 @@ public DatabricksResultSet(
       IDatabricksStatementInternal parentStatement,
       IDatabricksSession session)
       throws SQLException {
-    this.statementStatus = statementStatus;
+    this.executionStatus = new ExecutionStatus(statementStatus);
     this.statementId = statementId;
     if (resultsResp != null) {
       this.executionResult =
@@ -193,7 +194,7 @@ public DatabricksResultSet(
       int[] isNullables,
       Object[][] rows,
       StatementType statementType) {
-    this.statementStatus = statementStatus;
+    this.executionStatus = new ExecutionStatus(statementStatus);
     this.statementId = statementId;
     this.executionResult = ExecutionResultFactory.getResultSet(rows);
     this.resultSetMetaData =
@@ -223,7 +224,7 @@ public DatabricksResultSet(
       List<Nullable> columnNullables,
       List<List<Object>> rows,
       StatementType statementType) {
-    this.statementStatus = statementStatus;
+    this.executionStatus = new ExecutionStatus(statementStatus);
     this.statementId = statementId;
     this.executionResult = ExecutionResultFactory.getResultSet(rows);
     this.resultSetMetaData =
@@ -249,7 +250,7 @@ public DatabricksResultSet(
       List<ColumnMetadata> columnMetadataList,
       List<List<Object>> rows,
       StatementType statementType) {
-    this.statementStatus = statementStatus;
+    this.executionStatus = new ExecutionStatus(statementStatus);
     this.statementId = statementId;
     this.executionResult = ExecutionResultFactory.getResultSet(rows);
     this.resultSetMetaData =
@@ -1734,7 +1735,12 @@ public String getStatementId() {
 
   @Override
   public StatementStatus getStatementStatus() {
-    return statementStatus;
+    return executionStatus.getSdkStatus();
+  }
+
+  @Override
+  public IExecutionStatus getExecutionStatus() {
+    return executionStatus;
   }
 
   @Override
@@ -1837,7 +1843,7 @@ private BigDecimal applyScaleToBigDecimal(BigDecimal bigDecimal, int columnIndex
   @Override
   public String toString() {
     return (new ToStringer(DatabricksResultSet.class))
-        .add("statementStatus", this.statementStatus)
+        .add("statementStatus", this.executionStatus)
         .add("statementId", this.statementId)
         .add("statementType", this.statementType)
         .add("updateCount", this.updateCount)

src/main/java/com/databricks/jdbc/api/impl/EmptyResultSet.java

@@ -1,6 +1,7 @@
 package com.databricks.jdbc.api.impl;
 
 import com.databricks.jdbc.api.IDatabricksResultSet;
+import com.databricks.jdbc.api.IExecutionStatus;
 import com.databricks.jdbc.api.internal.IDatabricksResultSetInternal;
 import com.databricks.jdbc.exception.DatabricksSQLException;
 import com.databricks.jdbc.model.core.StatementStatus;
@@ -1127,6 +1128,11 @@ public StatementStatus getStatementStatus() {
     return new StatementStatus().setState(StatementState.SUCCEEDED);
   }
 
+  @Override
+  public IExecutionStatus getExecutionStatus() {
+    return new ExecutionStatus(new StatementStatus().setState(StatementState.SUCCEEDED));
+  }
+
   @Override
   public long getUpdateCount() throws SQLException {
     return 0;

src/main/java/com/databricks/jdbc/api/impl/ExecutionStatus.java

@@ -0,0 +1,68 @@
+package com.databricks.jdbc.api.impl;
+
+import com.databricks.jdbc.api.ExecutionState;
+import com.databricks.jdbc.api.IExecutionStatus;
+import com.databricks.jdbc.model.core.StatementStatus;
+import com.databricks.sdk.service.sql.StatementState;
+
+/**
+ * This class implements the IStatementStatus interface and provides a default implementation for
+ * the methods defined in the interface. It is used to represent the status of a SQL statement
+ * execution in Databricks.
+ */
+class ExecutionStatus implements IExecutionStatus {
+  private final ExecutionState state;
+  private final String errorMessage;
+  private final String sqlState;
+  private final StatementStatus sdkStatus;
+
+  public ExecutionStatus(StatementStatus status) {
+    this.state = getStateFromSdkState(status.getState());
+    this.errorMessage = status.getError() != null ? status.getError().getMessage() : null;
+    this.sqlState = status.getSqlState();
+    this.sdkStatus = status;
+  }
+
+  @Override
+  public String getErrorMessage() {
+    return errorMessage;
+  }
+
+  @Override
+  public String getSqlState() {
+    return sqlState;
+  }
+
+  @Override
+  public ExecutionState getExecutionState() {
+    return state;
+  }
+
+  StatementStatus getSdkStatus() {
+    return sdkStatus;
+  }
+
+  private ExecutionState getStateFromSdkState(StatementState state) {
+    if (state == null) {
+      return ExecutionState.PENDING;
+    }
+    // Map the SDK statement state to the JDBC statement state
+    switch (state) {
+      case PENDING:
+        return ExecutionState.PENDING;
+      case RUNNING:
+        return ExecutionState.RUNNING;
+      case SUCCEEDED:
+        return ExecutionState.SUCCEEDED;
+      case FAILED:
+        return ExecutionState.FAILED;
+      case CANCELED:
+        return ExecutionState.ABORTED;
+      case CLOSED:
+        return ExecutionState.CLOSED;
+        // should never reach here
+      default:
+        throw new IllegalArgumentException("Unknown statement execution state: " + state);
+    }
+  }
+}

src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java

@@ -8,7 +8,9 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.databricks.jdbc.api.ExecutionState;
 import com.databricks.jdbc.api.IDatabricksResultSet;
+import com.databricks.jdbc.api.IExecutionStatus;
 import com.databricks.jdbc.api.impl.volume.VolumeOperationResult;
 import com.databricks.jdbc.api.internal.IDatabricksResultSetInternal;
 import com.databricks.jdbc.api.internal.IDatabricksSession;
@@ -20,6 +22,7 @@
 import com.databricks.jdbc.exception.DatabricksSQLFeatureNotSupportedException;
 import com.databricks.jdbc.model.client.thrift.generated.*;
 import com.databricks.jdbc.model.core.StatementStatus;
+import com.databricks.sdk.service.sql.ServiceError;
 import com.databricks.sdk.service.sql.StatementState;
 import java.io.*;
 import java.math.BigDecimal;
@@ -63,6 +66,18 @@ private DatabricksResultSet getResultSet(
         false);
   }
 
+  private DatabricksResultSet getResultSet(
+      StatementStatus statementState, IDatabricksStatementInternal statement) {
+    return new DatabricksResultSet(
+        statementState,
+        STATEMENT_ID,
+        StatementType.METADATA,
+        statement,
+        mockedExecutionResult,
+        mockedResultSetMetadata,
+        false);
+  }
+
   private DatabricksResultSet getThriftResultSetMetadata() throws SQLException {
     TColumnValue columnValue = new TColumnValue();
     columnValue.setStringVal(new TStringValue().setValue("testString"));
@@ -118,8 +133,20 @@ void testThriftResultSet() throws SQLException {
 
   @Test
   void testGetStatementStatus() {
-    DatabricksResultSet resultSet = getResultSet(StatementState.PENDING, null);
-    assertEquals(StatementState.PENDING, resultSet.getStatementStatus().getState());
+    StatementStatus statementStatus =
+        new StatementStatus()
+            .setState(StatementState.FAILED)
+            .setError(new ServiceError().setMessage("error"))
+            .setSqlState("sqlState");
+    DatabricksResultSet resultSet = getResultSet(statementStatus, null);
+    assertEquals(STATEMENT_ID.toString(), resultSet.getStatementId());
+    assertEquals(StatementState.FAILED, resultSet.getStatementStatus().getState());
+    assertEquals(statementStatus, resultSet.getStatementStatus());
+
+    IExecutionStatus executionStatus = resultSet.getExecutionStatus();
+    assertEquals(ExecutionState.FAILED, executionStatus.getExecutionState());
+    assertEquals("error", executionStatus.getErrorMessage());
+    assertEquals("sqlState", executionStatus.getSqlState());
   }
 
   @Test