From a2dc9e5ae13df2cea9697a8122433d1ee63d1427 Mon Sep 17 00:00:00 2001
From: Karrot <karrot@dunamu.com>
Date: Tue, 10 Dec 2024 23:17:28 +0900
Subject: [PATCH] fix: Retry in POST method

---
 lib/trino/client/statement_client.rb | 46 ++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 10 deletions(-)

diff --git a/lib/trino/client/statement_client.rb b/lib/trino/client/statement_client.rb
index 1f2c94e6..699a9540 100644
--- a/lib/trino/client/statement_client.rb
+++ b/lib/trino/client/statement_client.rb
@@ -66,20 +66,46 @@ def init_request(req)
 
     def post_query_request!
       uri = "/v1/statement"
-      response = @faraday.post do |req|
-        req.url uri
 
-        req.body = @query
-        init_request(req)
-      end
+      start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      attempts = 0
 
-      # TODO error handling
-      if response.status != 200
-        exception! TrinoHttpError.new(response.status, "Failed to start query: #{response.body} (#{response.status})")
+      loop do
+        begin
+          response = @faraday.post do |req|
+            req.url uri
+            req.body = @query
+            init_request(req)
+          end
+        rescue Faraday::TimeoutError, Faraday::ConnectionFailed
+          # temporally error to retry
+          response = nil
+        rescue => e
+          exception! e
+        end
+
+        if response
+          if response.status == 200 && !response.body.to_s.empty?
+            @results_headers = response.headers
+            @results = decode_model(uri, parse_body(response), @models::QueryResults)
+            return
+          end
+          # retry if 502, 503, 504 according to the trino protocol
+          unless [502, 503, 504].include?(response.status)
+            # deterministic error
+            exception! TrinoHttpError.new(response.status, "Trino API error at #{uri} returned #{response.status}: #{response.body}")
+          end
+        end
+
+        raise_if_timeout!
+
+        attempts += 1
+        sleep attempts * 0.1
+
+        break unless (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) < @retry_timeout && !client_aborted?
       end
 
-      @results_headers = response.headers
-      @results = decode_model(uri, parse_body(response), @models::QueryResults)
+      exception! TrinoHttpError.new(408, "Trino API error due to timeout")
     end
 
     private :post_query_request!