goes/goes.go

426 lines
9.9 KiB
Go

// Copyright 2013 Belogik. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package goes provides an API to access Elasticsearch.
package goes
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io/ioutil"
"net/http"
"net/url"
"reflect"
"strconv"
"strings"
)
const (
BULK_COMMAND_INDEX = "index"
BULK_COMMAND_DELETE = "delete"
)
func (err *SearchError) Error() string {
return fmt.Sprintf("[%d] %s", err.StatusCode, err.Msg)
}
// NewConnection initiates a new Connection to an elasticsearch server
//
// This function is pretty useless for now but might be useful in a near future
// if wee need more features like connection pooling or load balancing.
func NewConnection(host string, port string) *Connection {
return &Connection{host, port, http.DefaultClient}
}
func (c *Connection) WithClient(cl *http.Client) *Connection {
c.Client = cl
return c
}
// CreateIndex creates a new index represented by a name and a mapping
func (c *Connection) CreateIndex(name string, mapping map[string]interface{}) (Response, error) {
r := Request{
Conn: c,
Query: mapping,
IndexList: []string{name},
method: "PUT",
}
return r.Run()
}
// DeleteIndex deletes an index represented by a name
func (c *Connection) DeleteIndex(name string) (Response, error) {
r := Request{
Conn: c,
IndexList: []string{name},
method: "DELETE",
}
return r.Run()
}
// RefreshIndex refreshes an index represented by a name
func (c *Connection) RefreshIndex(name string) (Response, error) {
r := Request{
Conn: c,
IndexList: []string{name},
method: "POST",
api: "_refresh",
}
return r.Run()
}
// Stats fetches statistics (_stats) for the current elasticsearch server
func (c *Connection) Stats(indexList []string, extraArgs url.Values) (Response, error) {
r := Request{
Conn: c,
IndexList: indexList,
ExtraArgs: extraArgs,
method: "GET",
api: "_stats",
}
return r.Run()
}
// IndexStatus fetches the status (_status) for the indices defined in
// indexList. Use _all in indexList to get stats for all indices
func (c *Connection) IndexStatus(indexList []string) (Response, error) {
r := Request{
Conn: c,
IndexList: indexList,
method: "GET",
api: "_status",
}
return r.Run()
}
// Bulk adds multiple documents in bulk mode
func (c *Connection) BulkSend(documents []Document) (Response, error) {
// We do not generate a traditional JSON here (often a one liner)
// Elasticsearch expects one line of JSON per line (EOL = \n)
// plus an extra \n at the very end of the document
//
// More informations about the Bulk JSON format for Elasticsearch:
//
// - http://www.elasticsearch.org/guide/reference/api/bulk.html
//
// This is quite annoying for us as we can not use the simple JSON
// Marshaler available in Run().
//
// We have to generate this special JSON by ourselves which leads to
// the code below.
//
// I know it is unreadable I must find an elegant way to fix this.
// len(documents) * 2 : action + optional_sources
// + 1 : room for the trailing \n
bulkData := make([][]byte, len(documents)*2+1)
i := 0
for _, doc := range documents {
action, err := json.Marshal(map[string]interface{}{
doc.BulkCommand: map[string]interface{}{
"_index": doc.Index,
"_type": doc.Type,
"_id": doc.Id,
},
})
if err != nil {
return Response{}, err
}
bulkData[i] = action
i++
if doc.Fields != nil {
if docFields, ok := doc.Fields.(map[string]interface{}); ok {
if len(docFields) == 0 {
continue
}
} else {
typeOfFields := reflect.TypeOf(doc.Fields)
if typeOfFields.Kind() == reflect.Ptr {
typeOfFields = typeOfFields.Elem()
}
if typeOfFields.Kind() != reflect.Struct {
return Response{}, fmt.Errorf("Document fields not in struct or map[string]interface{} format")
}
if typeOfFields.NumField() == 0 {
continue
}
}
sources, err := json.Marshal(doc.Fields)
if err != nil {
return Response{}, err
}
bulkData[i] = sources
i++
}
}
// forces an extra trailing \n absolutely necessary for elasticsearch
bulkData[len(bulkData)-1] = []byte(nil)
r := Request{
Conn: c,
method: "POST",
api: "_bulk",
bulkData: bytes.Join(bulkData, []byte("\n")),
}
return r.Run()
}
// Search executes a search query against an index
func (c *Connection) Search(query map[string]interface{}, indexList []string, typeList []string, extraArgs url.Values) (Response, error) {
r := Request{
Conn: c,
Query: query,
IndexList: indexList,
TypeList: typeList,
method: "POST",
api: "_search",
ExtraArgs: extraArgs,
}
return r.Run()
}
//Query runs a query against an index using the provided http method.
//This method can be used to execute a delete by query, just pass in "DELETE"
//for the HTTP method.
func (c *Connection) Query(query map[string]interface{}, indexList []string, typeList []string, httpMethod string, extraArgs url.Values) (Response, error) {
r := Request{
Conn: c,
Query: query,
IndexList: indexList,
TypeList: typeList,
method: httpMethod,
api: "_query",
ExtraArgs: extraArgs,
}
return r.Run()
}
// Scan starts scroll over an index
func (c *Connection) Scan(query map[string]interface{}, indexList []string, typeList []string, timeout string, size int) (Response, error) {
v := url.Values{}
v.Add("search_type", "scan")
v.Add("scroll", timeout)
v.Add("size", strconv.Itoa(size))
r := Request{
Conn: c,
Query: query,
IndexList: indexList,
TypeList: typeList,
method: "POST",
api: "_search",
ExtraArgs: v,
}
return r.Run()
}
// Scroll fetches data by scroll id
func (c *Connection) Scroll(scrollId string, timeout string) (Response, error) {
v := url.Values{}
v.Add("scroll", timeout)
r := Request{
Conn: c,
method: "POST",
api: "_search/scroll",
ExtraArgs: v,
Body: []byte(scrollId),
}
return r.Run()
}
// Get a typed document by its id
func (c *Connection) Get(index string, documentType string, id string, extraArgs url.Values) (Response, error) {
r := Request{
Conn: c,
IndexList: []string{index},
method: "GET",
api: documentType + "/" + id,
ExtraArgs: extraArgs,
}
return r.Run()
}
// Index indexes a Document
// The extraArgs is a list of url.Values that you can send to elasticsearch as
// URL arguments, for example, to control routing, ttl, version, op_type, etc.
func (c *Connection) Index(d Document, extraArgs url.Values) (Response, error) {
r := Request{
Conn: c,
Query: d.Fields,
IndexList: []string{d.Index.(string)},
TypeList: []string{d.Type},
ExtraArgs: extraArgs,
method: "POST",
}
if d.Id != nil {
r.method = "PUT"
r.id = d.Id.(string)
}
return r.Run()
}
// Delete deletes a Document d
// The extraArgs is a list of url.Values that you can send to elasticsearch as
// URL arguments, for example, to control routing.
func (c *Connection) Delete(d Document, extraArgs url.Values) (Response, error) {
r := Request{
Conn: c,
IndexList: []string{d.Index.(string)},
TypeList: []string{d.Type},
ExtraArgs: extraArgs,
method: "DELETE",
id: d.Id.(string),
}
return r.Run()
}
// Run executes an elasticsearch Request. It converts data to Json, sends the
// request and return the Response obtained
func (req *Request) Run() (Response, error) {
postData := []byte{}
// XXX : refactor this
if len(req.Body) > 0 {
postData = req.Body
} else if req.api == "_bulk" {
postData = req.bulkData
} else {
b, err := json.Marshal(req.Query)
if err != nil {
return Response{}, err
}
postData = b
}
reader := bytes.NewReader(postData)
newReq, err := http.NewRequest(req.method, req.Url(), reader)
if err != nil {
return Response{}, err
}
if req.method == "POST" || req.method == "PUT" {
newReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
}
resp, err := req.Conn.Client.Do(newReq)
if err != nil {
return Response{}, err
}
defer resp.Body.Close()
body, err := ioutil.ReadAll(resp.Body)
if err != nil {
return Response{}, err
}
if resp.StatusCode > 201 && resp.StatusCode < 400 {
return Response{}, errors.New(string(body))
}
esResp := new(Response)
err = json.Unmarshal(body, &esResp)
if err != nil {
return Response{}, err
}
if req.api == "_bulk" && esResp.Errors {
for _, item := range esResp.Items {
for _, i := range item {
if i.Error != "" {
return Response{}, &SearchError{i.Error, i.Status}
}
}
}
return Response{}, &SearchError{Msg: "Unknown error while bulk indexing"}
}
if esResp.Error != "" {
return Response{}, &SearchError{esResp.Error, esResp.Status}
}
return *esResp, nil
}
// Url builds a Request for a URL
func (r *Request) Url() string {
path := "/" + strings.Join(r.IndexList, ",")
if len(r.TypeList) > 0 {
path += "/" + strings.Join(r.TypeList, ",")
}
// XXX : for indexing documents using the normal (non bulk) API
if len(r.api) == 0 && len(r.id) > 0 {
path += "/" + r.id
}
path += "/" + r.api
u := url.URL{
Scheme: "http",
Host: fmt.Sprintf("%s:%s", r.Conn.Host, r.Conn.Port),
Path: path,
RawQuery: r.ExtraArgs.Encode(),
}
return u.String()
}
// Buckets returns list of buckets in aggregation
func (a Aggregation) Buckets() []Bucket {
result := []Bucket{}
if buckets, ok := a["buckets"]; ok {
for _, bucket := range buckets.([]interface{}) {
result = append(result, bucket.(map[string]interface{}))
}
}
return result
}
// Key returns key for aggregation bucket
func (b Bucket) Key() interface{} {
return b["key"]
}
// DocCount returns count of documents in this bucket
func (b Bucket) DocCount() uint64 {
return uint64(b["doc_count"].(float64))
}
// Aggregation returns aggregation by name from bucket
func (b Bucket) Aggregation(name string) Aggregation {
if agg, ok := b[name]; ok {
return agg.(map[string]interface{})
} else {
return Aggregation{}
}
}