wip: 重构基本完成, 还未测试

xes/es7/read.go

@@ -5,6 +5,9 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"github.com/elastic/go-elasticsearch/v7/esutil"
+	"github.com/loveuer/esgo2dump/internal/opt"
+	"github.com/loveuer/esgo2dump/pkg/log"
 	"time"
 
 	elastic "github.com/elastic/go-elasticsearch/v7"
@@ -22,7 +25,7 @@ type streamer struct {
 }
 
 // ReadData implements model.IO.
-func (s *streamer) ReadData(limit int, query map[string]any, fields []string, sort []string) ([]map[string]any, error) {
+func (s *streamer) ReadData(ctx context.Context, limit int, query map[string]any, fields []string, sort []string) ([]map[string]any, error) {
 	var (
 		err    error
 		qs     []func(*esapi.SearchRequest)
@@ -30,6 +33,10 @@ func (s *streamer) ReadData(limit int, query map[string]any, fields []string, so
 		result = new(model.ESResponseV7[map[string]any])
 	)
 
+	if limit == 0 {
+		return nil, nil
+	}
+
 	if s.scroll != "" {
 		if resp, err = s.client.Scroll(
 			s.client.Scroll.WithContext(tool.TimeoutCtx(s.ctx)),
@@ -82,17 +89,185 @@ HandleResp:
 
 	s.scroll = result.ScrollId
 
-	return lo.Map(
-		result.Hits.Hits,
-		func(item *model.ESSource[map[string]any], _ int) map[string]any {
-			return item.Content
-		},
+	return lo.Slice(
+		lo.Map(
+			result.Hits.Hits,
+			func(item *model.ESSource[map[string]any], _ int) map[string]any {
+				return item.Content
+			},
+		),
+		0,
+		limit,
 	), nil
 }
 
 // WriteData implements model.IO.
-func (s *streamer) WriteData([]map[string]any) (int, error) {
-	panic("unimplemented")
+func (s *streamer) WriteData(ctx context.Context, items []map[string]any) (int, error) {
+	var (
+		err     error
+		indexer esutil.BulkIndexer
+		total   int
+	)
+
+	if len(items) == 0 {
+		return 0, nil
+	}
+
+	count := 0
+
+	if indexer, err = esutil.NewBulkIndexer(esutil.BulkIndexerConfig{
+		NumWorkers:    0,
+		FlushBytes:    0,
+		FlushInterval: 0,
+		Client:        s.client,
+		Decoder:       nil,
+		OnError: func(ctx context.Context, err error) {
+			log.Error("es7.writer: on error log, err = %s", err.Error())
+		},
+		Index:               s.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 0, err
+	}
+
+	for _, item := range items {
+		var bs []byte
+
+		if bs, err = json.Marshal(item); err != nil {
+			return 0, err
+		}
+
+		if err = indexer.Add(context.Background(), esutil.BulkIndexerItem{
+			Action: "index",
+			Index:  s.index,
+			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 0, err
+		}
+
+		count++
+	}
+
+	total += count
+
+	if err = indexer.Close(ctx); err != nil {
+		return 0, err
+	}
+
+	stats := indexer.Stats()
+
+	return len(items) - int(stats.NumFailed), nil
+}
+
+func (s *streamer) ReadMapping(ctx context.Context) (map[string]any, error) {
+	r, err := s.client.Indices.GetMapping(
+		s.client.Indices.GetMapping.WithIndex(s.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 (s *streamer) WriteMapping(ctx context.Context, mapping map[string]any) error {
+	var (
+		err    error
+		bs     []byte
+		result *esapi.Response
+	)
+
+	for idxKey := range mapping {
+		if bs, err = json.Marshal(mapping[idxKey]); err != nil {
+			return err
+		}
+
+		if result, err = s.client.Indices.Create(
+			s.index,
+			s.client.Indices.Create.WithContext(tool.TimeoutCtx(ctx, opt.Timeout)),
+			s.client.Indices.Create.WithBody(bytes.NewReader(bs)),
+		); err != nil {
+			return err
+		}
+
+		if result.StatusCode != 200 {
+			return fmt.Errorf("status=%d, msg=%s", result.StatusCode, result.String())
+		}
+	}
+
+	return nil
+}
+
+func (s *streamer) ReadSetting(ctx context.Context) (map[string]any, error) {
+	r, err := s.client.Indices.GetSettings(
+		s.client.Indices.GetSettings.WithContext(tool.TimeoutCtx(ctx, opt.Timeout)),
+		s.client.Indices.GetSettings.WithIndex(s.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 (s *streamer) WriteSetting(ctx context.Context, setting map[string]any) error {
+	var (
+		err    error
+		bs     []byte
+		result *esapi.Response
+	)
+
+	if bs, err = json.Marshal(setting); err != nil {
+		return err
+	}
+
+	if result, err = s.client.Indices.PutSettings(
+		bytes.NewReader(bs),
+		s.client.Indices.PutSettings.WithContext(tool.TimeoutCtx(ctx, opt.Timeout)),
+	); err != nil {
+		return err
+	}
+
+	if result.StatusCode != 200 {
+		return fmt.Errorf("status=%d, msg=%s", result.StatusCode, result.String())
+	}
+
+	return nil
 }
 
 func NewStreamer(ctx context.Context, client *elastic.Client, index string) (model.IO[map[string]any], error) {