Fixed issue with 503 being not handled. Reorganized code handing status code to be more certain about codes

Author: chernser

client-v2/src/main/java/com/clickhouse/client/api/Client.java

@@ -1716,7 +1716,7 @@ public CompletableFuture<QueryResponse> query(String sqlQuery, Map<String, Objec
                         return new QueryResponse(transportResp, responseFormat, requestSettings, metrics);
 
                     } catch (Exception e) {
-                        ClientUtils.quiteClose(transportResp, LOG);
+                        ClientUtils.quietClose(transportResp, LOG);
                         String msg = requestExMsg("Query", (i + 1), durationSince(startTime).toMillis(), requestSettings.getQueryId());
                         lastException = httpClientHelper.wrapException(msg, e, requestSettings.getQueryId());
                         if (httpClientHelper.shouldRetry(e, requestSettings.getAllSettings())) {

client-v2/src/main/java/com/clickhouse/client/api/internal/ClientUtils.java

@@ -3,7 +3,6 @@
 import org.slf4j.Logger;
 
 import java.io.Closeable;
-import java.io.IOException;
 
 /**
  * Class containing utility methods used across the client.
@@ -20,7 +19,7 @@ public static boolean isBlank(String str) {
         return str == null || str.trim().isEmpty();
     }
 
-    public static void quiteClose(Closeable closeable, Logger log) {
+    public static void quietClose(Closeable closeable, Logger log) {
         if (closeable != null) {
             try {
                 closeable.close();

client-v2/src/main/java/com/clickhouse/client/api/internal/HttpAPIClientHelper.java

@@ -17,7 +17,6 @@
 import com.clickhouse.client.api.transport.Endpoint;
 import com.clickhouse.client.api.transport.internal.TransportRequest;
 import com.clickhouse.client.api.transport.internal.TransportResponse;
-import com.clickhouse.client.config.ClickHouseDefaultSslContextProvider;
 import com.clickhouse.data.ClickHouseFormat;
 import net.jpountz.lz4.LZ4Factory;
 import org.apache.commons.compress.compressors.CompressorStreamFactory;
@@ -46,6 +45,7 @@
 import org.apache.hc.core5.http.HttpHeaders;
 import org.apache.hc.core5.http.HttpHost;
 import org.apache.hc.core5.http.HttpRequest;
+import org.apache.hc.core5.http.HttpResponse;
 import org.apache.hc.core5.http.HttpStatus;
 import org.apache.hc.core5.http.NoHttpResponseException;
 import org.apache.hc.core5.http.ProtocolException;
@@ -369,10 +369,7 @@ public CloseableHttpClient createHttpClient(boolean initSslContext, Map<String,
      * @return exception object with server code
      */
     public Exception readError(HttpPost req, ClassicHttpResponse httpResponse) {
-        final Header serverQueryIdHeader = httpResponse.getFirstHeader(ClickHouseHttpProto.HEADER_QUERY_ID);
-        final Header clientQueryIdHeader = req.getFirstHeader(ClickHouseHttpProto.HEADER_QUERY_ID);
-        final Header queryHeader = Stream.of(serverQueryIdHeader, clientQueryIdHeader).filter(Objects::nonNull).findFirst().orElse(null);
-        final String queryId = queryHeader == null ? "" : queryHeader.getValue();
+        final String queryId = getQueryId(httpResponse, req);
         int serverCode = getHeaderInt(httpResponse.getFirstHeader(ClickHouseHttpProto.HEADER_EXCEPTION_CODE), 0);
         try {
             return serverCode > 0 ? readClickHouseError(httpResponse.getEntity(), serverCode, queryId, httpResponse.getCode()) :
@@ -608,11 +605,6 @@ private static final class TransportResponseImpl implements TransportResponse {
             this.delegate = delegate;
         }
 
-        @Override
-        public int getStatusCode() {
-            return 0;
-        }
-
         @Override
         public ClickHouseFormat getDataFormat() {
             Header formatHeader = delegate.getFirstHeader(ClickHouseHttpProto.HEADER_FORMAT);
@@ -632,6 +624,7 @@ public String getQueryId() {
         }
 
         @Override
+        @SuppressWarnings("unchecked")
         public <T> T getDelegate() {
             return (T) delegate;
         }
@@ -680,6 +673,7 @@ private ClassicHttpResponse doPostRequest(Map<String, Object> requestConfig, Htt
         doPoolVent();
 
         ClassicHttpResponse httpResponse = null;
+        boolean closeResponse = true;
         HttpContext context = createRequestHttpContext(requestConfig);
         try {
             httpResponse = httpClient.executeOpen(null, req, context);
@@ -688,35 +682,58 @@ private ClassicHttpResponse doPostRequest(Map<String, Object> requestConfig, Htt
                     httpResponse.getCode(),
                     requestConfig));
 
-            if (httpResponse.getCode() == HttpStatus.SC_PROXY_AUTHENTICATION_REQUIRED) {
-                throw new ClientMisconfigurationException("Proxy authentication required. Please check your proxy settings.");
-            } else if (httpResponse.getCode() == HttpStatus.SC_BAD_GATEWAY) {
-                httpResponse.close();
-                throw new ClientException("Server returned '502 Bad gateway'. Check network and proxy settings.");
-            } else if (httpResponse.getCode() >= HttpStatus.SC_BAD_REQUEST || httpResponse.containsHeader(ClickHouseHttpProto.HEADER_EXCEPTION_CODE)) {
-                try {
-                    throw readError(req, httpResponse);
-                } finally {
-                    httpResponse.close();
-                }
+            if (httpResponse.containsHeader(ClickHouseHttpProto.HEADER_EXCEPTION_CODE)) {
+                throw readError(req, httpResponse);
             }
-            return httpResponse;
 
+            int statusCode = httpResponse.getCode();
+            switch (statusCode) {
+                case HttpStatus.SC_OK:
+                    closeResponse = false;
+                    return httpResponse;
+                case HttpStatus.SC_PROXY_AUTHENTICATION_REQUIRED:
+                    throw new ClientMisconfigurationException("Proxy authentication required. Please check your proxy settings.");
+                case HttpStatus.SC_BAD_GATEWAY:
+                    throw new ClientException("Server returned '502 Bad gateway'. Check network and proxy settings.");
+                case HttpStatus.SC_SERVICE_UNAVAILABLE:
+                    throw new ServerException(0, "Server returned '503 Service Unavailable'. Check network settings.",
+                            HttpStatus.SC_SERVICE_UNAVAILABLE, getQueryId(httpResponse, req));
+                case HttpStatus.SC_BAD_REQUEST:
+                case HttpStatus.SC_UNAUTHORIZED:
+                case HttpStatus.SC_FORBIDDEN:
+                case HttpStatus.SC_SERVER_ERROR:
+                case HttpStatus.SC_NOT_FOUND:
+                    // ClickHouse usually uses SC_BAD_REQUEST and SC_SERVER_ERROR to return error.
+                    // SC_UNAUTHORIZED, SC_FORBIDDEN is for authentication
+                    // SC_NOT_FOUND can be returned by ClickHouse when path doesn't match database, but also by proxy
+                    // others we cannot handle properly
+                    throw readError(req, httpResponse);
+                default:
+                    throw new ClientException("Unexpected result status " + statusCode);
+            }
         } catch (UnknownHostException e) {
-            ClientUtils.quiteClose(httpResponse, LOG);
             LOG.warn("Host '{}' unknown", req.getAuthority());
             throw e;
         } catch (ConnectException | NoRouteToHostException e) {
-            ClientUtils.quiteClose(httpResponse, LOG);
             LOG.warn("Failed to connect to '{}': {}", req.getAuthority(), e.getMessage());
             throw e;
         } catch (Exception e) {
-            ClientUtils.quiteClose(httpResponse, LOG);
             LOG.debug("Failed to execute request to '{}': {}", req.getAuthority(), e.getMessage(), e);
             throw e;
+        } finally {
+            if (closeResponse) {
+                ClientUtils.quietClose(httpResponse, LOG);
+            }
         }
     }
 
+    private String getQueryId(HttpResponse httpResponse, HttpPost httpRequest) {
+        final Header serverQueryIdHeader = httpResponse.getFirstHeader(ClickHouseHttpProto.HEADER_QUERY_ID);
+        final Header clientQueryIdHeader = httpRequest.getFirstHeader(ClickHouseHttpProto.HEADER_QUERY_ID);
+        final Header queryHeader = Stream.of(serverQueryIdHeader, clientQueryIdHeader).filter(Objects::nonNull).findFirst().orElse(null);
+        return queryHeader == null ? "" : queryHeader.getValue();
+    }
+
     private static final ContentType CONTENT_TYPE = ContentType.create(ContentType.TEXT_PLAIN.getMimeType(), "UTF-8");
 
     private void addHeaders(HttpPost req, Map<String, Object> requestConfig) {

client-v2/src/main/java/com/clickhouse/client/api/transport/internal/TransportResponse.java

@@ -8,22 +8,6 @@
 
 public interface TransportResponse extends Closeable {
 
-    /**
-     * Transport status code translated to one of values:
-     * <ul>
-     *     <li>503 - for service unavailable</li>
-     *     <li>500 - server error</li>
-     *     <li>400 - user error</li>
-     *     <li>404 - endpoint not found</li>
-     *     <li>403 - access not granted</li>
-     *     <li>401 - no authentication information</li>
-     *     <li>200 - ok</li>
-     * </ul>
-     * @return integer value of status code
-     */
-    int getStatusCode();
-
-
     /**
      * Data format returned by server or calculated other way
      * @return data format

client-v2/src/test/java/com/clickhouse/client/HttpTransportTests.java

@@ -2114,6 +2114,102 @@ public static Object[][] testHttpStatusCompressedBodyDataProvider() {
         };
     }
 
+    @Test(groups = {"integration"})
+    public void test503ServiceUnavailableSurfacesAsServerException() throws Exception {
+        if (isCloud()) {
+            return; // mocked server
+        }
+
+        WireMockServer mockServer = new WireMockServer(WireMockConfiguration
+                .options().dynamicPort().notifier(new ConsoleNotifier(false)));
+        mockServer.start();
+
+        try {
+            mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
+                    .willReturn(WireMock.aResponse()
+                            .withStatus(HttpStatus.SC_SERVICE_UNAVAILABLE)
+                            .withBody("Service Unavailable"))
+                    .build());
+
+            try (Client client = new Client.Builder().addEndpoint(Protocol.HTTP, "localhost", mockServer.port(), false)
+                    .setUsername("default")
+                    .setPassword(ClickHouseServerForTest.getPassword())
+                    .compressServerResponse(false)
+                    .retryOnFailures(ClientFaultCause.None)
+                    .build()) {
+
+                Throwable thrown = Assert.expectThrows(Throwable.class,
+                        () -> client.query("SELECT 1").get(10, TimeUnit.SECONDS));
+
+                ServerException serverException = findServerException(thrown);
+                Assert.assertNotNull(serverException,
+                        "Expected 503 to be reported as a ServerException, but was: " + thrown);
+                Assert.assertEquals(serverException.getTransportProtocolCode(), HttpStatus.SC_SERVICE_UNAVAILABLE,
+                        "Expected transport protocol code 503, but was: " + serverException.getTransportProtocolCode());
+
+                Assert.assertTrue(containsMessageInCauseChain(thrown, "503 Service Unavailable"),
+                        "Expected '503 Service Unavailable' in failure message chain, but was: " + thrown);
+
+                Assert.assertNull(findCause(thrown, ConnectionInitiationException.class),
+                        "503 should not be reported as a ConnectionInitiationException, but was: " + thrown);
+            }
+        } finally {
+            mockServer.stop();
+        }
+    }
+
+    @Test(groups = {"integration"})
+    public void test503ServiceUnavailableIsRetried() throws Exception {
+        if (isCloud()) {
+            return; // mocked server
+        }
+
+        WireMockServer mockServer = new WireMockServer(WireMockConfiguration
+                .options().dynamicPort().notifier(new ConsoleNotifier(false)));
+        mockServer.start();
+
+        try {
+            // First request fails with a retryable 503 server error
+            mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
+                    .inScenario("ServiceUnavailable")
+                    .withRequestBody(WireMock.containing("SELECT 1"))
+                    .whenScenarioStateIs(STARTED)
+                    .willSetStateTo("Recovered")
+                    .willReturn(WireMock.aResponse()
+                            .withStatus(HttpStatus.SC_SERVICE_UNAVAILABLE)
+                            .withHeader("X-ClickHouse-Exception-Code", "202")
+                            .withBody("Code: 202. DB::Exception: Too many simultaneous queries. (TOO_MANY_SIMULTANEOUS_QUERIES)"))
+                    .build());
+
+            // Second request (retry) succeeds
+            mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
+                    .inScenario("ServiceUnavailable")
+                    .withRequestBody(WireMock.containing("SELECT 1"))
+                    .whenScenarioStateIs("Recovered")
+                    .willReturn(WireMock.aResponse()
+                            .withStatus(HttpStatus.SC_OK)
+                            .withHeader("X-ClickHouse-Summary", "{ \"read_bytes\": \"10\", \"read_rows\": \"1\"}"))
+                    .build());
+
+            try (Client client = new Client.Builder().addEndpoint(Protocol.HTTP, "localhost", mockServer.port(), false)
+                    .setUsername("default")
+                    .setPassword(ClickHouseServerForTest.getPassword())
+                    .compressServerResponse(false)
+                    .setMaxRetries(1)
+                    .retryOnFailures(ClientFaultCause.ServerRetryable)
+                    .build()) {
+
+                try (QueryResponse response = client.query("SELECT 1").get(10, TimeUnit.SECONDS)) {
+                    Assert.assertNotNull(response);
+                }
+            }
+
+            mockServer.verify(2, WireMock.postRequestedFor(WireMock.anyUrl()));
+        } finally {
+            mockServer.stop();
+        }
+    }
+
     private byte[] fetchBinaryPayload(String query, int networkBufferSize, int timeoutSec) throws Exception {
         QuerySettings querySettings = new QuerySettings().setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes);
         try (Client client = newClient()

client-v2/src/test/java/com/clickhouse/client/api/internal/ClientUtilsTest.java

@@ -11,36 +11,36 @@
 public class ClientUtilsTest {
 
     @Test(groups = {"unit"})
-    public void testQuiteCloseSwallowsExceptionAndLogs() throws IOException {
+    public void testQuietCloseSwallowsExceptionAndLogs() throws IOException {
         Logger log = Mockito.mock(Logger.class);
         IOException failure = new IOException("close failed");
         Closeable closeable = Mockito.mock(Closeable.class);
         Mockito.doThrow(failure).when(closeable).close();
 
         // Should not propagate the exception thrown by close()
-        ClientUtils.quiteClose(closeable, log);
+        ClientUtils.quietClose(closeable, log);
 
         Mockito.verify(closeable).close();
         Mockito.verify(log).warn(Mockito.contains("Failed to close object"), Mockito.eq(failure));
     }
 
     @Test(groups = {"unit"})
-    public void testQuiteCloseClosesSuccessfully() throws IOException {
+    public void testQuietCloseClosesSuccessfully() throws IOException {
         Logger log = Mockito.mock(Logger.class);
         Closeable closeable = Mockito.mock(Closeable.class);
 
-        ClientUtils.quiteClose(closeable, log);
+        ClientUtils.quietClose(closeable, log);
 
         Mockito.verify(closeable).close();
         Mockito.verifyNoInteractions(log);
     }
 
     @Test(groups = {"unit"})
-    public void testQuiteCloseWithNull() {
+    public void testQuietCloseWithNull() {
         Logger log = Mockito.mock(Logger.class);
 
         // Should be a no-op and not throw on a null closeable
-        ClientUtils.quiteClose(null, log);
+        ClientUtils.quietClose(null, log);
 
         Mockito.verifyNoInteractions(log);
         Assert.assertTrue(true);