[PECOBLR-1362] Include ExecuteStatement call duration in query timeout calculation

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

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);
+  }
 }

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,

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(),
         () -> {

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"));
+  }
 }