feat: 🎉 完成基本功能
This commit is contained in:
loveuer
295
internal/xes/xes.go
Normal file
295
internal/xes/xes.go
Normal file
@ -0,0 +1,295 @@
|
||||
package xes
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/json"
|
||||
"esgo2dump/internal/interfaces"
|
||||
"esgo2dump/internal/opt"
|
||||
"esgo2dump/internal/util"
|
||||
"fmt"
|
||||
elastic "github.com/elastic/go-elasticsearch/v7"
|
||||
"github.com/elastic/go-elasticsearch/v7/esapi"
|
||||
"github.com/elastic/go-elasticsearch/v7/esutil"
|
||||
"github.com/sirupsen/logrus"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func NewClient(url *url.URL, iot interfaces.IO, qm map[string]any) (interfaces.DumpIO, error) {
|
||||
|
||||
var (
|
||||
err error
|
||||
endpoint = fmt.Sprintf("%s://%s", url.Scheme, url.Host)
|
||||
c *elastic.Client
|
||||
infoResp *esapi.Response
|
||||
index = strings.TrimPrefix(url.Path, "/")
|
||||
username string
|
||||
password string
|
||||
)
|
||||
|
||||
if url.User != nil {
|
||||
username = url.User.Username()
|
||||
if p, ok := url.User.Password(); ok {
|
||||
password = p
|
||||
}
|
||||
}
|
||||
|
||||
logrus.Debugf("xes.NewClient: endpoint=%s index=%s (username=%s password=%s)", endpoint, index, username, password)
|
||||
|
||||
if index == "" {
|
||||
return nil, fmt.Errorf("please specify index name: (like => http://127.0.0.1:9200/my_index)")
|
||||
}
|
||||
|
||||
if c, err = elastic.NewClient(
|
||||
elastic.Config{
|
||||
Addresses: []string{endpoint},
|
||||
Username: username,
|
||||
Password: password,
|
||||
CACert: nil,
|
||||
RetryOnStatus: []int{429},
|
||||
MaxRetries: 3,
|
||||
RetryBackoff: nil,
|
||||
Transport: &http.Transport{
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
|
||||
},
|
||||
},
|
||||
); err != nil {
|
||||
logrus.Debugf("xes.NewClient: elastic new client with endpont=%s err=%v", endpoint, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if infoResp, err = c.Info(); err != nil {
|
||||
logrus.Debugf("xes.NewClient: ping err=%v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if infoResp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("info xes status=%d", infoResp.StatusCode)
|
||||
}
|
||||
|
||||
return &client{c: c, index: index, queryMap: qm, iot: iot}, nil
|
||||
}
|
||||
|
||||
type client struct {
|
||||
c *elastic.Client
|
||||
iot interfaces.IO
|
||||
index string
|
||||
from int
|
||||
scrollId string
|
||||
queryMap map[string]any
|
||||
}
|
||||
|
||||
func (c *client) checkResponse(r *esapi.Response) error {
|
||||
if r.StatusCode == 200 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("status=%d msg=%s", r.StatusCode, r.String())
|
||||
}
|
||||
|
||||
func (c *client) IOType() interfaces.IO {
|
||||
return c.iot
|
||||
}
|
||||
|
||||
func (c *client) IsFile() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func (c *client) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *client) WriteData(ctx context.Context, docs []*interfaces.ESSource) (int, error) {
|
||||
var (
|
||||
err error
|
||||
indexer esutil.BulkIndexer
|
||||
count int
|
||||
)
|
||||
if indexer, err = esutil.NewBulkIndexer(esutil.BulkIndexerConfig{
|
||||
Client: c.c,
|
||||
Index: c.index,
|
||||
Refresh: "",
|
||||
}); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
for _, doc := range docs {
|
||||
var bs []byte
|
||||
|
||||
if bs, err = json.Marshal(doc.Content); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
logrus.Debugf("xes.Write: doc content=%s", string(bs))
|
||||
|
||||
if err = indexer.Add(context.Background(), esutil.BulkIndexerItem{
|
||||
Action: "index",
|
||||
Index: c.index,
|
||||
DocumentID: doc.DocId,
|
||||
Body: bytes.NewReader(bs),
|
||||
}); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
count++
|
||||
}
|
||||
|
||||
if err = indexer.Close(util.TimeoutCtx(ctx, opt.Timeout)); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
stats := indexer.Stats()
|
||||
if stats.NumFailed > 0 {
|
||||
return count, fmt.Errorf("write to xes failed=%d", stats.NumFailed)
|
||||
}
|
||||
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func (c *client) ReadData(ctx context.Context, i int) ([]*interfaces.ESSource, error) {
|
||||
var (
|
||||
err error
|
||||
resp *esapi.Response
|
||||
result = new(interfaces.ESResponse)
|
||||
)
|
||||
|
||||
if c.scrollId == "" {
|
||||
qs := []func(*esapi.SearchRequest){
|
||||
c.c.Search.WithContext(util.TimeoutCtx(ctx, opt.Timeout)),
|
||||
c.c.Search.WithIndex(c.index),
|
||||
c.c.Search.WithSize(i),
|
||||
c.c.Search.WithFrom(0),
|
||||
c.c.Search.WithScroll(time.Duration(opt.ScrollDurationSeconds) * time.Second),
|
||||
}
|
||||
|
||||
if len(c.queryMap) > 0 {
|
||||
queryBs, _ := json.Marshal(map[string]any{"query": c.queryMap})
|
||||
qs = append(qs, c.c.Search.WithBody(bytes.NewReader(queryBs)))
|
||||
}
|
||||
|
||||
if resp, err = c.c.Search(qs...); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf(resp.String())
|
||||
}
|
||||
|
||||
decoder := json.NewDecoder(resp.Body)
|
||||
if err = decoder.Decode(result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
c.scrollId = result.ScrollId
|
||||
|
||||
return result.Hits.Hits, nil
|
||||
}
|
||||
|
||||
if resp, err = c.c.Scroll(
|
||||
c.c.Scroll.WithScrollID(c.scrollId),
|
||||
c.c.Scroll.WithScroll(time.Duration(opt.ScrollDurationSeconds)*time.Second),
|
||||
); err != nil {
|
||||
return result.Hits.Hits, nil
|
||||
}
|
||||
|
||||
decoder := json.NewDecoder(resp.Body)
|
||||
if err = decoder.Decode(result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return result.Hits.Hits, nil
|
||||
}
|
||||
|
||||
func (c *client) ReadMapping(ctx context.Context) (map[string]any, error) {
|
||||
r, err := c.c.Indices.GetMapping(
|
||||
c.c.Indices.GetMapping.WithIndex(c.index),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if r.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("status=%d, msg=%s", r.StatusCode, r.String())
|
||||
}
|
||||
|
||||
m := make(map[string]any)
|
||||
decoder := json.NewDecoder(r.Body)
|
||||
if err = decoder.Decode(&m); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return m, nil
|
||||
}
|
||||
func (c *client) WriteMapping(ctx context.Context, m map[string]any) error {
|
||||
var (
|
||||
err error
|
||||
bs []byte
|
||||
result *esapi.Response
|
||||
)
|
||||
|
||||
for idxKey := range m {
|
||||
if bs, err = json.Marshal(m[idxKey]); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if result, err = c.c.Indices.Create(
|
||||
c.index,
|
||||
c.c.Indices.Create.WithContext(util.TimeoutCtx(ctx, opt.Timeout)),
|
||||
c.c.Indices.Create.WithBody(bytes.NewReader(bs)),
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err = c.checkResponse(result); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *client) ReadSetting(ctx context.Context) (map[string]any, error) {
|
||||
r, err := c.c.Indices.GetSettings(
|
||||
c.c.Indices.GetSettings.WithContext(util.TimeoutCtx(ctx, opt.Timeout)),
|
||||
c.c.Indices.GetSettings.WithIndex(c.index),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if r.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("status=%d, msg=%s", r.StatusCode, r.String())
|
||||
}
|
||||
|
||||
m := make(map[string]any)
|
||||
decoder := json.NewDecoder(r.Body)
|
||||
if err = decoder.Decode(&m); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (c *client) WriteSetting(ctx context.Context, m map[string]any) error {
|
||||
var (
|
||||
err error
|
||||
bs []byte
|
||||
result *esapi.Response
|
||||
)
|
||||
|
||||
if bs, err = json.Marshal(m); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if result, err = c.c.Indices.PutSettings(
|
||||
bytes.NewReader(bs),
|
||||
c.c.Indices.PutSettings.WithContext(util.TimeoutCtx(ctx, opt.Timeout)),
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return c.checkResponse(result)
|
||||
}
|
||||
38
internal/xes/xes_test.go
Normal file
38
internal/xes/xes_test.go
Normal file
@ -0,0 +1,38 @@
|
||||
package xes
|
||||
|
||||
import (
|
||||
"esgo2dump/internal/util"
|
||||
elastic "github.com/elastic/go-elasticsearch/v7"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGetESMapping(t *testing.T) {
|
||||
endpoint := "http://127.0.0.1:9200"
|
||||
index := "some_index"
|
||||
|
||||
cli, err := elastic.NewClient(elastic.Config{
|
||||
Addresses: []string{endpoint},
|
||||
})
|
||||
if err != nil {
|
||||
t.Error(1, err)
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := cli.Info(cli.Info.WithContext(util.Timeout(5)))
|
||||
if err != nil {
|
||||
t.Error(2, err)
|
||||
return
|
||||
}
|
||||
|
||||
t.Log("info:", resp.String())
|
||||
|
||||
r, err := cli.Indices.GetMapping(
|
||||
cli.Indices.GetMapping.WithIndex(index),
|
||||
)
|
||||
if err != nil {
|
||||
t.Error(3, err)
|
||||
return
|
||||
}
|
||||
|
||||
t.Log("get source:", r.String())
|
||||
}
|
||||
Reference in New Issue
Block a user