refactor: project code arch; feat: cancel operation anytime

internal/core/io.go

@@ -0,0 +1,81 @@
+package core
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	elastic7 "github.com/elastic/go-elasticsearch/v7"
+	"github.com/go-resty/resty/v2"
+	"github.com/loveuer/esgo2dump/internal/opt"
+	"github.com/loveuer/esgo2dump/internal/tool"
+	"github.com/loveuer/esgo2dump/internal/xfile"
+	"github.com/loveuer/esgo2dump/pkg/log"
+	"github.com/loveuer/esgo2dump/pkg/model"
+	"github.com/loveuer/esgo2dump/xes/es7"
+	"net/url"
+	"strings"
+)
+
+func NewIO(ctx context.Context, uri string, ioType model.IOType) (model.IO[map[string]any], error) {
+	type Version struct {
+		Name    string
+		Version struct {
+			Number string `json:"number"`
+		} `json:"version"`
+	}
+
+	var (
+		err    error
+		target *url.URL
+		rr     *resty.Response
+		v      Version
+	)
+
+	if target, err = url.Parse(uri); err != nil {
+		log.Debug("parse uri failed, type = %s, uri = %s, err = %s", ioType, uri, err.Error())
+		return xfile.NewClient(uri, ioType)
+	}
+
+	if err = tool.ValidScheme(target.Scheme); err != nil {
+		log.Debug("uri scheme check failed, type = %s, uri = %s", ioType, uri)
+		return xfile.NewClient(uri, ioType)
+	}
+
+	// elastic uri
+	index := strings.TrimPrefix(target.Path, "/")
+	if index == "" {
+		return nil, fmt.Errorf("uri invalid without index(path)")
+	}
+
+	log.Debug("%s uri es index = %s", ioType, index)
+
+	versionURL := fmt.Sprintf("%s://%s", target.Scheme, strings.Split(target.Host, ",")[0])
+	log.Debug("%s version url = %s", ioType, versionURL)
+	if rr, err = opt.HttpClient.R().Get(versionURL); err != nil {
+		log.Debug("get uri es version failed, type = %s, uri = %s, version_url = %s, err = %s", ioType, uri, versionURL, err.Error())
+	}
+
+	if err = json.Unmarshal(rr.Body(), &v); err != nil {
+		log.Debug("decode uri es version failed, type = %s, uri = %s, version_url = %s, err = %s", ioType, uri, versionURL, err.Error())
+		return nil, err
+	}
+
+	log.Debug("%s uri es version = %s", ioType, v.Version.Number)
+
+	mainVersion := strings.Split(v.Version.Number, ".")[0]
+	switch mainVersion {
+	case "8":
+	case "7":
+		var client *elastic7.Client
+		if client, err = es7.NewClient(ctx, uri); err != nil {
+			return nil, err
+		}
+
+		return es7.NewStreamer(ctx, client, index)
+	case "6":
+	default:
+		return nil, fmt.Errorf("es version not supported yet: %s", mainVersion)
+	}
+
+	return nil, nil
+}

internal/core/run.data.go

@@ -0,0 +1,125 @@
+package core
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"github.com/loveuer/esgo2dump/internal/opt"
+	"github.com/loveuer/esgo2dump/internal/tool"
+	"github.com/loveuer/esgo2dump/pkg/log"
+	"github.com/loveuer/esgo2dump/pkg/model"
+	"github.com/samber/lo"
+	"github.com/spf13/cobra"
+	"os"
+	"strings"
+	"sync"
+)
+
+func RunData(cmd *cobra.Command, input, output model.IO[map[string]any]) error {
+	var (
+		err error
+		// query chan
+		qc = make(chan map[string]any)
+		// error chan
+		ec = make(chan error)
+		// done chan
+		wc    = &sync.WaitGroup{}
+		total = 0
+	)
+
+	wc.Add(1)
+
+	go func() {
+		var (
+			wroteCount = 0
+			items      []map[string]any
+		)
+
+		defer wc.Done()
+
+		for query := range qc {
+			for {
+				limit := tool.CalculateLimit(opt.Cfg.Args.Limit, total, opt.Cfg.Args.Max)
+				log.Debug("one-step dump: arg.limit = %d, total = %d, arg.max = %d, calculate.limit = %d", opt.Cfg.Args.Limit, total, opt.Cfg.Args.Max, limit)
+				if limit == 0 {
+					break
+				}
+
+				if items, err = input.ReadData(
+					cmd.Context(),
+					limit,
+					query,
+					lo.Filter(strings.Split(opt.Cfg.Args.Field, ","), func(x string, _ int) bool { return x != "" }),
+					lo.Filter(strings.Split(opt.Cfg.Args.Sort, ","), func(x string, _ int) bool { return x != "" }),
+				); err != nil {
+					ec <- err
+					return
+				}
+
+				if len(items) == 0 {
+					break
+				}
+
+				if wroteCount, err = output.WriteData(cmd.Context(), items); err != nil {
+					ec <- err
+					return
+				}
+
+				total += wroteCount
+
+				if wroteCount != len(items) {
+					ec <- fmt.Errorf("got items %d, but wrote %d", len(items), wroteCount)
+					return
+				}
+
+				log.Info("Dump: dump data success = %d total = %d", wroteCount, total)
+			}
+		}
+	}()
+
+	switch {
+	case opt.Cfg.Args.QueryFile != "":
+		var (
+			// query file
+			qf         *os.File
+			queryCount = 0
+		)
+		if qf, err = os.Open(opt.Cfg.Args.QueryFile); err != nil {
+			return err
+		}
+
+		scanner := bufio.NewScanner(qf)
+		for scanner.Scan() {
+			queryCount++
+			qm := make(map[string]any)
+			if err = json.Unmarshal(scanner.Bytes(), &qm); err != nil {
+				return err
+			}
+
+			qc <- qm
+
+			log.Debug("Dump: queries[%06d]", queryCount)
+		}
+	case opt.Cfg.Args.Query != "":
+		var (
+			qm = make(map[string]any)
+		)
+
+		if err = json.Unmarshal([]byte(opt.Cfg.Args.Query), &qm); err != nil {
+			return err
+		}
+
+		qc <- qm
+	default:
+		qc <- nil
+	}
+
+	// close query chan to stop trans_io_goroutine
+	close(qc)
+
+	wc.Wait()
+
+	log.Info("Dump: dump all data success, total = %d", total)
+
+	return nil
+}

internal/core/run.mapping.go

@@ -0,0 +1,19 @@
+package core
+
+import (
+	"github.com/loveuer/esgo2dump/pkg/model"
+	"github.com/spf13/cobra"
+)
+
+func RunMapping(cmd *cobra.Command, input model.IO[map[string]any], output model.IO[map[string]any]) error {
+	mapping, err := input.ReadMapping(cmd.Context())
+	if err != nil {
+		return err
+	}
+
+	if err = output.WriteMapping(cmd.Context(), mapping); err != nil {
+		return err
+	}
+
+	return nil
+}

internal/core/run.setting.go

@@ -0,0 +1,19 @@
+package core
+
+import (
+	"github.com/loveuer/esgo2dump/pkg/model"
+	"github.com/spf13/cobra"
+)
+
+func RunSetting(cmd *cobra.Command, input model.IO[map[string]any], output model.IO[map[string]any]) error {
+	setting, err := input.ReadSetting(cmd.Context())
+	if err != nil {
+		return err
+	}
+
+	if err = output.WriteSetting(cmd.Context(), setting); err != nil {
+		return err
+	}
+
+	return nil
+}