From 66376b8c4ea59febab94917e5a1f9037d6cd0d71 Mon Sep 17 00:00:00 2001
From: Naoki Takezoe <takezoe@gmail.com>
Date: Sat, 7 Dec 2024 02:05:20 +0900
Subject: [PATCH] Make sure sending cancel request in StatementClient.close

---
 lib/trino/client/statement_client.rb |  6 +++--
 spec/statement_client_spec.rb        | 33 ++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/lib/trino/client/statement_client.rb b/lib/trino/client/statement_client.rb
index 1f2c94e6..6a2e059d 100644
--- a/lib/trino/client/statement_client.rb
+++ b/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
diff --git a/spec/statement_client_spec.rb b/spec/statement_client_spec.rb
index 17805d23..854fb169 100644
--- a/spec/statement_client_spec.rb
+++ b/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