[client-v2] Request cancellation

client-v2/src/main/java/com/clickhouse/client/api/insert/InsertResponse.java

@@ -3,21 +3,17 @@
 import com.clickhouse.client.api.http.ClickHouseHttpProto;
 import com.clickhouse.client.api.metrics.OperationMetrics;
 import com.clickhouse.client.api.metrics.ServerMetrics;
+import com.clickhouse.client.api.transport.internal.TransportResponse;
 
-import java.util.Collections;
 import java.util.Map;
 
 public class InsertResponse implements AutoCloseable {
     private OperationMetrics operationMetrics;
     private final Map<String, String> responseHeaders;
 
-    public InsertResponse(OperationMetrics metrics) {
-        this(metrics, Collections.emptyMap());
-    }
-
-    public InsertResponse(OperationMetrics metrics, Map<String, String> responseHeaders) {
+    public InsertResponse(TransportResponse transportResponse, OperationMetrics metrics) {
         this.operationMetrics = metrics;
-        this.responseHeaders = responseHeaders;
+        this.responseHeaders = transportResponse.getHeaders();
     }
 
     @Override

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

@@ -1,5 +1,9 @@
 package com.clickhouse.client.api.internal;
 
+import org.slf4j.Logger;
+
+import java.io.Closeable;
+
 /**
  * Class containing utility methods used across the client.
  */
@@ -14,4 +18,14 @@ public static boolean isNotBlank(String str) {
     public static boolean isBlank(String str) {
         return str == null || str.trim().isEmpty();
     }
+
+    public static void quietClose(Closeable closeable, Logger log) {
+        if (closeable != null) {
+            try {
+                closeable.close();
+            } catch (Exception e) {
+                log.warn("Failed to close object " + closeable, e);
+            }
+        }
+    }
 }

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

@@ -15,6 +15,8 @@
 import com.clickhouse.client.api.enums.SSLMode;
 import com.clickhouse.client.api.http.ClickHouseHttpProto;
 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.data.ClickHouseFormat;
 import net.jpountz.lz4.LZ4Factory;
 import org.apache.commons.compress.compressors.CompressorStreamFactory;
@@ -25,6 +27,7 @@
 import org.apache.hc.client5.http.entity.mime.MultipartEntityBuilder;
 import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
 import org.apache.hc.client5.http.impl.classic.HttpClientBuilder;
+import org.apache.hc.client5.http.impl.classic.RequestFailedException;
 import org.apache.hc.client5.http.impl.io.BasicHttpClientConnectionManager;
 import org.apache.hc.client5.http.impl.io.ManagedHttpClientConnectionFactory;
 import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManager;
@@ -43,8 +46,10 @@
 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;
 import org.apache.hc.core5.http.URIScheme;
 import org.apache.hc.core5.http.config.CharCodingConfig;
 import org.apache.hc.core5.http.config.Http1Config;
@@ -99,6 +104,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
@@ -365,10 +371,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()) :
@@ -529,9 +532,44 @@ private HttpPost createPostRequest(URI uri, Map<String, Object> requestConfig) {
         return req;
     }
 
-    public ClassicHttpResponse executeRequest(Endpoint server, Map<String, Object> requestConfig,
-                                              String body) throws Exception {
+    private static final class TransportRequestImpl implements TransportRequest {
+        private final HttpPost delegate;
+        private final Map<String, Object> config;
+        private final AtomicBoolean cancelled = new AtomicBoolean(false);
+
+        TransportRequestImpl(HttpPost delegate, Map<String, Object> config) {
+            this.delegate = delegate;
+            this.config = config;
+        }
+
+        @Override
+        public boolean cancel() {
+            cancelled.set(true);
+            if (delegate.isCancelled()) {
+                return true;
+            }
+            return delegate.cancel();
+        }
+
+        @Override
+        public boolean isCancelled() {
+            return cancelled.get();
+        }
+
+        @Override
+        public Map<String, Object> getConfig() {
+            return config;
+        }
+
+        @Override
+        @SuppressWarnings("unchecked")
+        public <T> T getDelegate() {
+            return (T) delegate;
+        }
+    }
 
+    public TransportRequest createRequest(Endpoint server, Map<String, Object> requestConfig,
+                                          String body) {
         boolean useMultipart = ClientConfigProperties.HTTP_SEND_PARAMS_IN_BODY.<Boolean>getOrDefault(requestConfig) &&
                 requestConfig.containsKey(HttpAPIClientHelper.KEY_STATEMENT_PARAMS);
 
@@ -554,34 +592,97 @@ public ClassicHttpResponse executeRequest(Endpoint server, Map<String, Object> r
             req.setEntity(wrapRequestEntity(httpEntity, requestConfig));
 
         } else {
-            final String contentEncoding = req.containsHeader(HttpHeaders.CONTENT_ENCODING) ? req.getHeader(HttpHeaders.CONTENT_ENCODING).getValue() : null;
-
-            HttpEntity httpEntity = new ByteArrayEntity(body.getBytes(StandardCharsets.UTF_8.name()), CONTENT_TYPE, contentEncoding);
+            final HttpEntity httpEntity;
+            try {
+                final String contentEncoding = req.containsHeader(HttpHeaders.CONTENT_ENCODING) ? req.getHeader(HttpHeaders.CONTENT_ENCODING).getValue() : null;
+                httpEntity = new ByteArrayEntity(body.getBytes(StandardCharsets.UTF_8.name()), CONTENT_TYPE, contentEncoding);
+            } catch (UnsupportedEncodingException | ProtocolException e) {
+                throw new ClientException("failed to create request body entity", e);
+            }
             req.setEntity(wrapRequestEntity(httpEntity, requestConfig));
         }
 
-        // execute
-        return doPostRequest(requestConfig, req);
+        return new TransportRequestImpl(req, requestConfig);
+    }
+
+
+    private static final class TransportResponseImpl implements TransportResponse {
+
+        private final ClassicHttpResponse delegate;
+
+        TransportResponseImpl(ClassicHttpResponse delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public ClickHouseFormat getDataFormat() {
+            Header formatHeader = delegate.getFirstHeader(ClickHouseHttpProto.HEADER_FORMAT);
+            return formatHeader == null ? null : ClickHouseFormat.valueOf(formatHeader.getValue());
+        }
+
+        @Override
+        public String getSummaryJson() {
+            return HttpAPIClientHelper.getHeaderVal(delegate
+                    .getFirstHeader(ClickHouseHttpProto.HEADER_SRV_SUMMARY), "{}");
+        }
+
+        @Override
+        public String getQueryId() {
+            return HttpAPIClientHelper.getHeaderVal(delegate
+                    .getFirstHeader(ClickHouseHttpProto.HEADER_QUERY_ID), null);
+        }
+
+        @Override
+        @SuppressWarnings("unchecked")
+        public <T> T getDelegate() {
+            return (T) delegate;
+        }
+
+        @Override
+        public Map<String, String> getHeaders() {
+            return HttpAPIClientHelper.collectResponseHeaders(delegate);
+        }
+
+        @Override
+        public void close() throws IOException {
+            delegate.close();
+        }
+
+        @Override
+        public InputStream createDataInputStream() {
+            try {
+                return delegate.getEntity().getContent();
+            } catch (Exception e) {
+                throw new ClientException("Failed to construct input stream", e);
+            }
+        }
     }
 
-    public ClassicHttpResponse executeRequest(Endpoint server, Map<String, Object> requestConfig,
-                                              IOCallback<OutputStream> writeCallback) throws Exception {
+    public TransportResponse executeRequest(TransportRequest transportRequest) throws Exception {
+        return new TransportResponseImpl(doPostRequest(transportRequest.getConfig(), transportRequest.getDelegate()));
+    }
 
+    public TransportRequest createRequest(Endpoint server, Map<String, Object> requestConfig, IOCallback<OutputStream> writeCallback) {
         final URI uri = createRequestURI(server, requestConfig, true);
         final HttpPost req = createPostRequest(uri, requestConfig);
-        String contentEncoding = req.containsHeader(HttpHeaders.CONTENT_ENCODING) ? req.getHeader(HttpHeaders.CONTENT_ENCODING).getValue() : null;
-        req.setEntity(wrapRequestEntity(
-                new EntityTemplate(-1, CONTENT_TYPE, contentEncoding , writeCallback),
-                requestConfig));
+        try {
+            String contentEncoding = req.containsHeader(HttpHeaders.CONTENT_ENCODING) ? req.getHeader(HttpHeaders.CONTENT_ENCODING).getValue() : null;
+            req.setEntity(wrapRequestEntity(
+                    new EntityTemplate(-1, CONTENT_TYPE, contentEncoding, writeCallback),
+                    requestConfig));
+        } catch (ProtocolException e) {
+            throw new ClientException("failed to create request body entity", e);
+        }
 
-        return doPostRequest(requestConfig, req);
+        return new TransportRequestImpl(req, requestConfig);
     }
 
     private ClassicHttpResponse doPostRequest(Map<String, Object> requestConfig, HttpPost req) throws Exception {
 
         doPoolVent();
 
         ClassicHttpResponse httpResponse = null;
+        boolean closeResponse = true;
         HttpContext context = createRequestHttpContext(requestConfig);
         try {
             httpResponse = httpClient.executeOpen(null, req, context);
@@ -590,43 +691,59 @@ 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) {
-            closeQuietly(httpResponse);
             LOG.warn("Host '{}' unknown", req.getAuthority());
             throw e;
         } catch (ConnectException | NoRouteToHostException e) {
-            closeQuietly(httpResponse);
             LOG.warn("Failed to connect to '{}': {}", req.getAuthority(), e.getMessage());
             throw e;
         } catch (Exception e) {
-            closeQuietly(httpResponse);
             LOG.debug("Failed to execute request to '{}': {}", req.getAuthority(), e.getMessage(), e);
+            if (e instanceof RequestFailedException && req.isCancelled()) {
+                throw new TransportException("Request was cancelled on client side", e, getQueryId(httpResponse, req));
+            }
             throw e;
+        } finally {
+            if (closeResponse) {
+                ClientUtils.quietClose(httpResponse, LOG);
+            }
         }
     }
 
-    public static void closeQuietly(ClassicHttpResponse httpResponse) {
-        if (httpResponse != null) {
-            try {
-                httpResponse.close();
-            } catch (IOException e) {
-                LOG.warn("Failed to close response");
-            }
-        }
+    private String getQueryId(HttpResponse httpResponse, HttpPost httpRequest) {
+        final Header serverQueryIdHeader = httpResponse == null ? null : httpResponse.getFirstHeader(ClickHouseHttpProto.HEADER_QUERY_ID);
+        final Header clientQueryIdHeader = httpResponse == null ? null : 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");
@@ -820,7 +937,10 @@ public static int getHeaderInt(Header header, int defaultValue) {
             ClickHouseHttpProto.HEADER_SRV_SUMMARY,
             ClickHouseHttpProto.HEADER_SRV_DISPLAY_NAME,
             ClickHouseHttpProto.HEADER_DATABASE,
-            ClickHouseHttpProto.HEADER_DB_USER
+            ClickHouseHttpProto.HEADER_DB_USER,
+            ClickHouseHttpProto.HEADER_TIMEZONE,
+            ClickHouseHttpProto.HEADER_FORMAT,
+            ClickHouseHttpProto.HEADER_PROGRESS
     ));
 
     /**