Added custom header tests

clickhouse-http-client/pom.xml

@@ -87,7 +87,7 @@
         <dependency>
             <groupId>org.wiremock</groupId>
             <artifactId>wiremock-standalone</artifactId>
-            <version>3.13.0</version>
+            <version>${wiremock.version}</version>
             <scope>test</scope>
         </dependency>
 

clickhouse-http-client/src/main/java/com/clickhouse/client/http/config/ClickHouseHttpOption.java

@@ -116,6 +116,12 @@ public enum ClickHouseHttpOption implements ClickHouseOption {
      */
     USE_BASIC_AUTHENTICATION("http_use_basic_auth", true, "Whether to use basic authentication.");
 
+    /**
+     * Replica tag header used by a proxy to route a request to a specific replica.
+     * ClickHouse Cloud feature only.
+     */
+    public static final String HEADER_REPLICA_TAG = "X-ClickHouse-Replica-Tag";
+
     private final String key;
     private final Serializable defaultValue;
     private final Class<? extends Serializable> clazz;

clickhouse-http-client/src/test/java/com/clickhouse/client/http/ApacheHttpConnectionImplTest.java

@@ -14,6 +14,7 @@
 import com.clickhouse.client.http.config.ClickHouseHttpOption;
 import com.clickhouse.client.http.config.HttpConnectionProvider;
 import com.clickhouse.config.ClickHouseOption;
+import com.clickhouse.data.ClickHouseFormat;
 import com.clickhouse.data.ClickHouseUtils;
 import com.github.tomakehurst.wiremock.WireMockServer;
 import com.github.tomakehurst.wiremock.client.WireMock;
@@ -77,6 +78,57 @@ protected Map<ClickHouseOption, Serializable> getClientOptions() {
                 HttpConnectionProvider.APACHE_HTTP_CLIENT);
     }
 
+    @Test(groups = { "unit" }, dataProvider = "replicaTags")
+    public void testCustomHeadersRouteToReplica(String replicaTag) throws Exception {
+        String host = "replica-router.clickhouse.test";
+        String expectedReplica = "replica-for-requested-tag";
+        String otherReplicaTag = "other-" + replicaTag;
+        WireMockServer mockServer = new WireMockServer(WireMockConfiguration.options().dynamicPort());
+        mockServer.start();
+        try {
+            mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
+                    .withHeader(ClickHouseHttpOption.HEADER_REPLICA_TAG, WireMock.equalTo(replicaTag))
+                    .withHeader("Host", WireMock.equalTo(host))
+                    .withRequestBody(WireMock.matching("(?is)select\\s+hostname\\(\\).*"))
+                    .willReturn(WireMock.ok("hostname()\nString\n" + expectedReplica + "\n"))
+                    .build());
+            mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
+                    .withHeader(ClickHouseHttpOption.HEADER_REPLICA_TAG, WireMock.equalTo(otherReplicaTag))
+                    .withHeader("Host", WireMock.equalTo(host))
+                    .withRequestBody(WireMock.matching("(?is)select\\s+hostname\\(\\).*"))
+                    .willReturn(WireMock.ok("hostname()\nString\nother-replica\n"))
+                    .build());
+
+            Map<ClickHouseOption, Serializable> options = new HashMap<>();
+            options.put(ClickHouseHttpOption.CONNECTION_PROVIDER, HttpConnectionProvider.APACHE_HTTP_CLIENT);
+            options.put(ClickHouseClientOption.COMPRESS, false);
+            options.put(ClickHouseHttpOption.CUSTOM_HEADERS,
+                    ClickHouseHttpOption.HEADER_REPLICA_TAG + "=" + replicaTag + ",Host=" + host);
+
+            try (ClickHouseClient client = ClickHouseClient.builder().config(new ClickHouseConfig(options)).build();
+                    ClickHouseResponse response = client.read("http://localhost:" + mockServer.port())
+                            .format(ClickHouseFormat.TabSeparatedWithNamesAndTypes)
+                            .query("select hostname()").executeAndWait()) {
+                Assert.assertEquals(response.firstRecord().getValue(0).asString(), expectedReplica);
+            }
+
+            mockServer.verify(WireMock.postRequestedFor(WireMock.anyUrl())
+                    .withHeader(ClickHouseHttpOption.HEADER_REPLICA_TAG, WireMock.equalTo(replicaTag))
+                    .withHeader("Host", WireMock.equalTo(host)));
+        } finally {
+            mockServer.stop();
+        }
+    }
+
+    @DataProvider(name = "replicaTags")
+    public static Object[][] replicaTags() {
+        return new Object[][] {
+                { "550e8400-e29b-41d4-a716-446655440000" },
+                { "replica-primary" },
+                { "replica=primary" }
+        };
+    }
+
     @Test(groups = { "integration" })
     public void testConnection() throws Exception {
         ClickHouseNode server = getServer(ClickHouseProtocol.HTTP);

clickhouse-jdbc/pom.xml

@@ -96,6 +96,12 @@
             <artifactId>testng</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.wiremock</groupId>
+            <artifactId>wiremock-standalone</artifactId>
+            <version>${wiremock.version}</version>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>com.mysql</groupId>
             <artifactId>mysql-connector-j</artifactId>

clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHouseConnectionTest.java

@@ -13,14 +13,20 @@
 import com.clickhouse.client.ClickHouseRequest;
 import com.clickhouse.client.ClickHouseServerForTest;
 import com.clickhouse.client.config.ClickHouseClientOption;
+import com.clickhouse.client.http.config.ClickHouseHttpOption;
 import com.clickhouse.data.ClickHouseCompression;
+import com.clickhouse.data.ClickHouseFormat;
 import com.clickhouse.data.ClickHouseUtils;
 import com.clickhouse.data.value.UnsignedByte;
+import com.github.tomakehurst.wiremock.WireMockServer;
+import com.github.tomakehurst.wiremock.client.WireMock;
+import com.github.tomakehurst.wiremock.core.WireMockConfiguration;
 
 import org.testng.Assert;
 import org.testng.SkipException;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 public class ClickHouseConnectionTest extends JdbcIntegrationTest {
@@ -37,6 +43,62 @@ public ClickHouseConnection newConnection(Properties properties) throws SQLExcep
         return (ClickHouseConnection) newDataSource(properties).getConnection();
     }
 
+    @Test(groups = "unit", dataProvider = "replicaTags")
+    public void testCustomHeadersRouteToReplica(String replicaTag) throws Exception {
+        String host = "replica-router.clickhouse.test";
+        String expectedReplica = "replica-for-requested-tag";
+        String otherReplicaTag = "other-" + replicaTag;
+        WireMockServer mockServer = new WireMockServer(WireMockConfiguration.options().dynamicPort());
+        mockServer.start();
+        try {
+            mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
+                    .withHeader(ClickHouseHttpOption.HEADER_REPLICA_TAG, WireMock.equalTo(replicaTag))
+                    .withHeader("Host", WireMock.equalTo(host))
+                    .withRequestBody(WireMock.matching("(?is)select\\s+hostname\\(\\).*"))
+                    .willReturn(WireMock.ok("hostname()\nString\n" + expectedReplica + "\n"))
+                    .build());
+            mockServer.addStubMapping(WireMock.post(WireMock.anyUrl())
+                    .withHeader(ClickHouseHttpOption.HEADER_REPLICA_TAG, WireMock.equalTo(otherReplicaTag))
+                    .withHeader("Host", WireMock.equalTo(host))
+                    .withRequestBody(WireMock.matching("(?is)select\\s+hostname\\(\\).*"))
+                    .willReturn(WireMock.ok("hostname()\nString\nother-replica\n"))
+                    .build());
+
+            Properties properties = new Properties();
+            properties.setProperty(ClickHouseClientOption.SERVER_TIME_ZONE.getKey(), "UTC");
+            properties.setProperty(ClickHouseClientOption.SERVER_VERSION.getKey(), "25.8");
+            properties.setProperty(ClickHouseClientOption.COMPRESS.getKey(), Boolean.FALSE.toString());
+            properties.setProperty(ClickHouseClientOption.FORMAT.getKey(),
+                    ClickHouseFormat.TabSeparatedWithNamesAndTypes.name());
+            properties.setProperty(ClickHouseHttpOption.CUSTOM_HEADERS.getKey(),
+                    ClickHouseHttpOption.HEADER_REPLICA_TAG + "=" + replicaTag + ",Host=" + host);
+
+            String url = "jdbc:clickhouse:http://localhost:" + mockServer.port() + "?clickhouse.jdbc.v1=true";
+            try (Connection connection = new ClickHouseDataSource(url, properties).getConnection();
+                    Statement statement = connection.createStatement();
+                    ResultSet resultSet = statement.executeQuery("select hostname()")) {
+                Assert.assertTrue(resultSet.next());
+                Assert.assertEquals(resultSet.getString(1), expectedReplica);
+                Assert.assertFalse(resultSet.next());
+            }
+
+            mockServer.verify(WireMock.postRequestedFor(WireMock.anyUrl())
+                    .withHeader(ClickHouseHttpOption.HEADER_REPLICA_TAG, WireMock.equalTo(replicaTag))
+                    .withHeader("Host", WireMock.equalTo(host)));
+        } finally {
+            mockServer.stop();
+        }
+    }
+
+    @DataProvider(name = "replicaTags")
+    public static Object[][] replicaTags() {
+        return new Object[][] {
+                { "550e8400-e29b-41d4-a716-446655440000" },
+                { "replica-primary" },
+                { "replica=primary" }
+        };
+    }
+
     @Test(groups = "integration")
     public void testCentralizedConfiguration() throws SQLException {
         if (isCloud()) return; //TODO: testCentralizedConfiguration - Revisit, see: https://github.com/ClickHouse/clickhouse-java/issues/1747

client-v2/pom.xml

@@ -137,7 +137,7 @@
         <dependency>
             <groupId>org.wiremock</groupId>
             <artifactId>wiremock-standalone</artifactId>
-            <version>3.13.0</version>
+            <version>${wiremock.version}</version>
             <scope>test</scope>
         </dependency>
 

client-v2/src/main/java/com/clickhouse/client/api/http/ClickHouseHttpProto.java

@@ -52,6 +52,11 @@ public class ClickHouseHttpProto {
 
     public static final String HEADER_SSL_CERT_AUTH = "x-clickhouse-ssl-certificate-auth";
 
+    /**
+     * Replica tag used by a proxy to route a request to a specific replica.
+     */
+    public static final String HEADER_REPLICA_TAG = "X-ClickHouse-Replica-Tag";
+
     /**
      * Query parameter to specify the query ID.
      */