diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 434c4a84f..420d88fd3 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -18,6 +18,7 @@ - **Enhanced `enableMultipleCatalogSupport` behavior**: When this parameter is disabled (`enableMultipleCatalogSupport=0`), metadata operations (such as `getSchemas()`, `getTables()`, `getColumns()`, etc.) now return results only when the catalog parameter is either `null` or matches the current catalog. For any other catalog name, an empty result set is returned. This ensures metadata queries are restricted to the current catalog context. When enabled (`enableMultipleCatalogSupport=1`), metadata operations continue to work across all accessible catalogs. ### Fixed +- Fixed query timeout to account for time spent in the initial `ExecuteStatement` HTTP call. Previously, the timeout timer started only after the call returned, so that duration was not counted toward the timeout. This applies to both Thrift and SEA flows. - Fixed timeout exception handling to throw `SQLTimeoutException` instead of `DatabricksHttpException` when queries timeout during result fetching phase. This completes the timeout exception fix to handle both query execution polling and result fetching phases. - Fixed `getTypeInfo()` and `getClientInfoProperties()` to return fresh ResultSet instances on each call instead of shared static instances. This resolves issues where calling these methods multiple times would fail due to exhausted cursor state (Issue #1178). - Fixed complex data type metadata support when retrieving 0 rows in Arrow format diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java index edc85e454..51c0507ea 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java @@ -32,7 +32,30 @@ public TimeoutHandler( String operationDescription, Runnable onTimeoutAction, DatabricksDriverErrorCode internalErrorCode) { - this.startTimeMillis = System.currentTimeMillis(); + this( + System.currentTimeMillis(), + timeoutSeconds, + operationDescription, + onTimeoutAction, + internalErrorCode); + } + + /** + * Creates a new timeout handler with an explicit start time. + * + * @param startTimeMillis The start time in milliseconds (from System.currentTimeMillis()) + * @param timeoutSeconds Timeout in seconds, 0 means no timeout + * @param operationDescription Description of the operation for logging + * @param onTimeoutAction Runnable to call when a timeout occurs + * @param internalErrorCode Internal driver error code to annotate timeout exceptions + */ + public TimeoutHandler( + long startTimeMillis, + int timeoutSeconds, + String operationDescription, + Runnable onTimeoutAction, + DatabricksDriverErrorCode internalErrorCode) { + this.startTimeMillis = startTimeMillis; this.timeoutSeconds = timeoutSeconds; this.operationDescription = operationDescription; this.onTimeoutAction = onTimeoutAction; @@ -100,4 +123,29 @@ public static TimeoutHandler forStatement( }, internalErrorCode); } + + /** + * Factory method with an explicit start time, to account for time already spent before the + * handler is created (e.g., the initial execute HTTP call). + */ + public static TimeoutHandler forStatement( + long startTimeMillis, + int timeoutSeconds, + StatementId statementId, + IDatabricksClient client, + DatabricksDriverErrorCode internalErrorCode) { + + return new TimeoutHandler( + startTimeMillis, + timeoutSeconds, + "Statement ID: " + statementId, + () -> { + try { + client.cancelStatement(statementId); + } catch (Exception e) { + LOGGER.warn("Cancel statement on timeout failed: " + e.getMessage()); + } + }, + internalErrorCode); + } } 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 dd3911740..3745bae03 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 @@ -237,9 +237,10 @@ public DatabricksResultSet executeStatement( int timeoutInSeconds = parentStatement != null ? parentStatement.getStatement().getQueryTimeout() : 0; - // Create timeout handler + // Create timeout handler with start time from before the execute statement call TimeoutHandler timeoutHandler = TimeoutHandler.forStatement( + executionStartTime, timeoutInSeconds, typedStatementId, this, 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 14af79bcb..57f4271f0 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 @@ -197,6 +197,8 @@ DatabricksResultSet execute( StatementType statementType) throws SQLException { + long executeStartTime = System.currentTimeMillis(); + try { // Set direct result configuration if (enableDirectResults) { @@ -232,7 +234,8 @@ DatabricksResultSet execute( TGetOperationStatusResp statusResp = pollTillOperationFinished( - response, parentStatement, session, statementId, sessionDebugInfo); + response, parentStatement, session, statementId, sessionDebugInfo, executeStartTime); + if (hasResultDataInDirectResults(response)) { // The first response has result data // There is no polling in this case as status was already finished @@ -279,7 +282,8 @@ private TGetOperationStatusResp pollTillOperationFinished( IDatabricksStatementInternal parentStatement, IDatabricksSession session, StatementId statementId, - String sessionDebugInfo) + String sessionDebugInfo, + long executeStartTimeMillis) throws SQLException, TException { int timeoutInSeconds = (parentStatement == null) ? 0 : parentStatement.getStatement().getQueryTimeout(); @@ -297,7 +301,10 @@ private TGetOperationStatusResp pollTillOperationFinished( TimeoutHandler timeoutHandler = getTimeoutHandler( - response, timeoutInSeconds, DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); + response, + timeoutInSeconds, + DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT, + executeStartTimeMillis); // Polling until query operation state is finished long pollingStartTime = System.nanoTime(); @@ -819,10 +826,12 @@ void setServerProtocolVersion(TProtocolVersion protocolVersion) { private TimeoutHandler getTimeoutHandler( TExecuteStatementResp response, int timeoutInSeconds, - DatabricksDriverErrorCode internalErrorCode) { + DatabricksDriverErrorCode internalErrorCode, + long startTimeMillis) { final TOperationHandle operationHandle = response.getOperationHandle(); return new TimeoutHandler( + startTimeMillis, timeoutInSeconds, "Thrift Operation Handle: " + operationHandle.toString(), () -> { diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java index ae75478fc..7aad3a2fe 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java @@ -195,4 +195,55 @@ void testForStatementCancelFailure() throws Exception { // Verify DatabricksTimeoutException is still thrown assertTrue(exception.getMessage().contains("timed-out after 5 seconds")); } + + @Test + void testConstructorWithExplicitStartTime() { + // Create handler with startTimeMillis set to 3 seconds in the past + long pastStartTime = System.currentTimeMillis() - TimeUnit.SECONDS.toMillis(3); + TimeoutHandler handler = + new TimeoutHandler( + pastStartTime, + 2, + "Test operation", + mockTimeoutAction, + DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); + + // This should immediately throw a DatabricksTimeoutException since 3s > 2s timeout + DatabricksTimeoutException exception = + assertThrows(DatabricksTimeoutException.class, handler::checkTimeout); + + // Verify timeout action was called + verify(mockTimeoutAction, times(1)).run(); + + // Verify exception message + assertTrue(exception.getMessage().contains("timed-out after 2 seconds")); + assertTrue(exception.getMessage().contains("Test operation")); + } + + @Test + void testForStatementFactoryWithExplicitStartTime() throws Exception { + when(mockStatementId.toString()).thenReturn("test-statement-id"); + + // Create handler with factory method using startTimeMillis set to 6 seconds in the past + long pastStartTime = System.currentTimeMillis() - TimeUnit.SECONDS.toMillis(6); + TimeoutHandler handler = + TimeoutHandler.forStatement( + pastStartTime, + 5, + mockStatementId, + mockClient, + DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); + + // Verify handler was created with correct start time (no reflection needed) + // This should immediately throw since 6s > 5s timeout + DatabricksTimeoutException exception = + assertThrows(DatabricksTimeoutException.class, handler::checkTimeout); + + // Verify client.cancelStatement was called + verify(mockClient, times(1)).cancelStatement(mockStatementId); + + // Verify exception message + assertTrue(exception.getMessage().contains("timed-out after 5 seconds")); + assertTrue(exception.getMessage().contains("test-statement-id")); + } }