From 7fa767fc28e316a6ee9e0794d868a3b5473e6e9e Mon Sep 17 00:00:00 2001
From: Ido Perlmuter <ido@ido50.net>
Date: Mon, 6 Apr 2020 09:24:44 +0000
Subject: [PATCH] Add support for Count requests (#8)

This commit adds initial support for Count requests, which are simple
search requests asking to get the number of matches for a query.

The functionality is provided by the Count() function, which accepts a
query object (implementing the Mappable interface), and can be executed
just like a search request.

    res, err := esquery.
        Count(esquery.Match("user", "someone")).
        Run(es)
---
 count.go      | 62 +++++++++++++++++++++++++++++++++++++++++++++++++++
 count_test.go | 17 ++++++++++++++
 2 files changed, 79 insertions(+)
 create mode 100644 count.go
 create mode 100644 count_test.go

diff --git a/count.go b/count.go
new file mode 100644
index 0000000..2e01b5c
--- /dev/null
+++ b/count.go
@@ -0,0 +1,62 @@
+package esquery
+
+import (
+	"bytes"
+	"encoding/json"
+
+	"github.com/elastic/go-elasticsearch/v7"
+	"github.com/elastic/go-elasticsearch/v7/esapi"
+)
+
+// CountRequest represents a request to get the number of matches for a search
+// query, as described in:
+// https://www.elastic.co/guide/en/elasticsearch/reference/current/search-count.html
+type CountRequest struct {
+	query Mappable
+}
+
+// Count creates a new count request with the provided query.
+func Count(q Mappable) *CountRequest {
+	return &CountRequest{
+		query: q,
+	}
+}
+
+// Map returns a map representation of the request, thus implementing the
+// Mappable interface.
+func (req *CountRequest) Map() map[string]interface{} {
+	return map[string]interface{}{
+		"query": req.query.Map(),
+	}
+}
+
+// Run executes the request using the provided ElasticCount client. Zero or
+// more search options can be provided as well. It returns the standard Response
+// type of the official Go client.
+func (req *CountRequest) Run(
+	api *elasticsearch.Client,
+	o ...func(*esapi.CountRequest),
+) (res *esapi.Response, err error) {
+	return req.RunCount(api.Count, o...)
+}
+
+// RunCount is the same as the Run method, except that it accepts a value of
+// type esapi.Count (usually this is the Count field of an elasticsearch.Client
+// object). Since the ElasticCount client does not provide an interface type
+// for its API (which would allow implementation of mock clients), this provides
+// a workaround. The Count function in the ES client is actually a field of a
+// function type.
+func (req *CountRequest) RunCount(
+	count esapi.Count,
+	o ...func(*esapi.CountRequest),
+) (res *esapi.Response, err error) {
+	var b bytes.Buffer
+	err = json.NewEncoder(&b).Encode(req.Map())
+	if err != nil {
+		return nil, err
+	}
+
+	opts := append([]func(*esapi.CountRequest){count.WithBody(&b)}, o...)
+
+	return count(opts...)
+}
diff --git a/count_test.go b/count_test.go
new file mode 100644
index 0000000..359e7d5
--- /dev/null
+++ b/count_test.go
@@ -0,0 +1,17 @@
+package esquery
+
+import "testing"
+
+func TestCount(t *testing.T) {
+	runMapTests(t, []mapTest{
+		{
+			"a simple count request",
+			Count(MatchAll()),
+			map[string]interface{}{
+				"query": map[string]interface{}{
+					"match_all": map[string]interface{}{},
+				},
+			},
+		},
+	})
+}