Support _source, sort and post_filter in search requests (#10)

* Add support for post filter in Search

Search requests can now accept a post filter applied after all hits were
returned from the database. For example:

    esquery.Search().
        Query(...).
        Aggs(...).
        PostFilter(esquery.Range(field).Gt(0)).
        Run(...)

* Add support for _source and sort in search requests

This commit adds the ability to use the "sort" and "_source" attributes
in search requests, via the new methods Sort, SourceIncludes and
SourceExcludes.

common.go

@@ -4,13 +4,19 @@ package esquery
 // queries. Currently, only the "includes" option is supported.
 type Source struct {
 	includes []string
+	excludes []string
 }
 
 // Map returns a map representation of the Source object.
 func (source Source) Map() map[string]interface{} {
-	return map[string]interface{}{
+	m := make(map[string]interface{})
-		"includes": source.includes,
+	if len(source.includes) > 0 {
+		m["includes"] = source.includes
 	}
+	if len(source.excludes) > 0 {
+		m["excludes"] = source.excludes
+	}
+	return m
 }
 
 // Sort represents a list of keys to sort by.

search.go

@@ -15,12 +15,15 @@ import (
 // Not all features of the search API are currently supported, but a request can
 // currently include a query, aggregations, and more.
 type SearchRequest struct {
-	query   Mappable
+	query      Mappable
-	aggs    []Aggregation
+	aggs       []Aggregation
-	from    *uint64
+	postFilter Mappable
-	size    *uint64
+	from       *uint64
-	explain *bool
+	size       *uint64
-	timeout *time.Duration
+	explain    *bool
+	timeout    *time.Duration
+	source     Source
+	sort       Sort
 }
 
 // Search creates a new SearchRequest object, to be filled via method chaining.
@@ -36,7 +39,13 @@ func (req *SearchRequest) Query(q Mappable) *SearchRequest {
 
 // Aggs sets one or more aggregations for the request.
 func (req *SearchRequest) Aggs(aggs ...Aggregation) *SearchRequest {
-	req.aggs = aggs
+	req.aggs = append(req.aggs, aggs...)
+	return req
+}
+
+// PostFilter sets a post_filter for the request.
+func (req *SearchRequest) PostFilter(filter Mappable) *SearchRequest {
+	req.postFilter = filter
 	return req
 }
 
@@ -53,6 +62,17 @@ func (req *SearchRequest) Size(size uint64) *SearchRequest {
 	return req
 }
 
+// Sort sets how the results should be sorted.
+func (req *SearchRequest) Sort(name string, order Order) *SearchRequest {
+	req.sort = append(req.sort, map[string]interface{}{
+		name: map[string]interface{}{
+			"order": order,
+		},
+	})
+
+	return req
+}
+
 // Explain sets whether the ElasticSearch API should return an explanation for
 // how each hit's score was calculated.
 func (req *SearchRequest) Explain(b bool) *SearchRequest {
@@ -66,6 +86,18 @@ func (req *SearchRequest) Timeout(dur time.Duration) *SearchRequest {
 	return req
 }
 
+// SourceIncludes sets the keys to return from the matching documents.
+func (req *SearchRequest) SourceIncludes(keys ...string) *SearchRequest {
+	req.source.includes = keys
+	return req
+}
+
+// SourceExcludes sets the keys to not return from the matching documents.
+func (req *SearchRequest) SourceExcludes(keys ...string) *SearchRequest {
+	req.source.excludes = keys
+	return req
+}
+
 // Map implements the Mappable interface. It converts the request to into a
 // nested map[string]interface{}, as expected by the go-elasticsearch library.
 func (req *SearchRequest) Map() map[string]interface{} {
@@ -81,9 +113,15 @@ func (req *SearchRequest) Map() map[string]interface{} {
 
 		m["aggs"] = aggs
 	}
+	if req.postFilter != nil {
+		m["post_filter"] = req.postFilter.Map()
+	}
 	if req.size != nil {
 		m["size"] = *req.size
 	}
+	if len(req.sort) > 0 {
+		m["sort"] = req.sort
+	}
 	if req.from != nil {
 		m["from"] = *req.from
 	}
@@ -94,6 +132,11 @@ func (req *SearchRequest) Map() map[string]interface{} {
 		m["timeout"] = fmt.Sprintf("%.0fs", req.timeout.Seconds())
 	}
 
+	source := req.source.Map()
+	if len(source) > 0 {
+		m["_source"] = source
+	}
+
 	return m
 }
 

search_test.go

@@ -42,9 +42,14 @@ func TestSearchMaps(t *testing.T) {
 					StringStats("tag_stats", "tags").
 						ShowDistribution(true),
 				).
+				PostFilter(Range("score").Gt(0)).
 				Size(30).
 				From(5).
 				Explain(true).
+				Sort("field_1", OrderDesc).
+				Sort("field_2", OrderAsc).
+				SourceIncludes("field_1", "field_2").
+				SourceExcludes("field_3").
 				Timeout(time.Duration(20000000000)),
 			map[string]interface{}{
 				"query": map[string]interface{}{
@@ -87,10 +92,25 @@ func TestSearchMaps(t *testing.T) {
 						},
 					},
 				},
+				"post_filter": map[string]interface{}{
+					"range": map[string]interface{}{
+						"score": map[string]interface{}{
+							"gt": 0,
+						},
+					},
+				},
 				"size":    30,
 				"from":    5,
 				"explain": true,
 				"timeout": "20s",
+				"sort": []map[string]interface{}{
+					{"field_1": map[string]interface{}{"order": "desc"}},
+					{"field_2": map[string]interface{}{"order": "asc"}},
+				},
+				"_source": map[string]interface{}{
+					"includes": []string{"field_1", "field_2"},
+					"excludes": []string{"field_3"},
+				},
 			},
 		},
 	})