Add RegisterDial function: #372

Author: kshvakov

clickhouse_test.go

@@ -1774,7 +1774,7 @@ func Test_Tuple(t *testing.T) {
  			    ?,
  				?,
  				?,
- 				?,     
+ 				?,
  				?,
  				?,
  				?,
@@ -2061,7 +2061,7 @@ func Test_ReadHistogram(t *testing.T) {
  			    ?,
  				?,
  				?,
- 				?,     
+ 				?,
  				?,
  				?,
  				?,
@@ -2173,12 +2173,12 @@ func Test_ReadHistogram(t *testing.T) {
 func Test_ReadArrayArrayTuple(t *testing.T) {
 	const (
 		query = `
- 			select 
+ 			select
  			       [
- 			           [(1.0, 2.0, 3.0)], 
- 			           [(4.0, 5.0, 6.0), (7.0, 8.0, 9.0)], 
+ 			           [(1.0, 2.0, 3.0)],
+ 			           [(4.0, 5.0, 6.0), (7.0, 8.0, 9.0)],
  			           [(10.0, 11.0, 12.0), (13.0, 14.0, 15.0), (16.0, 17.0, 18.0), (19.0, 20.0, 21.0), (22.0, 23.0, 24.0)]
-				   ], 
+				   ],
  			       number
 			from numbers(2)
 			group by number;
@@ -2229,3 +2229,13 @@ func Test_ReadArrayArrayTuple(t *testing.T) {
 		}
 	}
 }
+
+func Test_RegisterDial(t *testing.T) {
+	clickhouse.RegisterDial(func(network, address string, timeout time.Duration, config *tls.Config) (net.Conn, error) {
+		return net.DialTimeout(network, address, timeout)
+	})
+	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?debug=true"); assert.NoError(t, err) {
+		assert.NoError(t, connect.Ping())
+	}
+	clickhouse.DeregisterDial()
+}

connect.go

@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"database/sql/driver"
 	"net"
+	"sync"
 	"sync/atomic"
 	"time"
 )
@@ -39,6 +40,28 @@ type connOptions struct {
 	logf                                   func(string, ...interface{})
 }
 
+// DialFunc is a function which can be used to establish the network connection.
+// Custom dial functions must be registered with RegisterDial
+type DialFunc func(network, address string, timeout time.Duration, config *tls.Config) (net.Conn, error)
+
+var (
+	customDialLock sync.RWMutex
+	customDial     DialFunc
+)
+
+// RegisterDial registers a custom dial function.
+func RegisterDial(dial DialFunc) {
+	customDialLock.Lock()
+	customDial = dial
+	customDialLock.Unlock()
+}
+
+// DeregisterDial deregisters the custom dial function.
+func DeregisterDial() {
+	customDialLock.Lock()
+	customDial = nil
+	customDialLock.Unlock()
+}
 func dial(options connOptions) (*connect, error) {
 	var (
 		err error
@@ -74,18 +97,29 @@ func dial(options connOptions) (*connect, error) {
 			}
 			checkedHosts[num] = struct{}{}
 		}
+		customDialLock.RLock()
+		cd := customDial
+		customDialLock.RUnlock()
 		switch {
 		case options.secure:
-			conn, err = tls.DialWithDialer(
-				&net.Dialer{
-					Timeout: options.connTimeout,
-				},
-				"tcp",
-				options.hosts[num],
-				tlsConfig,
-			)
+			if cd != nil {
+				conn, err = cd("tcp", options.hosts[num], options.connTimeout, tlsConfig)
+			} else {
+				conn, err = tls.DialWithDialer(
+					&net.Dialer{
+						Timeout: options.connTimeout,
+					},
+					"tcp",
+					options.hosts[num],
+					tlsConfig,
+				)
+			}
 		default:
-			conn, err = net.DialTimeout("tcp", options.hosts[num], options.connTimeout)
+			if cd != nil {
+				conn, err = cd("tcp", options.hosts[num], options.connTimeout, nil)
+			} else {
+				conn, err = net.DialTimeout("tcp", options.hosts[num], options.connTimeout)
+			}
 		}
 		if err == nil {
 			options.logf(