diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 671f893bf8..012b71d3cd 100644 --- a/NEXT_CHANGELOG.md +++ b/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 - diff --git a/src/main/java/com/databricks/jdbc/api/ExecutionState.java b/src/main/java/com/databricks/jdbc/api/ExecutionState.java new file mode 100644 index 0000000000..c7bc4cd73b --- /dev/null +++ b/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; +} diff --git a/src/main/java/com/databricks/jdbc/api/IDatabricksResultSet.java b/src/main/java/com/databricks/jdbc/api/IDatabricksResultSet.java index d9e6197f34..b085b94bd6 100644 --- a/src/main/java/com/databricks/jdbc/api/IDatabricksResultSet.java +++ b/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. diff --git a/src/main/java/com/databricks/jdbc/api/IExecutionStatus.java b/src/main/java/com/databricks/jdbc/api/IExecutionStatus.java new file mode 100644 index 0000000000..b79a9ee295 --- /dev/null +++ b/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(); +} diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java index 9facf35fcf..bb9c1161de 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/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 columnNullables, List> 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 columnMetadataList, List> 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) diff --git a/src/main/java/com/databricks/jdbc/api/impl/EmptyResultSet.java b/src/main/java/com/databricks/jdbc/api/impl/EmptyResultSet.java index 6dc477bd64..181f575eea 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/EmptyResultSet.java +++ b/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; diff --git a/src/main/java/com/databricks/jdbc/api/impl/ExecutionStatus.java b/src/main/java/com/databricks/jdbc/api/impl/ExecutionStatus.java new file mode 100644 index 0000000000..9d392b6578 --- /dev/null +++ b/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); + } + } +} diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java index 590d161382..075a4ae4cc 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java +++ b/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