Add gzip/compression support

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
+				}
 			}
 		}
 	}

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

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 {

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)
+}