Added field, which specify alternative unmarshal behaviour

This commit is contained in:
atyron 2016-10-05 09:30:45 +03:00
parent d5bd899311
commit 658df7cd94
3 changed files with 60 additions and 67 deletions

28
goes.go
View File

@ -32,7 +32,7 @@ func (err *SearchError) Error() string {
// This function is pretty useless for now but might be useful in a near future // This function is pretty useless for now but might be useful in a near future
// if wee need more features like connection pooling or load balancing. // if wee need more features like connection pooling or load balancing.
func NewConnection(host string, port string) *Connection { func NewConnection(host string, port string) *Connection {
return &Connection{host, port, http.DefaultClient} return &Connection{host, port, http.DefaultClient, false}
} }
func (c *Connection) WithClient(cl *http.Client) *Connection { func (c *Connection) WithClient(cl *http.Client) *Connection {
@ -342,6 +342,27 @@ func (c *Connection) Delete(d Document, extraArgs url.Values) (*Response, error)
return r.Run() return r.Run()
} }
// Enable use number behaviour
func (c *Connection) UseNumberEnable() {
c.UseNumber = true
}
func (req *Request) unmarshalBody(body []byte, esResp *Response) error {
if !req.Conn.UseNumber {
err := json.Unmarshal(body, esResp)
if err != nil {
return err
}
} else {
dec := json.NewDecoder(bytes.NewReader(body))
dec.UseNumber()
if err := dec.Decode(esResp); err != nil {
return err
}
}
return nil
}
// Run executes an elasticsearch Request. It converts data to Json, sends the // Run executes an elasticsearch Request. It converts data to Json, sends the
// request and returns the Response obtained // request and returns the Response obtained
func (req *Request) Run() (*Response, error) { func (req *Request) Run() (*Response, error) {
@ -353,9 +374,8 @@ func (req *Request) Run() (*Response, error) {
} }
if req.method != "HEAD" { if req.method != "HEAD" {
dec := json.NewDecoder(bytes.NewReader(body)) err = req.unmarshalBody(body, esResp)
dec.UseNumber() if err != nil {
if err := dec.Decode(&esResp); err != nil {
return esResp, err return esResp, err
} }
err = json.Unmarshal(body, &esResp.Raw) err = json.Unmarshal(body, &esResp.Raw)

View File

@ -13,7 +13,6 @@ import (
"time" "time"
. "github.com/go-check/check" . "github.com/go-check/check"
"encoding/json"
) )
var ( var (
@ -22,9 +21,7 @@ var (
) )
// Hook up gocheck into the gotest runner. // Hook up gocheck into the gotest runner.
func Test(t *testing.T) { func Test(t *testing.T) { TestingT(t) }
TestingT(t)
}
type GoesTestSuite struct{} type GoesTestSuite struct{}
@ -44,7 +41,7 @@ func (s *GoesTestSuite) SetUpTest(c *C) {
func (s *GoesTestSuite) TestNewConnection(c *C) { func (s *GoesTestSuite) TestNewConnection(c *C) {
conn := NewConnection(ES_HOST, ES_PORT) conn := NewConnection(ES_HOST, ES_PORT)
c.Assert(conn, DeepEquals, &Connection{ES_HOST, ES_PORT, http.DefaultClient}) c.Assert(conn, DeepEquals, &Connection{ES_HOST, ES_PORT, http.DefaultClient, false})
} }
func (s *GoesTestSuite) TestWithClient(c *C) { func (s *GoesTestSuite) TestWithClient(c *C) {
@ -57,7 +54,7 @@ func (s *GoesTestSuite) TestWithClient(c *C) {
} }
conn := NewConnection(ES_HOST, ES_PORT).WithClient(cl) conn := NewConnection(ES_HOST, ES_PORT).WithClient(cl)
c.Assert(conn, DeepEquals, &Connection{ES_HOST, ES_PORT, cl}) c.Assert(conn, DeepEquals, &Connection{ES_HOST, ES_PORT, cl, false})
c.Assert(conn.Client.Transport.(*http.Transport).DisableCompression, Equals, true) c.Assert(conn.Client.Transport.(*http.Transport).DisableCompression, Equals, true)
c.Assert(conn.Client.Transport.(*http.Transport).ResponseHeaderTimeout, Equals, 1*time.Second) c.Assert(conn.Client.Transport.(*http.Transport).ResponseHeaderTimeout, Equals, 1*time.Second)
} }
@ -738,13 +735,6 @@ func (s *GoesTestSuite) TestSearch(c *C) {
} }
response, err := conn.Search(query, []string{indexName}, []string{docType}, url.Values{}) response, err := conn.Search(query, []string{indexName}, []string{docType}, url.Values{})
//Type of response.Hits.MaxScore is json.Number, and DeepEquals will return false,
// so we need to change it manualy
jsonNumberMaxScore, _ := response.Hits.MaxScore.(json.Number)
maxScore, _ := jsonNumberMaxScore.Float64()
response.Hits.MaxScore = maxScore
expectedHits := Hits{ expectedHits := Hits{
Total: 1, Total: 1,
MaxScore: 1.0, MaxScore: 1.0,
@ -835,25 +825,19 @@ func (s *GoesTestSuite) TestIndexStatus(c *C) {
expectedShards := Shard{Total: 2, Successful: 1, Failed: 0} expectedShards := Shard{Total: 2, Successful: 1, Failed: 0}
c.Assert(response.Shards, Equals, expectedShards) c.Assert(response.Shards, Equals, expectedShards)
jsonPrimarySizeInBytes, _ := response.Indices[indexName].Index["primary_size_in_bytes"].(json.Number) primarySizeInBytes := response.Indices[indexName].Index["primary_size_in_bytes"].(float64)
primarySizeInBytes, _ := jsonPrimarySizeInBytes.Float64() sizeInBytes := response.Indices[indexName].Index["size_in_bytes"].(float64)
jsonSizeInBytes, _ := response.Indices[indexName].Index["size_in_bytes"].(json.Number) refreshTotal := response.Indices[indexName].Refresh["total"].(float64)
sizeInBytes, _ := jsonSizeInBytes.Float64()
jsonRefreshTotal, _ := response.Indices[indexName].Refresh["total"].(json.Number)
refreshTotal, _ := jsonRefreshTotal.Float64()
c.Assert(primarySizeInBytes > 0, Equals, true) c.Assert(primarySizeInBytes > 0, Equals, true)
c.Assert(sizeInBytes > 0, Equals, true) c.Assert(sizeInBytes > 0, Equals, true)
c.Assert(refreshTotal > 0, Equals, true) c.Assert(refreshTotal > 0, Equals, true)
var mockNullValue json.Number
mockNullValue = "0"
expectedIndices := map[string]IndexStatus{ expectedIndices := map[string]IndexStatus{
indexName: { indexName: {
Index: map[string]interface{}{ Index: map[string]interface{}{
"primary_size_in_bytes": jsonPrimarySizeInBytes, "primary_size_in_bytes": primarySizeInBytes,
"size_in_bytes": jsonSizeInBytes, "size_in_bytes": sizeInBytes,
}, },
Translog: map[string]uint64{ Translog: map[string]uint64{
"operations": 0, "operations": 0,
@ -864,21 +848,21 @@ func (s *GoesTestSuite) TestIndexStatus(c *C) {
"deleted_docs": 0, "deleted_docs": 0,
}, },
Merges: map[string]interface{}{ Merges: map[string]interface{}{
"current": mockNullValue, "current": float64(0),
"current_docs": mockNullValue, "current_docs": float64(0),
"current_size_in_bytes": mockNullValue, "current_size_in_bytes": float64(0),
"total": mockNullValue, "total": float64(0),
"total_time_in_millis": mockNullValue, "total_time_in_millis": float64(0),
"total_docs": mockNullValue, "total_docs": float64(0),
"total_size_in_bytes": mockNullValue, "total_size_in_bytes": float64(0),
}, },
Refresh: map[string]interface{}{ Refresh: map[string]interface{}{
"total": jsonRefreshTotal, "total": refreshTotal,
"total_time_in_millis": mockNullValue, "total_time_in_millis": float64(0),
}, },
Flush: map[string]interface{}{ Flush: map[string]interface{}{
"total": mockNullValue, "total": float64(0),
"total_time_in_millis": mockNullValue, "total_time_in_millis": float64(0),
}, },
}, },
} }
@ -1082,30 +1066,18 @@ func (s *GoesTestSuite) TestAggregations(c *C) {
c.Assert(user.Buckets()[1].Key(), Equals, "foo") c.Assert(user.Buckets()[1].Key(), Equals, "foo")
barAge := user.Buckets()[0].Aggregation("age") barAge := user.Buckets()[0].Aggregation("age")
jCount, _ := barAge["count"].(json.Number) c.Assert(barAge["count"], Equals, 1.0)
countInt64, _ := jCount.Int64() c.Assert(barAge["sum"], Equals, 30.0)
c.Assert(countInt64, Equals, int64(1))
jSum, _ := barAge["sum"].(json.Number)
sumFloat64, _ := jSum.Float64()
c.Assert(sumFloat64, Equals, 30.0)
fooAge := user.Buckets()[1].Aggregation("age") fooAge := user.Buckets()[1].Aggregation("age")
jCount, _ = fooAge["count"].(json.Number) c.Assert(fooAge["count"], Equals, 1.0)
countInt64, _ = jCount.Int64() c.Assert(fooAge["sum"], Equals, 25.0)
c.Assert(countInt64, Equals, int64(1))
jSum, _ = fooAge["sum"].(json.Number)
sumFloat64, _ = jSum.Float64()
c.Assert(sumFloat64, Equals, 25.0)
age, ok := resp.Aggregations["age"] age, ok := resp.Aggregations["age"]
c.Assert(ok, Equals, true) c.Assert(ok, Equals, true)
jCount, _ = age["count"].(json.Number) c.Assert(age["count"], Equals, 2.0)
countInt64, _ = jCount.Int64() c.Assert(age["sum"], Equals, 25.0+30.0)
c.Assert(countInt64, Equals, int64(2))
jSum, _ = age["sum"].(json.Number)
sumFloat64, _ = jSum.Float64()
c.Assert(sumFloat64, Equals, 25.0 + 30.0)
} }
func (s *GoesTestSuite) TestPutMapping(c *C) { func (s *GoesTestSuite) TestPutMapping(c *C) {
@ -1238,9 +1210,7 @@ func (s *GoesTestSuite) TestUpdate(c *C) {
response, err = conn.Get(indexName, docType, docId, url.Values{}) response, err = conn.Get(indexName, docType, docId, url.Values{})
c.Assert(err, Equals, nil) c.Assert(err, Equals, nil)
jsonCounter, _ := response.Source["counter"].(json.Number) c.Assert(response.Source["counter"], Equals, float64(6))
counter, _ := jsonCounter.Float64()
c.Assert(counter, Equals, float64(6))
c.Assert(response.Source["user"], Equals, "foo") c.Assert(response.Source["user"], Equals, "foo")
c.Assert(response.Source["message"], Equals, "bar") c.Assert(response.Source["message"], Equals, "bar")

View File

@ -20,6 +20,9 @@ type Connection struct {
// Client is the http client used to make requests, allowing settings things // Client is the http client used to make requests, allowing settings things
// such as timeouts etc // such as timeouts etc
Client *http.Client Client *http.Client
// Specify to use numbers when unmarshal
UseNumber bool
} }
// Represents a Request to elasticsearch // Represents a Request to elasticsearch