Extend the list of status codes retried in client

client/trino-client/src/main/java/io/trino/client/HttpStatusCodes.java

@@ -0,0 +1,35 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.client;
+
+import static java.net.HttpURLConnection.HTTP_BAD_GATEWAY;
+import static java.net.HttpURLConnection.HTTP_GATEWAY_TIMEOUT;
+import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;
+
+public final class HttpStatusCodes
+{
+    private HttpStatusCodes() {}
+
+    public static boolean shouldRetry(int statusCode)
+    {
+        switch (statusCode) {
+            case HTTP_BAD_GATEWAY:
+            case HTTP_UNAVAILABLE:
+            case HTTP_GATEWAY_TIMEOUT:
+                return true;
+            default:
+                return false;
+        }
+    }
+}

client/trino-client/src/main/java/io/trino/client/StatementClientV1.java

@@ -51,12 +51,12 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.net.HttpHeaders.ACCEPT_ENCODING;
 import static com.google.common.net.HttpHeaders.USER_AGENT;
+import static io.trino.client.HttpStatusCodes.shouldRetry;
 import static io.trino.client.JsonCodec.jsonCodec;
 import static io.trino.client.ProtocolHeaders.TRINO_HEADERS;
 import static java.lang.String.format;
 import static java.net.HttpURLConnection.HTTP_OK;
 import static java.net.HttpURLConnection.HTTP_UNAUTHORIZED;
-import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;
 import static java.util.Arrays.stream;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
@@ -391,7 +391,7 @@ public boolean advance()
                 return true;
             }
 
-            if (response.getStatusCode() != HTTP_UNAVAILABLE) {
+            if (!shouldRetry(response.getStatusCode())) {
                 state.compareAndSet(State.RUNNING, State.CLIENT_ERROR);
                 throw requestFailedException("fetching next", request, response);
             }

client/trino-client/src/main/java/io/trino/client/auth/external/HttpTokenPoller.java

@@ -35,12 +35,12 @@
 
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.net.HttpHeaders.USER_AGENT;
+import static io.trino.client.HttpStatusCodes.shouldRetry;
 import static io.trino.client.JsonCodec.jsonCodec;
 import static io.trino.client.JsonResponse.execute;
 import static java.lang.String.format;
 import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
 import static java.net.HttpURLConnection.HTTP_OK;
-import static java.net.HttpURLConnection.HTTP_UNAVAILABLE;
 import static java.time.temporal.ChronoUnit.MILLIS;
 import static java.util.Objects.requireNonNull;
 
@@ -142,7 +142,7 @@ private TokenPollResult executePoll(Request request)
 
         String message = format("Request to %s failed: %s [Error: %s]", request.url(), response, response.getResponseBody().orElse("<Response Too Large>"));
 
-        if (response.getStatusCode() == HTTP_UNAVAILABLE) {
+        if (shouldRetry(response.getStatusCode())) {
             throw new IOException(message);
         }
 

docs/src/main/sphinx/develop/client-protocol.rst

@@ -26,9 +26,12 @@ The caller may set various :ref:`client-request-headers`. The headers are
 only required on the initial ``POST`` request, and not when following the
 ``nextUri`` links.
 
-If the client request returns an HTTP 503, that means the server was busy,
-and the client should try again in 50-100 milliseconds.  Any HTTP status other
-than 503 or 200 means that query processing has failed.
+If the client request returns an HTTP 502, 503 or 504, that means there was
+intermittent problem processing request and the client should try again
+in 50-100 milliseconds. Trino does not generate those codes by itself
+but those can be generated by gateways/load balancers in front of Trino.
+Any HTTP status other than 502, 503, 504 or 200 means that query processing
+has failed.
 
 The ``/v1/statement`` ``POST`` request returns a JSON document of type
 ``QueryResults``, as well as a collection of response headers.  The