Add DialContext option (#487)

* Add DialContext option

It allows to using a cancellable context for a connection

* Remove previous Dial-option

Author: slaskis

clickhouse.go

@@ -181,13 +181,13 @@ func (ch *clickhouse) Stats() driver.Stats {
 	}
 }
 
-func (ch *clickhouse) dial() (conn *connect, err error) {
+func (ch *clickhouse) dial(ctx context.Context) (conn *connect, err error) {
 	connID := int(atomic.AddInt64(&ch.connID, 1))
 	for num := range ch.opt.Addr {
 		if ch.opt.ConnOpenStrategy == ConnOpenRoundRobin {
 			num = int(connID) % len(ch.opt.Addr)
 		}
-		if conn, err = dial(ch.opt.Addr[num], connID, ch.opt); err == nil {
+		if conn, err = dial(ctx, ch.opt.Addr[num], connID, ch.opt); err == nil {
 			return conn, nil
 		}
 	}
@@ -213,7 +213,7 @@ func (ch *clickhouse) acquire(ctx context.Context) (conn *connect, err error) {
 	case conn := <-ch.idle:
 		if conn.isBad() {
 			conn.close()
-			if conn, err = ch.dial(); err != nil {
+			if conn, err = ch.dial(ctx); err != nil {
 				select {
 				case <-ch.open:
 				default:
@@ -225,7 +225,7 @@ func (ch *clickhouse) acquire(ctx context.Context) (conn *connect, err error) {
 		return conn, nil
 	default:
 	}
-	if conn, err = ch.dial(); err != nil {
+	if conn, err = ch.dial(ctx); err != nil {
 		select {
 		case <-ch.open:
 		default:

clickhouse_options.go

@@ -18,6 +18,7 @@
 package clickhouse
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"net"
@@ -29,9 +30,7 @@ import (
 	"github.com/ClickHouse/clickhouse-go/v2/lib/compress"
 )
 
-var (
-	CompressionLZ4 compress.Method = compress.LZ4
-)
+var CompressionLZ4 compress.Method = compress.LZ4
 
 type Auth struct { // has_control_character
 	Database string
@@ -62,7 +61,7 @@ type Options struct {
 	TLS              *tls.Config
 	Addr             []string
 	Auth             Auth
-	Dial             func(addr string) (net.Conn, error)
+	DialContext      func(ctx context.Context, addr string) (net.Conn, error)
 	Debug            bool
 	Settings         Settings
 	Compression      *Compression

clickhouse_std.go

@@ -42,7 +42,7 @@ func (o *stdConnOpener) Driver() driver.Driver {
 	return &stdDriver{}
 }
 
-func (o *stdConnOpener) Connect(context.Context) (_ driver.Conn, err error) {
+func (o *stdConnOpener) Connect(ctx context.Context) (_ driver.Conn, err error) {
 	if o.err != nil {
 		return nil, o.err
 	}
@@ -54,7 +54,7 @@ func (o *stdConnOpener) Connect(context.Context) (_ driver.Conn, err error) {
 		if o.opt.ConnOpenStrategy == ConnOpenRoundRobin {
 			num = int(connID) % len(o.opt.Addr)
 		}
-		if conn, err = dial(o.opt.Addr[num], connID, o.opt); err == nil {
+		if conn, err = dial(ctx, o.opt.Addr[num], connID, o.opt); err == nil {
 			return &stdDriver{
 				conn: conn,
 			}, nil

conn.go

@@ -32,15 +32,15 @@ import (
 	"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
 )
 
-func dial(addr string, num int, opt *Options) (*connect, error) {
+func dial(ctx context.Context, addr string, num int, opt *Options) (*connect, error) {
 	var (
 		err    error
 		conn   net.Conn
 		debugf = func(format string, v ...interface{}) {}
 	)
 	switch {
-	case opt.Dial != nil:
-		conn, err = opt.Dial(addr)
+	case opt.DialContext != nil:
+		conn, err = opt.DialContext(ctx, addr)
 	default:
 		switch {
 		case opt.TLS != nil:

tests/custom_dial_test.go

@@ -26,7 +26,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestCustomDial(t *testing.T) {
+func TestCustomDialContext(t *testing.T) {
 	var (
 		dialCount int
 		conn, err = clickhouse.Open(&clickhouse.Options{
@@ -36,16 +36,39 @@ func TestCustomDial(t *testing.T) {
 				Username: "default",
 				Password: "",
 			},
-			Dial: func(addr string) (net.Conn, error) {
+			DialContext: func(ctx context.Context, addr string) (net.Conn, error) {
 				dialCount++
-				return net.Dial("tcp", addr)
+				var d net.Dialer
+				return d.DialContext(ctx, "tcp", addr)
 			},
 		})
 	)
 	if !assert.NoError(t, err) {
 		return
 	}
-	if err := conn.Ping(context.Background()); assert.NoError(t, err) {
+	ctx := context.Background()
+	if err := conn.Ping(ctx); assert.NoError(t, err) {
 		assert.Equal(t, 1, dialCount)
 	}
+
+	ctx1, cancel1 := context.WithCancel(ctx)
+	go func() {
+		cancel1()
+	}()
+
+	ctx2, cancel2 := context.WithCancel(ctx)
+	defer cancel2()
+
+	// query is cancelled with context
+	err = conn.QueryRow(ctx1, "SELECT sleep(3)").Scan()
+	if assert.Error(t, err, "context cancelled") {
+		assert.Equal(t, 1, dialCount)
+	}
+
+	// uncancelled context still works (new connection is acquired)
+	var i uint8
+	err = conn.QueryRow(ctx2, "SELECT 1").Scan(&i)
+	if assert.NoError(t, err) {
+		assert.Equal(t, 2, dialCount)
+	}
 }