feat: make es read/write can be imported
This commit is contained in:
loveuer
85
xes/es6/client.go
Normal file
85
xes/es6/client.go
Normal file
@ -0,0 +1,85 @@
|
||||
package es6
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
elastic "github.com/elastic/go-elasticsearch/v6"
|
||||
"github.com/elastic/go-elasticsearch/v6/esapi"
|
||||
"github.com/loveuer/esgo2dump/internal/util"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"time"
|
||||
)
|
||||
|
||||
func NewClient(ctx context.Context, url *url.URL) (*elastic.Client, error) {
|
||||
var (
|
||||
err error
|
||||
urlUsername string
|
||||
urlPassword string
|
||||
client *elastic.Client
|
||||
errCh = make(chan error)
|
||||
cliCh = make(chan *elastic.Client)
|
||||
address = fmt.Sprintf("%s://%s", url.Scheme, url.Host)
|
||||
)
|
||||
|
||||
if url.User != nil {
|
||||
urlUsername = url.User.Username()
|
||||
if p, ok := url.User.Password(); ok {
|
||||
urlPassword = p
|
||||
}
|
||||
}
|
||||
|
||||
ncFunc := func(endpoints []string, username, password string) {
|
||||
var (
|
||||
err error
|
||||
cli *elastic.Client
|
||||
infoResp *esapi.Response
|
||||
)
|
||||
|
||||
if cli, err = elastic.NewClient(
|
||||
elastic.Config{
|
||||
Addresses: endpoints,
|
||||
Username: username,
|
||||
Password: password,
|
||||
CACert: nil,
|
||||
RetryOnStatus: []int{429},
|
||||
MaxRetries: 3,
|
||||
RetryBackoff: nil,
|
||||
Transport: &http.Transport{
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
|
||||
DialContext: (&net.Dialer{Timeout: 10 * time.Second}).DialContext,
|
||||
},
|
||||
},
|
||||
); err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
|
||||
if infoResp, err = cli.Info(); err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
|
||||
if infoResp.StatusCode != 200 {
|
||||
err = fmt.Errorf("info es7 status=%d", infoResp.StatusCode)
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
|
||||
cliCh <- cli
|
||||
}
|
||||
|
||||
go ncFunc([]string{address}, urlUsername, urlPassword)
|
||||
timeout := util.TimeoutCtx(ctx, 10)
|
||||
|
||||
select {
|
||||
case <-timeout.Done():
|
||||
return nil, fmt.Errorf("dial es=%s err=%v", address, context.DeadlineExceeded)
|
||||
case client = <-cliCh:
|
||||
return client, nil
|
||||
case err = <-errCh:
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
136
xes/es6/read.go
Normal file
136
xes/es6/read.go
Normal file
@ -0,0 +1,136 @@
|
||||
package es6
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
elastic "github.com/elastic/go-elasticsearch/v6"
|
||||
"github.com/elastic/go-elasticsearch/v6/esapi"
|
||||
"github.com/loveuer/esgo2dump/internal/util"
|
||||
"github.com/loveuer/esgo2dump/log"
|
||||
"github.com/loveuer/esgo2dump/model"
|
||||
"time"
|
||||
)
|
||||
|
||||
func ReadData(ctx context.Context, client *elastic.Client, index string, size, max int, query map[string]any, source []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.ESResponse)
|
||||
scrollId string
|
||||
total int
|
||||
)
|
||||
|
||||
defer func() {
|
||||
close(dataCh)
|
||||
close(errCh)
|
||||
|
||||
if scrollId != "" {
|
||||
bs, _ := json.Marshal(map[string]string{
|
||||
"scroll_id": scrollId,
|
||||
})
|
||||
|
||||
var (
|
||||
rr *esapi.Response
|
||||
)
|
||||
|
||||
if rr, err = client.ClearScroll(
|
||||
client.ClearScroll.WithContext(util.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(util.TimeoutCtx(ctx, 20)),
|
||||
client.Search.WithIndex(index),
|
||||
client.Search.WithSize(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 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.ESResponse)
|
||||
|
||||
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
|
||||
}
|
||||
77
xes/es6/write.go
Normal file
77
xes/es6/write.go
Normal file
@ -0,0 +1,77 @@
|
||||
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"
|
||||
)
|
||||
|
||||
func WriteData(ctx context.Context, client *elastic.Client, index string, docsCh <-chan []*model.ESSource) error {
|
||||
var (
|
||||
err error
|
||||
indexer esutil.BulkIndexer
|
||||
)
|
||||
|
||||
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++
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user