wip: refactory
This commit is contained in:
loveuer
145
xes/es6/read.go
145
xes/es6/read.go
@ -1,145 +0,0 @@
|
||||
package es6
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
elastic "github.com/elastic/go-elasticsearch/v6"
|
||||
"github.com/elastic/go-elasticsearch/v6/esapi"
|
||||
"github.com/loveuer/esgo2dump/internal/tool"
|
||||
"github.com/loveuer/esgo2dump/model"
|
||||
"github.com/loveuer/nf/nft/log"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
func ReadData(ctx context.Context, client *elastic.Client, index string, size, max int, query map[string]any, source []string, sort []string) (<-chan []*model.ESSource, <-chan error) {
|
||||
var (
|
||||
dataCh = make(chan []*model.ESSource)
|
||||
errCh = make(chan error)
|
||||
)
|
||||
|
||||
go func() {
|
||||
var (
|
||||
err error
|
||||
resp *esapi.Response
|
||||
result = new(model.ESResponseV6)
|
||||
scrollId string
|
||||
total int
|
||||
)
|
||||
|
||||
defer func() {
|
||||
close(dataCh)
|
||||
|
||||
if scrollId != "" {
|
||||
bs, _ := json.Marshal(map[string]string{
|
||||
"scroll_id": scrollId,
|
||||
})
|
||||
|
||||
var rr *esapi.Response
|
||||
|
||||
if rr, err = client.ClearScroll(
|
||||
client.ClearScroll.WithContext(tool.Timeout(3)),
|
||||
client.ClearScroll.WithBody(bytes.NewReader(bs)),
|
||||
); err != nil {
|
||||
log.Warn("clear scroll id=%s err=%v", scrollId, err)
|
||||
return
|
||||
}
|
||||
|
||||
if rr.StatusCode != 200 {
|
||||
log.Warn("clear scroll id=%s status=%d msg=%s", scrollId, rr.StatusCode, rr.String())
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
if client == nil {
|
||||
errCh <- fmt.Errorf("client is nil")
|
||||
}
|
||||
|
||||
qs := []func(*esapi.SearchRequest){
|
||||
client.Search.WithContext(tool.TimeoutCtx(ctx, 20)),
|
||||
client.Search.WithIndex(index),
|
||||
client.Search.WithSize(int(size)),
|
||||
client.Search.WithFrom(0),
|
||||
client.Search.WithScroll(time.Duration(120) * time.Second),
|
||||
}
|
||||
|
||||
if len(source) > 0 {
|
||||
qs = append(qs, client.Search.WithSourceIncludes(source...))
|
||||
}
|
||||
|
||||
if len(sort) > 0 {
|
||||
sorts := lo.Filter(sort, func(item string, index int) bool {
|
||||
return item != ""
|
||||
})
|
||||
|
||||
if len(sorts) > 0 {
|
||||
qs = append(qs, client.Search.WithSort(sorts...))
|
||||
}
|
||||
}
|
||||
|
||||
if query != nil && len(query) > 0 {
|
||||
queryBs, _ := json.Marshal(map[string]any{"query": query})
|
||||
qs = append(qs, client.Search.WithBody(bytes.NewReader(queryBs)))
|
||||
}
|
||||
|
||||
if resp, err = client.Search(qs...); err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
errCh <- fmt.Errorf("resp status=%d, resp=%s", resp.StatusCode, resp.String())
|
||||
return
|
||||
}
|
||||
|
||||
decoder := json.NewDecoder(resp.Body)
|
||||
if err = decoder.Decode(result); err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
|
||||
scrollId = result.ScrollId
|
||||
|
||||
dataCh <- result.Hits.Hits
|
||||
total += len(result.Hits.Hits)
|
||||
|
||||
if len(result.Hits.Hits) < size || (max > 0 && total >= max) {
|
||||
return
|
||||
}
|
||||
|
||||
for {
|
||||
if resp, err = client.Scroll(
|
||||
client.Scroll.WithScrollID(scrollId),
|
||||
client.Scroll.WithScroll(time.Duration(120)*time.Second),
|
||||
); err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
|
||||
result = new(model.ESResponseV6)
|
||||
|
||||
decoder = json.NewDecoder(resp.Body)
|
||||
if err = decoder.Decode(result); err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
errCh <- fmt.Errorf("resp status=%d, resp=%s", resp.StatusCode, resp.String())
|
||||
return
|
||||
}
|
||||
|
||||
dataCh <- result.Hits.Hits
|
||||
total += len(result.Hits.Hits)
|
||||
|
||||
if len(result.Hits.Hits) < size || (max > 0 && total >= max) {
|
||||
break
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return dataCh, errCh
|
||||
}
|
||||
@ -1,85 +0,0 @@
|
||||
package es6
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
elastic "github.com/elastic/go-elasticsearch/v6"
|
||||
"github.com/elastic/go-elasticsearch/v6/esutil"
|
||||
"github.com/loveuer/esgo2dump/model"
|
||||
"github.com/loveuer/nf/nft/log"
|
||||
)
|
||||
|
||||
func WriteData(ctx context.Context, client *elastic.Client, index string, docsCh <-chan []*model.ESSource, logs ...log.WroteLogger) error {
|
||||
var (
|
||||
err error
|
||||
indexer esutil.BulkIndexer
|
||||
total = 0
|
||||
)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case docs, ok := <-docsCh:
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
if len(docs) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
count := 0
|
||||
|
||||
if indexer, err = esutil.NewBulkIndexer(esutil.BulkIndexerConfig{
|
||||
Client: client,
|
||||
Index: index,
|
||||
ErrorTrace: true,
|
||||
OnError: func(ctx context.Context, err error) {
|
||||
},
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, doc := range docs {
|
||||
var bs []byte
|
||||
|
||||
if bs, err = json.Marshal(doc.Content); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err = indexer.Add(context.Background(), esutil.BulkIndexerItem{
|
||||
Action: "index",
|
||||
Index: index,
|
||||
DocumentID: doc.DocId,
|
||||
DocumentType: "_doc",
|
||||
Body: bytes.NewReader(bs),
|
||||
OnFailure: func(ctx context.Context, item esutil.BulkIndexerItem, item2 esutil.BulkIndexerResponseItem, bulkErr error) {
|
||||
},
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
count++
|
||||
}
|
||||
|
||||
total += count
|
||||
|
||||
if err = indexer.Close(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
stats := indexer.Stats()
|
||||
if stats.NumFailed > 0 {
|
||||
return fmt.Errorf("write to es failed_count=%d bulk_count=%d", stats.NumFailed, count)
|
||||
}
|
||||
|
||||
if len(logs) > 0 && logs[0] != nil {
|
||||
logs[0].Info("Dump: succeed=%d total=%d docs succeed!!!", count, total)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user