Merge pull request #1864 from ClickHouse/v2_log_comment

[client-v2] Added setting log_comment
ClickHouse · Oct 29, 2024 · db1c498 · db1c498
2 parents 3b831cd + a5b3277
commit db1c498
Show file tree

Hide file tree

Showing 6 changed files with 102 additions and 14 deletions.
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/ClientSettings.java b/client-v2/src/main/java/com/clickhouse/client/api/ClientSettings.java
@@ -35,4 +35,6 @@ public static List<String> valuesFromCommaSeparated(String value) {
     }
 
     public static final String SESSION_DB_ROLES = "session_db_roles";
+
+    public static final String SETTING_LOG_COMMENT = SERVER_SETTING_PREFIX + "log_comment";
 }
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java b/client-v2/src/main/java/com/clickhouse/client/api/insert/InsertSettings.java
@@ -224,4 +224,23 @@ public InsertSettings setDBRoles(Collection<String> dbRoles) {
     public Collection<String> getDBRoles() {
         return (Collection<String>) rawSettings.get(ClientSettings.SESSION_DB_ROLES);
     }
+
+    /**
+     * Sets the comment that will be added to the query log record associated with the query.
+     * @param logComment - comment to be added to the log
+     * @return same instance of the builder
+     */
+    public InsertSettings logComment(String logComment) {
+        this.logComment = logComment;
+        if (logComment != null && !logComment.isEmpty()) {
+            rawSettings.put(ClientSettings.SETTING_LOG_COMMENT, logComment);
+        }
+        return this;
+    }
+
+    private String logComment = null;
+
+    public String getLogComment() {
+        return logComment;
+    }
 }
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/internal/HttpAPIClientHelper.java b/client-v2/src/main/java/com/clickhouse/client/api/internal/HttpAPIClientHelper.java
@@ -54,29 +54,22 @@
 
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLException;
-import java.io.BufferedReader;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.InputStreamReader;
 import java.io.OutputStream;
 import java.net.ConnectException;
 import java.net.InetSocketAddress;
 import java.net.NoRouteToHostException;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.UnknownHostException;
-import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.security.NoSuchAlgorithmException;
 import java.util.Base64;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.EnumSet;
-import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
-import java.util.StringTokenizer;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 

diff --git a/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java b/client-v2/src/main/java/com/clickhouse/client/api/query/QuerySettings.java
@@ -240,4 +240,23 @@ public QuerySettings setDBRoles(Collection<String> dbRoles) {
     public Collection<String> getDBRoles() {
         return (Collection<String>) rawSettings.get(ClientSettings.SESSION_DB_ROLES);
     }
+
+    /**
+     * Sets the comment that will be added to the query log record associated with the query.
+     * @param logComment - comment to be added to the log
+     * @return same instance of the builder
+     */
+    public QuerySettings logComment(String logComment) {
+        this.logComment = logComment;
+        if (logComment != null && !logComment.isEmpty()) {
+            rawSettings.put(ClientSettings.SETTING_LOG_COMMENT, logComment);
+        }
+        return this;
+    }
+
+    private String logComment = null;
+
+    public String getLogComment() {
+        return logComment;
+    }
 }
diff --git a/client-v2/src/test/java/com/clickhouse/client/insert/InsertTests.java b/client-v2/src/test/java/com/clickhouse/client/insert/InsertTests.java
@@ -9,6 +9,7 @@
 import com.clickhouse.client.ClickHouseProtocol;
 import com.clickhouse.client.api.Client;
 import com.clickhouse.client.api.ClientException;
+import com.clickhouse.client.api.command.CommandResponse;
 import com.clickhouse.client.api.data_formats.ClickHouseBinaryFormatReader;
 import com.clickhouse.client.api.enums.Protocol;
 import com.clickhouse.client.api.insert.InsertResponse;
@@ -18,17 +19,12 @@
 import com.clickhouse.client.api.metrics.ServerMetrics;
 import com.clickhouse.client.api.query.GenericRecord;
 import com.clickhouse.client.api.query.QueryResponse;
-import com.clickhouse.client.config.ClickHouseClientOption;
 import com.clickhouse.data.ClickHouseFormat;
-import com.github.tomakehurst.wiremock.WireMockServer;
-import com.github.tomakehurst.wiremock.client.WireMock;
-import com.github.tomakehurst.wiremock.common.ConsoleNotifier;
-import com.github.tomakehurst.wiremock.core.WireMockConfiguration;
-import com.github.tomakehurst.wiremock.http.Fault;
 import org.testcontainers.shaded.org.apache.commons.lang3.RandomStringUtils;
 import org.testng.Assert;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 import java.io.ByteArrayInputStream;
@@ -41,7 +37,6 @@
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
-import static com.github.tomakehurst.wiremock.stubbing.Scenario.STARTED;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
@@ -294,4 +289,44 @@ public void testInsertSettingsAddDatabase() throws Exception {
         List<GenericRecord> records = client.queryAll("SELECT * FROM " + new_database + "." + tableName);
         assertEquals(records.size(), 1000);
     }
+
+    @Test(groups = {"integration"}, dataProviderClass = InsertTests.class, dataProvider = "logCommentDataProvider")
+    public void testLogComment(String logComment) throws Exception {
+
+        InsertSettings settings = new InsertSettings()
+                .setQueryId(UUID.randomUUID().toString())
+                .logComment(logComment);
+
+        final String tableName = "single_pojo_table";
+        final String createSQL = SamplePOJO.generateTableCreateSQL(tableName);
+        final SamplePOJO pojo = new SamplePOJO();
+
+        dropTable(tableName);
+        createTable(createSQL);
+        client.register(SamplePOJO.class, client.getTableSchema(tableName, "default"));
+
+        try (InsertResponse response = client.insert(tableName, Collections.singletonList(pojo), settings).get(30, TimeUnit.SECONDS)) {
+            Assert.assertEquals(response.getWrittenRows(), 1);
+        }
+
+        try (CommandResponse resp = client.execute("SYSTEM FLUSH LOGS").get()) {
+        }
+
+        List<GenericRecord> logRecords = client.queryAll("SELECT query_id, log_comment FROM system.query_log WHERE query_id = '" + settings.getQueryId() + "'");
+        Assert.assertEquals(logRecords.get(0).getString("query_id"), settings.getQueryId());
+        Assert.assertEquals(logRecords.get(0).getString("log_comment"), logComment == null ? "" : logComment);
+    }
+
+    @DataProvider( name = "logCommentDataProvider")
+    public static Object[] logCommentDataProvider() {
+        return new Object[][] {
+                { "Test log comment" },
+                { "Another log comment?" },
+                { "Log comment with special characters: !@#$%^&*()" },
+                { "Log comment with unicode: 你好" },
+                { "", },
+                { "               "},
+                { null }
+        };
+    }
 }
diff --git a/client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java b/client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java
@@ -1743,6 +1743,26 @@ public void testClientCustomRoles(String[] roles) throws Exception {
     }
 
 
+    @Test(groups = {"integration"})
+    public void testLogComment() throws Exception {
+
+        String logComment = "Test log comment";
+        QuerySettings settings = new QuerySettings()
+                .setQueryId(UUID.randomUUID().toString())
+                .logComment(logComment);
+        try (QueryResponse response = client.query("SELECT 1", settings).get()) {
+            Assert.assertNotNull(response.getQueryId());
+            Assert.assertTrue(response.getQueryId().startsWith(settings.getQueryId()));
+        }
+
+        try (CommandResponse resp = client.execute("SYSTEM FLUSH LOGS").get()) {
+        }
+
+        List<GenericRecord> logRecords = client.queryAll("SELECT query_id, log_comment FROM system.query_log WHERE query_id = '" + settings.getQueryId() + "'");
+        Assert.assertEquals(logRecords.get(0).getString("query_id"), settings.getQueryId());
+        Assert.assertEquals(logRecords.get(0).getString("log_comment"), logComment);
+    }
+
     protected Client.Builder newClient() {
         ClickHouseNode node = getServer(ClickHouseProtocol.HTTP);
         return new Client.Builder()