diff --git a/conn.go b/conn.go index 6c59406..363f4e4 100644 --- a/conn.go +++ b/conn.go @@ -1,6 +1,7 @@ package clickhouse import ( + "compress/gzip" "context" "database/sql" "database/sql/driver" @@ -13,7 +14,6 @@ import ( "net/http" "net/url" "os" - "strings" "sync/atomic" "time" @@ -302,7 +302,21 @@ func (c *conn) buildRequest(ctx context.Context, query string, params []driver.V method = http.MethodPost } c.log("query: ", query) - req, err := http.NewRequest(method, c.url.String(), strings.NewReader(query)) + + bodyReader, bodyWriter := io.Pipe() + go func() { + if c.useGzipCompression { + gz := gzip.NewWriter(bodyWriter) + _, err := gz.Write([]byte(query)) + gz.Close() + bodyWriter.CloseWithError(err) + } else { + bodyWriter.Write([]byte(query)) + bodyWriter.Close() + } + }() + + req, err := http.NewRequest(method, c.url.String(), bodyReader) if err != nil { return nil, err } @@ -335,6 +349,10 @@ func (c *conn) buildRequest(ctx context.Context, query string, params []driver.V } req.URL.RawQuery = reqQuery.Encode() + if c.useGzipCompression { + req.Header.Set("Content-Encoding", "gzip") + } + return req, nil } diff --git a/conn_test.go b/conn_test.go index 21edb4d..54a637e 100644 --- a/conn_test.go +++ b/conn_test.go @@ -1,6 +1,7 @@ package clickhouse import ( + "compress/gzip" "context" "database/sql" "database/sql/driver" @@ -380,6 +381,24 @@ func (s *connSuite) TestBuildRequestParamsInterpolation() { } } +func (s *connSuite) TestRequestBodyGzipCompression() { + query := `INSERT INTO test (str) VALUES ("Question?")` + cn := newConn(NewConfig()) + cn.useGzipCompression = true + req, err := cn.buildRequest(context.Background(), query, make([]driver.Value, 0), false) + if s.NoError(err) { + s.Contains(req.Header, "Content-Encoding") + gz, err := gzip.NewReader(req.Body) + if s.NoError(err) { + defer gz.Close() + body, e := ioutil.ReadAll(gz) + if s.NoError(e) { + s.Equal(query, string(body)) + } + } + } +} + func TestConn(t *testing.T) { suite.Run(t, new(connSuite)) }