Close StatementClient when hitting client-side timeout (#136)

treasure-data · Feb 6, 2025 · b3261b1 · b3261b1
1 parent da87c84
commit b3261b1
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/lib/trino/client/statement_client.rb b/lib/trino/client/statement_client.rb
@@ -233,13 +233,15 @@ def raise_if_timeout!
         elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - @started_at
 
         if @query_timeout && elapsed > @query_timeout
+          close
           raise_timeout_error!
         end
 
         if @plan_timeout && (@results == nil || @results.columns == nil) &&
             elapsed > @plan_timeout
           # @results is not set (even first faraday_get_with_retry isn't called yet) or
           # result from Trino doesn't include result schema. Query planning isn't done yet.
+          close
           raise_timeout_error!
         end
       end

diff --git a/spec/statement_client_spec.rb b/spec/statement_client_spec.rb
@@ -540,12 +540,21 @@
         client.advance
 
         sleep 1
+
         stub_request(:get, "localhost/v1/next_uri").
           with(headers: headers).
           to_return(body: planning_response.to_json)
+
+        cancel = stub_request(:delete, "localhost/v1/next_uri").
+          with(headers: headers).
+          to_return(status: 204) # NoContent
+
         expect do
           client.advance
         end.to raise_error(Trino::Client::TrinoQueryTimeoutError, "Query queryid timed out")
+
+        expect(cancel).to have_been_requested
+        expect(client.client_error?).to eq true
       end
 
       it "raises TrinoQueryTimeoutError if timeout during initial resuming" do
@@ -602,12 +611,21 @@
       client.advance
 
       sleep 1
+
       stub_request(:get, "localhost/v1/next_uri").
         with(headers: headers).
         to_return(body: late_running_response.to_json)
+
+      cancel = stub_request(:delete, "localhost/v1/next_uri").
+        with(headers: headers).
+        to_return(status: 204) # NoContent
+
       expect do
         client.advance
       end.to raise_error(Trino::Client::TrinoQueryTimeoutError, "Query queryid timed out")
+
+      expect(cancel).to have_been_requested
+      expect(client.client_error?).to eq true
     end
 
     it "doesn't raise errors if query is done" do