Merge pull request #1846 from ClickHouse/clientv2_fix_statistics

reimplemented in way ClientStats is always not null

Author: chernser

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

@@ -1083,7 +1083,7 @@ public CompletableFuture<InsertResponse> insert(String tableName, List<?> data,
         }
 
 
-        String operationId = startOperation();
+        String operationId = registerOperationMetrics();
         settings.setOperationId(operationId);
         if (useNewImplementation) {
             globalClientStats.get(operationId).start(ClientMetrics.OP_DURATION);
@@ -1225,13 +1225,18 @@ public CompletableFuture<InsertResponse> insert(String tableName,
                                      InputStream data,
                                      ClickHouseFormat format,
                                      InsertSettings settings) {
+
         String operationId = (String) settings.getOperationId();
-        if (operationId == null) {
-            operationId = startOperation();
-            settings.setOperationId(operationId);
+        ClientStatisticsHolder clientStats = null;
+        if (operationId != null) {
+            clientStats = globalClientStats.remove(operationId);
+        }
+
+        if (clientStats == null) {
+            clientStats = new ClientStatisticsHolder();
         }
-        ClientStatisticsHolder clientStats = globalClientStats.remove(operationId);
         clientStats.start(ClientMetrics.OP_DURATION);
+        final ClientStatisticsHolder finalClientStats = clientStats;
 
         Supplier<InsertResponse> responseSupplier;
         if (useNewImplementation) {
@@ -1280,7 +1285,7 @@ public CompletableFuture<InsertResponse> insert(String tableName,
                             continue;
                         }
 
-                        OperationMetrics metrics = new OperationMetrics(clientStats);
+                        OperationMetrics metrics = new OperationMetrics(finalClientStats);
                         String summary = HttpAPIClientHelper.getHeaderVal(httpResponse.getFirstHeader(ClickHouseHttpProto.HEADER_SRV_SUMMARY), "{}");
                         ProcessParser.parseSummary(summary, metrics);
                         String queryId =  HttpAPIClientHelper.getHeaderVal(httpResponse.getFirstHeader(ClickHouseHttpProto.HEADER_QUERY_ID), finalSettings.getQueryId(), String::valueOf);
@@ -1333,7 +1338,7 @@ public CompletableFuture<InsertResponse> insert(String tableName,
                     } else {
                         clickHouseResponse = future.get();
                     }
-                    InsertResponse response = new InsertResponse(clickHouseResponse, clientStats);
+                    InsertResponse response = new InsertResponse(clickHouseResponse, finalClientStats);
                     return response;
                 } catch (ExecutionException e) {
                     throw  new ClientException("Failed to get insert response", e.getCause());
@@ -1762,7 +1767,7 @@ public ClickHouseBinaryFormatReader newBinaryFormatReader(QueryResponse response
         return  newBinaryFormatReader(response, null);
     }
 
-    private String startOperation() {
+    private String registerOperationMetrics() {
         String operationId = UUID.randomUUID().toString();
         globalClientStats.put(operationId, new ClientStatisticsHolder());
         return operationId;

client-v2/src/test/java/com/clickhouse/client/insert/InsertTests.java

@@ -230,4 +230,31 @@ public void insertRawDataSimple(int numberOfRecords) throws Exception {
         OperationMetrics metrics = response.getMetrics();
         assertEquals((int)response.getWrittenRows(), numberOfRecords );
     }
+
+    @Test(groups = { "integration" })
+    public void testInsertMetricsOperationId() throws Exception {
+        final String tableName = "insert_metrics_test";
+        final String createSQL = "CREATE TABLE " + tableName +
+                " (Id UInt32, event_ts Timestamp, name String, p1 Int64, p2 String) ENGINE = MergeTree() ORDER BY ()";
+        dropTable(tableName);
+        createTable(createSQL);
+
+        ByteArrayOutputStream data = new ByteArrayOutputStream();
+        PrintWriter writer = new PrintWriter(data);
+        int numberOfRecords = 3;
+        for (int i = 0; i < numberOfRecords; i++) {
+            writer.printf("%d\t%s\t%s\t%d\t%s\n", i, "2021-01-01 00:00:00", "name" + i, i, "p2");
+        }
+        writer.flush();
+
+        InsertSettings settings = new InsertSettings()
+                .setQueryId(String.valueOf(UUID.randomUUID()))
+                .setOperationId(UUID.randomUUID().toString());
+        InsertResponse response = client.insert(tableName, new ByteArrayInputStream(data.toByteArray()),
+                ClickHouseFormat.TSV, settings).get(30, TimeUnit.SECONDS);
+        OperationMetrics metrics = response.getMetrics();
+        assertEquals((int)response.getWrittenRows(), numberOfRecords );
+        assertEquals(metrics.getQueryId(), settings.getQueryId());
+        assertTrue(metrics.getMetric(ClientMetrics.OP_DURATION).getLong() > 0);
+    }
 }