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{}{
-		"includes": source.includes,
+	m := make(map[string]interface{})
+	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

@@ -22,6 +22,8 @@ type SearchRequest struct {
 	size       *uint64
 	explain    *bool
 	timeout    *time.Duration
+	source     Source
+	sort       Sort
 }
 
 // Search creates a new SearchRequest object, to be filled via method chaining.
@@ -37,7 +39,7 @@ 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
 }
 
@@ -60,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 {
@@ -73,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{} {
@@ -94,6 +119,9 @@ func (req *SearchRequest) Map() map[string]interface{} {
 	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
 	}
@@ -104,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

@@ -46,6 +46,10 @@ func TestSearchMaps(t *testing.T) {
 				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{}{
@@ -99,6 +103,14 @@ func TestSearchMaps(t *testing.T) {
 				"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"},
+				},
 			},
 		},
 	})