From 21ab94d0cabcf90ab5aae2731625f9a3bcb63ddc Mon Sep 17 00:00:00 2001
From: Stephen Asbury <sasbury@sasbury.com>
Date: Mon, 26 Jan 2015 17:43:50 -0800
Subject: [PATCH] count: support for count requests

Added count method to connection
Added Count field to response
Added test for count method
---
 goes.go      | 19 +++++++++++++++++--
 goes_test.go | 48 +++++++++++++++++++++++++++++++++++++++++++++++-
 structs.go   |  1 +
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/goes.go b/goes.go
index a20de12..a85cd7f 100644
--- a/goes.go
+++ b/goes.go
@@ -210,6 +210,21 @@ func (c *Connection) Search(query map[string]interface{}, indexList []string, ty
 	return r.Run()
 }
 
+// Count executes a count query against an index, use the Count field in the response for the result
+func (c *Connection) Count(query map[string]interface{}, indexList []string, typeList []string, extraArgs url.Values) (Response, error) {
+	r := Request{
+		Conn:      c,
+		Query:     query,
+		IndexList: indexList,
+		TypeList:  typeList,
+		method:    "POST",
+		api:       "_count",
+		ExtraArgs: extraArgs,
+	}
+
+	return r.Run()
+}
+
 //Query runs a query against an index using the provided http method.
 //This method can be used to execute a delete by query, just pass in "DELETE"
 //for the HTTP method.
@@ -494,9 +509,9 @@ func (c *Connection) Update(d Document, query map[string]interface{}, extraArgs
 		api:       "_update",
 	}
 
-    if d.Id != nil {
+	if d.Id != nil {
 		r.id = d.Id.(string)
-    }
+	}
 
 	return r.Run()
 }
diff --git a/goes_test.go b/goes_test.go
index 49f3d5b..fdb6023 100644
--- a/goes_test.go
+++ b/goes_test.go
@@ -201,7 +201,7 @@ func (s *GoesTestSuite) TestOptimize(c *C) {
 	// we must wait for a bit otherwise ES crashes
 	time.Sleep(1 * time.Second)
 
-	response, err := conn.Optimize([]string{indexName}, url.Values{"max_num_segments" : []string{"1"}})
+	response, err := conn.Optimize([]string{indexName}, url.Values{"max_num_segments": []string{"1"}})
 	c.Assert(err, IsNil)
 
 	c.Assert(response.All.Indices[indexName].Primaries["docs"].Count, Equals, 0)
@@ -722,6 +722,52 @@ func (s *GoesTestSuite) TestSearch(c *C) {
 	c.Assert(response.Hits, DeepEquals, expectedHits)
 }
 
+func (s *GoesTestSuite) TestCount(c *C) {
+	indexName := "testcount"
+	docType := "tweet"
+	docId := "1234"
+	source := map[string]interface{}{
+		"user":    "foo",
+		"message": "bar",
+	}
+
+	conn := NewConnection(ES_HOST, ES_PORT)
+	conn.DeleteIndex(indexName)
+
+	_, err := conn.CreateIndex(indexName, map[string]interface{}{})
+	c.Assert(err, IsNil)
+	defer conn.DeleteIndex(indexName)
+
+	d := Document{
+		Index:  indexName,
+		Type:   docType,
+		Id:     docId,
+		Fields: source,
+	}
+
+	_, err = conn.Index(d, url.Values{})
+	c.Assert(err, IsNil)
+
+	_, err = conn.RefreshIndex(indexName)
+	c.Assert(err, IsNil)
+
+	// I can feel my eyes bleeding
+	query := map[string]interface{}{
+		"query": map[string]interface{}{
+			"bool": map[string]interface{}{
+				"must": []map[string]interface{}{
+					{
+						"match_all": map[string]interface{}{},
+					},
+				},
+			},
+		},
+	}
+	response, err := conn.Count(query, []string{indexName}, []string{docType}, url.Values{})
+
+	c.Assert(response.Count, Equals, 1)
+}
+
 func (s *GoesTestSuite) TestIndexStatus(c *C) {
 	indexName := "testindexstatus"
 	conn := NewConnection(ES_HOST, ES_PORT)
diff --git a/structs.go b/structs.go
index e397a9e..ea29437 100644
--- a/structs.go
+++ b/structs.go
@@ -70,6 +70,7 @@ type Response struct {
 	Type         string `json:"_type"`
 	Version      int    `json:"_version"`
 	Found        bool
+	Count        int
 
 	// Used by the _stats API
 	All All `json:"_all"`