Fix: try cancelling statement in all cases before throing an error#1202
Fix: try cancelling statement in all cases before throing an error#1202zashroof wants to merge 1 commit intodatabricks:mainfrom
Conversation
There was a problem hiding this comment.
Pull request overview
Ensure server-side operations/statements are cancelled when polling or thread interruption errors occur, without surfacing cancellation failures as additional exceptions.
Changes:
- Add “best-effort” cancellation helpers (
tryCancelOperation/tryCancelStatement) that swallow cancellation failures - Invoke cancellation on polling errors and thread interruptions before rethrowing
- Add tests validating cancellation is attempted when polling fails
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java | Adds tryCancelOperation and uses it on polling errors/interruption during async status polling |
| src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java | Adds tryCancelStatement and uses it on polling errors/interruption during statement execution polling |
| src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java | Adds test ensuring CancelOperation is invoked when polling throws |
| src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java | Adds test ensuring statement cancel endpoint is called when polling throws |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| try { | ||
| Thread.sleep(connectionContext.getAsyncExecPollInterval()); | ||
| } catch (InterruptedException e) { | ||
| // Cancel the statement on the server before throwing - it may still be running | ||
| tryCancelStatement(typedStatementId, "thread interruption"); | ||
| String timeoutErrorMessage = | ||
| String.format( | ||
| "Thread interrupted due to statement timeout. StatementID %s", statementId); |
There was a problem hiding this comment.
The InterruptedException catch block does not restore the thread interrupt flag. Please call Thread.currentThread().interrupt() before proceeding (even if you cancel and throw), so upstream code can observe the interruption consistently.
| "Failed to cancel statement {} after {}: {}", | ||
| statementId.toSQLExecStatementId(), | ||
| reason, | ||
| e.getMessage()); |
There was a problem hiding this comment.
The warning log drops the exception stack trace by logging only e.getMessage(). Please include the throwable in the log call (e.g., as the last argument) so cancellation failures can be debugged when needed.
| e.getMessage()); | |
| e.getMessage(), | |
| e); |
| "Failed to cancel operation for statement {} after {}: {}", | ||
| statementId.toSQLExecStatementId(), | ||
| reason, | ||
| e.getMessage()); |
There was a problem hiding this comment.
Same logging issue here: only e.getMessage() is logged, which loses the stack trace/context. Please pass the exception into the logger so the full failure cause is captured.
| "Failed to cancel operation for statement {} after {}: {}", | |
| statementId.toSQLExecStatementId(), | |
| reason, | |
| e.getMessage()); | |
| "Failed to cancel operation for statement {} after {}", | |
| statementId.toSQLExecStatementId(), | |
| reason, | |
| e); |
|
|
||
| // Mock CancelOperation to succeed | ||
| when(thriftClient.CancelOperation(any(TCancelOperationReq.class))) | ||
| .thenReturn(new TCancelOperationResp()); |
There was a problem hiding this comment.
This stub returns an empty TCancelOperationResp. If production code inspects the response status (common in Thrift patterns), an empty response can cause unexpected exceptions and make the test pass for the wrong reason (because tryCancelOperation swallows exceptions). Consider returning a response populated with a SUCCESS status to better reflect a real successful cancellation.
| .thenReturn(new TCancelOperationResp()); | |
| .thenReturn( | |
| new TCancelOperationResp() | |
| .setStatus(new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS))); |
| try { | ||
| Thread.sleep(connectionContext.getAsyncExecPollInterval()); | ||
| } catch (InterruptedException e) { | ||
| // Cancel the statement on the server before throwing - it may still be running |
There was a problem hiding this comment.
We have cancel being fired from checkTimeout above also. We can keep a check that if cancel was already attempted, don't do cancel again.
| @@ -320,8 +357,7 @@ private TGetOperationStatusResp pollTillOperationFinished( | |||
| TimeUnit.MILLISECONDS.sleep(asyncPollIntervalMillis); | |||
| } catch (InterruptedException e) { | |||
| Thread.currentThread().interrupt(); // Restore interrupt flag | |||
There was a problem hiding this comment.
restore interrupt flag -> can we add this to sdkClient handling as well?
2 participants
Description
If thread is inturrpted or a statment timesout, an exception is thrown while sometimes the statement continues to execute on the server. Also we see lots of exceptions that cancelling a statement failed (404).
Adding a tryCancelStatement to always try to cancel the statement on the server even if the chances the cancelStatement succeeds are slim. And not throwing an exception if canceling a statement fails.
Testing
Added tests
Additional Notes to the Reviewer
Cursor helped me create this PR. This is my first PR in this repo as well so probably I got a few things wrong.