fetcher/i2es/elastic.go

package i2es

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"strings"

	log "github.com/sirupsen/logrus"
)

type topicid struct {
	es ESConf
}

// doPostRequest ...
func doRequest(uri, data, method string) ([]byte, error) {
	req, err := http.NewRequest(method, uri, bytes.NewBuffer([]byte(data)))
	req.Header.Add("Content-Type", "application/json")
	if err != nil {
		return []byte(""), err
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return []byte(""), err
	}

	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return []byte(""), err
	}

	return body, nil
}

// PutToIndex ...
func (es ESConf) PutToIndex(msg ESDoc) (ESDoc, error) {
	putURI := strings.Join([]string{es.Host, es.Index, es.Type, msg.MsgID}, "/")
	log.Print(putURI)

	// Assign topicID for top message
	var t topicid
	t.es = es
	t.getOrCreate(&msg)

	doc, err := json.Marshal(msg)
	if err != nil {
		return msg, err
	}

	req, err := http.NewRequest("PUT", putURI, bytes.NewBuffer(doc))
	req.Header.Add("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return msg, err
	}

	defer resp.Body.Close()

	return msg, err
}

// CheckID ...
func (es ESConf) CheckID(id string) (bool, error) {
	searchURI := strings.Join([]string{es.Host, es.Index, es.Type, "_search"}, "/")

	searchQ := strings.Join([]string{`{"query": {"match": {"_id": "`, id, `"}}}`}, "")

	body, err := doRequest(searchURI, searchQ, "POST")
	if err != nil {
		return false, err
	}

	var e ESRes
	err = json.Unmarshal(body, &e)
	if err != nil {
		return false, err
	}

	if e.Hits.Total > 0 {
		err = errors.New(strings.Join([]string{"Message ", id, " already in index"}, ""))
		return false, err
	}

	if e.TimedOut {
		err = errors.New("Request time out")
		return false, err
	}

	return false, nil
}