wip: refactory

xes/es7/client.go

@@ -7,7 +7,6 @@ import (
 	"net"
 	"net/http"
 	"net/url"
-	"strconv"
 	"strings"
 	"time"
 
@@ -17,22 +16,12 @@ import (
 	"github.com/samber/lo"
 )
 
-// Deprecated. use uri query: http://<username>:<password>@example.com:port?ping=false&...
-type Config struct {
-	DisablePing bool
-}
-
-type UriConfig struct {
-	Ping  bool `json:"ping"`
-	Sniff bool `json:"sniff"`
-}
-
 // NewClient
 // new esv7 client
 // uri example:
 //   - http://127.0.0.1:9200
 //   - https://<username>:<password>@node1.dev:9200,node2.dev:19200,node3.dev:29200
-func NewClient(ctx context.Context, uri string, configs ...Config) (*elastic.Client, error) {
+func NewClient(ctx context.Context, uri string) (*elastic.Client, error) {
 	var (
 		err      error
 		username string
@@ -45,11 +34,6 @@ func NewClient(ctx context.Context, uri string, configs ...Config) (*elastic.Cli
 		return nil, err
 	}
 
-	cfg := Config{}
-	if len(configs) > 0 {
-		cfg = configs[0]
-	}
-
 	endpoints := lo.Map(
 		strings.Split(ins.Host, ","),
 		func(item string, index int) string {
@@ -64,10 +48,6 @@ func NewClient(ctx context.Context, uri string, configs ...Config) (*elastic.Cli
 
 	query := ins.Query()
 
-	cfg2 := &UriConfig{}
-	cfg2.Ping, _ = strconv.ParseBool(query.Get("ping"))
-	cfg2.Sniff, _ = strconv.ParseBool(query.Get("sniff"))
-
 	if client, err = elastic.NewClient(
 		elastic.Config{
 			Addresses:     endpoints,
@@ -81,15 +61,13 @@ func NewClient(ctx context.Context, uri string, configs ...Config) (*elastic.Cli
 				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
 				DialContext:     (&net.Dialer{Timeout: 10 * time.Second}).DialContext,
 			},
-			DiscoverNodesOnStart: cfg2.Sniff,
+			DiscoverNodesOnStart: lo.If(query.Get("sniff") == "true", true).Else(false),
 		},
 	); err != nil {
 		return nil, err
 	}
 
-	// Deprecated.
-	cfg.DisablePing = cfg.DisablePing || cfg2.Ping
-	if cfg.DisablePing {
+	if query.Get("ping") != "false" {
 		var res *esapi.Response
 		if res, err = client.Ping(client.Ping.WithContext(tool.TimeoutCtx(ctx, 5))); err != nil {
 			return nil, err

xes/es7/read.go

@@ -10,255 +10,92 @@ import (
 	elastic "github.com/elastic/go-elasticsearch/v7"
 	"github.com/elastic/go-elasticsearch/v7/esapi"
 	"github.com/loveuer/esgo2dump/internal/tool"
-	"github.com/loveuer/esgo2dump/model"
-	"github.com/loveuer/nf/nft/log"
+	"github.com/loveuer/esgo2dump/pkg/model"
 	"github.com/samber/lo"
 )
 
-// ReadData
-// @param[source]: a list of include fields to extract and return from the _source field.
-// @param[sort]:   a list of <field>:<direction> pairs.
-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.ESResponseV7)
-			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(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(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.ESResponseV7)
-
-			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
+type streamer struct {
+	ctx    context.Context
+	client *elastic.Client
+	index  string
+	scroll string
 }
 
-// ReadDataV2 es7 read data
-// Deprecated: bug, when can't sort by _id
-/*
-	- @param[source]: a list of include fields to extract and return from the _source field.
-    - @param[sort]:   a list of <field>:<direction> pairs.
-*/
-func ReadDataV2(
-	ctx context.Context,
-	client *elastic.Client,
-	index string,
-	size, max int,
-	query map[string]any,
-	source []string,
-	sort []string,
-) (<-chan []*model.ESSource, <-chan error) {
+// ReadData implements model.IO.
+func (s *streamer) ReadData(limit int, query map[string]any, fields []string, sort []string) ([]map[string]any, error) {
 	var (
-		dataCh = make(chan []*model.ESSource)
-		errCh  = make(chan error)
+		err    error
+		qs     []func(*esapi.SearchRequest)
+		resp   *esapi.Response
+		result = new(model.ESResponseV7[map[string]any])
 	)
 
-	log.Debug("es7.ReadDataV2: arg.index = %s, arg.size = %d, arg.max = %d", index, size, max)
-
-	go func() {
-		var (
-			err         error
-			bs          []byte
-			resp        *esapi.Response
-			searchAfter     = make([]any, 0)
-			total       int = 0
-			body            = make(map[string]any)
-			qs          []func(request *esapi.SearchRequest)
-		)
-
-		if sort == nil {
-			sort = []string{}
+	if s.scroll != "" {
+		if resp, err = s.client.Scroll(
+			s.client.Scroll.WithContext(tool.TimeoutCtx(s.ctx)),
+			s.client.Scroll.WithScrollID(s.scroll),
+			s.client.Scroll.WithScroll(35*time.Second),
+		); err != nil {
+			return nil, err
 		}
 
-		if len(query) > 0 {
-			body["query"] = query
+		goto HandleResp
+	}
+
+	qs = []func(*esapi.SearchRequest){
+		s.client.Search.WithContext(tool.TimeoutCtx(s.ctx)),
+		s.client.Search.WithIndex(s.index),
+		s.client.Search.WithSize(limit),
+		s.client.Search.WithScroll(35 * time.Second),
+	}
+
+	if len(fields) > 0 {
+		qs = append(qs, s.client.Search.WithSourceIncludes(fields...))
+	}
+
+	if len(sort) > 0 {
+		qs = append(qs, s.client.Search.WithSort(sort...))
+	}
+
+	if len(query) > 0 {
+		queryBs, err := json.Marshal(map[string]any{"query": query})
+		if err != nil {
+			return nil, err
 		}
 
-		sort = append(sort, "_id:ASC")
+		qs = append(qs, s.client.Search.WithBody(bytes.NewReader(queryBs)))
+	}
 
-		sorts := lo.Filter(sort, func(item string, index int) bool {
-			return item != ""
-		})
+	if resp, err = s.client.Search(qs...); err != nil {
+		return nil, err
+	}
 
-		defer func() {
-			close(dataCh)
-			close(errCh)
-		}()
+HandleResp:
 
-		for {
-			finaSize := tool.CalcSize(size, max, total)
-			qs = []func(*esapi.SearchRequest){
-				client.Search.WithContext(tool.TimeoutCtx(ctx, 30)),
-				client.Search.WithIndex(index),
-				client.Search.WithSize(finaSize),
-				client.Search.WithSort(sorts...),
-			}
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("resp status=%d, resp=%s", resp.StatusCode, resp.String())
+	}
 
-			if len(source) > 0 {
-				qs = append(qs, client.Search.WithSourceIncludes(source...))
-			}
+	if err = json.NewDecoder(resp.Body).Decode(result); err != nil {
+		return nil, err
+	}
 
-			delete(body, "search_after")
-			if len(searchAfter) > 0 {
-				body["search_after"] = searchAfter
-			}
+	s.scroll = result.ScrollId
 
-			if bs, err = json.Marshal(body); err != nil {
-				errCh <- err
-				return
-			}
-
-			log.Debug("es7.ReadDataV2: search request size = %d, body = %s", finaSize, string(bs))
-
-			qs = append(qs, client.Search.WithBody(bytes.NewReader(bs)))
-			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
-			}
-
-			result := new(model.ESResponseV7)
-			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
-			log.Debug("es7.ReadDataV2: search response hits = %d", len(result.Hits.Hits))
-			total += len(result.Hits.Hits)
-
-			if len(result.Hits.Hits) < size || (max > 0 && total >= max) {
-				break
-			}
-
-			searchAfter = result.Hits.Hits[len(result.Hits.Hits)-1].Sort
-		}
-	}()
-
-	return dataCh, errCh
+	return lo.Map(
+		result.Hits.Hits,
+		func(item *model.ESSource[map[string]any], _ int) map[string]any {
+			return item.Content
+		},
+	), nil
+}
+
+// WriteData implements model.IO.
+func (s *streamer) WriteData([]map[string]any) (int, error) {
+	panic("unimplemented")
+}
+
+func NewStreamer(ctx context.Context, client *elastic.Client, index string) (model.IO[map[string]any], error) {
+	s := &streamer{ctx: ctx, client: client, index: index}
+	return s, nil
 }

xes/es7/write.go

@@ -8,94 +8,82 @@ import (
 
 	elastic "github.com/elastic/go-elasticsearch/v7"
 	"github.com/elastic/go-elasticsearch/v7/esutil"
-	"github.com/loveuer/esgo2dump/model"
-	"github.com/loveuer/nf/nft/log"
+	"github.com/loveuer/esgo2dump/pkg/log"
+	"github.com/loveuer/esgo2dump/pkg/model"
 )
 
-func WriteData(ctx context.Context, client *elastic.Client, index string, docsCh <-chan []*model.ESSource, logs ...log.WroteLogger) error {
+func WriteData[T any](ctx context.Context, client *elastic.Client, index string, docs ...*model.ESSource[T]) error {
 	var (
 		err     error
 		indexer esutil.BulkIndexer
 		total   int
 	)
 
-	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{
-				NumWorkers:    0,
-				FlushBytes:    0,
-				FlushInterval: 0,
-				Client:        client,
-				Decoder:       nil,
-				OnError: func(ctx context.Context, err error) {
-					log.Error("es7.writer: on error log, err = %s", err.Error())
-				},
-				Index:               index,
-				ErrorTrace:          true,
-				FilterPath:          []string{},
-				Header:              map[string][]string{},
-				Human:               false,
-				Pipeline:            "",
-				Pretty:              false,
-				Refresh:             "",
-				Routing:             "",
-				Source:              []string{},
-				SourceExcludes:      []string{},
-				SourceIncludes:      []string{},
-				Timeout:             0,
-				WaitForActiveShards: "",
-			}); 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,
-					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)
-			}
-		}
+	if len(docs) == 0 {
+		return nil
 	}
+
+	count := 0
+
+	if indexer, err = esutil.NewBulkIndexer(esutil.BulkIndexerConfig{
+		NumWorkers:    0,
+		FlushBytes:    0,
+		FlushInterval: 0,
+		Client:        client,
+		Decoder:       nil,
+		OnError: func(ctx context.Context, err error) {
+			log.Error("es7.writer: on error log, err = %s", err.Error())
+		},
+		Index:               index,
+		ErrorTrace:          true,
+		FilterPath:          []string{},
+		Header:              map[string][]string{},
+		Human:               false,
+		Pipeline:            "",
+		Pretty:              false,
+		Refresh:             "",
+		Routing:             "",
+		Source:              []string{},
+		SourceExcludes:      []string{},
+		SourceIncludes:      []string{},
+		Timeout:             0,
+		WaitForActiveShards: "",
+	}); 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,
+			Body:       bytes.NewReader(bs),
+			OnFailure: func(ctx context.Context, item esutil.BulkIndexerItem, item2 esutil.BulkIndexerResponseItem, bulkErr error) {
+				log.Error("es7.writer: on failure err log, err = %s", 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)
+	}
+
+	return nil
 }