esquery/aggregations.go

package esquery

import (
	"bytes"
	"encoding/json"

	"github.com/elastic/go-elasticsearch/v7"
	"github.com/elastic/go-elasticsearch/v7/esapi"
)

// AggregationRequest represents a complete request of type "aggregations"
// (a.k.a "aggs") to ElasticSearch's search API. It simply wraps a map of named
// aggregations, which are values of a type that implements the Mappable
// interface.
type AggregationRequest struct {
	Aggs map[string]Mappable
}

// Aggregation is an interface that each aggregation type must implement. It
// is simply an extension of the Mappable interface to include a Named function,
// which returns the name of the aggregation.
type Aggregation interface {
	Mappable
	Name() string
}

// Aggregate generates a search request of type "aggs", represented by a
// *AggregationRequest object. It receives a variadic amount of values that
// implement the Aggregation interface, whether provided internally by the
// library or custom aggregations provided by consuming code.
func Aggregate(aggs ...Aggregation) *AggregationRequest {
	req := &AggregationRequest{
		Aggs: make(map[string]Mappable),
	}
	for _, agg := range aggs {
		req.Aggs[agg.Name()] = agg
	}

	return req
}

// Map implements the Mappable interface. It converts the "aggs" request into a
// (potentially nested) map[string]interface{}.
func (req *AggregationRequest) Map() map[string]interface{} {
	m := make(map[string]interface{})

	for name, agg := range req.Aggs {
		m[name] = agg.Map()
	}

	return map[string]interface{}{
		"aggs": m,
	}
}

// MarshalJSON implements the json.Marshaler interface, it simply encodes the
// map representation of the request (provided by the Map method) as JSON.
func (req *AggregationRequest) MarshalJSON() ([]byte, error) {
	return json.Marshal(req.Map())
}

// Run executes the request using the provided ElasticSearch client. Zero or
// more search options can be provided as well. It returns the standard Response
// type of the official Go client.
func (req *AggregationRequest) Run(
	api *elasticsearch.Client,
	o ...func(*esapi.SearchRequest),
) (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.SearchRequest){api.Search.WithBody(&b)}, o...)

	return api.Search(opts...)
}

// RunSearch is the same as the Run method, except that it accepts a value of
// type esapi.Search (usually this is the Search field of an elasticsearch.Client
// object). Since the ElasticSearch client does not provide an interface type
// for its API (which would allow implementation of mock clients), this provides
// a workaround. The Search function in the ES client is actually a field of a
// function type.
func (req *AggregationRequest) RunSearch(
	search esapi.Search,
	o ...func(*esapi.SearchRequest),
) (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.SearchRequest){search.WithBody(&b)}, o...)

	return search(opts...)
}