implemented configurable operation timeout

chernser · chernser · commit 4e3bf9f9b010 · 2024-07-03T12:56:15.000-07:00
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/Client.java b/client-v2/src/main/java/com/clickhouse/client/api/Client.java
@@ -29,6 +29,7 @@
 import com.clickhouse.client.api.query.QuerySettings;
 import com.clickhouse.client.api.query.Records;
 import com.clickhouse.client.config.ClickHouseClientOption;
+import com.clickhouse.client.config.ClickHouseDefaults;
 import com.clickhouse.data.ClickHouseColumn;
 import com.clickhouse.data.ClickHouseDataStreamFactory;
 import com.clickhouse.data.ClickHouseFormat;
@@ -44,7 +45,6 @@
 import java.io.InputStream;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
-import java.net.URISyntaxException;
 import java.net.URL;
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
@@ -96,18 +96,13 @@
  *
  */
 public class Client {
-    private static final long TIMEOUT = TimeUnit.SECONDS.toMillis(30);
-
-    private static final String DEFAULT_DB_NAME = "default";
-
-    private static final String DEFAULT_THREADS_PER_CLIENT = "10";
 
     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<?>, List<POJOSerializer>> serializers; //Order is important to preserve for RowBinary
     private Map<Class<?>, Map<String, Method>> getterMethods;
-    private Map<Class<?>, Boolean> hasDefaults;
+    private Map<Class<?>, Boolean> hasDefaults; // Whether the POJO has defaults
     private static final Logger LOG = LoggerFactory.getLogger(Client.class);
     private ExecutorService sharedOperationExecutor;
 
@@ -123,9 +118,7 @@ private Client(Set<String> endpoints, Map<String,String> configuration) {
         this.getterMethods = new HashMap<>();
         this.hasDefaults = new HashMap<>();
 
-        final int numThreads = Integer.parseInt(configuration.getOrDefault(
-                ClickHouseClientOption.MAX_THREADS_PER_CLIENT.getKey(), DEFAULT_THREADS_PER_CLIENT));
-
+        final int numThreads = Integer.parseInt(configuration.get(ClickHouseClientOption.MAX_THREADS_PER_CLIENT.getKey()));
         this.sharedOperationExecutor = Executors.newFixedThreadPool(numThreads, new DefaultThreadFactory("chc-operation"));
         LOG.debug("Query executor created with {} threads", numThreads);
     }
@@ -141,6 +134,8 @@ public String getDefaultDatabase() {
 
     public static class Builder {
         private Set<String> endpoints;
+
+        // Read-only configuration
         private Map<String, String> configuration;
 
         public Builder() {
@@ -396,6 +391,17 @@ public Builder addProxy(ProxyType type, String host, int port) {
             return this;
         }
 
+        /**
+         * Sets the maximum time for operation to complete. By default, it is set to 3 hours.
+         * @param timeout
+         * @param timeUnit
+         * @return
+         */
+        public Builder setExecutionTimeout(long timeout, TimeUnit timeUnit) {
+            this.configuration.put(ClickHouseClientOption.MAX_EXECUTION_TIME.getKey(), String.valueOf(timeUnit.toMillis(timeout)));
+            return this;
+        }
+
         public Client build() {
             // check if endpoint are empty. so can not initiate client
             if (this.endpoints.isEmpty()) {
@@ -405,11 +411,30 @@ public Client build() {
             if (!this.configuration.containsKey("access_token") && (!this.configuration.containsKey("user") || !this.configuration.containsKey("password"))) {
                 throw new IllegalArgumentException("Username and password are required");
             }
+
+            this.configuration = setDefaults(this.configuration);
+
+            return new Client(this.endpoints, this.configuration);
+        }
+
+        private Map<String, String> setDefaults(Map<String, String> userConfig) {
+
             // set default database name if not specified
-            if (!this.configuration.containsKey("database")) {
-                this.configuration.put("database", DEFAULT_DB_NAME);
+            if (!userConfig.containsKey("database")) {
+                userConfig.put("database", (String) ClickHouseDefaults.DATABASE.getDefaultValue());
             }
-            return new Client(this.endpoints, this.configuration);
+
+            if (!userConfig.containsKey(ClickHouseClientOption.MAX_EXECUTION_TIME.getKey())) {
+                userConfig.put(ClickHouseClientOption.MAX_EXECUTION_TIME.getKey(),
+                        String.valueOf(ClickHouseClientOption.MAX_EXECUTION_TIME.getDefaultValue()));
+            }
+
+            if (!userConfig.containsKey(ClickHouseClientOption.MAX_THREADS_PER_CLIENT.getKey())) {
+                userConfig.put(ClickHouseClientOption.MAX_THREADS_PER_CLIENT.getKey(),
+                        String.valueOf(ClickHouseClientOption.MAX_THREADS_PER_CLIENT.getDefaultValue()));
+            }
+
+            return userConfig;
         }
     }
 
@@ -423,7 +448,7 @@ private ClickHouseNode getServerNode() {
      * @return true if the server is alive, false otherwise
      */
     public boolean ping() {
-        return ping(Client.TIMEOUT);
+        return ping(getOperationTimeout());
     }
 
     /**
@@ -641,7 +666,14 @@ public CompletableFuture<InsertResponse> insert(String tableName,
 
             if (!responseFuture.isCompletedExceptionally()) {
                 try {
-                    InsertResponse response = new InsertResponse(client, future.get(TIMEOUT, TimeUnit.MILLISECONDS), clientStats);
+                    int operationTimeout = getOperationTimeout();
+                    ClickHouseResponse clickHouseResponse;
+                    if (operationTimeout > 0) {
+                        clickHouseResponse = future.get(operationTimeout, TimeUnit.MILLISECONDS);
+                    } else {
+                        clickHouseResponse = future.get();
+                    }
+                    InsertResponse response = new InsertResponse(client, clickHouseResponse, clientStats);
                     responseFuture.complete(response);
                 } catch (ExecutionException e) {
                     responseFuture.completeExceptionally(new ClientException("Failed to get insert response", e.getCause()));
@@ -729,9 +761,16 @@ public CompletableFuture<QueryResponse> query(String sqlQuery, Map<String, Objec
         CompletableFuture<QueryResponse> future = CompletableFuture.supplyAsync(() -> {
             LOG.trace("Executing request: {}", request);
             try {
-                QueryResponse queryResponse = new QueryResponse(client, request.execute(), finalSettings, format, clientStats);
-                queryResponse.ensureDone();
-                return queryResponse;
+
+                int operationTimeout = getOperationTimeout();
+                ClickHouseResponse clickHouseResponse;
+                if (operationTimeout > 0) {
+                    clickHouseResponse = request.execute().get(operationTimeout, TimeUnit.MILLISECONDS);
+                } else {
+                    clickHouseResponse = request.execute().get();
+                }
+
+                return new QueryResponse(client, clickHouseResponse, finalSettings, format, clientStats);
             } catch (ClientException e) {
                 throw e;
             } catch (Exception e) {
@@ -782,8 +821,15 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
         CompletableFuture<Records> future = CompletableFuture.supplyAsync(() -> {
             LOG.trace("Executing request: {}", request);
             try {
-                QueryResponse queryResponse = new QueryResponse(client, request.execute(), finalSettings, format, clientStats);
-                queryResponse.ensureDone();
+                int operationTimeout = getOperationTimeout();
+                ClickHouseResponse clickHouseResponse;
+                if (operationTimeout > 0) {
+                    clickHouseResponse = request.execute().get(operationTimeout, TimeUnit.MILLISECONDS);
+                } else {
+                    clickHouseResponse = request.execute().get();
+                }
+
+                QueryResponse queryResponse = new QueryResponse(client, clickHouseResponse, finalSettings, format, clientStats);
                 return new Records(queryResponse, finalSettings);
             } catch (ClientException e) {
                 throw e;
@@ -804,7 +850,10 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
      */
     public List<GenericRecord> queryAll(String sqlQuery) {
         try {
-            try (QueryResponse response = query(sqlQuery).get(TIMEOUT, TimeUnit.MILLISECONDS)) {
+            int operationTimeout = getOperationTimeout();
+            QueryResponse response = operationTimeout == 0 ? query(sqlQuery).get() :
+                    query(sqlQuery).get(operationTimeout, TimeUnit.MILLISECONDS);
+            try (response) {
                 List<GenericRecord> records = new ArrayList<>();
                 if (response.getResultRows() > 0) {
                     ClickHouseBinaryFormatReader reader = new RowBinaryWithNamesAndTypesFormatReader(response.getInputStream());
@@ -914,6 +963,10 @@ public Map<String, String> getConfiguration() {
         return Collections.unmodifiableMap(configuration);
     }
 
+    protected int getOperationTimeout() {
+        return Integer.parseInt(configuration.get(ClickHouseClientOption.MAX_EXECUTION_TIME.getKey()));
+    }
+
     /**
      * Returns unmodifiable set of endpoints.
      * @return - set of endpoints
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/query/QueryResponse.java b/client-v2/src/main/java/com/clickhouse/client/api/query/QueryResponse.java
@@ -10,7 +10,6 @@
 import com.clickhouse.data.ClickHouseInputStream;
 
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
@@ -29,7 +28,7 @@
  */
 public class QueryResponse implements AutoCloseable {
 
-    private final Future<ClickHouseResponse> responseRef;
+    private final ClickHouseResponse clickHouseResponse;
     private final ClickHouseFormat format;
 
     private long completeTimeout = TimeUnit.MINUTES.toMillis(1);
@@ -42,43 +41,19 @@ public class QueryResponse implements AutoCloseable {
 
     private volatile boolean completed = false;
 
-    public QueryResponse(ClickHouseClient client, Future<ClickHouseResponse> responseRef,
+    public QueryResponse(ClickHouseClient client, ClickHouseResponse clickHouseResponse,
                          QuerySettings settings, ClickHouseFormat format,
                          ClientStatisticsHolder clientStatisticsHolder) {
         this.client = client;
-        this.responseRef = responseRef;
+        this.clickHouseResponse = clickHouseResponse;
         this.format = format;
         this.settings = settings;
         this.operationMetrics = new OperationMetrics(clientStatisticsHolder);
     }
 
-    /**
-     * Called internally to finalize the query execution.
-     * Do not call this method directly.
-     */
-    public void ensureDone() {
-        if (!completed) {
-            // TODO: thread-safety
-            makeComplete();
-        }
-    }
-
-    private void makeComplete() {
-        try {
-            ClickHouseResponse response = responseRef.get(completeTimeout, TimeUnit.MILLISECONDS);
-            completed = true;
-            operationMetrics.operationComplete(response.getSummary());
-        } catch (ExecutionException e) {
-            throw new ClientException("Failed to get command response", e.getCause());
-        } catch (TimeoutException | InterruptedException e) {
-            throw new ClientException("Query request failed", e);
-        }
-    }
-
     public ClickHouseInputStream getInputStream() {
-        ensureDone();
         try {
-            return responseRef.get().getInputStream();
+            return clickHouseResponse.getInputStream();
         } catch (Exception e) {
             throw new RuntimeException(e); // TODO: handle exception
         }
@@ -87,7 +62,7 @@ public ClickHouseInputStream getInputStream() {
     @Override
     public void close() throws Exception {
         try {
-            responseRef.get(completeTimeout, TimeUnit.MILLISECONDS ).close();
+            clickHouseResponse.close();
         } catch (Exception e) {
             throw new ClientException("Failed to close response", e);
         }
@@ -109,7 +84,6 @@ public ClickHouseFormat getFormat() {
      * @return metrics of this operation
      */
     public OperationMetrics getMetrics() {
-        ensureDone();
         return operationMetrics;
     }
 
