From 4184d6b95a2011874567d1ac7f6b841256b0d93b Mon Sep 17 00:00:00 2001
From: Zhichun Wu <zhicwu@gmail.com>
Date: Sun, 11 Jul 2021 18:20:29 +0800
Subject: [PATCH] Enable ipv6 support in BalancedClickHouseDataSource

---
 .../clickhouse/BalancedClickhouseDataSource.java   |  2 +-
 .../BalancedClickhouseDataSourceTest.java          | 14 ++++++++++++++
 .../clickhouse/ClickHouseContainerForTest.java     |  6 ++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/BalancedClickhouseDataSource.java b/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/BalancedClickhouseDataSource.java
index 8b1c00c42..3bf9a4b36 100644
--- a/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/BalancedClickhouseDataSource.java
+++ b/clickhouse-jdbc/src/main/java/ru/yandex/clickhouse/BalancedClickhouseDataSource.java
@@ -27,7 +27,7 @@
 public class BalancedClickhouseDataSource implements DataSource {
     private static final org.slf4j.Logger log = LoggerFactory.getLogger(BalancedClickhouseDataSource.class);
     private static final Pattern URL_TEMPLATE = Pattern.compile(JDBC_CLICKHOUSE_PREFIX + "" +
-            "//([a-zA-Z0-9_:,.-]+)" +
+            "//([a-zA-Z0-9_\\[\\]:,.-]+)" +
             "(/[a-zA-Z0-9_]+" +
             "([?][a-zA-Z0-9_]+[=][a-zA-Z0-9_]+([&][a-zA-Z0-9_]+[=][a-zA-Z0-9_]+)*)?" +
             ")?");
diff --git a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/BalancedClickhouseDataSourceTest.java b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/BalancedClickhouseDataSourceTest.java
index cbbab493a..d321f1bbd 100644
--- a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/BalancedClickhouseDataSourceTest.java
+++ b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/BalancedClickhouseDataSourceTest.java
@@ -313,4 +313,18 @@ public void testConnectionWithAuth() throws SQLException {
         }
     }
 
+    @Test
+    public void testIPv6() throws Exception {
+        // dedup is not supported at all :<
+        assertEquals(Arrays.asList("jdbc:clickhouse://[::1]:12345", "jdbc:clickhouse://[0:0:0:0:0:0:0:1]:12345"),
+                BalancedClickhouseDataSource.splitUrl("jdbc:clickhouse://[::1]:12345,[0:0:0:0:0:0:0:1]:12345"));
+        assertEquals(Arrays.asList("jdbc:clickhouse://[192:168:0:0:0:0:0:1]:12345", "jdbc:clickhouse://[192:168:0:0:0:0:0:2]:12345"),
+                BalancedClickhouseDataSource.splitUrl("jdbc:clickhouse://[192:168:0:0:0:0:0:1]:12345,[192:168:0:0:0:0:0:2]:12345"));
+        
+        ClickHouseProperties properties = new ClickHouseProperties();
+        String hostAddr = ClickHouseContainerForTest.getClickHouseHttpAddress();
+        String ipAddr = ClickHouseContainerForTest.getClickHouseHttpAddress("[::1]");
+        assertEquals(ClickHouseContainerForTest.newBalancedDataSource(properties, ipAddr).getConnection().getServerVersion(),
+            ClickHouseContainerForTest.newBalancedDataSource(properties, hostAddr).getConnection().getServerVersion());
+    }
 }
diff --git a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/ClickHouseContainerForTest.java b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/ClickHouseContainerForTest.java
index 70a696f63..9aae3ebcc 100644
--- a/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/ClickHouseContainerForTest.java
+++ b/clickhouse-jdbc/src/test/java/ru/yandex/clickhouse/ClickHouseContainerForTest.java
@@ -75,6 +75,12 @@ public static String getClickHouseHttpAddress(boolean useIPaddress) {
                 .append(':').append(clickhouseContainer.getMappedPort(HTTP_PORT)).toString();
     }
 
+    public static String getClickHouseHttpAddress(String customHostOrIp) {
+        return new StringBuilder()
+                .append(customHostOrIp == null || customHostOrIp.isEmpty() ? clickhouseContainer.getContainerIpAddress() : customHostOrIp)
+                .append(':').append(clickhouseContainer.getMappedPort(HTTP_PORT)).toString();
+    }
+
     public static ClickHouseDataSource newDataSource() {
         return newDataSource(new ClickHouseProperties());
     }