From 0966f69a3abe427a9c4037b20e5ddd42a22b26f9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=AB=98=E9=93=81?= <shaojin.wensj@alibaba-inc.com>
Date: Wed, 17 Feb 2021 17:21:31 +0800
Subject: [PATCH] support BalancedClickhouseDataSource, fix #3984

---
 .../alibaba/druid/pool/DruidDataSource.java   |  7 +++
 .../clickhouse/BalancedClickhouseDriver.java  | 59 +++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 src/main/java/com/alibaba/druid/support/clickhouse/BalancedClickhouseDriver.java

diff --git a/src/main/java/com/alibaba/druid/pool/DruidDataSource.java b/src/main/java/com/alibaba/druid/pool/DruidDataSource.java
index c1546e3b46..b4bb94d47d 100644
--- a/src/main/java/com/alibaba/druid/pool/DruidDataSource.java
+++ b/src/main/java/com/alibaba/druid/pool/DruidDataSource.java
@@ -82,6 +82,7 @@ import com.alibaba.druid.stat.DruidDataSourceStatManager;
 import com.alibaba.druid.stat.JdbcDataSourceStat;
 import com.alibaba.druid.stat.JdbcSqlStat;
 import com.alibaba.druid.stat.JdbcSqlStatValue;
+import com.alibaba.druid.support.clickhouse.BalancedClickhouseDriver;
 import com.alibaba.druid.support.logging.Log;
 import com.alibaba.druid.support.logging.LogFactory;
 import com.alibaba.druid.util.JMXUtils;
@@ -1196,6 +1197,12 @@ public class DruidDataSource extends DruidAbstractDataSource implements DruidDat
 
             if (MockDriver.class.getName().equals(driverClass)) {
                 driver = MockDriver.instance;
+            } else if ("com.alibaba.druid.support.clickhouse.BalancedClickhouseDriver".equals(driverClass)) {
+                Properties info = new Properties();
+                info.put("user", username);
+                info.put("password", password);
+                info.putAll(connectProperties);
+                driver = new BalancedClickhouseDriver(jdbcUrl, info);
             } else {
                 if (jdbcUrl == null && (driverClass == null || driverClass.length() == 0)) {
                     throw new SQLException("url not set");
diff --git a/src/main/java/com/alibaba/druid/support/clickhouse/BalancedClickhouseDriver.java b/src/main/java/com/alibaba/druid/support/clickhouse/BalancedClickhouseDriver.java
new file mode 100644
index 0000000000..40fb10913a
--- /dev/null
+++ b/src/main/java/com/alibaba/druid/support/clickhouse/BalancedClickhouseDriver.java
@@ -0,0 +1,59 @@
+package com.alibaba.druid.support.clickhouse;
+
+import ru.yandex.clickhouse.BalancedClickhouseDataSource;
+
+import java.sql.Connection;
+import java.sql.DriverPropertyInfo;
+import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
+import java.util.Properties;
+import java.util.logging.Logger;
+
+public class BalancedClickhouseDriver implements java.sql.Driver {
+    private final String url;
+    private BalancedClickhouseDataSource dataSource;
+
+    public BalancedClickhouseDriver(final String url, Properties properties) {
+        this.url = url;
+        this.dataSource = new BalancedClickhouseDataSource(url, properties);
+    }
+
+    @Override
+    public Connection connect(String url, Properties info) throws SQLException {
+        if (!acceptsURL(url)) {
+            throw new SQLException("TODO");
+        }
+
+        return dataSource.getConnection();
+    }
+
+    @Override
+    public boolean acceptsURL(String url) throws SQLException {
+        return this.url.equals(url);
+    }
+
+    @Override
+    public DriverPropertyInfo[] getPropertyInfo(String url, Properties info) throws SQLException {
+        return new DriverPropertyInfo[0];
+    }
+
+    @Override
+    public int getMajorVersion() {
+        return 0;
+    }
+
+    @Override
+    public int getMinorVersion() {
+        return 0;
+    }
+
+    @Override
+    public boolean jdbcCompliant() {
+        return false;
+    }
+
+    @Override
+    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
+        return null;
+    }
+}