fetcher/i2es/elastic.go

100 lines
1.8 KiB
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
}