Shared cookie store

clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/settings/ClickHouseConnectionSettings.java

@@ -42,6 +42,7 @@ public enum ClickHouseConnectionSettings implements DriverPropertyCreator {
      * additional
      */
     USE_OBJECTS_IN_ARRAYS("use_objects_in_arrays", false, "Whether Object[] should be used instead primitive arrays."),
+    USE_SHARED_COOKIE_STORE("useSharedCookieStore", false, "Whether to use shared cookie to store among all http clients of db are not"),
     MAX_COMPRESS_BUFFER_SIZE("maxCompressBufferSize", 1024*1024, ""),
 
     USE_SERVER_TIME_ZONE("use_server_time_zone", true, "Whether to use timezone from server. On connection init select timezone() will be executed"),

clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/settings/ClickHouseProperties.java

@@ -56,6 +56,8 @@ public class ClickHouseProperties {
     private String useTimeZone;
     private boolean useServerTimeZoneForDates;
     private boolean useObjectsInArrays;
+    // the shared cookie store is scoped to a database
+    private boolean useSharedCookieStore;
 
     // queries settings
     private Integer maxParallelReplicas;
@@ -130,6 +132,7 @@ public ClickHouseProperties(Properties info) {
         this.useTimeZone = (String)getSetting(info, ClickHouseConnectionSettings.USE_TIME_ZONE);
         this.useServerTimeZoneForDates = (Boolean)getSetting(info, ClickHouseConnectionSettings.USE_SERVER_TIME_ZONE_FOR_DATES);
         this.useObjectsInArrays = (Boolean)getSetting(info, ClickHouseConnectionSettings.USE_OBJECTS_IN_ARRAYS);
+        this.useSharedCookieStore = (Boolean)getSetting(info, ClickHouseConnectionSettings.USE_SHARED_COOKIE_STORE);
         this.clientName = (String)getSetting(info, ClickHouseConnectionSettings.CLIENT_NAME);
         this.useNewParser = (Boolean)getSetting(info, ClickHouseConnectionSettings.USE_NEW_PARSER);
 
@@ -199,6 +202,7 @@ public Properties asProperties() {
         ret.put(ClickHouseConnectionSettings.USE_TIME_ZONE.getKey(), String.valueOf(useTimeZone));
         ret.put(ClickHouseConnectionSettings.USE_SERVER_TIME_ZONE_FOR_DATES.getKey(), String.valueOf(useServerTimeZoneForDates));
         ret.put(ClickHouseConnectionSettings.USE_OBJECTS_IN_ARRAYS.getKey(), String.valueOf(useObjectsInArrays));
+        ret.put(ClickHouseConnectionSettings.USE_SHARED_COOKIE_STORE.getKey(), String.valueOf(useSharedCookieStore));
         ret.put(ClickHouseConnectionSettings.CLIENT_NAME.getKey(), String.valueOf(clientName));
         ret.put(ClickHouseConnectionSettings.USE_NEW_PARSER.getKey(), String.valueOf(useNewParser));
 
@@ -271,6 +275,7 @@ public ClickHouseProperties(ClickHouseProperties properties) {
         setUseTimeZone(properties.useTimeZone);
         setUseServerTimeZoneForDates(properties.useServerTimeZoneForDates);
         setUseObjectsInArrays(properties.useObjectsInArrays);
+        setUseSharedCookieStore(properties.useSharedCookieStore);
         setClientName(properties.clientName);
         setUseNewParser(properties.useNewParser);
         setMaxParallelReplicas(properties.maxParallelReplicas);
@@ -688,6 +693,14 @@ public void setUseObjectsInArrays(boolean useObjectsInArrays) {
         this.useObjectsInArrays = useObjectsInArrays;
     }
 
+    public boolean isUseSharedCookieStore() {
+        return useSharedCookieStore;
+    }
+
+    public void setUseSharedCookieStore(boolean useSharedCookieStore) {
+        this.useSharedCookieStore = useSharedCookieStore;
+    }
+
     public String getClientName() {
         return this.clientName;
     }

clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/util/ClickHouseCookieStoreProvider.java

@@ -0,0 +1,29 @@
+package ru.yandex.clickhouse.util;
+
+import org.apache.http.client.CookieStore;
+import org.apache.http.impl.client.BasicCookieStore;
+import ru.yandex.clickhouse.settings.ClickHouseProperties;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+public class ClickHouseCookieStoreProvider {
+    private static final Map<String, CookieStore> cookieStoreMap = new ConcurrentHashMap<>();
+
+    public CookieStore getCookieStore(ClickHouseProperties properties) {
+        return hasValidProperties(properties) && properties.isUseSharedCookieStore() ?
+                cookieStoreMap.computeIfAbsent(getCookieStoreKey(properties), k -> new BasicCookieStore()) :
+                null;
+    }
+
+    private boolean hasValidProperties(ClickHouseProperties properties) {
+        return properties != null
+                && !Utils.isNullOrEmptyString(properties.getHost())
+                && properties.getPort() > 0
+                && !Utils.isNullOrEmptyString(properties.getDatabase());
+    }
+
+    private String getCookieStoreKey(ClickHouseProperties properties) {
+        return String.format("%s:%s/%s", properties.getHost(), properties.getPort(), properties.getDatabase());
+    }
+}

clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/util/ClickHouseHttpClientBuilder.java

@@ -36,6 +36,7 @@
 import org.apache.http.client.AuthCache;
 import org.apache.http.client.CredentialsProvider;
 import org.apache.http.client.HttpRequestRetryHandler;
+import org.apache.http.client.config.CookieSpecs;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.protocol.HttpClientContext;
 import org.apache.http.config.ConnectionConfig;
@@ -60,6 +61,7 @@
 
 public class ClickHouseHttpClientBuilder {
 
+    private static final ClickHouseCookieStoreProvider cookieStoreProvider = new ClickHouseCookieStoreProvider();
     private final ClickHouseProperties properties;
 
     public ClickHouseHttpClientBuilder(ClickHouseProperties properties) {
@@ -76,6 +78,7 @@ public CloseableHttpClient buildClient() throws Exception {
                 .setDefaultHeaders(getDefaultHeaders())
                 .setDefaultCredentialsProvider(getDefaultCredentialsProvider())
                 .disableContentCompression() // gzip is not needed. Use lz4 when compress=1
+                .setDefaultCookieStore(cookieStoreProvider.getCookieStore(properties))
                 .disableRedirectHandling();
 
         String clientName = properties != null ? properties.getClientName() : null;
@@ -112,6 +115,7 @@ public static HttpClientContext createClientContext(ClickHouseProperties props)
         authCache.put(getTargetHost(props), basicAuth);
         HttpClientContext ctx = HttpClientContext.create();
         ctx.setAuthCache(authCache);
+        ctx.setCookieStore(cookieStoreProvider.getCookieStore(props));
         return ctx;
     }
 
@@ -164,6 +168,7 @@ private RequestConfig getRequestConfig() {
                 .setSocketTimeout(properties.getSocketTimeout())
                 .setConnectTimeout(properties.getConnectionTimeout())
                 .setConnectionRequestTimeout(properties.getConnectionTimeout())
+                .setCookieSpec(CookieSpecs.STANDARD)
                 .build();
     }
 

clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/util/ClickHouseCookieStoreProviderTest.java

@@ -0,0 +1,74 @@
+package ru.yandex.clickhouse.util;
+
+import org.testng.annotations.Test;
+import ru.yandex.clickhouse.settings.ClickHouseProperties;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+
+public class ClickHouseCookieStoreProviderTest {
+    ClickHouseCookieStoreProvider cookieStoreProvider = new ClickHouseCookieStoreProvider();
+
+    @Test
+    public void testCookieStoreProviderWithNullHost() {
+        ClickHouseProperties props = new ClickHouseProperties();
+        props.setUseSharedCookieStore(true);
+        props.setPort(8080);
+        props.setDatabase("default");
+        assertNull(cookieStoreProvider.getCookieStore(props));
+    }
+
+    @Test
+    public void testCookieStoreProviderWithInvalidPort() {
+        ClickHouseProperties props = new ClickHouseProperties();
+        props.setUseSharedCookieStore(true);
+        props.setHost("127.0.0.1");
+        props.setPort(0);
+        props.setDatabase("default");
+        assertNull(cookieStoreProvider.getCookieStore(props));
+    }
+
+    @Test
+    public void testCookieStoreProviderWithNullDBName() {
+        ClickHouseProperties props = new ClickHouseProperties();
+        props.setUseSharedCookieStore(true);
+        props.setHost("127.0.0.1");
+        props.setPort(0);
+        assertNull(cookieStoreProvider.getCookieStore(props));
+    }
+
+    @Test
+    public void testCookieStoreProviderWithSameDBAndSharedCookieStore() {
+        ClickHouseProperties props = new ClickHouseProperties();
+        props.setUseSharedCookieStore(true);
+        props.setHost("127.0.0.1");
+        props.setPort(8080);
+        props.setDatabase("default");
+        assertNotNull(cookieStoreProvider.getCookieStore(props));
+        assertEquals(cookieStoreProvider.getCookieStore(props), cookieStoreProvider.getCookieStore(props));
+    }
+
+    @Test
+    public void testCookieStoreProviderWithPrivateCookieStore() {
+        ClickHouseProperties props = new ClickHouseProperties();
+        props.setUseSharedCookieStore(false);
+        props.setHost("127.0.0.1");
+        props.setPort(8080);
+        props.setDatabase("default");
+        assertNull(cookieStoreProvider.getCookieStore(props));
+    }
+
+    @Test
+    public void testCookieStoreProviderWithDiffDB() {
+        ClickHouseProperties props1 = new ClickHouseProperties();
+        props1.setUseSharedCookieStore(true);
+        props1.setHost("127.0.0.1");
+        props1.setPort(8080);
+        props1.setDatabase("default1");
+        ClickHouseProperties props2 = new ClickHouseProperties(props1);
+        props2.setDatabase("default2");
+        assertNotEquals(cookieStoreProvider.getCookieStore(props1), cookieStoreProvider.getCookieStore(props2));
+    }
+}

clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/util/ClickHouseHttpClientBuilderTest.java

@@ -2,6 +2,7 @@
 
 import org.apache.http.HttpHost;
 import org.apache.http.NoHttpResponseException;
+import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.conn.HttpHostConnectException;
 import org.apache.http.impl.client.CloseableHttpClient;
@@ -19,6 +20,9 @@
 
 import ru.yandex.clickhouse.settings.ClickHouseProperties;
 
+import java.time.Instant;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNull;
 
@@ -101,6 +105,31 @@ public void testCreateClientContextOnlyPass() {
             "basic");
     }
 
+    @Test
+    public void testHttpClientsWithSharedCookie() throws Exception {
+        ClickHouseProperties props = new ClickHouseProperties();
+        props.setHost("localhost");
+        props.setPort(mockServer.port());
+        props.setDatabase("default");
+        props.setUseSharedCookieStore(true);
+        CloseableHttpClient client = new ClickHouseHttpClientBuilder(props).buildClient();
+        String cookie = "AWS-ALB=random-value-" + Instant.now().toEpochMilli();
+        mockServer.stubFor(WireMock.get(WireMock.urlPathMatching("/cookie/get"))
+                                   .willReturn(WireMock.aResponse()
+                                                       .withStatus(200)
+                                                       .withHeader("set-cookie", cookie)
+                                                       .withBody("OK")));
+        HttpGet getCookie = new HttpGet(mockServer.baseUrl() + "/cookie/get");
+        client.execute(getCookie);
+        CloseableHttpClient clientWithSharedCookieStore = new ClickHouseHttpClientBuilder(props).buildClient();
+        props.setUseSharedCookieStore(false);
+        CloseableHttpClient clientWithPrivateCookieStore = new ClickHouseHttpClientBuilder(props).buildClient();
+        HttpGet checkCookie = new HttpGet(mockServer.baseUrl() + "/cookie/check");
+        clientWithPrivateCookieStore.execute(checkCookie);
+        mockServer.verify(getRequestedFor(WireMock.urlEqualTo("/cookie/check")).withoutHeader("cookie"));
+        clientWithSharedCookieStore.execute(checkCookie);
+        mockServer.verify(getRequestedFor(WireMock.urlEqualTo("/cookie/check")).withHeader("cookie", equalTo(cookie)));
+    }
 
     @Test(dataProvider = "authUserPassword")
     public void testHttpAuthParametersCombination(String authorization, String user,