Merge pull request #1577 from ClickHouse/open_http

Support HTTP connnections in Native Go interface

Author: SpencerTorres

benchmark/json_test.go

@@ -18,7 +18,7 @@ func TestMain(m *testing.M) {
 }
 
 func GetNativeConnection(settings clickhouse.Settings, tlsConfig *tls.Config, compression *clickhouse.Compression) (driver.Conn, error) {
-	return clickhouse_tests.GetConnection(testSet, settings, tlsConfig, compression)
+	return clickhouse_tests.GetConnectionTCP(testSet, settings, tlsConfig, compression)
 }
 
 func prepareJSONTest(ctx context.Context, b *testing.B) driver.Conn {

clickhouse.go

@@ -27,7 +27,6 @@ import (
 
 	_ "time/tzdata"
 
-	chproto "github.com/ClickHouse/ch-go/proto"
 	"github.com/ClickHouse/clickhouse-go/v2/contributors"
 	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
@@ -83,19 +82,43 @@ func Open(opt *Options) (driver.Conn, error) {
 		opt = &Options{}
 	}
 	o := opt.setDefaults()
+
 	conn := &clickhouse{
 		opt:  o,
-		idle: make(chan *connect, o.MaxIdleConns),
+		idle: make(chan nativeTransport, o.MaxIdleConns),
 		open: make(chan struct{}, o.MaxOpenConns),
 		exit: make(chan struct{}),
 	}
 	go conn.startAutoCloseIdleConnections()
 	return conn, nil
 }
 
+// nativeTransport represents an implementation (TCP or HTTP) that can be pooled by the main clickhouse struct.
+// Implementations are not expected to be thread safe, which is why we provide acquire/release functions.
+type nativeTransport interface {
+	serverVersion() (*ServerVersion, error)
+	query(ctx context.Context, release nativeTransportRelease, query string, args ...any) (*rows, error)
+	queryRow(ctx context.Context, release nativeTransportRelease, query string, args ...any) *row
+	prepareBatch(ctx context.Context, release nativeTransportRelease, acquire nativeTransportAcquire, query string, opts driver.PrepareBatchOptions) (driver.Batch, error)
+	exec(ctx context.Context, query string, args ...any) error
+	asyncInsert(ctx context.Context, query string, wait bool, args ...any) error
+	ping(context.Context) error
+	isBad() bool
+	connID() int
+	connectedAtTime() time.Time
+	isReleased() bool
+	setReleased(released bool)
+	debugf(format string, v ...any)
+	// freeBuffer is called if Options.FreeBufOnConnRelease is set
+	freeBuffer()
+	close() error
+}
+type nativeTransportAcquire func(context.Context) (nativeTransport, error)
+type nativeTransportRelease func(nativeTransport, error)
+
 type clickhouse struct {
 	opt    *Options
-	idle   chan *connect
+	idle   chan nativeTransport
 	open   chan struct{}
 	exit   chan struct{}
 	connID int64
@@ -118,27 +141,27 @@ func (ch *clickhouse) ServerVersion() (*driver.ServerVersion, error) {
 	if err != nil {
 		return nil, err
 	}
-	ch.release(conn, nil)
-	return &conn.server, nil
+	defer ch.release(conn, nil)
+	return conn.serverVersion()
 }
 
 func (ch *clickhouse) Query(ctx context.Context, query string, args ...any) (rows driver.Rows, err error) {
 	conn, err := ch.acquire(ctx)
 	if err != nil {
 		return nil, err
 	}
-	conn.debugf("[acquired] connection [%d]", conn.id)
+	conn.debugf("[acquired] connection [%d]", conn.connID())
 	return conn.query(ctx, ch.release, query, args...)
 }
 
-func (ch *clickhouse) QueryRow(ctx context.Context, query string, args ...any) (rows driver.Row) {
+func (ch *clickhouse) QueryRow(ctx context.Context, query string, args ...any) driver.Row {
 	conn, err := ch.acquire(ctx)
 	if err != nil {
 		return &row{
 			err: err,
 		}
 	}
-	conn.debugf("[acquired] connection [%d]", conn.id)
+	conn.debugf("[acquired] connection [%d]", conn.connID())
 	return conn.queryRow(ctx, ch.release, query, args...)
 }
 
@@ -147,7 +170,7 @@ func (ch *clickhouse) Exec(ctx context.Context, query string, args ...any) error
 	if err != nil {
 		return err
 	}
-	conn.debugf("[acquired] connection [%d]", conn.id)
+	conn.debugf("[acquired] connection [%d]", conn.connID())
 
 	if err := conn.exec(ctx, query, args...); err != nil {
 		ch.release(conn, err)
@@ -162,7 +185,7 @@ func (ch *clickhouse) PrepareBatch(ctx context.Context, query string, opts ...dr
 	if err != nil {
 		return nil, err
 	}
-	batch, err := conn.prepareBatch(ctx, query, getPrepareBatchOptions(opts...), ch.release, ch.acquire)
+	batch, err := conn.prepareBatch(ctx, ch.release, ch.acquire, query, getPrepareBatchOptions(opts...))
 	if err != nil {
 		return nil, err
 	}
@@ -214,11 +237,18 @@ func (ch *clickhouse) Stats() driver.Stats {
 	}
 }
 
-func (ch *clickhouse) dial(ctx context.Context) (conn *connect, err error) {
+func (ch *clickhouse) dial(ctx context.Context) (conn nativeTransport, err error) {
 	connID := int(atomic.AddInt64(&ch.connID, 1))
 
 	dialFunc := func(ctx context.Context, addr string, opt *Options) (DialResult, error) {
-		conn, err := dial(ctx, addr, connID, opt)
+		var conn nativeTransport
+		var err error
+		switch opt.Protocol {
+		case HTTP:
+			conn, err = dialHttp(ctx, addr, connID, opt)
+		default:
+			conn, err = dial(ctx, addr, connID, opt)
+		}
 
 		return DialResult{conn}, err
 	}
@@ -260,7 +290,7 @@ func DefaultDialStrategy(ctx context.Context, connID int, opt *Options, dial Dia
 	return r, err
 }
 
-func (ch *clickhouse) acquire(ctx context.Context) (conn *connect, err error) {
+func (ch *clickhouse) acquire(ctx context.Context) (conn nativeTransport, err error) {
 	timer := time.NewTimer(ch.opt.DialTimeout)
 	defer timer.Stop()
 	select {
@@ -293,7 +323,7 @@ func (ch *clickhouse) acquire(ctx context.Context) (conn *connect, err error) {
 				return nil, err
 			}
 		}
-		conn.released = false
+		conn.setReleased(false)
 		return conn, nil
 	default:
 	}
@@ -326,7 +356,7 @@ func (ch *clickhouse) closeIdleExpired() {
 	for {
 		select {
 		case conn := <-ch.idle:
-			if conn.connectedAt.Before(cutoff) {
+			if conn.connectedAtTime().Before(cutoff) {
 				conn.close()
 			} else {
 				select {
@@ -342,24 +372,23 @@ func (ch *clickhouse) closeIdleExpired() {
 	}
 }
 
-func (ch *clickhouse) release(conn *connect, err error) {
-	if conn.released {
+func (ch *clickhouse) release(conn nativeTransport, err error) {
+	if conn.isReleased() {
 		return
 	}
-	conn.released = true
-	conn.debugf("[released] connection [%d]", conn.id)
+	conn.setReleased(true)
+	conn.debugf("[released] connection [%d]", conn.connID())
 
 	select {
 	case <-ch.open:
 	default:
 	}
-	if err != nil || time.Since(conn.connectedAt) >= ch.opt.ConnMaxLifetime {
+	if err != nil || time.Since(conn.connectedAtTime()) >= ch.opt.ConnMaxLifetime {
 		conn.close()
 		return
 	}
 	if ch.opt.FreeBufOnConnRelease {
-		conn.buffer = new(chproto.Buffer)
-		conn.compressor.Data = nil
+		conn.freeBuffer()
 	}
 	select {
 	case ch.idle <- conn:
@@ -374,7 +403,12 @@ func (ch *clickhouse) Close() error {
 		case c := <-ch.idle:
 			c.close()
 		default:
-			ch.exit <- struct{}{}
+			// In rare cases, close may be called multiple times, don't block
+			//TODO: add proper close flag to indicate this pool is unusable after Close
+			select {
+			case ch.exit <- struct{}{}:
+			default:
+			}
 			return nil
 		}
 	}

clickhouse_options.go

@@ -124,7 +124,7 @@ func ParseDSN(dsn string) (*Options, error) {
 
 type Dial func(ctx context.Context, addr string, opt *Options) (DialResult, error)
 type DialResult struct {
-	conn *connect
+	conn nativeTransport
 }
 
 type HTTPProxy func(*http.Request) (*url.URL, error)

clickhouse_std.go

@@ -34,7 +34,7 @@ import (
 	"syscall"
 
 	"github.com/ClickHouse/clickhouse-go/v2/lib/column"
-	ldriver "github.com/ClickHouse/clickhouse-go/v2/lib/driver"
+	chdriver "github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 )
 
 var globalConnID int64
@@ -196,10 +196,10 @@ func OpenDB(opt *Options) *sql.DB {
 type stdConnect interface {
 	isBad() bool
 	close() error
-	query(ctx context.Context, release func(*connect, error), query string, args ...any) (*rows, error)
+	query(ctx context.Context, release nativeTransportRelease, query string, args ...any) (*rows, error)
 	exec(ctx context.Context, query string, args ...any) error
 	ping(ctx context.Context) (err error)
-	prepareBatch(ctx context.Context, query string, options ldriver.PrepareBatchOptions, release func(*connect, error), acquire func(context.Context) (*connect, error)) (ldriver.Batch, error)
+	prepareBatch(ctx context.Context, release nativeTransportRelease, acquire nativeTransportAcquire, query string, options chdriver.PrepareBatchOptions) (chdriver.Batch, error)
 	asyncInsert(ctx context.Context, query string, wait bool, args ...any) error
 }
 
@@ -333,7 +333,7 @@ func (std *stdDriver) QueryContext(ctx context.Context, query string, args []dri
 		return nil, driver.ErrBadConn
 	}
 
-	r, err := std.conn.query(ctx, func(*connect, error) {}, query, rebind(args)...)
+	r, err := std.conn.query(ctx, func(nativeTransport, error) {}, query, rebind(args)...)
 	if isConnBrokenError(err) {
 		std.debugf("QueryContext got a fatal error, resetting connection: %v\n", err)
 		return nil, driver.ErrBadConn
@@ -358,7 +358,7 @@ func (std *stdDriver) PrepareContext(ctx context.Context, query string) (driver.
 		return nil, driver.ErrBadConn
 	}
 
-	batch, err := std.conn.prepareBatch(ctx, query, ldriver.PrepareBatchOptions{}, func(*connect, error) {}, func(context.Context) (*connect, error) { return nil, nil })
+	batch, err := std.conn.prepareBatch(ctx, func(nativeTransport, error) {}, func(context.Context) (nativeTransport, error) { return nil, nil }, query, chdriver.PrepareBatchOptions{})
 	if err != nil {
 		if isConnBrokenError(err) {
 			std.debugf("PrepareContext got a fatal error, resetting connection: %v\n", err)
@@ -387,7 +387,7 @@ func (std *stdDriver) Close() error {
 }
 
 type stdBatch struct {
-	batch  ldriver.Batch
+	batch  chdriver.Batch
 	debugf func(format string, v ...any)
 }
 

conn.go

@@ -96,7 +96,7 @@ func dial(ctx context.Context, addr string, num int, opt *Options) (*connect, er
 			id:                   num,
 			opt:                  opt,
 			conn:                 conn,
-			debugf:               debugf,
+			debugfFunc:           debugf,
 			buffer:               new(chproto.Buffer),
 			reader:               chproto.NewReader(conn),
 			revision:             ClientTCPProtocolVersion,
@@ -144,7 +144,7 @@ type connect struct {
 	id                   int
 	opt                  *Options
 	conn                 net.Conn
-	debugf               func(format string, v ...any)
+	debugfFunc           func(format string, v ...any)
 	server               ServerVersion
 	closed               bool
 	buffer               *chproto.Buffer
@@ -162,6 +162,22 @@ type connect struct {
 	closeMutex           sync.Mutex
 }
 
+func (c *connect) debugf(format string, v ...any) {
+	c.debugfFunc(format, v...)
+}
+
+func (c *connect) connID() int {
+	return c.id
+}
+
+func (c *connect) connectedAtTime() time.Time {
+	return c.connectedAt
+}
+
+func (c *connect) serverVersion() (*ServerVersion, error) {
+	return &c.server, nil
+}
+
 func (c *connect) settings(querySettings Settings) []proto.Setting {
 	settingToProtoSetting := func(k string, v any) proto.Setting {
 		isCustom := false
@@ -206,6 +222,14 @@ func (c *connect) isBad() bool {
 	return false
 }
 
+func (c *connect) isReleased() bool {
+	return c.released
+}
+
+func (c *connect) setReleased(released bool) {
+	c.released = released
+}
+
 func (c *connect) isClosed() bool {
 	c.closeMutex.Lock()
 	defer c.closeMutex.Unlock()
@@ -372,6 +396,11 @@ func (c *connect) readData(ctx context.Context, packet byte, compressible bool)
 	return &block, nil
 }
 
+func (c *connect) freeBuffer() {
+	c.buffer = new(chproto.Buffer)
+	c.compressor.Data = nil
+}
+
 func (c *connect) flush() error {
 	if len(c.buffer.Buf) == 0 {
 		// Nothing to flush.

conn_batch.go

@@ -35,7 +35,7 @@ import (
 var insertMatch = regexp.MustCompile(`(?i)(INSERT\s+INTO\s+[^( ]+(?:\s*\([^()]*(?:\([^()]*\)[^()]*)*\))?)(?:\s*VALUES)?`)
 var columnMatch = regexp.MustCompile(`INSERT INTO .+\s\((?P<Columns>.+)\)$`)
 
-func (c *connect) prepareBatch(ctx context.Context, query string, opts driver.PrepareBatchOptions, release func(*connect, error), acquire func(context.Context) (*connect, error)) (driver.Batch, error) {
+func (c *connect) prepareBatch(ctx context.Context, release nativeTransportRelease, acquire nativeTransportAcquire, query string, opts driver.PrepareBatchOptions) (driver.Batch, error) {
 	query, _, queryColumns, verr := extractNormalizedInsertQueryAndColumns(query)
 	if verr != nil {
 		return nil, verr
@@ -60,17 +60,29 @@ func (c *connect) prepareBatch(ctx context.Context, query string, opts driver.Pr
 	}
 	// resort batch to specified columns
 	if err = block.SortColumns(queryColumns); err != nil {
+		release(c, err)
 		return nil, err
 	}
 
+	connRelease := func(conn *connect, err error) {
+		release(conn, err)
+	}
+	connAcquire := func(ctx context.Context) (*connect, error) {
+		conn, err := acquire(ctx)
+		if err != nil {
+			return nil, err
+		}
+		return conn.(*connect), nil
+	}
+
 	b := &batch{
 		ctx:          ctx,
 		query:        query,
 		conn:         c,
 		block:        block,
 		released:     false,
-		connRelease:  release,
-		connAcquire:  acquire,
+		connRelease:  connRelease,
+		connAcquire:  connAcquire,
 		onProcess:    onProcess,
 		closeOnFlush: opts.CloseOnFlush,
 	}