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 index f7cc187d4..a51c4abe7 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/Client.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/Client.java @@ -139,20 +139,21 @@ public class Client implements AutoCloseable { private String dbUser; private String serverVersion; private Object metricsRegistry; - private int retries; + private final int retries; private LZ4Factory lz4Factory = null; + private final Supplier queryIdGenerator; private Client(Set endpoints, Map configuration, - ExecutorService sharedOperationExecutor, ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy) { - this(endpoints, configuration, sharedOperationExecutor, columnToMethodMatchingStrategy, null); + ExecutorService sharedOperationExecutor, ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy, Object metricRegistry, Supplier queryIdGenerator) { + this(endpoints, configuration, sharedOperationExecutor, columnToMethodMatchingStrategy, null, metricRegistry, queryIdGenerator); } private Client(Set endpoints, Map configuration, - ExecutorService sharedOperationExecutor, ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy, Object metricsRegistry) { - // Simple initialization + ExecutorService sharedOperationExecutor, ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy, Object metricsRegistry, Object metricRegistry, Supplier queryIdGenerator) { this.configuration = ClientConfigProperties.parseConfigMap(configuration); this.readOnlyConfig = Collections.unmodifiableMap(configuration); this.metricsRegistry = metricsRegistry; + this.queryIdGenerator = queryIdGenerator; // Serialization this.pojoSerDe = new POJOSerDe(columnToMethodMatchingStrategy); @@ -267,6 +268,8 @@ public static class Builder { private ExecutorService sharedOperationExecutor = null; private ColumnToMethodMatchingStrategy columnToMethodMatchingStrategy; private Object metricRegistry = null; + private Supplier queryIdGenerator; + public Builder() { this.endpoints = new HashSet<>(); this.configuration = new HashMap<>(); @@ -1048,6 +1051,16 @@ public Builder sslSocketSNI(String sni) { return this; } + /** + * Sets query id generator. Will be used when operation settings (InsertSettings, QuerySettings) do not have query id set. + * @param supplier + * @return + */ + public Builder queryIdGenerator(Supplier supplier) { + this.queryIdGenerator = supplier; + return this; + } + public Client build() { // check if endpoint are empty. so can not initiate client if (this.endpoints.isEmpty()) { @@ -1106,7 +1119,7 @@ public Client build() { } return new Client(this.endpoints, this.configuration, this.sharedOperationExecutor, - this.columnToMethodMatchingStrategy, this.metricRegistry); + this.columnToMethodMatchingStrategy, this.metricRegistry, this.queryIdGenerator); } } @@ -1245,6 +1258,9 @@ public CompletableFuture insert(String tableName, List data, final int maxRetries = retry == null ? 0 : retry; requestSettings.setOption(ClientConfigProperties.INPUT_OUTPUT_FORMAT.getKey(), format); + if (requestSettings.getQueryId() == null && queryIdGenerator != null) { + requestSettings.setQueryId(queryIdGenerator.get()); + } Supplier supplier = () -> { long startTime = System.nanoTime(); // Selecting some node @@ -1462,6 +1478,9 @@ public CompletableFuture insert(String tableName, } sqlStmt.append(" FORMAT ").append(format.name()); requestSettings.serverSetting(ClickHouseHttpProto.QPARAM_QUERY_STMT, sqlStmt.toString()); + if (requestSettings.getQueryId() == null && queryIdGenerator != null) { + requestSettings.setQueryId(queryIdGenerator.get()); + } responseSupplier = () -> { long startTime = System.nanoTime(); // Selecting some node @@ -1587,6 +1606,9 @@ public CompletableFuture query(String sqlQuery, Map { long startTime = System.nanoTime(); // Selecting some node diff --git a/client-v2/src/main/java/com/clickhouse/client/api/query/Records.java b/client-v2/src/main/java/com/clickhouse/client/api/query/Records.java index 46d0e827d..f8eeba40c 100644 --- a/client-v2/src/main/java/com/clickhouse/client/api/query/Records.java +++ b/client-v2/src/main/java/com/clickhouse/client/api/query/Records.java @@ -125,6 +125,14 @@ public long getResultRows() { return response.getMetrics().getMetric(ServerMetrics.RESULT_ROWS).getLong(); } + /** + * Returns response query id + * @return query id of the request + */ + public String getQueryId() { + return response.getQueryId(); + } + @Override public void close() throws Exception { response.close();