Make sure sending cancel request in StatementClient.close

lib/trino/client/statement_client.rb

@@ -262,9 +262,11 @@ def cancel_leaf_stage
     end
 
     def close
-      return unless running?
+      return if finished? || query_failed? || client_aborted?
 
-      @state = :client_aborted
+      if running?
+        @state = :client_aborted
+      end
 
       begin
         if uri = @results.next_uri

spec/statement_client_spec.rb

@@ -608,6 +608,13 @@
       expect do
         client.advance
       end.to raise_error(Trino::Client::TrinoQueryTimeoutError, "Query queryid timed out")
+
+      # close sends a cancel request even after TrinoQueryTimeoutError
+      expect do
+        client.close
+      end.to raise_error(WebMock::NetConnectNotAllowedError, /.*Unregistered request: DELETE.*/)
+
+      expect(client.client_error?).to eq true
     end
 
     it "doesn't raise errors if query is done" do
@@ -629,6 +636,32 @@
 
       sleep 1
       client.advance # set finished
+
+      # close doesn't send a cancel request if query has finished
+      client.close
+      expect(client.finished?).to eq true
+    end
+
+    it "sends a cancel request when close is called during query execution" do
+      stub_request(:post, "localhost/v1/statement").
+        with(body: query, headers: headers).
+        to_return(body: planning_response.to_json)
+
+      client = StatementClient.new(faraday, query, options)
+
+      stub_request(:get, "localhost/v1/next_uri").
+        with(headers: headers).
+        to_return(body: planning_response.to_json)
+      client.advance
+
+      expect(client.running?).to eq true
+
+      # close sends a cancel request if query is executing
+      expect do
+        client.close
+      end.to raise_error(WebMock::NetConnectNotAllowedError, /.*Unregistered request: DELETE.*/)
+
+      expect(client.client_aborted?).to eq true
     end
   end
 end