Merge pull request #656 from ClickHouse/compress_enhance

Clean up compression + add zstd

Author: gingerwizard

README.md

@@ -1,29 +1,29 @@
 # ClickHouse [![run-tests](https://github.com/ClickHouse/clickhouse-go/actions/workflows/run-tests.yml/badge.svg?branch=v2)](https://github.com/ClickHouse/clickhouse-go/actions/workflows/run-tests.yml) [![Go Reference](https://pkg.go.dev/badge/github.com/ClickHouse/clickhouse-go/v2.svg)](https://pkg.go.dev/github.com/ClickHouse/clickhouse-go/v2)
 
-Golang SQL database driver for [ClickHouse](https://clickhouse.com/).
+Golang SQL database client for [ClickHouse](https://clickhouse.com/).
 
 ## Versions
 
-There are two version of this driver, v1 and v2, available as separate branches. 
+There are two version of this client, v1 and v2, available as separate branches. 
 
 **v1 is now in a state of a maintenance - we will only accept PRs for bug and security fixes.**
 
 Users should use v2 which is production ready and [significantly faster than v1](#benchmark).
 
 ## Supported ClickHouse Versions
 
-The driver is tested against the currently [supported versions](https://github.com/ClickHouse/ClickHouse/blob/master/SECURITY.md) of ClickHouse
+The client is tested against the currently [supported versions](https://github.com/ClickHouse/ClickHouse/blob/master/SECURITY.md) of ClickHouse
 
 ## Supported Golang Versions
 
-| Driver Version | Golang Versions |
+| Client Version | Golang Versions |
 |----------------|-----------------|
 | => 2.0 <= 2.2  | 1.17, 1.18      |
 | >= 2.3         | 1.18            |
 
 ## Key features
 
-* Uses ClickHouse native format for optimal performance. Utilises low level [ch-go](https://github.com/ClickHouse/ch-go) driver for encoding/decoding and compression (versions >= 2.3.0).
+* Uses ClickHouse native format for optimal performance. Utilises low level [ch-go](https://github.com/ClickHouse/ch-go) client for encoding/decoding and compression (versions >= 2.3.0).
 * Supports native ClickHouse TCP client-server protocol
 * Compatibility with [`database/sql`](#std-databasesql-interface) ([slower](#benchmark) than [native interface](#native-interface)!)
 * [`database/sql`](#std-databasesql-interface) supports http protocol for transport. (Experimental)
@@ -34,7 +34,7 @@ The driver is tested against the currently [supported versions](https://github.c
 * [Bulk write support](examples/native/batch/main.go) (for `database/sql` [use](examples/std/batch/main.go) `begin->prepare->(in loop exec)->commit`)
 * [AsyncInsert](benchmark/v2/write-async/main.go)
 * Named and numeric placeholders support
-* LZ4 compression support
+* LZ4/ZSTD compression support
 * External data
 
 Support for the ClickHouse protocol advanced features using `Context`:
@@ -133,13 +133,15 @@ conn := clickhouse.OpenDB(&clickhouse.Options{
 
 ## Compression
 
-Compression is supported over http and native. This is performed at a block level and is only used for inserts.
+ZSTD/LZ4 compression is supported over native and http. This is performed at a block level and is only used for inserts.
+
+If using `Open` via the std interface and specifying a DSN, compression can be enabled via the `compress` flag. Currently, this is a boolean flag which enables `LZ4` compression.
 
 Other compression methods will be added in future PRs.
 
 ## TLS/SSL
 
-At a low level all driver connect methods (DSN/OpenDB/Open) will use the [Go tls package](https://pkg.go.dev/crypto/tls) to establish a secure connection. The driver knows to use TLS if the Options struct contains a non-nil tls.Config pointer.
+At a low level all client connect methods (DSN/OpenDB/Open) will use the [Go tls package](https://pkg.go.dev/crypto/tls) to establish a secure connection. The client knows to use TLS if the Options struct contains a non-nil tls.Config pointer.
 
 Setting secure in the DSN creates a minimal tls.Config struct with only the InsecureSkipVerify field set (either true or false).  It is equivalent to this code:
 
@@ -220,13 +222,13 @@ go get -u github.com/ClickHouse/clickhouse-go/v2
 
 ## ClickHouse alternatives
 
-Versions of this driver >=2.3.x utilise [ch-go](https://github.com/ClickHouse/ch-go) for their low level encoding/decoding. This low lever driver provides a high performance columnar interface and should be used in performance critical use cases. This driver provides more familar row orientated and `database/sql` semantics at the cost of some performance.
+Versions of this client >=2.3.x utilise [ch-go](https://github.com/ClickHouse/ch-go) for their low level encoding/decoding. This low lever client provides a high performance columnar interface and should be used in performance critical use cases. This client provides more familar row orientated and `database/sql` semantics at the cost of some performance.
 
-Both drivers are supported by ClickHouse.
+Both clients are supported by ClickHouse.
 
 ## Third-party alternatives
 
-* Database drivers:
+* Database client/clients:
 	* [mailru/go-clickhouse](https://github.com/mailru/go-clickhouse) (uses the HTTP protocol)
 	* [uptrace/go-clickhouse](https://github.com/uptrace/go-clickhouse) (uses the native TCP protocol with `database/sql`-like API)
 	* Drivers with columnar interface:

clickhouse_options.go

@@ -21,16 +21,21 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"github.com/ClickHouse/ch-go/compress"
 	"net"
 	"net/url"
 	"strconv"
 	"strings"
 	"time"
-
-	"github.com/ClickHouse/clickhouse-go/v2/lib/compress"
 )
 
-var CompressionLZ4 compress.Method = compress.LZ4
+type CompressionMethod compress.Method
+
+const (
+	CompressionNone = CompressionMethod(compress.None)
+	CompressionLZ4  = CompressionMethod(compress.LZ4)
+	CompressionZSTD = CompressionMethod(compress.ZSTD)
+)
 
 type Auth struct { // has_control_character
 	Database string
@@ -39,7 +44,7 @@ type Auth struct { // has_control_character
 }
 
 type Compression struct {
-	Method compress.Method
+	Method CompressionMethod
 }
 
 type ConnOpenStrategy uint8

conn.go

@@ -59,9 +59,9 @@ func dial(ctx context.Context, addr string, num int, opt *Options) (*connect, er
 			debugf = log.New(os.Stdout, fmt.Sprintf("[clickhouse][conn=%d][%s]", num, conn.RemoteAddr()), 0).Printf
 		}
 	}
-	var compression bool
+	compression := CompressionNone
 	if opt.Compression != nil {
-		compression = opt.Compression.Method == CompressionLZ4
+		compression = opt.Compression.Method
 	}
 	var (
 		connect = &connect{
@@ -95,7 +95,7 @@ type connect struct {
 	released    bool
 	revision    uint64
 	structMap   *structMap
-	compression bool
+	compression CompressionMethod
 	// lastUsedIn  time.Time
 	connectedAt time.Time
 	compressor  *compress.Writer
@@ -169,12 +169,10 @@ func (c *connect) sendData(block *proto.Block, name string) error {
 	if err := block.Encode(c.buffer, c.revision); err != nil {
 		return err
 	}
-	if c.compression {
-		// Performing compression.
-		//
-		// Note: only blocks are compressed.
+	if c.compression != CompressionNone {
+		// Performing compression. Note: only blocks are compressed.
 		data := c.buffer.Buf[start:]
-		if err := c.compressor.Compress(compress.LZ4, data); err != nil {
+		if err := c.compressor.Compress(compress.Method(c.compression), data); err != nil {
 			return errors.Wrap(err, "compress")
 		}
 		c.buffer.Buf = append(c.buffer.Buf[:start], c.compressor.Data...)
@@ -193,7 +191,7 @@ func (c *connect) readData(packet byte, compressible bool) (*proto.Block, error)
 	if _, err := c.reader.Str(); err != nil {
 		return nil, err
 	}
-	if compressible && c.compression {
+	if compressible && c.compression != CompressionNone {
 		c.reader.EnableCompression()
 		defer c.reader.DisableCompression()
 	}

conn_http.go

@@ -70,9 +70,9 @@ func dialHttp(ctx context.Context, addr string, num int, opt *Options) (*httpCon
 		query.Set(k, fmt.Sprint(v))
 	}
 	query.Set("default_format", "Native")
-	var compression bool
+	compression := CompressionNone
 	if opt.Compression != nil {
-		compression = opt.Compression.Method == CompressionLZ4
+		compression = opt.Compression.Method
 	}
 	u.RawQuery = query.Encode()
 
@@ -120,7 +120,7 @@ type httpConnect struct {
 	client      *http.Client
 	location    *time.Location
 	buffer      *chproto.Buffer
-	compression bool
+	compression CompressionMethod
 	compressor  *compress.Writer
 }
 
@@ -137,10 +137,10 @@ func (h *httpConnect) writeData(block *proto.Block) error {
 	if err := block.Encode(h.buffer, 0); err != nil {
 		return err
 	}
-	if h.compression {
+	if h.compression != CompressionNone {
 		// Performing compression. Supported and requires
 		data := h.buffer.Buf[start:]
-		if err := h.compressor.Compress(compress.LZ4, data); err != nil {
+		if err := h.compressor.Compress(compress.Method(h.compression), data); err != nil {
 			return errors.Wrap(err, "compress")
 		}
 		h.buffer.Buf = append(h.buffer.Buf[:start], h.compressor.Data...)
@@ -150,6 +150,10 @@ func (h *httpConnect) writeData(block *proto.Block) error {
 
 func (h *httpConnect) readData(reader *chproto.Reader) (*proto.Block, error) {
 	var block proto.Block
+	if h.compression != CompressionNone {
+		reader.EnableCompression()
+		defer reader.DisableCompression()
+	}
 	if err := block.Decode(reader, 0); err != nil {
 		return nil, err
 	}

conn_http_batch.go

@@ -170,7 +170,7 @@ func (b *httpBatch) Send() (err error) {
 
 	options := queryOptions(b.ctx)
 	// only compress blocks
-	if b.conn.compression {
+	if b.conn.compression != CompressionNone {
 		options.settings["decompress"] = "1"
 	}
 	options.settings["query"] = b.query

conn_http_query.go

@@ -34,8 +34,10 @@ func (h *httpConnect) query(ctx context.Context, release func(*connect, error),
 	if err != nil {
 		return nil, err
 	}
-
 	options := queryOptions(ctx)
+	if h.compression != CompressionNone {
+		options.settings["compress"] = "1"
+	}
 	res, err := h.sendQuery(ctx, strings.NewReader(query), &options, nil)
 	if err != nil {
 		return nil, err

conn_send_query.go

@@ -31,7 +31,7 @@ func (c *connect) sendQuery(body string, o *QueryOptions) error {
 		Body:           body,
 		Span:           o.span,
 		QuotaKey:       o.quotaKey,
-		Compression:    c.compression,
+		Compression:    c.compression != CompressionNone,
 		InitialAddress: c.conn.LocalAddr().String(),
 		Settings:       c.settings(o.settings),
 	}