feat: dump from es with query, limit

internal/xes/xes.go

@@ -0,0 +1,196 @@
+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, 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}, nil
+}
+
+type client struct {
+	c        *elastic.Client
+	index    string
+	from     int
+	scrollId string
+	queryMap map[string]any
+}
+
+func (c *client) IsInput() bool {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (c *client) IsFile() bool {
+	return false
+}
+
+func (c *client) Close() error {
+	return nil
+}
+
+func (c *client) Write(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.Timeout(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) Read(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.Timeout(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
+}