diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index c72052218..a5085bd6c 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -10,6 +10,7 @@ - Added support for using SQL SHOW commands for Thrift-mode metadata operations (`getTables`, `getColumns`, `getSchemas`, `getFunctions`, `getPrimaryKeys`, `getImportedKeys`, `getCrossReference`). Enable by setting `UseQueryForMetadata=1`. This aligns Thrift metadata behavior with Statement Execution API (SEA) mode. ### Fixed +- Improved error messages for cancelled statements: operations cancelled via `Statement.cancel()` or closed connections now return SQL state `HY008` (operation cancelled) instead of generic error codes, making it easier for applications to detect and handle cancellations. - Fixed race condition between chunk download error handling and result set close that could cause invalid state transition warnings (`CHUNK_RELEASED -> DOWNLOAD_FAILED`) during Arrow Cloud Fetch operations in resource-constrained environments. - Fixed `EnableBatchedInserts` silently falling back to individual execution when table or schema names contain special characters (e.g., hyphens) inside backtick-quoted identifiers. Added a warn log when the fallback occurs. - Fixed `IntervalConverter` crash (`IllegalArgumentException: Invalid interval metadata`) when INTERVAL columns are returned via CloudFetch. Arrow metadata from CloudFetch uses underscored format (`INTERVAL_YEAR_MONTH`, `INTERVAL_DAY_TIME`) which the driver's regex did not accept. diff --git a/src/main/java/com/databricks/jdbc/api/impl/ExecutionStatus.java b/src/main/java/com/databricks/jdbc/api/impl/ExecutionStatus.java index 9d392b657..d0b8c6d29 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/ExecutionStatus.java +++ b/src/main/java/com/databricks/jdbc/api/impl/ExecutionStatus.java @@ -18,7 +18,10 @@ class ExecutionStatus implements IExecutionStatus { public ExecutionStatus(StatementStatus status) { this.state = getStateFromSdkState(status.getState()); - this.errorMessage = status.getError() != null ? status.getError().getMessage() : null; + this.errorMessage = + status.getError() != null && status.getError().getMessage() != null + ? status.getError().getMessage() + : (this.state == ExecutionState.ABORTED ? "Statement was cancelled" : null); this.sqlState = status.getSqlState(); this.sdkStatus = status; } diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java index 7849411f5..58cea1adb 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java @@ -109,6 +109,10 @@ public final class DatabricksJdbcConstants { public static final String GCP_GOOGLE_ID_AUTH_TYPE = "google-id"; public static final String DEFAULT_HTTP_EXCEPTION_SQLSTATE = "08000"; public static final String QUERY_EXECUTION_TIMEOUT_SQLSTATE = "57KD0"; + + /** Standard SQL state for operation cancelled (SQLSTATE HY008). */ + public static final String OPERATION_CANCELLED_SQLSTATE = "HY008"; + public static final int TEMPORARY_REDIRECT_STATUS_CODE = 307; public static final String REDACTED_TOKEN = "****"; public static final String QUERY_TAGS = "query_tags"; diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java index f97e47327..edda15a2b 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java @@ -1,6 +1,7 @@ package com.databricks.jdbc.dbclient.impl.sqlexec; import static com.databricks.jdbc.common.DatabricksJdbcConstants.JSON_HTTP_HEADERS; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.OPERATION_CANCELLED_SQLSTATE; import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE; import static com.databricks.jdbc.common.DatabricksJdbcConstants.TEMPORARY_REDIRECT_STATUS_CODE; import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_RESULT_ROW_LIMIT; @@ -709,6 +710,17 @@ void handleFailedExecution( ExecuteStatementResponse response, String statementId, String statement) throws SQLException { StatementState statementState = response.getStatus().getState(); ServiceError error = response.getStatus().getError(); + + // Distinguish cancellation from failure + if (statementState == StatementState.CANCELED) { + String cancelMessage = String.format("Statement [%s] was cancelled", statementId); + LOGGER.info(cancelMessage); + throw new DatabricksSQLException( + cancelMessage, + OPERATION_CANCELLED_SQLSTATE, + DatabricksDriverErrorCode.EXECUTE_STATEMENT_CANCELLED); + } + String errorMessage = String.format( "Statement execution failed %s -> %s\n%s.", statementId, statement, statementState); diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index c816fbdde..6de0a624f 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -1,5 +1,6 @@ package com.databricks.jdbc.dbclient.impl.thrift; +import static com.databricks.jdbc.common.DatabricksJdbcConstants.OPERATION_CANCELLED_SQLSTATE; import static com.databricks.jdbc.common.DatabricksJdbcConstants.QUERY_EXECUTION_TIMEOUT_SQLSTATE; import static com.databricks.jdbc.common.EnvironmentVariables.*; import static com.databricks.jdbc.common.util.DatabricksThriftUtil.*; @@ -422,6 +423,9 @@ DatabricksResultSet getStatementResult( try { response = getOperationStatus(request, statementId); TOperationState operationState = response.getOperationState(); + if (operationState == TOperationState.CANCELED_STATE) { + throw cancelledStatementException(statementId.toSQLExecStatementId()); + } if (operationState == TOperationState.FINISHED_STATE) { verifySuccessStatus( response.getStatus(), "getStatementResult", statementId.toSQLExecStatementId()); @@ -828,6 +832,11 @@ private void checkOperationStatusForErrors(TGetOperationStatusResp statusResp, S errorMsg, statusResp.isSetSqlState() ? statusResp.getSqlState() : null); } + if (statusResp.isSetOperationState() + && statusResp.getOperationState() == TOperationState.CANCELED_STATE) { + throw cancelledStatementException(statementId); + } + if (statusResp.isSetOperationState() && isErrorOperationState(statusResp.getOperationState())) { String errorMsg = String.format( @@ -863,6 +872,13 @@ private , F extends TFieldIdEnum> boolean hasResultDataInD return directResults.isSetResultSet() && directResults.isSetResultSetMetadata(); } + private DatabricksSQLException cancelledStatementException(String statementId) { + String msg = String.format("Statement [%s] was cancelled", statementId); + LOGGER.info(msg); + return new DatabricksSQLException( + msg, OPERATION_CANCELLED_SQLSTATE, DatabricksDriverErrorCode.EXECUTE_STATEMENT_CANCELLED); + } + private boolean isErrorStatusCode(TStatus status) { if (status == null || !status.isSetStatusCode()) { LOGGER.error("Status code is not set, marking the response as failed"); diff --git a/src/test/java/com/databricks/jdbc/api/impl/ExecutionStatusTest.java b/src/test/java/com/databricks/jdbc/api/impl/ExecutionStatusTest.java new file mode 100644 index 000000000..aa54af8ff --- /dev/null +++ b/src/test/java/com/databricks/jdbc/api/impl/ExecutionStatusTest.java @@ -0,0 +1,35 @@ +package com.databricks.jdbc.api.impl; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.jdbc.model.core.StatementStatus; +import com.databricks.sdk.service.sql.StatementState; +import org.junit.jupiter.api.Test; + +class ExecutionStatusTest { + + @Test + void cancelledStateWithNullError_hasMeaningfulMessage() { + StatementStatus cancelledStatus = + new StatementStatus().setState(StatementState.CANCELED).setError(null); + ExecutionStatus status = new ExecutionStatus(cancelledStatus); + assertNotNull(status.getErrorMessage(), "Cancelled state should have non-null error message"); + assertTrue(status.getErrorMessage().contains("cancelled")); + } + + @Test + void failedStateWithNullError_hasNullMessage() { + StatementStatus failedStatus = + new StatementStatus().setState(StatementState.FAILED).setError(null); + ExecutionStatus status = new ExecutionStatus(failedStatus); + assertNull(status.getErrorMessage(), "Failed state with null error should have null message"); + } + + @Test + void succeededStateWithNullError_hasNullMessage() { + StatementStatus succeededStatus = + new StatementStatus().setState(StatementState.SUCCEEDED).setError(null); + ExecutionStatus status = new ExecutionStatus(succeededStatus); + assertNull(status.getErrorMessage()); + } +} diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java index c484e6b21..3c8aafda6 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java @@ -290,6 +290,54 @@ public void testCancelStatement() throws DatabricksSQLException, IOException { eq(Void.class)); } + @Test + public void testHandleFailedExecution_CancelledState_ThrowsWithHY008() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksSdkClient databricksSdkClient = + new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); + + StatementStatus cancelledStatus = new StatementStatus().setState(StatementState.CANCELED); + ExecuteStatementResponse response = + new ExecuteStatementResponse() + .setStatementId(STATEMENT_ID.toSQLExecStatementId()) + .setStatus(cancelledStatus); + + DatabricksSQLException exception = + assertThrows( + DatabricksSQLException.class, + () -> + databricksSdkClient.handleFailedExecution( + response, STATEMENT_ID.toSQLExecStatementId(), STATEMENT)); + + assertEquals("HY008", exception.getSQLState()); + assertTrue(exception.getMessage().contains("was cancelled")); + } + + @Test + public void testHandleFailedExecution_FailedState_ThrowsWithoutHY008() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksSdkClient databricksSdkClient = + new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); + + StatementStatus failedStatus = new StatementStatus().setState(StatementState.FAILED); + ExecuteStatementResponse response = + new ExecuteStatementResponse() + .setStatementId(STATEMENT_ID.toSQLExecStatementId()) + .setStatus(failedStatus); + + DatabricksSQLException exception = + assertThrows( + DatabricksSQLException.class, + () -> + databricksSdkClient.handleFailedExecution( + response, STATEMENT_ID.toSQLExecStatementId(), STATEMENT)); + + assertNotEquals("HY008", exception.getSQLState()); + assertTrue(exception.getMessage().contains("execution failed")); + } + @Test public void testDisposition_arrowAndCloudFetchEnabled_usesExternalLinks() throws Exception { setupClientMocks(true, false); diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java index b1998dbea..ce2d981f7 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java @@ -449,6 +449,29 @@ void testGetStatementResult_pending() throws Exception { assertNull(resultSet.getMetaData()); } + @Test + void testGetStatementResult_cancelled_throwsWithHY008() throws Exception { + when(connectionContext.getDirectResultMode()).thenReturn(false); + accessor = spy(new DatabricksThriftAccessor(connectionContext)); + doReturn(thriftClient).when(accessor).getThriftClient(); + + // Server returns CANCELED_STATE with OK_STATUS and null errorMessage + TGetOperationStatusResp cancelledResp = + new TGetOperationStatusResp() + .setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS)) + .setOperationState(TOperationState.CANCELED_STATE); + when(thriftClient.GetOperationStatus(any(TGetOperationStatusReq.class))) + .thenReturn(cancelledResp); + + DatabricksSQLException exception = + assertThrows( + DatabricksSQLException.class, + () -> accessor.getStatementResult(tOperationHandle, null, session)); + + assertEquals("HY008", exception.getSQLState()); + assertTrue(exception.getMessage().contains("was cancelled")); + } + @Test void testListPrimaryKeys() throws TException, SQLException, DatabricksValidationException { setup(false);