Skip to content

Commit

Permalink
Added support for gzip compression of request body
Browse files Browse the repository at this point in the history
Usefull for massive data load to ClickHouse
  • Loading branch information
Fiery-Fenix committed Feb 17, 2021
1 parent 828678e commit b741c8f
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 2 deletions.
22 changes: 20 additions & 2 deletions conn.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package clickhouse

import (
"compress/gzip"
"context"
"database/sql"
"database/sql/driver"
Expand All @@ -13,7 +14,6 @@ import (
"net/http"
"net/url"
"os"
"strings"
"sync/atomic"
"time"

Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}

Expand Down
19 changes: 19 additions & 0 deletions conn_test.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package clickhouse

import (
"compress/gzip"
"context"
"database/sql"
"database/sql/driver"
Expand Down Expand Up @@ -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))
}

0 comments on commit b741c8f

Please sign in to comment.