2024-04-01 18:13:10 +08:00

192 lines
4.0 KiB
Go

package es7
import (
"bytes"
"context"
"encoding/json"
"fmt"
elastic "github.com/elastic/go-elasticsearch/v7"
"github.com/elastic/go-elasticsearch/v7/esapi"
"github.com/loveuer/nfflow/internal/model"
"github.com/loveuer/nfflow/internal/opt"
"github.com/loveuer/nfflow/internal/util"
"github.com/sirupsen/logrus"
)
type ES7 struct {
cli *elastic.Client
scroll string
cfg struct {
Endpoints []string
Username string
Password string
Size int
Query map[string]any
Source []string
}
}
func (e *ES7) init() error {
var (
err error
cfg = elastic.Config{
Addresses: e.cfg.Endpoints,
Username: e.cfg.Username,
Password: e.cfg.Password,
RetryOnStatus: []int{429},
}
info *esapi.Response
)
if e.cli, err = elastic.NewClient(cfg); err != nil {
return err
}
if info, err = e.cli.Info(e.cli.Info.WithContext(util.Timeout(5))); err != nil {
return err
}
if info.StatusCode != 200 {
return fmt.Errorf("status=%d msg=%s", info.StatusCode, info.String())
}
return nil
}
func (e *ES7) Start(ctx context.Context, task *model.Task, rowCh chan<- model.TaskRow, errCh chan<- error) error {
var (
err error
result *esapi.Response
ready = make(chan bool)
decoder *json.Decoder
hits = new(model.ESResponse)
)
if err = e.init(); err != nil {
logrus.Debugf("ES7.Start: init err=%v", err)
return err
}
qs := []func(*esapi.SearchRequest){
e.cli.Search.WithContext(util.TimeoutCtx(ctx, opt.ES7OperationTimeout)),
e.cli.Search.WithScroll(opt.ScrollTimeout),
e.cli.Search.WithSize(e.cfg.Size),
}
if e.cfg.Query != nil && len(e.cfg.Query) > 0 {
var bs []byte
if bs, err = json.Marshal(e.cfg.Query); err != nil {
logrus.Debugf("ES7.Start: marshal query err=%v", err)
return err
}
qs = append(qs, e.cli.Search.WithBody(bytes.NewReader(bs)))
}
go func() {
defer func() {
if e.scroll != "" {
var csr *esapi.Response
if csr, err = e.cli.ClearScroll(
e.cli.ClearScroll.WithContext(util.TimeoutCtx(ctx, 5)),
e.cli.ClearScroll.WithScrollID(e.scroll),
); err != nil {
logrus.Warnf("ES7.Start: clear scroll=%s err=%v", e.scroll, err)
} else {
if csr.StatusCode != 200 {
logrus.Warnf("ES7.Start: clear scroll=%s status=%d msg=%s", e.scroll, csr.StatusCode, csr.String())
}
}
}
close(rowCh)
close(errCh)
}()
ready <- true
if result, err = e.cli.Search(qs...); err != nil {
logrus.Debugf("ES7.Start: search err=%v", err)
errCh <- err
return
}
if err = util.CheckES7Response(result); err != nil {
logrus.Debugf("ES7.Start: search resp err=%v", err)
errCh <- err
return
}
decoder = json.NewDecoder(result.Body)
if err = decoder.Decode(hits); err != nil {
logrus.Debugf("ES7.Start: decode err=%v", err)
errCh <- err
return
}
if hits.TimedOut {
err = fmt.Errorf("timeout")
logrus.Debugf("ES7.Start: search timeout")
errCh <- err
return
}
e.scroll = hits.ScrollId
for idx := range hits.Hits.Hits {
rowCh <- hits.Hits.Hits[idx]
}
if len(hits.Hits.Hits) < e.cfg.Size {
return
}
for {
if result, err = e.cli.Scroll(
e.cli.Scroll.WithContext(util.TimeoutCtx(ctx, opt.ES7OperationTimeout)),
e.cli.Scroll.WithScrollID(e.scroll),
); err != nil {
logrus.Debugf("ES7.Start: search err=%v", err)
errCh <- err
return
}
if err = util.CheckES7Response(result); err != nil {
logrus.Debugf("ES7.Start: search resp err=%v", err)
errCh <- err
return
}
decoder = json.NewDecoder(result.Body)
hits = new(model.ESResponse)
if err = decoder.Decode(hits); err != nil {
logrus.Debugf("ES7.Start: decode err=%v", err)
errCh <- err
return
}
if hits.TimedOut {
err = fmt.Errorf("timeout")
logrus.Debugf("ES7.Start: search timeout")
errCh <- err
return
}
for idx := range hits.Hits.Hits {
rowCh <- hits.Hits.Hits[idx]
}
if len(hits.Hits.Hits) < e.cfg.Size {
return
}
}
}()
<-ready
return nil
}