loveuer/nfflow
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

wip: flow task start

Browse Source
This commit is contained in:
loveuer
2024-04-03 16:46:25 +08:00
parent 4718532458
commit bb56271348
11 changed files with 381 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
internal/controller/input/es7/es7.go
View File

@ -7,8 +7,10 @@ import (
"fmt"
elastic "github.com/elastic/go-elasticsearch/v7"
"github.com/elastic/go-elasticsearch/v7/esapi"
"github.com/loveuer/nfflow/internal/interfaces"
"github.com/loveuer/nfflow/internal/model"
"github.com/loveuer/nfflow/internal/opt"
"github.com/loveuer/nfflow/internal/sqlType"
"github.com/loveuer/nfflow/internal/util"
"github.com/sirupsen/logrus"
)
@ -16,23 +18,52 @@ import (
type ES7 struct {
cli *elastic.Client
scroll string
cfg struct {
Endpoints []string
Username string
Password string
Size int
Query map[string]any
Source []string
}
query map[string]any
size int
max int
source []string
}
func (e *ES7) init() error {
func NewInput(cfg sqlType.JSONB) (interfaces.Input, error) {
type config struct {
Endpoints []string `json:"endpoints"`
Username string `json:"username"`
Password string `json:"password"`
Size int `json:"size"`
Max int `json:"max"`
Query map[string]any `json:"query"`
Source []string `json:"source"`
}
var (
err error
c = new(config)
ins = &ES7{}
)
if err = cfg.Bind(c); err != nil {
return nil, err
}
if err = ins.Init(c.Endpoints, c.Username, c.Password); err != nil {
return nil, err
}
ins.query = c.Query
ins.size = c.Size
ins.max = c.Max
ins.source = c.Source
return ins, nil
}
func (e *ES7) Init(endpoints []string, username, password string) error {
var (
err error
cfg = elastic.Config{
Addresses: e.cfg.Endpoints,
Username: e.cfg.Username,
Password: e.cfg.Password,
Addresses: endpoints,
Username: username,
Password: password,
RetryOnStatus: []int{429},
}
info *esapi.Response
@ -53,7 +84,7 @@ func (e *ES7) init() error {
return nil
}
func (e *ES7) Start(ctx context.Context, task *model.Task, rowCh chan<- model.TaskRow, errCh chan<- error) error {
func (e *ES7) Start(ctx context.Context, rowCh chan<- interfaces.Row, errCh chan<- error) error {
var (
err error
result *esapi.Response
@ -63,20 +94,15 @@ func (e *ES7) Start(ctx context.Context, task *model.Task, rowCh chan<- model.Ta
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),
e.cli.Search.WithSize(e.size),
}
if e.cfg.Query != nil && len(e.cfg.Query) > 0 {
if e.query != nil && len(e.query) > 0 {
var bs []byte
if bs, err = json.Marshal(e.cfg.Query); err != nil {
if bs, err = json.Marshal(map[string]any{"query": e.query}); err != nil {
logrus.Debugf("ES7.Start: marshal query err=%v", err)
return err
}
@ -139,7 +165,7 @@ func (e *ES7) Start(ctx context.Context, task *model.Task, rowCh chan<- model.Ta
rowCh <- hits.Hits.Hits[idx]
}
if len(hits.Hits.Hits) < e.cfg.Size {
if len(hits.Hits.Hits) < e.size {
return
}
@ -179,7 +205,7 @@ func (e *ES7) Start(ctx context.Context, task *model.Task, rowCh chan<- model.Ta
rowCh <- hits.Hits.Hits[idx]
}
if len(hits.Hits.Hits) < e.cfg.Size {
if len(hits.Hits.Hits) < e.size {
return
}
}
@ -189,3 +215,5 @@ func (e *ES7) Start(ctx context.Context, task *model.Task, rowCh chan<- model.Ta
return nil
}
func (e *ES7) Close() {}