diff --git a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/settings/ClickHouseConnectionSettings.java b/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/settings/ClickHouseConnectionSettings.java index ed9bedadb..05028647d 100644 --- a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/settings/ClickHouseConnectionSettings.java +++ b/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"), diff --git a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/settings/ClickHouseProperties.java b/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/settings/ClickHouseProperties.java index b99b476e1..435a4366f 100644 --- a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/settings/ClickHouseProperties.java +++ b/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; } diff --git a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/util/ClickHouseCookieStoreProvider.java b/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/util/ClickHouseCookieStoreProvider.java new file mode 100644 index 000000000..921226f66 --- /dev/null +++ b/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 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()); + } +} diff --git a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/util/ClickHouseHttpClientBuilder.java b/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/util/ClickHouseHttpClientBuilder.java index 9dbb8af6b..a7fd7956d 100644 --- a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/util/ClickHouseHttpClientBuilder.java +++ b/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(); } diff --git a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/util/ClickHouseCookieStoreProviderTest.java b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/util/ClickHouseCookieStoreProviderTest.java new file mode 100644 index 000000000..41500a9ae --- /dev/null +++ b/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)); + } +} diff --git a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/util/ClickHouseHttpClientBuilderTest.java b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/util/ClickHouseHttpClientBuilderTest.java index 41502c0ee..8ecbcd1f8 100644 --- a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/util/ClickHouseHttpClientBuilderTest.java +++ b/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,