Merge pull request #1699 from ClickHouse/add-query-id-summary

Adding query_id to V1 client
ClickHouse · Jun 20, 2024 · c30c346 · c30c346
2 parents a593b52 + b061b48
commit c30c346
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 5 deletions.
diff --git a/clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseResponseSummary.java b/clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseResponseSummary.java
@@ -20,7 +20,7 @@ public class ClickHouseResponseSummary implements Serializable {
     public static final class Progress implements Serializable {
         private static final long serialVersionUID = -1447066780591278108L;
 
-        static final Progress EMPTY = new Progress(0L, 0L, 0L, 0L, 0L, 0L, 0L);
+        static final Progress EMPTY = new Progress(0L, 0L, 0L, 0L, 0L, 0L, 0L, "");
 
         private final long read_rows;
         private final long read_bytes;
@@ -29,6 +29,7 @@ public static final class Progress implements Serializable {
         private final long written_bytes;
         private final long elapsed_time;
         private final long result_rows;
+        private final String query_id;
 
         /**
          * Default constructor.
@@ -39,16 +40,34 @@ public static final class Progress implements Serializable {
          * @param written_rows       Number of rows written
          * @param written_bytes      Volume of data written in bytes
          * @param elapsed_time       Query processing time in (ns)
+         * @param result_rows        Number of rows in the result
          */
         public Progress(long read_rows, long read_bytes, long total_rows_to_read, long written_rows,
-                long written_bytes, long elapsed_time, long result_rows) {
+                        long written_bytes, long elapsed_time, long result_rows) {
+            this(read_rows, read_bytes, total_rows_to_read, written_rows, written_bytes, elapsed_time, result_rows, "");
+        }
+        /**
+         * Default constructor.
+         *
+         * @param read_rows          Number of rows read
+         * @param read_bytes         Volume of data read in bytes
+         * @param total_rows_to_read Total number of rows to be read
+         * @param written_rows       Number of rows written
+         * @param written_bytes      Volume of data written in bytes
+         * @param elapsed_time       Query processing time in (ns)
+         * @param result_rows        Number of rows in the result
+         * @param query_id           Query ID
+         */
+        public Progress(long read_rows, long read_bytes, long total_rows_to_read, long written_rows,
+                long written_bytes, long elapsed_time, long result_rows, String query_id) {
             this.read_rows = read_rows;
             this.read_bytes = read_bytes;
             this.total_rows_to_read = total_rows_to_read;
             this.written_rows = written_rows;
             this.written_bytes = written_bytes;
             this.elapsed_time = elapsed_time;
             this.result_rows = result_rows;
+            this.query_id = query_id;
         }
 
         public long getReadRows() {
@@ -78,6 +97,10 @@ public long getElapsedTime() {
         public long getResultRows() {
             return result_rows;
         }
+
+        public String getQueryId() {
+            return query_id;
+        }
         public Progress add(Progress progress) {
             if (progress == null) {
                 return this;
@@ -86,7 +109,7 @@ public Progress add(Progress progress) {
             return new Progress(read_rows + progress.read_rows, read_bytes + progress.read_bytes,
                     total_rows_to_read + progress.total_rows_to_read, written_rows + progress.written_rows,
                     written_bytes + progress.written_bytes,elapsed_time + progress.elapsed_time,
-                    result_rows + progress.result_rows);
+                    result_rows + progress.result_rows, query_id + ", " + progress.query_id);
         }
 
         public boolean isEmpty() {
@@ -322,6 +345,10 @@ public long getResultRows() {
         return progress.get().getResultRows();
     }
 
+    public String getQueryId() {
+        return progress.get().getQueryId();
+    }
+
     public boolean isEmpty() {
         return progress.get().isEmpty() && stats.get().isEmpty();
     }

diff --git a/clickhouse-client/src/test/java/com/clickhouse/client/ClickHouseResponseSummaryTest.java b/clickhouse-client/src/test/java/com/clickhouse/client/ClickHouseResponseSummaryTest.java
@@ -8,7 +8,7 @@
 
 public class ClickHouseResponseSummaryTest {
     @Test(groups = { "unit" })
-    public void testConsutrctor() {
+    public void testConstructor() {
         ClickHouseResponseSummary summary = new ClickHouseResponseSummary(null, null);
         Assert.assertNotNull(summary.getProgress());
         Assert.assertNotNull(summary.getStatistics());

diff --git a/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpResponse.java b/clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpResponse.java
@@ -14,6 +14,7 @@
 import com.clickhouse.data.ClickHouseFormat;
 import com.clickhouse.data.ClickHouseInputStream;
 import com.clickhouse.data.ClickHouseUtils;
+import com.clickhouse.logging.Logger;
 
 public class ClickHouseHttpResponse {
     private static long getLongValue(Map<String, String> map, String key) {
@@ -70,7 +71,7 @@ public ClickHouseHttpResponse(ClickHouseHttpConnection connection, ClickHouseInp
                 new ClickHouseResponseSummary.Progress(getLongValue(map, "read_rows"), getLongValue(map, "read_bytes"),
                         getLongValue(map, "total_rows_to_read"), getLongValue(map, "written_rows"),
                         getLongValue(map, "written_bytes"), getLongValue(map, "elapsed_ns"),
-                getLongValue(map, "result_rows")),
+                        getLongValue(map, "result_rows"), this.queryId),
                 null);
 
         this.format = format != null ? format : connection.config.getFormat();

diff --git a/...khouse-http-client/src/test/java/com/clickhouse/client/http/ClickHouseHttpClientTest.java b/...khouse-http-client/src/test/java/com/clickhouse/client/http/ClickHouseHttpClientTest.java
@@ -414,6 +414,21 @@ public void testPost() throws ClickHouseException {
         }
     }
 
+    @Test(groups = {"integration"})
+    public void testQueryId() throws ClickHouseException {
+        ClickHouseNode server = getServer(ClickHouseProtocol.HTTP);
+        String uuid = UUID.randomUUID().toString();
+
+        try (ClickHouseClient client = ClickHouseClient.builder().options(getClientOptions())
+                .defaultCredentials(ClickHouseCredentials.fromUserAndPassword("foo", "bar")).build()) {
+            try (ClickHouseResponse resp = newRequest(client, server).compressServerResponse(false)
+                    .format(ClickHouseFormat.RowBinaryWithNamesAndTypes)
+                    .query("select 1,2", uuid).executeAndWait()) {
+                Assert.assertEquals(resp.getSummary().getQueryId(), uuid);
+            }
+        }
+    }
+
     @Test(groups = {"integration"})
     public void testProxyConnection() throws ClickHouseException, IOException {
         ToxiproxyContainer toxiproxy = null;