From b3f495a5f074e74032aac3fab5a45496868fa39b Mon Sep 17 00:00:00 2001 From: Travis Jeffery Date: Mon, 14 Mar 2016 13:41:35 -0500 Subject: [PATCH 1/2] add gzip support --- lib/baserequest.go | 26 +++++++++++++++----------- lib/connection.go | 4 +++- lib/request.go | 38 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 56 insertions(+), 12 deletions(-) diff --git a/lib/baserequest.go b/lib/baserequest.go index f9a9dcbf..b77b2418 100644 --- a/lib/baserequest.go +++ b/lib/baserequest.go @@ -36,17 +36,21 @@ func (c *Conn) DoCommand(method string, url string, args map[string]interface{}, } if data != nil { - switch v := data.(type) { - case string: - req.SetBodyString(v) - case io.Reader: - req.SetBody(v) - case []byte: - req.SetBodyBytes(v) - default: - err = req.SetBodyJson(v) - if err != nil { - return body, err + if c.Gzip { + req.SetBodyGzip(data) + } else { + switch v := data.(type) { + case string: + req.SetBodyString(v) + case io.Reader: + req.SetBody(v) + case []byte: + req.SetBodyBytes(v) + default: + err = req.SetBodyJson(v) + if err != nil { + return body, err + } } } } diff --git a/lib/connection.go b/lib/connection.go index 895eb4fc..db70d1c7 100644 --- a/lib/connection.go +++ b/lib/connection.go @@ -14,13 +14,14 @@ package elastigo import ( "errors" "fmt" - hostpool "github.com/bitly/go-hostpool" "net/http" "net/url" "runtime" "strings" "sync" "time" + + hostpool "github.com/bitly/go-hostpool" ) const ( @@ -41,6 +42,7 @@ type Conn struct { Username string Password string Hosts []string + Gzip bool RequestTracer func(method, url, body string) hp hostpool.HostPool once sync.Once diff --git a/lib/request.go b/lib/request.go index 4f72955e..385faa10 100644 --- a/lib/request.go +++ b/lib/request.go @@ -13,6 +13,7 @@ package elastigo import ( "bytes" + "compress/gzip" "encoding/json" "fmt" "io" @@ -32,6 +33,43 @@ type Request struct { hostResponse hostpool.HostPoolResponse } +func (r *Request) SetBodyGzip(data interface{}) error { + buf := new(bytes.Buffer) + gw := gzip.NewWriter(buf) + + switch v := data.(type) { + case string: + if _, err := gw.Write([]byte(v)); err != nil { + return err + } + case []byte: + if _, err := gw.Write([]byte(v)); err != nil { + return err + } + case io.Reader: + if _, err := io.Copy(gw, v); err != nil { + return err + } + default: + b, err := json.Marshal(data) + if err != nil { + return err + } + if _, err := gw.Write(b); err != nil { + return err + } + } + + if err := gw.Close(); err != nil { + return err + } + r.SetBody(bytes.NewReader(buf.Bytes())) + r.ContentLength = int64(len(buf.Bytes())) + r.Header.Add("Accept-Charset", "utf-8") + r.Header.Set("Content-Encoding", "gzip") + return nil +} + func (r *Request) SetBodyJson(data interface{}) error { body, err := json.Marshal(data) if err != nil { From f0cfaf3bc672f99d3649980305e9e986b2de5e46 Mon Sep 17 00:00:00 2001 From: Travis Jeffery Date: Thu, 17 Mar 2016 10:36:42 -0500 Subject: [PATCH 2/2] Add tests --- lib/request_test.go | 61 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/lib/request_test.go b/lib/request_test.go index fa173d6e..b286f9ec 100644 --- a/lib/request_test.go +++ b/lib/request_test.go @@ -11,7 +11,12 @@ package elastigo import ( + "bytes" + "compress/gzip" "fmt" + "io/ioutil" + "net/http" + "strings" "testing" "github.com/bmizerany/assert" @@ -72,3 +77,59 @@ func TestQueryString(t *testing.T) { s, err = Escape(map[string]interface{}{"foo": []int{}}) assert.T(t, err != nil, fmt.Sprintf("Expected err to not be nil")) } + +func TestSetBodyGzip(t *testing.T) { + s := "foo" + + // test []byte + expB := []byte(s) + actB, err := gzipHelper(t, expB) + assert.T(t, err == nil, fmt.Sprintf("Expected err to be nil")) + assert.T(t, bytes.Compare(actB, expB) == 0, fmt.Sprintf("Expected: %s, got: %s", expB, actB)) + + // test string + expS := s + actS, err := gzipHelper(t, expS) + assert.T(t, err == nil, fmt.Sprintf("Expected err to be nil")) + assert.T(t, string(actS) == expS, fmt.Sprintf("Expected: %s, got: %s", expS, actS)) + + // test io.Reader + expR := strings.NewReader(s) + actR, err := gzipHelper(t, expR) + assert.T(t, err == nil, fmt.Sprintf("Expected err to be nil")) + assert.T(t, bytes.Compare([]byte(s), actR) == 0, fmt.Sprintf("Expected: %s, got: %s", s, actR)) + + // test other + expO := testStruct{Name: "Travis"} + actO, err := gzipHelper(t, expO) + assert.T(t, err == nil, fmt.Sprintf("Expected err to not be nil")) + assert.T(t, bytes.Compare([]byte(`{"name":"Travis"}`), actO) == 0, fmt.Sprintf("Expected: %s, got: %s", s, actO)) +} + +type testStruct struct { + Name string `json:"name"` +} + +func gzipHelper(t *testing.T, data interface{}) ([]byte, error) { + r, err := http.NewRequest("GET", "http://google.com", nil) + if err != nil { + return nil, err + } + + // test string + req := &Request{ + Request: r, + } + + err = req.SetBodyGzip(data) + if err != nil { + return nil, err + } + + gr, err := gzip.NewReader(req.Body) + if err != nil { + return nil, err + } + + return ioutil.ReadAll(gr) +}