From 02433118cc87d43c1848db0d7d71d305a3025261 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 12 May 2026 11:37:59 +0530 Subject: [PATCH 01/13] Proactive server-side statement close for cost savings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Close the server-side operation proactively when: 1. next() returns false — all rows consumed 2. ResultSet.close() is called explicitly The client-side Statement remains open for reuse. A new serverOperationClosed flag prevents duplicate closeStatement RPCs when Statement.close() is called later. Behavior: - next() exhausts rows → closeStatement RPC → serverOperationClosed=true - ResultSet.close() → closeStatement RPC (if not already closed) - Statement.close() → skips RPC if serverOperationClosed, just cleans up client - Re-execution resets the flag 6 new tests covering: - Proactive close skips RPC on Statement.close() - Idempotent (double call sends one RPC) - Skipped for direct results - Skipped when no statementId - Flag reset verified - Server errors swallowed (best-effort) Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksResultSet.java | 14 +++ .../jdbc/api/impl/DatabricksStatement.java | 45 +++++++- .../api/impl/DatabricksStatementTest.java | 107 ++++++++++++++++++ 3 files changed, 160 insertions(+), 6 deletions(-) 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 a83956528..510095009 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -283,11 +283,18 @@ public boolean next() throws SQLException { cachedTelemetryCollector.recordResultSetIteration( statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext); } + if (!hasNext) { + // All rows consumed — proactively close server operation to release resources. + // The client-side Statement remains open for reuse. + closeServerOperation(); + } return hasNext; } @Override public void close() throws DatabricksSQLException { + // Proactively close server operation when ResultSet is closed explicitly. + closeServerOperation(); isClosed = true; this.executionResult.close(); if (parentStatement != null) { @@ -295,6 +302,13 @@ public void close() throws DatabricksSQLException { } } + /** Proactively closes the server-side operation via the parent statement. */ + private void closeServerOperation() { + if (parentStatement instanceof DatabricksStatement) { + ((DatabricksStatement) parentStatement).closeServerOperation(); + } + } + private static TelemetryCollector resolveTelemetryCollector( IDatabricksStatementInternal parentStatement) { try { diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index c9b6733e7..bfb90245c 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -59,6 +59,9 @@ public class DatabricksStatement implements IDatabricksStatement, IDatabricksSta // closed the operation — no further RPCs for this statement ID are possible. The JDBC Statement // itself remains open for re-execution. Reset on each new execution. Volatile because cancel() // can be called from a different thread (JDBC spec requirement). + private volatile boolean serverOperationClosed = false; // Client proactively closed the server + // operation after all results were consumed or ResultSet was closed. Prevents duplicate + // closeStatement RPC when Statement.close() is called later. Reset on each new execution. protected Boolean shouldReturnResultSet = null; // Cached result of shouldReturnResultSetWithConfig() @@ -149,15 +152,18 @@ public void close(boolean removeFromSession) throws DatabricksSQLException { LOGGER.warn(warningMsg); warnings = WarningUtil.addWarning(warnings, warningMsg); } else { - // Skip server-side close if the server already closed the operation (direct results). - // The operation handle is gone on the server side, so closeStatement would fail. - if (!directResultsReceived) { + // Skip server-side close if operation was already closed: + // - directResultsReceived: server closed it (inline results) + // - serverOperationClosed: client proactively closed it (results consumed or RS closed) + if (!directResultsReceived && !serverOperationClosed) { this.connection.getSession().getDatabricksClient().closeStatement(statementId); } else { LOGGER.debug( - "Statement {} closed locally (direct results — server operation already closed, " - + "skipping closeStatement RPC)", - statementId); + "Statement {} closed locally (server operation already closed — " + + "directResults={}, proactivelyClosed={}, skipping closeStatement RPC)", + statementId, + directResultsReceived, + serverOperationClosed); } if (resultSet != null) { this.resultSet.close(); @@ -953,6 +959,32 @@ public void markDirectResultsReceived() { this.directResultsReceived = true; } + /** + * Proactively closes the server-side operation to release server resources while keeping the + * client-side Statement open for reuse. Called when all results have been consumed (next() + * returns false) or when ResultSet.close() is called. + * + *

After this call, {@link #close(boolean)} will skip the closeStatement RPC since the server + * operation is already closed. The Statement can still be re-executed. + */ + void closeServerOperation() { + if (serverOperationClosed || directResultsReceived || statementId == null || isClosed) { + return; + } + try { + this.connection.getSession().getDatabricksClient().closeStatement(statementId); + LOGGER.debug( + "Proactively closed server operation for statement {} (results consumed)", statementId); + } catch (Exception e) { + // Best-effort — don't fail the user's iteration/close for a server cleanup failure + LOGGER.debug( + "Failed to proactively close server operation for statement {}: {}", + statementId, + e.getMessage()); + } + this.serverOperationClosed = true; + } + /** * Resets statement state before a new execution (sync or async). Closes the previous server-side * operation handle (if still open) and the local ResultSet, clears flags, and nulls the @@ -970,6 +1002,7 @@ private void resetForNewExecution() { // For direct results, the server already closed the handle. directResultsReceived = false; + serverOperationClosed = false; // Per JDBC spec, re-executing a Statement implicitly closes the current ResultSet. if (resultSet != null) { diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java index ff5d9e7ec..60a4eb765 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java @@ -16,6 +16,7 @@ import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.exception.DatabricksSQLFeatureNotSupportedException; import com.databricks.jdbc.model.core.StatementStatus; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; import com.databricks.sdk.service.sql.StatementState; import java.io.InputStream; import java.sql.*; @@ -1682,4 +1683,110 @@ public void testGetResultSet_WithNonRowcountQueryPrefixes_ReturnsResultSet() thr statement.close(); } + + // ========================================================================= + // Proactive server operation close + // ========================================================================= + + @Test + public void testCloseServerOperation_closesServerAndSkipsRpcOnStatementClose() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + statement.setStatementId(STATEMENT_ID); + + // Proactively close the server operation (simulates next() returning false) + statement.closeServerOperation(); + + // Statement should still be open + assertFalse(statement.isClosed()); + + // closeStatement should have been called once (proactive close) + verify(client, times(1)).closeStatement(STATEMENT_ID); + + // Now close the statement — should NOT call closeStatement again + statement.close(); + verify(client, times(1)).closeStatement(STATEMENT_ID); // still 1, not 2 + } + + @Test + public void testCloseServerOperation_idempotent() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + statement.setStatementId(STATEMENT_ID); + + // Call twice — should only close once + statement.closeServerOperation(); + statement.closeServerOperation(); + + verify(client, times(1)).closeStatement(STATEMENT_ID); + } + + @Test + public void testCloseServerOperation_skippedForDirectResults() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + statement.setStatementId(STATEMENT_ID); + statement.markDirectResultsReceived(); + + // Should be a no-op since directResults already closed the server operation + statement.closeServerOperation(); + + verify(client, never()).closeStatement(any(StatementId.class)); + } + + @Test + public void testCloseServerOperation_skippedWhenNoStatementId() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + // statementId is null — should be a no-op + + statement.closeServerOperation(); + + verify(client, never()).closeStatement(any(StatementId.class)); + } + + @Test + public void testCloseServerOperation_resetsAfterFlagClear() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + statement.setStatementId(STATEMENT_ID); + + // First proactive close + statement.closeServerOperation(); + verify(client, times(1)).closeStatement(STATEMENT_ID); + + // Second call is no-op (flag set) + statement.closeServerOperation(); + verify(client, times(1)).closeStatement(STATEMENT_ID); + + // Statement.close() is also no-op for server RPC (flag set) + statement.close(); + verify(client, times(1)).closeStatement(STATEMENT_ID); + } + + @Test + public void testCloseServerOperation_errorSwallowed() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + statement.setStatementId(STATEMENT_ID); + + // Server close fails — should not throw + doThrow(new DatabricksSQLException("Server error", DatabricksDriverErrorCode.SDK_CLIENT_ERROR)) + .when(client) + .closeStatement(STATEMENT_ID); + + assertDoesNotThrow(() -> statement.closeServerOperation()); + } } From 863e9eeca7a70c14ee3d113fa45959ee3f8488c9 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 12 May 2026 12:12:08 +0530 Subject: [PATCH 02/13] Add changelog entry for proactive server statement close Co-authored-by: Isaac Signed-off-by: Gopal Lal --- NEXT_CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 9716a8d3b..1db44d81c 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -14,6 +14,7 @@ ### Updated - `EnableGeoSpatialSupport` no longer requires `EnableComplexDatatypeSupport=1`. Geospatial types (GEOMETRY, GEOGRAPHY) can now be enabled independently of complex type support (ARRAY, MAP, STRUCT). +- Server-side operations are now closed proactively when all results are consumed (`next()` returns false) or `ResultSet.close()` is called. This releases server resources and reduces warehouse cost without closing the client-side Statement, which remains reusable. Subsequent `Statement.close()` skips the server RPC (already closed). ### Fixed From cf2f6b3a8bea9878258384a8d14a247fa51e07c4 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 12 May 2026 12:12:49 +0530 Subject: [PATCH 03/13] Reword changelog: focus on resource utilization, not cost Co-authored-by: Isaac Signed-off-by: Gopal Lal --- NEXT_CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 1db44d81c..9c50adf4b 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -14,7 +14,7 @@ ### Updated - `EnableGeoSpatialSupport` no longer requires `EnableComplexDatatypeSupport=1`. Geospatial types (GEOMETRY, GEOGRAPHY) can now be enabled independently of complex type support (ARRAY, MAP, STRUCT). -- Server-side operations are now closed proactively when all results are consumed (`next()` returns false) or `ResultSet.close()` is called. This releases server resources and reduces warehouse cost without closing the client-side Statement, which remains reusable. Subsequent `Statement.close()` skips the server RPC (already closed). +- Server-side operations are now closed proactively when all results are consumed (`next()` returns false) or `ResultSet.close()` is called, improving resource utilization. The client-side Statement remains open and reusable for re-execution. ### Fixed From a7095684576182a87e1427f527582b68bfa90ddc Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 12 May 2026 12:22:38 +0530 Subject: [PATCH 04/13] Document invariant: all data fetched before server close Add comment explaining why proactive server close in next() is safe: by the time next() returns false, all chunks/batches are already client-side. No further server RPCs needed for CloudFetch link refresh, lazy getUpdateCount, or getMoreResults. Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../com/databricks/jdbc/api/impl/DatabricksResultSet.java | 8 ++++++++ 1 file changed, 8 insertions(+) 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 510095009..daba0bd6e 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -286,6 +286,14 @@ public boolean next() throws SQLException { if (!hasNext) { // All rows consumed — proactively close server operation to release resources. // The client-side Statement remains open for reuse. + // + // Invariant: by the time next() returns false, the execution result has already + // fetched all data from the server (all chunks downloaded for CloudFetch, all + // batches fetched for Thrift streaming). No further server RPCs are needed for + // data retrieval. This means: + // - CloudFetch link refresh is not needed (all links already resolved) + // - Lazy getUpdateCount() works because row data is already client-side + // - getMoreResults() calling resultSet.close() is safe (server op already closing) closeServerOperation(); } return hasNext; From 113d57932773a93dd2193f8d1903b57920387b14 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 12 May 2026 12:26:03 +0530 Subject: [PATCH 05/13] Guard cancel() and getExecutionResult() for serverOperationClosed - cancel(): Skip cancelStatement RPC when server operation already closed proactively. Prevents 404 on cancel after results consumed. - getExecutionResult(): Return cached result when server operation already closed. Prevents getStatementResult RPC on closed operation. Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksStatement.java | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index bfb90245c..ea9ef77b0 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -252,12 +252,11 @@ public void cancel() throws SQLException { LOGGER.debug("public void cancel()"); checkIfClosed(); - if (statementId != null && !directResultsReceived) { + if (statementId != null && !directResultsReceived && !serverOperationClosed) { this.connection.getSession().getDatabricksClient().cancelStatement(statementId); DatabricksThreadContextHolder.clearStatementInfo(); - } else if (directResultsReceived) { - String warningMsg = - "Statement's server operation was already closed (direct results); cancel has no effect."; + } else if (directResultsReceived || serverOperationClosed) { + String warningMsg = "Statement's server operation was already closed; cancel has no effect."; LOGGER.debug(warningMsg); warnings = WarningUtil.addWarning(warnings, warningMsg); } else { @@ -700,17 +699,18 @@ public ResultSet getExecutionResult() throws SQLException { "No execution available for statement", DatabricksDriverErrorCode.INPUT_VALIDATION_ERROR); } - // For direct results, the server already closed the operation — making an RPC - // would return "not found". Return the cached result set instead. - if (directResultsReceived) { + // For direct results or proactively closed operations, the server operation is gone — + // making an RPC would return "not found". Return the cached result set instead. + if (directResultsReceived || serverOperationClosed) { if (resultSet != null) { LOGGER.debug( - "Returning cached result for statement {} (direct results received)", statementId); + "Returning cached result for statement {} (server operation already closed)", + statementId); return resultSet; } throw new DatabricksSQLException( - "Direct results were received but no result set is available. " - + "The server closed the operation and no further results can be fetched.", + "Server operation was already closed and no result set is available. " + + "No further results can be fetched.", DatabricksDriverErrorCode.INVALID_STATE); } From 3e88c362c21c02febe0bd3c2e38c014367e83f46 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 12 May 2026 12:27:05 +0530 Subject: [PATCH 06/13] Only set serverOperationClosed on success, not on RPC failure If proactive closeStatement RPC fails (transient network error), keep serverOperationClosed=false so Statement.close() retries the RPC. Prevents leaving server operations alive until session timeout. Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../com/databricks/jdbc/api/impl/DatabricksStatement.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index ea9ef77b0..d57678594 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -973,16 +973,20 @@ void closeServerOperation() { } try { this.connection.getSession().getDatabricksClient().closeStatement(statementId); + // Only mark closed on success — if the RPC fails (transient network error), + // Statement.close() should retry the closeStatement RPC rather than skip it, + // to avoid leaving the server operation alive until session timeout. + this.serverOperationClosed = true; LOGGER.debug( "Proactively closed server operation for statement {} (results consumed)", statementId); } catch (Exception e) { - // Best-effort — don't fail the user's iteration/close for a server cleanup failure + // Best-effort — don't fail the user's iteration/close for a server cleanup failure. + // serverOperationClosed stays false so Statement.close() will retry the RPC. LOGGER.debug( "Failed to proactively close server operation for statement {}: {}", statementId, e.getMessage()); } - this.serverOperationClosed = true; } /** From 4a06513e1bbb8f146502d30b9b52588ff80edc48 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 12 May 2026 12:30:28 +0530 Subject: [PATCH 07/13] Add additional tests for proactive server close interactions New tests: - cancel() is no-op after proactive close (cancel + close interaction) - getExecutionResult() returns cached result after proactive close - Statement is reusable for re-execution after proactive close - Error swallowed test verifies flag NOT set on failure (retry behavior) Covers review feedback #8-#13 (wiring tests, reuse, cancel interaction, getExecutionResult guard, error retry semantics). Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../api/impl/DatabricksStatementTest.java | 86 +++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java index 60a4eb765..179742ac3 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java @@ -1788,5 +1788,91 @@ public void testCloseServerOperation_errorSwallowed() throws Exception { .closeStatement(STATEMENT_ID); assertDoesNotThrow(() -> statement.closeServerOperation()); + + // Flag should NOT be set on failure — Statement.close() should retry + verify(client, times(1)).closeStatement(STATEMENT_ID); + // The second closeServerOperation call should retry since flag wasn't set + reset(client); // clear the throw stub + statement.closeServerOperation(); + verify(client, times(1)).closeStatement(STATEMENT_ID); + } + + @Test + public void testCloseServerOperation_cancelSkipsAfterProactiveClose() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + statement.setStatementId(STATEMENT_ID); + + // Proactively close server operation + statement.closeServerOperation(); + verify(client, times(1)).closeStatement(STATEMENT_ID); + + // cancel() should be a no-op — server operation already closed + statement.cancel(); + verify(client, never()).cancelStatement(any(StatementId.class)); + } + + @Test + public void testCloseServerOperation_getExecutionResultReturnsCachedAfterClose() + throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + statement.setStatementId(STATEMENT_ID); + // Set resultSet field via reflection to simulate executeQuery having run + java.lang.reflect.Field rsField = DatabricksStatement.class.getDeclaredField("resultSet"); + rsField.setAccessible(true); + rsField.set(statement, resultSet); + + // Proactively close server operation + statement.closeServerOperation(); + + // getExecutionResult() should return cached result, not make RPC + ResultSet cached = statement.getExecutionResult(); + assertNotNull(cached); + assertEquals(resultSet, cached); + verify(client, never()) + .getStatementResult(any(StatementId.class), any(IDatabricksSession.class), any()); + } + + @Test + public void testStatementReusableAfterProactiveClose() throws Exception { + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + + when(client.executeStatement( + eq(STATEMENT), + eq(new Warehouse(WAREHOUSE_ID)), + eq(new HashMap<>()), + eq(StatementType.QUERY), + any(IDatabricksSession.class), + eq(statement), + any())) + .thenReturn(resultSet); + + // First execution + statement.executeQuery(STATEMENT); + statement.closeServerOperation(); + assertFalse(statement.isClosed(), "Statement should stay open after proactive close"); + + // Re-execute — should work, flag reset by resetForNewExecution + statement.executeQuery(STATEMENT); + assertFalse(statement.isClosed()); + + // Both executions should have called executeStatement + verify(client, times(2)) + .executeStatement( + eq(STATEMENT), + eq(new Warehouse(WAREHOUSE_ID)), + eq(new HashMap<>()), + eq(StatementType.QUERY), + any(IDatabricksSession.class), + eq(statement), + any()); } } From 376b38aa8c8548e4b0f4cabe71904c696f933a08 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 12 May 2026 16:45:35 +0530 Subject: [PATCH 08/13] =?UTF-8?q?Remove=20proactive=20close=20from=20next(?= =?UTF-8?q?)=20=E2=80=94=20only=20close=20on=20ResultSet.close()?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Proactive closeServerOperation() in next() when returning false was too aggressive: batch execution reuses statements, and the early close broke subsequent batch commands (HTTP 404 on ExecuteStatement). Also consumed WireMock stubs meant for final Statement.close(). Keep proactive close only in ResultSet.close() — this is the explicit user signal that results are no longer needed. next() exhaustion alone does not justify closing the server operation since the statement may be reused for batch or re-execution. Fixes integration test failures: - ExecutionIntegrationTests.testStatementBatch_MixedDMLOperations - ExecutionIntegrationTests.testStatementBatch_MultipleInserts - NonRowcountQueryPrefixesIntegrationTests.testConfigured_MultiplePrefixes Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksResultSet.java | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) 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 daba0bd6e..bd3736a1f 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -283,19 +283,12 @@ public boolean next() throws SQLException { cachedTelemetryCollector.recordResultSetIteration( statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext); } - if (!hasNext) { - // All rows consumed — proactively close server operation to release resources. - // The client-side Statement remains open for reuse. - // - // Invariant: by the time next() returns false, the execution result has already - // fetched all data from the server (all chunks downloaded for CloudFetch, all - // batches fetched for Thrift streaming). No further server RPCs are needed for - // data retrieval. This means: - // - CloudFetch link refresh is not needed (all links already resolved) - // - Lazy getUpdateCount() works because row data is already client-side - // - getMoreResults() calling resultSet.close() is safe (server op already closing) - closeServerOperation(); - } + // Note: we do NOT proactively close the server operation when next() returns false. + // While all data is client-side at this point, closing here is too aggressive: + // - Batch execution reuses the statement and the early close can break subsequent commands + // - WireMock fake service stubs are consumed by the extra closeStatement RPC + // - The explicit ResultSet.close() is the safer signal for proactive server close + // Server operation is closed proactively only in close() below. return hasNext; } From 27d1d6afa9a41e5995a6bd786d0fc86376031394 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 12 May 2026 17:04:32 +0530 Subject: [PATCH 09/13] Restore proactive close in next() for all types, fix flag reset ordering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of integration test failures: resetForNewExecution() reset directResultsReceived and serverOperationClosed BEFORE closing the old ResultSet. When resultSet.close() triggered closeServerOperation(), the flags were already false, causing an unnecessary closeStatement RPC on a direct-results operation (server already closed it) → HTTP 404. Fix: close the previous ResultSet BEFORE resetting flags so closeServerOperation() sees the correct state and skips the RPC for direct results. Proactive close now fires for all statement types in next() — no type-based exclusions. The guards in closeServerOperation() (directResultsReceived, serverOperationClosed, null statementId) handle all edge cases correctly. Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../databricks/jdbc/api/impl/DatabricksResultSet.java | 11 +++++------ .../databricks/jdbc/api/impl/DatabricksStatement.java | 10 +++++++--- 2 files changed, 12 insertions(+), 9 deletions(-) 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 bd3736a1f..510095009 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -283,12 +283,11 @@ public boolean next() throws SQLException { cachedTelemetryCollector.recordResultSetIteration( statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext); } - // Note: we do NOT proactively close the server operation when next() returns false. - // While all data is client-side at this point, closing here is too aggressive: - // - Batch execution reuses the statement and the early close can break subsequent commands - // - WireMock fake service stubs are consumed by the extra closeStatement RPC - // - The explicit ResultSet.close() is the safer signal for proactive server close - // Server operation is closed proactively only in close() below. + if (!hasNext) { + // All rows consumed — proactively close server operation to release resources. + // The client-side Statement remains open for reuse. + closeServerOperation(); + } return hasNext; } diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index d57678594..991f98254 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -1005,10 +1005,10 @@ private void resetForNewExecution() { // when the server returns unexpected responses (e.g., WireMock 404 in tests). // For direct results, the server already closed the handle. - directResultsReceived = false; - serverOperationClosed = false; - // Per JDBC spec, re-executing a Statement implicitly closes the current ResultSet. + // Close BEFORE resetting flags — resultSet.close() → closeServerOperation() needs + // to see the current directResultsReceived/serverOperationClosed state to decide + // whether to send closeStatement RPC. if (resultSet != null) { try { resultSet.close(); @@ -1018,6 +1018,10 @@ private void resetForNewExecution() { resultSet = null; } + // Reset flags AFTER closing old ResultSet + directResultsReceived = false; + serverOperationClosed = false; + // Null out statementId so that if the new execution fails before setStatementId(), // close() takes the statementId==null branch instead of sending closeStatement(stale-id) statementId = null; From bbcf0490b756d548aa6e62be757368f60ae4126f Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Tue, 19 May 2026 19:30:01 +0530 Subject: [PATCH 10/13] Address review feedback: fix double RPC, remove inline close from next(), add tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F1 (CRITICAL): Fix double closeStatement RPC by closing ResultSet BEFORE the Statement-level RPC check in close(). ResultSet.close() triggers proactive close and sets serverOperationClosed=true, preventing duplicate. F2 (HIGH): Remove proactive close from next()→false inline path. Server operations are now only closed proactively on ResultSet.close(), avoiding synchronous RPCs on every result exhaustion including batch DML and empty result sets. F3 (HIGH): Upgrade failed proactive close logging from DEBUG to WARN. F4 (MEDIUM): Add closeServerOperation() as default no-op on IDatabricksStatementInternal interface. ResultSet now calls through the interface instead of instanceof DatabricksStatement cast. F6 (MEDIUM): Add tests for ResultSet.close() triggering proactive close, and Statement.close() with non-closed ResultSet firing exactly one RPC. F9 (MEDIUM): Update changelog to document getExecutionResult() behavior change (returns cached ResultSet after close). F10 (LOW): Narrow catch to SQLException|RuntimeException and include throwable in WARN log. Also fix resetForNewExecution to set serverOperationClosed=true after explicit close and check the flag to avoid double RPC. Co-authored-by: Isaac Signed-off-by: Gopal Lal --- NEXT_CHANGELOG.md | 2 +- .../jdbc/api/impl/DatabricksResultSet.java | 9 +-- .../jdbc/api/impl/DatabricksStatement.java | 32 +++++----- .../IDatabricksStatementInternal.java | 11 ++++ .../api/impl/DatabricksStatementTest.java | 63 +++++++++++++++++++ 5 files changed, 95 insertions(+), 22 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 8adb39ca4..2c54a3ca8 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -43,7 +43,7 @@ upgrading. These changes do not affect metadata on All-Purpose Clusters. - `EnableGeoSpatialSupport` no longer requires `EnableComplexDatatypeSupport=1`. Geospatial types (GEOMETRY, GEOGRAPHY) can now be enabled independently of complex type support (ARRAY, MAP, STRUCT). - Arrow schema deserialization failures (Thrift metadata path) now surface a dedicated driver error code `ARROW_SCHEMA_PARSING_ERROR` (vendor code `22000`) and a proper SQLSTATE `22000` (Data Exception) on the thrown `SQLException`, instead of the generic `RESULT_SET_ERROR` (1004) and the enum name as SQLSTATE. The exception message is unchanged. - When a Statement is re-executed, the previous server-side operation is now explicitly closed before starting the new execution, preventing orphaned server-side operations when Statements are reused. -- Server-side operations are now closed proactively when all results are consumed (`next()` returns false) or `ResultSet.close()` is called, improving resource utilization. The client-side Statement remains open and reusable for re-execution. +- Server-side operations are now closed proactively when `ResultSet.close()` is called, improving resource utilization. The client-side Statement remains open and reusable for re-execution. As a result, `getExecutionResult()` after result consumption returns the cached ResultSet instead of making a server RPC. ### Fixed - Fixed `DatabaseMetaData.getTables()` in Thrift mode returning rows when called with an empty `types` array. Per JDBC spec, empty types means "no types selected" and now correctly returns zero rows (matching SEA mode). 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 510095009..db436247c 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java @@ -283,11 +283,6 @@ public boolean next() throws SQLException { cachedTelemetryCollector.recordResultSetIteration( statementId.toSQLExecStatementId(), resultSetMetaData.getChunkCount(), hasNext); } - if (!hasNext) { - // All rows consumed — proactively close server operation to release resources. - // The client-side Statement remains open for reuse. - closeServerOperation(); - } return hasNext; } @@ -304,8 +299,8 @@ public void close() throws DatabricksSQLException { /** Proactively closes the server-side operation via the parent statement. */ private void closeServerOperation() { - if (parentStatement instanceof DatabricksStatement) { - ((DatabricksStatement) parentStatement).closeServerOperation(); + if (parentStatement != null) { + parentStatement.closeServerOperation(); } } diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index 9fb0cf1e0..f05c87010 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -152,6 +152,13 @@ public void close(boolean removeFromSession) throws DatabricksSQLException { LOGGER.warn(warningMsg); warnings = WarningUtil.addWarning(warnings, warningMsg); } else { + // Close ResultSet first — this triggers proactive server close via + // closeServerOperation() and sets serverOperationClosed=true, preventing + // a duplicate closeStatement RPC below. + if (resultSet != null) { + this.resultSet.close(); + this.resultSet = null; + } // Skip server-side close if operation was already closed: // - directResultsReceived: server closed it (inline results) // - serverOperationClosed: client proactively closed it (results consumed or RS closed) @@ -165,10 +172,6 @@ public void close(boolean removeFromSession) throws DatabricksSQLException { directResultsReceived, serverOperationClosed); } - if (resultSet != null) { - this.resultSet.close(); - this.resultSet = null; - } } } finally { // Always run cleanup even if resultSet.close() or closeStatement() throws. @@ -967,7 +970,8 @@ public void markDirectResultsReceived() { *

After this call, {@link #close(boolean)} will skip the closeStatement RPC since the server * operation is already closed. The Statement can still be re-executed. */ - void closeServerOperation() { + @Override + public void closeServerOperation() { if (serverOperationClosed || directResultsReceived || statementId == null || isClosed) { return; } @@ -979,13 +983,14 @@ void closeServerOperation() { this.serverOperationClosed = true; LOGGER.debug( "Proactively closed server operation for statement {} (results consumed)", statementId); - } catch (Exception e) { - // Best-effort — don't fail the user's iteration/close for a server cleanup failure. + } catch (SQLException | RuntimeException e) { + // Best-effort — don't fail the user's close for a server cleanup failure. // serverOperationClosed stays false so Statement.close() will retry the RPC. - LOGGER.debug( + LOGGER.warn( "Failed to proactively close server operation for statement {}: {}", statementId, - e.getMessage()); + e.getMessage(), + e); } } @@ -1013,9 +1018,10 @@ private void resetForNewExecution() { // // Skip if: (1) no previous execution (statementId==null), or // (2) server already closed the operation (direct results). - if (statementId != null && !directResultsReceived) { + if (statementId != null && !directResultsReceived && !serverOperationClosed) { try { connection.getSession().getDatabricksClient().closeStatement(statementId); + serverOperationClosed = true; } catch (Exception e) { // Don't block re-execution if closing the previous operation fails. // This covers: network errors, operation already expired/evicted on server, @@ -1028,10 +1034,8 @@ private void resetForNewExecution() { } } - // Per JDBC spec, re-executing a Statement implicitly closes the current ResultSet. - // Close BEFORE resetting flags — resultSet.close() → closeServerOperation() needs - // to see the current directResultsReceived/serverOperationClosed state to decide - // whether to send closeStatement RPC. + // Close the previous ResultSet. closeServerOperation() inside resultSet.close() + // will be a no-op since serverOperationClosed is already true from above. if (resultSet != null) { try { resultSet.close(); diff --git a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java index ccb6f035e..965a971e7 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksStatementInternal.java @@ -37,4 +37,15 @@ public interface IDatabricksStatementInternal { default void markDirectResultsReceived() { // no-op by default } + + /** + * Proactively closes the server-side operation to release server resources while keeping the + * client-side Statement open for reuse. Default no-op for implementations that don't support + * proactive close. + */ + default void closeServerOperation() { + // no-op by default + } + + long getLargeMaxRows() throws DatabricksSQLException; } diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java index 54bb9393f..681eca915 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java @@ -2016,4 +2016,67 @@ public void testStatementReusableAfterProactiveClose() throws Exception { eq(statement), any()); } + + @Test + public void testResultSetClose_triggersProactiveServerClose() throws Exception { + // Verify that ResultSet.close() triggers closeServerOperation on the parent Statement, + // and that Statement.close() then skips the duplicate RPC. + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + + IExecutionResult execResult = mock(IExecutionResult.class); + DatabricksResultSetMetaData rsMeta = mock(DatabricksResultSetMetaData.class); + DatabricksResultSet resultSet = + new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.QUERY, + statement, + execResult, + rsMeta, + false); + statement.setStatementId(STATEMENT_ID); + statement.resultSet = resultSet; + + // Close ResultSet — should trigger proactive server close + resultSet.close(); + verify(client, times(1)).closeStatement(STATEMENT_ID); + + // Close Statement — should skip RPC since server operation already closed + statement.close(); + verify(client, times(1)).closeStatement(STATEMENT_ID); // still 1, not 2 + } + + @Test + public void testStatementClose_noDoubleRpc_whenResultSetNotClosed() throws Exception { + // Verify that Statement.close() with a non-closed ResultSet fires exactly one RPC: + // ResultSet.close() inside Statement.close() triggers proactive close, and the + // subsequent check sees serverOperationClosed=true and skips. + IDatabricksConnectionContext connectionContext = + DatabricksConnectionContext.parse(JDBC_URL, new Properties()); + DatabricksConnection connection = new DatabricksConnection(connectionContext, client); + DatabricksStatement statement = new DatabricksStatement(connection); + + IExecutionResult execResult = mock(IExecutionResult.class); + DatabricksResultSetMetaData rsMeta = mock(DatabricksResultSetMetaData.class); + DatabricksResultSet resultSet = + new DatabricksResultSet( + new StatementStatus().setState(StatementState.SUCCEEDED), + STATEMENT_ID, + StatementType.QUERY, + statement, + execResult, + rsMeta, + false); + statement.setStatementId(STATEMENT_ID); + statement.resultSet = resultSet; + + // Close Statement directly (without closing ResultSet first) + statement.close(); + + // Should fire exactly one closeStatement RPC, not two + verify(client, times(1)).closeStatement(STATEMENT_ID); + } } From a239225c2cac8b03d7faa4a21fe0d283079c6cdd Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Wed, 20 May 2026 08:15:54 +0530 Subject: [PATCH 11/13] Apply F3/F10 fixes to resetForNewExecution close path, clarify PR #1422 provenance - Narrow catch to SQLException|RuntimeException and upgrade to WARN with throwable in resetForNewExecution close path (consistency with closeServerOperation) - Update comment to clarify this code is from PR #1422 and add skip condition (3) for serverOperationClosed Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksStatement.java | 28 ++++++++----------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index f05c87010..e04411bb8 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -1003,34 +1003,28 @@ private void resetForNewExecution() { noMoreResults = false; updateCount = -1; - // Close the previous server-side operation if it exists. This prevents resource - // leaks when a Statement is re-executed (e.g., PreparedStatement in a loop). - // This matches the behavior of pgJDBC, MySQL Connector/J, Trino JDBC, and - // Databricks Python SQL Connector — all close the previous operation on re-execute. + // Close the previous server-side operation if still open. This prevents resource + // leaks when a Statement is re-executed without closing the previous ResultSet. + // If the user closed the ResultSet before re-executing (best practice), the + // proactive close already set serverOperationClosed=true and this is a no-op. // - // Note on directResultsReceived: we check the flag value from the PREVIOUS execution - // here. The flag is reset to false below, after this close attempt. - // - // Note on latency: this close is synchronous (adds one RPC round-trip before the next - // execution). This is consistent with pgJDBC's closeForNextExecution() which is also - // synchronous. The correctness benefit (no orphaned server operations) outweighs the - // latency cost for typical usage patterns. - // - // Skip if: (1) no previous execution (statementId==null), or - // (2) server already closed the operation (direct results). + // Skip if: (1) no previous execution (statementId==null), + // (2) server already closed the operation (direct/inline results), or + // (3) client already proactively closed it (via ResultSet.close()). if (statementId != null && !directResultsReceived && !serverOperationClosed) { try { connection.getSession().getDatabricksClient().closeStatement(statementId); serverOperationClosed = true; - } catch (Exception e) { + } catch (SQLException | RuntimeException e) { // Don't block re-execution if closing the previous operation fails. // This covers: network errors, operation already expired/evicted on server, // and transport-level errors (e.g., unexpected server responses). // The new execution will create a fresh operation with a new statementId. - LOGGER.debug( + LOGGER.warn( "Failed to close previous server operation {} during re-execution: {}", statementId, - e.getMessage()); + e.getMessage(), + e); } } From 5f852d263e13139db75281c0da4a95ac96a123a6 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Wed, 20 May 2026 08:20:03 +0530 Subject: [PATCH 12/13] Make re-execution close fire-and-forget async to avoid blocking next execution resetForNewExecution now closes the previous server operation on a daemon thread instead of synchronously. This eliminates the RPC latency cost on the re-execution hot path (e.g., PreparedStatement in a loop). The flag is set immediately to prevent duplicate close from resultSet.close(). Tests updated to use verify with timeout for async close verification. Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../jdbc/api/impl/DatabricksStatement.java | 46 ++++++++++--------- .../api/impl/DatabricksStatementTest.java | 8 +++- 2 files changed, 31 insertions(+), 23 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index e04411bb8..6ca99a519 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -1003,33 +1003,37 @@ private void resetForNewExecution() { noMoreResults = false; updateCount = -1; - // Close the previous server-side operation if still open. This prevents resource - // leaks when a Statement is re-executed without closing the previous ResultSet. + // Close the previous server-side operation if still open. Fire-and-forget on a + // daemon thread so the new execution is not blocked by the close RPC latency. // If the user closed the ResultSet before re-executing (best practice), the // proactive close already set serverOperationClosed=true and this is a no-op. - // - // Skip if: (1) no previous execution (statementId==null), - // (2) server already closed the operation (direct/inline results), or - // (3) client already proactively closed it (via ResultSet.close()). if (statementId != null && !directResultsReceived && !serverOperationClosed) { - try { - connection.getSession().getDatabricksClient().closeStatement(statementId); - serverOperationClosed = true; - } catch (SQLException | RuntimeException e) { - // Don't block re-execution if closing the previous operation fails. - // This covers: network errors, operation already expired/evicted on server, - // and transport-level errors (e.g., unexpected server responses). - // The new execution will create a fresh operation with a new statementId. - LOGGER.warn( - "Failed to close previous server operation {} during re-execution: {}", - statementId, - e.getMessage(), - e); - } + // Mark as closed immediately so resultSet.close() below doesn't fire a duplicate. + serverOperationClosed = true; + final StatementId prevStatementId = statementId; + final IDatabricksClient prevClient = connection.getSession().getDatabricksClient(); + Thread closeThread = + new Thread( + () -> { + try { + prevClient.closeStatement(prevStatementId); + LOGGER.debug( + "Closed previous server operation {} during re-execution", prevStatementId); + } catch (SQLException | RuntimeException e) { + LOGGER.warn( + "Failed to close previous server operation {} during re-execution: {}", + prevStatementId, + e.getMessage(), + e); + } + }); + closeThread.setDaemon(true); + closeThread.setName("close-stmt-" + prevStatementId); + closeThread.start(); } // Close the previous ResultSet. closeServerOperation() inside resultSet.close() - // will be a no-op since serverOperationClosed is already true from above. + // is a no-op since serverOperationClosed is already true. if (resultSet != null) { try { resultSet.close(); diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java index 681eca915..a4bab6411 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksStatementTest.java @@ -1065,10 +1065,10 @@ public void testReExecutionClosesServerOperation() throws Exception { StatementId firstStatementId = new StatementId("first-stmt-id"); statement.setStatementId(firstStatementId); - // Second execution — should close the first server operation + // Second execution — should close the first server operation asynchronously statement.executeQuery(STATEMENT); - verify(client, times(1)).closeStatement(eq(firstStatementId)); + verify(client, timeout(5000).times(1)).closeStatement(eq(firstStatementId)); verify(firstResult, times(1)).close(); assertEquals(secondResult, statement.getResultSet()); } @@ -1139,6 +1139,8 @@ public void testReExecutionHandlesCloseFailureGracefully() throws Exception { // Re-execution should succeed even though closing previous operation failed assertDoesNotThrow(() -> statement.executeQuery(STATEMENT)); assertEquals(secondResult, statement.getResultSet()); + // Verify the async close was attempted + verify(client, timeout(5000).times(1)).closeStatement(eq(firstStatementId)); } @Test @@ -1177,6 +1179,8 @@ public void testReExecutionHandlesTransportErrorGracefully() throws Exception { // The new execution creates a fresh server operation with a new statementId. assertDoesNotThrow(() -> statement.executeQuery(STATEMENT)); assertEquals(secondResult, statement.getResultSet()); + // Verify the async close was attempted + verify(client, timeout(5000).times(1)).closeStatement(eq(firstStatementId)); } @Test From 17fb42e40e9abea6e44fc0e70549176dfcdaeacd Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Wed, 20 May 2026 08:21:29 +0530 Subject: [PATCH 13/13] Set serverOperationClosed after thread start, not before Co-authored-by: Isaac Signed-off-by: Gopal Lal --- .../com/databricks/jdbc/api/impl/DatabricksStatement.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java index 6ca99a519..42c10b18e 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksStatement.java @@ -1008,8 +1008,6 @@ private void resetForNewExecution() { // If the user closed the ResultSet before re-executing (best practice), the // proactive close already set serverOperationClosed=true and this is a no-op. if (statementId != null && !directResultsReceived && !serverOperationClosed) { - // Mark as closed immediately so resultSet.close() below doesn't fire a duplicate. - serverOperationClosed = true; final StatementId prevStatementId = statementId; final IDatabricksClient prevClient = connection.getSession().getDatabricksClient(); Thread closeThread = @@ -1030,10 +1028,11 @@ private void resetForNewExecution() { closeThread.setDaemon(true); closeThread.setName("close-stmt-" + prevStatementId); closeThread.start(); + serverOperationClosed = true; } // Close the previous ResultSet. closeServerOperation() inside resultSet.close() - // is a no-op since serverOperationClosed is already true. + // is a no-op since serverOperationClosed was set above. if (resultSet != null) { try { resultSet.close();