added async config and executor

Author: chernser

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

@@ -40,9 +40,7 @@
 import com.clickhouse.client.config.ClickHouseDefaults;
 import com.clickhouse.client.http.ClickHouseHttpProto;
 import com.clickhouse.data.ClickHouseColumn;
-import com.clickhouse.data.ClickHouseDataStreamFactory;
 import com.clickhouse.data.ClickHouseFormat;
-import com.clickhouse.data.ClickHousePipedOutputStream;
 import com.clickhouse.data.format.BinaryStreamUtils;
 import org.apache.hc.core5.concurrent.DefaultThreadFactory;
 import org.apache.hc.core5.http.ClassicHttpResponse;
@@ -114,22 +112,23 @@
 public class Client implements AutoCloseable {
     private HttpAPIClientHelper httpClientHelper = null;
 
-    private Set<String> endpoints;
-    private Map<String, String> configuration;
-    private List<ClickHouseNode> serverNodes = new ArrayList<>();
-    private Map<Class<?>, List<POJOSerializer>> serializers; //Order is important to preserve for RowBinary
-    private Map<Class<?>, Map<String, Method>> getterMethods;
-    private Map<Class<?>, Boolean> hasDefaults; // Whether the POJO has defaults
+    private final Set<String> endpoints;
+    private final Map<String, String> configuration;
+    private final List<ClickHouseNode> serverNodes = new ArrayList<>();
+    private final Map<Class<?>, List<POJOSerializer>> serializers; //Order is important to preserve for RowBinary
+    private final Map<Class<?>, Map<String, Method>> getterMethods;
+    private final Map<Class<?>, Boolean> hasDefaults; // Whether the POJO has defaults
     private static final Logger LOG = LoggerFactory.getLogger(Client.class);
-    private ExecutorService sharedOperationExecutor;
+    private final ExecutorService sharedOperationExecutor;
 
-    private Map<String, ClientStatisticsHolder> globalClientStats = new ConcurrentHashMap<>();
+    private final Map<String, ClientStatisticsHolder> globalClientStats = new ConcurrentHashMap<>();
 
     private boolean useNewImplementation = false;
 
     private ClickHouseClient oldClient = null;
 
-    private Client(Set<String> endpoints, Map<String,String> configuration, boolean useNewImplementation) {
+    private Client(Set<String> endpoints, Map<String,String> configuration, boolean useNewImplementation,
+                   ExecutorService sharedOperationExecutor) {
         this.endpoints = endpoints;
         this.configuration = configuration;
         this.endpoints.forEach(endpoint -> {
@@ -139,7 +138,12 @@ private Client(Set<String> endpoints, Map<String,String> configuration, boolean
         this.getterMethods = new HashMap<>();
         this.hasDefaults = new HashMap<>();
 
-        this.sharedOperationExecutor = Executors.newCachedThreadPool(new DefaultThreadFactory("chc-operation"));
+        boolean isAsyncEnabled = MapUtils.getFlag(this.configuration, ClickHouseClientOption.ASYNC.getKey());
+        if (isAsyncEnabled && sharedOperationExecutor == null) {
+            this.sharedOperationExecutor = Executors.newCachedThreadPool(new DefaultThreadFactory("chc-operation"));
+        } else {
+            this.sharedOperationExecutor = sharedOperationExecutor;
+        }
         this.useNewImplementation = useNewImplementation;
         if (useNewImplementation) {
             this.httpClientHelper = new HttpAPIClientHelper(configuration);
@@ -169,7 +173,7 @@ public String getDefaultDatabase() {
     @Override
     public void close() {
         try {
-            if (!sharedOperationExecutor.isShutdown()) {
+            if (sharedOperationExecutor != null && !sharedOperationExecutor.isShutdown()) {
                 this.sharedOperationExecutor.shutdownNow();
             }
         } catch (Exception e) {
@@ -188,6 +192,8 @@ public static class Builder {
         private Map<String, String> configuration;
         private boolean useNewImplementation = false;
 
+        private ExecutorService sharedOperationExecutor = null;
+
         public Builder() {
             this.endpoints = new HashSet<>();
             this.configuration = new HashMap<String, String>();
@@ -603,6 +609,34 @@ public Builder setServerTimeZone(String timeZone) {
             return this;
         }
 
+        /**
+         * Configures client to execute requests in a separate thread. By default, operations (query, insert)
+         * are executed in the same thread as the caller.
+         * It is possible to set a shared executor for all operations. See {@link #setSharedOperationExecutor(ExecutorService)}
+         *
+         * Note: Async operations a using executor what expects having a queue of tasks for a pool of executors.
+         * The queue size limit is small it may quickly become a problem for scheduling new tasks.
+         *
+         * @param async - if to use async requests
+         * @return
+         */
+        public Builder useAsyncRequests(boolean async) {
+            this.configuration.put(ClickHouseClientOption.ASYNC.getKey(), String.valueOf(async));
+            return this;
+        }
+
+        /**
+         * Sets an executor for running operations. If async operations are enabled and no executor is specified
+         * client will create a default executor.
+         *
+         * @param executorService - executor service for async operations
+         * @return
+         */
+        public Builder setSharedOperationExecutor(ExecutorService executorService) {
+            this.sharedOperationExecutor = executorService;
+            return this;
+        }
+
         public Client build() {
             this.configuration = setDefaults(this.configuration);
 
@@ -648,7 +682,7 @@ public Client build() {
                 throw new IllegalArgumentException("Nor server timezone nor specific timezone is set");
             }
 
-            return new Client(this.endpoints, this.configuration, this.useNewImplementation);
+            return new Client(this.endpoints, this.configuration, this.useNewImplementation, this.sharedOperationExecutor);
         }
 
         private Map<String, String> setDefaults(Map<String, String> userConfig) {
@@ -908,8 +942,8 @@ public CompletableFuture<InsertResponse> insert(String tableName, List<?> data,
                 }
                 throw new ClientException("Failed to get table schema: too many retries");
             };
-            boolean isAsync = MapUtils.getFlag(configuration, settings.getAllSettings(), ClickHouseClientOption.ASYNC.getKey());
-            return isAsync ? CompletableFuture.supplyAsync(supplier, sharedOperationExecutor) : CompletableFuture.completedFuture(supplier.get());
+
+            return runAsyncOperation(supplier, settings.getAllSettings());
         } else {
             //Create an output stream to write the data to
             ByteArrayOutputStream stream = new ByteArrayOutputStream();
@@ -964,6 +998,7 @@ public CompletableFuture<InsertResponse> insert(String tableName,
         ClientStatisticsHolder clientStats = globalClientStats.remove(operationId);
         clientStats.start(ClientMetrics.OP_DURATION);
 
+        Supplier<InsertResponse> responseSupplier;
         if (useNewImplementation) {
 
             String retry = configuration.get(ClickHouseClientOption.RETRY.getKey());
@@ -978,7 +1013,7 @@ public CompletableFuture<InsertResponse> insert(String tableName,
 
             settings.setOption(ClickHouseClientOption.FORMAT.getKey(), format.name());
             final InsertSettings finalSettings = settings;
-            Supplier<InsertResponse> supplier = () -> {
+            responseSupplier = () -> {
                 // Selecting some node
                 ClickHouseNode selectedNode = getNextAliveNode();
 
@@ -1036,29 +1071,22 @@ public CompletableFuture<InsertResponse> insert(String tableName,
                 }
                 throw new ClientException("Failed to insert data: too many retries");
             };
-            boolean isAsync = MapUtils.getFlag(configuration, settings.getAllSettings(), ClickHouseClientOption.ASYNC.getKey());
-            return isAsync ? CompletableFuture.supplyAsync(supplier, sharedOperationExecutor) : CompletableFuture.completedFuture(supplier.get());
         } else {
-            CompletableFuture<InsertResponse> responseFuture = new CompletableFuture<>();
-
-            ClickHouseRequest.Mutation request = ClientV1AdaptorHelper
-                    .createMutationRequest(oldClient.write(getServerNode()), tableName, settings, configuration).format(format);
-
-            CompletableFuture<ClickHouseResponse> future = null;
-            try (ClickHousePipedOutputStream stream = ClickHouseDataStreamFactory.getInstance().createPipedOutputStream(request.getConfig())) {
-                future = request.data(stream.getInputStream()).execute();
-
-                //Copy the data from the input stream to the output stream
-                byte[] buffer = new byte[settings.getInputStreamCopyBufferSize()];
-                int bytesRead;
-                while ((bytesRead = data.read(buffer)) != -1) {
-                    stream.write(buffer, 0, bytesRead);
-                }
-            } catch (IOException e) {
-                responseFuture.completeExceptionally(new ClientException("Failed to write data to the output stream", e));
-            }
+            responseSupplier = () -> {
+                ClickHouseRequest.Mutation request = ClientV1AdaptorHelper
+                        .createMutationRequest(oldClient.write(getServerNode()), tableName, settings, configuration).format(format);
+
+                CompletableFuture<ClickHouseResponse> future = null;
+                future = request.data(output -> {
+                    //Copy the data from the input stream to the output stream
+                    byte[] buffer = new byte[settings.getInputStreamCopyBufferSize()];
+                    int bytesRead;
+                    while ((bytesRead = data.read(buffer)) != -1) {
+                        output.write(buffer, 0, bytesRead);
+                    }
+                    output.close();
+                }).option(ClickHouseClientOption.ASYNC, false).execute();
 
-            if (!responseFuture.isCompletedExceptionally()) {
                 try {
                     int operationTimeout = getOperationTimeout();
                     ClickHouseResponse clickHouseResponse;
@@ -1068,17 +1096,17 @@ public CompletableFuture<InsertResponse> insert(String tableName,
                         clickHouseResponse = future.get();
                     }
                     InsertResponse response = new InsertResponse(clickHouseResponse, clientStats);
-                    responseFuture.complete(response);
+                    LOG.debug("Total insert (InputStream) time: {}", clientStats.getElapsedTime("insert"));
+                    return response;
                 } catch (ExecutionException e) {
-                    responseFuture.completeExceptionally(new ClientException("Failed to get insert response", e.getCause()));
+                    throw  new ClientException("Failed to get insert response", e.getCause());
                 } catch (InterruptedException | TimeoutException e) {
-                    responseFuture.completeExceptionally(new ClientException("Operation has likely timed out.", e));
+                    throw  new ClientException("Operation has likely timed out.", e);
                 }
-            }
-            LOG.debug("Total insert (InputStream) time: {}", clientStats.getElapsedTime("insert"));
-
-            return responseFuture;
+            };
         }
+
+        return runAsyncOperation(responseSupplier, settings.getAllSettings());
     }
 
     /**
@@ -1144,6 +1172,8 @@ public CompletableFuture<QueryResponse> query(String sqlQuery, Map<String, Objec
         clientStats.start(ClientMetrics.OP_DURATION);
         applyDefaults(settings);
 
+        Supplier<QueryResponse> responseSupplier;
+
         if (useNewImplementation) {
             String retry = configuration.get(ClickHouseClientOption.RETRY.getKey());
             final int maxRetries = retry == null ? (int) ClickHouseClientOption.RETRY.getDefaultValue() : Integer.parseInt(retry);
@@ -1152,7 +1182,7 @@ public CompletableFuture<QueryResponse> query(String sqlQuery, Map<String, Objec
                 settings.setOption("statement_params", queryParams);
             }
             final QuerySettings finalSettings = settings;
-            Supplier<QueryResponse> supplier = () -> {
+            responseSupplier = () -> {
                 // Selecting some node
                 ClickHouseNode selectedNode = getNextAliveNode();
                 for (int i = 0; i <= maxRetries; i++) {
@@ -1188,8 +1218,6 @@ public CompletableFuture<QueryResponse> query(String sqlQuery, Map<String, Objec
                 }
                 throw new ClientException("Failed to get table schema: too many retries");
             };
-            boolean isAsync = MapUtils.getFlag(configuration, settings.getAllSettings(), ClickHouseClientOption.ASYNC.getKey());
-            return isAsync ? CompletableFuture.supplyAsync(supplier, sharedOperationExecutor) : CompletableFuture.completedFuture(supplier.get());
         } else {
             ClickHouseRequest<?> request = oldClient.read(getServerNode());
             request.options(SettingsConverter.toRequestOptions(settings.getAllSettings()));
@@ -1200,7 +1228,7 @@ public CompletableFuture<QueryResponse> query(String sqlQuery, Map<String, Objec
             request.format(format);
 
             final QuerySettings finalSettings = settings;
-            CompletableFuture<QueryResponse> future = CompletableFuture.supplyAsync(() -> {
+            responseSupplier = () -> {
                 LOG.trace("Executing request: {}", request);
                 try {
 
@@ -1218,9 +1246,10 @@ public CompletableFuture<QueryResponse> query(String sqlQuery, Map<String, Objec
                 } catch (Exception e) {
                     throw new ClientException("Failed to get query response", e);
                 }
-            }, sharedOperationExecutor);
-            return future;
+            };
         }
+
+        return runAsyncOperation(responseSupplier, settings.getAllSettings());
     }
 
     /**
@@ -1252,13 +1281,13 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
         settings.waitEndOfQuery(true); // we rely on the summery
 
         final QuerySettings finalSettings = settings;
-        return query(sqlQuery, settings).thenApplyAsync(response -> {
+        return query(sqlQuery, settings).thenApply(response -> {
             try {
                 return new Records(response, finalSettings);
             } catch (Exception e) {
                 throw new ClientException("Failed to get query response", e);
             }
-        }, sharedOperationExecutor);
+        });
     }
 
     /**
@@ -1360,13 +1389,13 @@ public CompletableFuture<CommandResponse> execute(String sql, CommandSettings se
      */
     public CompletableFuture<CommandResponse> execute(String sql) {
         return query(sql)
-                .thenApplyAsync(response -> {
+                .thenApply(response -> {
                     try {
                         return new CommandResponse(response);
                     } catch (Exception e) {
                         throw new ClientException("Failed to get command response", e);
                     }
-                }, sharedOperationExecutor);
+                });
     }
 
     /**
@@ -1427,6 +1456,11 @@ private void applyDefaults(QuerySettings settings) {
         }
     }
 
+    private <T> CompletableFuture<T> runAsyncOperation(Supplier<T> resultSupplier, Map<String, Object> requestSettings) {
+        boolean isAsync = MapUtils.getFlag(configuration, requestSettings, ClickHouseClientOption.ASYNC.getKey());
+        return isAsync ? CompletableFuture.supplyAsync(resultSupplier, sharedOperationExecutor) : CompletableFuture.completedFuture(resultSupplier.get());
+    }
+
     public String toString() {
         return "Client{" +
                 "endpoints=" + endpoints +

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

@@ -10,9 +10,7 @@
 import com.clickhouse.client.insert.SamplePOJO;
 import com.github.tomakehurst.wiremock.WireMockServer;
 import com.github.tomakehurst.wiremock.client.WireMock;
-import com.github.tomakehurst.wiremock.common.Slf4jNotifier;
 import com.github.tomakehurst.wiremock.core.WireMockConfiguration;
-import com.github.tomakehurst.wiremock.http.Fault;
 import com.github.tomakehurst.wiremock.stubbing.Scenario;
 import org.apache.hc.core5.http.HttpHeaders;
 import org.apache.hc.core5.http.HttpStatus;
@@ -28,6 +26,7 @@
 import java.util.concurrent.TimeUnit;
 
 import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.anyUrl;
 import static com.github.tomakehurst.wiremock.client.WireMock.equalTo;
 import static com.github.tomakehurst.wiremock.client.WireMock.post;
 import static com.github.tomakehurst.wiremock.client.WireMock.urlMatching;
@@ -88,9 +87,12 @@ public void testPrivateProxyWithoutAuth() {
         try {
             client.get().execute("select 1").get();
             Assert.fail("Should have thrown exception.");
+        } catch (ExecutionException e) {
+            Assert.assertTrue(e.getCause() instanceof ClientException);
+        } catch (ClientMisconfigurationException e) {
+            Assert.assertTrue(e.getMessage().contains("Proxy authentication required"));
         } catch (Exception e) {
-            e.printStackTrace();
-            Assert.assertTrue(e.getCause() instanceof ClientMisconfigurationException);
+            Assert.fail("Should have thrown exception.", e);
         }
     }
 
@@ -138,18 +140,25 @@ public void testProxyWithDisabledCookies() {
         client.set(clientBuilder(initProxy(), true).setHttpCookiesEnabled(false).build());
         final int targetPort = getServer(ClickHouseProtocol.HTTP).getPort();
 
-        proxy.get().addStubMapping(post(urlMatching("/.*"))
+        proxy.get().addStubMapping(post(anyUrl())
                 .inScenario("routeCookies")
                 .whenScenarioStateIs(Scenario.STARTED)
                 .willReturn(aResponse().withHeader(HttpHeaders.SET_COOKIE, "routeName=routeA")
-                        .proxiedFrom("http://localhost:" + targetPort)).willSetStateTo("cookies").build());
+                        .proxiedFrom("http://localhost:" + targetPort))
+                .willSetStateTo("cookies").build());
 
-        proxy.get().addStubMapping(post(urlMatching("/.*"))
+        proxy.get().addStubMapping(post(anyUrl())
                 .inScenario("routeCookies")
                 .whenScenarioStateIs("cookies")
                 .withHeader(HttpHeaders.COOKIE, equalTo("routeName=routeA"))
                 .willReturn(aResponse().proxiedFrom("http://localhost:" + targetPort)).build());
 
+        proxy.get().addStubMapping(post(anyUrl())
+                .inScenario("routeCookies")
+                .whenScenarioStateIs("cookies")
+                .withHeader(HttpHeaders.COOKIE, WireMock.absent())
+                .willReturn(aResponse().withStatus(HttpStatus.SC_BAD_GATEWAY)).build());
+
         try {
             client.get().execute("select 1").get();
         } catch (Exception e) {
@@ -159,8 +168,10 @@ public void testProxyWithDisabledCookies() {
             client.get().execute("select 1").get();
         } catch (ExecutionException e) {
             Assert.assertTrue(e.getCause() instanceof ClientException);
+        } catch (ClientException e) {
+            Assert.assertTrue(e.getMessage().contains("Server returned '502 Bad gateway'"));
         } catch (Exception e) {
-            fail("Should have thrown exception.", e);
+            Assert.fail("Should have thrown exception.", e);
         }
     }
 
@@ -175,7 +186,9 @@ private Client.Builder clientBuilder(int proxyPort, boolean onlyNewImplementatio
     }
 
     private int initProxy() {
-        WireMockServer wireMock = new WireMockServer(WireMockConfiguration.options().notifier(new Slf4jNotifier(true)));
+        WireMockServer wireMock = new WireMockServer(WireMockConfiguration.options()
+//                .notifier(new Slf4jNotifier(true))
+        );
         wireMock.start();
         proxy.set(wireMock);
         return wireMock.port();