From 48ad6f919c5171884faf69bfc88884cd5d73ca90 Mon Sep 17 00:00:00 2001
From: Ido Perlmuter <ido@ido50.net>
Date: Thu, 20 Feb 2020 16:45:08 +0200
Subject: [PATCH] Bugfix: Run() fails for queries, add MarshalJSON()

The `Run()` method on the `QueryRequest` type would fail, since it would
encode the _inner_ body of the query to JSON and not the complete, outer
body (i.e. including the "body: {}" portion).

The commit also adds `MarshalJSON()` methods to both `QueryRequest` and
`AggregationRequest`, allowing them to implement the `json.Marshaller`
interface, and providing easier debugging of the library. A test
skeleton for this is also added.
---
 aggregations.go |  4 ++++
 es_test.go      | 19 +++++++++++++++++++
 go.mod          |  2 ++
 go.sum          |  4 ++++
 queries.go      |  6 +++++-
 queries_test.go | 16 +++++++++++++++-
 6 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/aggregations.go b/aggregations.go
index 254f3fe..f1d7479 100644
--- a/aggregations.go
+++ b/aggregations.go
@@ -40,6 +40,10 @@ func (req *AggregationRequest) Map() map[string]interface{} {
 	}
 }
 
+func (req *AggregationRequest) MarshalJSON() ([]byte, error) {
+	return json.Marshal(req.Map())
+}
+
 func (req *AggregationRequest) Run(
 	api *elasticsearch.Client,
 	o ...func(*esapi.SearchRequest),
diff --git a/es_test.go b/es_test.go
index 346330a..85f9a4a 100644
--- a/es_test.go
+++ b/es_test.go
@@ -4,6 +4,8 @@ import (
 	"encoding/json"
 	"reflect"
 	"testing"
+
+	"github.com/jgroeneveld/trial/assert"
 )
 
 type mapTest struct {
@@ -38,3 +40,20 @@ func sameJSON(a, b map[string]interface{}) (aJSON, bJSON []byte, ok bool) {
 	ok = reflect.DeepEqual(aJSON, bJSON)
 	return aJSON, bJSON, ok
 }
+
+type jsonTest struct {
+	name    string
+	q       json.Marshaler
+	expJSON string
+	expErr  error
+}
+
+func runJSONTests(t *testing.T, tests []jsonTest) {
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			b, err := test.q.MarshalJSON()
+			assert.Equal(t, test.expErr, err)
+			assert.Equal(t, test.expJSON, string(b))
+		})
+	}
+}
diff --git a/go.mod b/go.mod
index 79ae130..d17acbf 100644
--- a/go.mod
+++ b/go.mod
@@ -5,4 +5,6 @@ go 1.13
 require (
 	github.com/elastic/go-elasticsearch/v7 v7.6.0
 	github.com/fatih/structs v1.1.0
+	github.com/jgroeneveld/schema v1.0.0 // indirect
+	github.com/jgroeneveld/trial v2.0.0+incompatible
 )
diff --git a/go.sum b/go.sum
index 680eeb3..4108027 100644
--- a/go.sum
+++ b/go.sum
@@ -6,3 +6,7 @@ github.com/elastic/go-elasticsearch/v8 v8.0.0-20200210103600-aff00e5adfde h1:Y9S
 github.com/elastic/go-elasticsearch/v8 v8.0.0-20200210103600-aff00e5adfde/go.mod h1:xe9a/L2aeOgFKKgrO3ibQTnMdpAeL0GC+5/HpGScSa4=
 github.com/fatih/structs v1.1.0 h1:Q7juDM0QtcnhCpeyLGQKyg4TOIghuNXrkL32pHAUMxo=
 github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga6PJ7M=
+github.com/jgroeneveld/schema v1.0.0 h1:J0E10CrOkiSEsw6dfb1IfrDJD14pf6QLVJ3tRPl/syI=
+github.com/jgroeneveld/schema v1.0.0/go.mod h1:M14lv7sNMtGvo3ops1MwslaSYgDYxrSmbzWIQ0Mr5rs=
+github.com/jgroeneveld/trial v2.0.0+incompatible h1:d59ctdgor+VqdZCAiUfVN8K13s0ALDioG5DWwZNtRuQ=
+github.com/jgroeneveld/trial v2.0.0+incompatible/go.mod h1:I6INLW96EN8WysNBXUFI3M4RIC8ePg9ntAc/Wy+U/+M=
diff --git a/queries.go b/queries.go
index 25a4fdd..e4edfa7 100644
--- a/queries.go
+++ b/queries.go
@@ -22,12 +22,16 @@ func (req *QueryRequest) Map() map[string]interface{} {
 	}
 }
 
+func (req *QueryRequest) MarshalJSON() ([]byte, error) {
+	return json.Marshal(req.Map())
+}
+
 func (req *QueryRequest) Run(
 	api *elasticsearch.Client,
 	o ...func(*esapi.SearchRequest),
 ) (res *esapi.Response, err error) {
 	var b bytes.Buffer
-	err = json.NewEncoder(&b).Encode(req.Query.Map())
+	err = json.NewEncoder(&b).Encode(req.Map())
 	if err != nil {
 		return nil, err
 	}
diff --git a/queries_test.go b/queries_test.go
index 64b4d60..029f042 100644
--- a/queries_test.go
+++ b/queries_test.go
@@ -4,7 +4,7 @@ import (
 	"testing"
 )
 
-func TestQueries(t *testing.T) {
+func TestQueryMaps(t *testing.T) {
 	runMapTests(t, []mapTest{
 		{
 			"a simple match_all query",
@@ -66,3 +66,17 @@ func TestQueries(t *testing.T) {
 		},
 	})
 }
+
+func TestQueryJSONs(t *testing.T) {
+	runJSONTests(t, []jsonTest{
+		{
+			"simple query",
+			Query(
+				Bool().
+					Must(Term("account_id", "bla")),
+			),
+			`{"query":{"bool":{"must":[{"term":{"account_id":{"value":"bla"}}}]}}}`,
+			nil,
+		},
+	})
+}