esgo2dump/internal/cmd/run.go

package cmd

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"strings"

	elastic "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"
	"github.com/samber/lo"
	"github.com/spf13/cobra"
)

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 *elastic.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
}

func run(cmd *cobra.Command, args []string) error {
	var (
		err    error
		input  model.IO[map[string]any]
		output model.IO[map[string]any]
	)

	if input, err = newIO(cmd.Context(), opt.Cfg.Args.Input, model.Input); err != nil {
		return err
	}

	if output, err = newIO(cmd.Context(), opt.Cfg.Args.Output, model.Output); err != nil {
		return err
	}

	go func() {
		<-cmd.Context().Done()
		os.Exit(1)
	}()

	if opt.Cfg.Args.QueryFile != "" {
		// query file
		var (
			items []map[string]any
			qf    *os.File
			// wrote count
			wc int
		)
		if qf, err = os.Open(opt.Cfg.Args.QueryFile); err != nil {
			return err
		}

		scanner := bufio.NewScanner(qf)
		// query count
		qc := 0
		for scanner.Scan() {
			qc++
			qm := make(map[string]any)
			if err = json.Unmarshal(scanner.Bytes(), &qm); err != nil {
				return err
			}

			for {
				if items, err = input.ReadData(
					opt.Cfg.Args.Limit,
					qm,
					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 {
					return err
				}

				if len(items) == 0 {
					break
				}

				if wc, err = output.WriteData(items); err != nil {
					return err
				}

				if wc != len(items) {
					return fmt.Errorf("got items %d, but wrote %d", len(items), wc)
				}

				log.Info("Dump: query_file[%06d] dump success = %d", qc, wc)
			}
		}
	}

	if opt.Cfg.Args.Query != "" {
		var (
			items []map[string]any
			qm    = make(map[string]any)
			wc    int
		)

		if err = json.Unmarshal([]byte(opt.Cfg.Args.Query), &qm); err != nil {
			return err
		}

		for {
			if items, err = input.ReadData(
				opt.Cfg.Args.Limit,
				qm,
				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 {
				return err
			}

			if len(items) == 0 {
				break
			}

			if wc, err = output.WriteData(items); err != nil {
				return err
			}

			if wc != len(items) {
				return fmt.Errorf("got items %d, but wrote %d", len(items), wc)
			}

			log.Info("Dump: query dump success = %d", wc)
		}
	}

	var (
		items []map[string]any
		wc    int
	)

	for {
		if items, err = input.ReadData(
			opt.Cfg.Args.Limit,
			nil,
			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 {
			return err
		}

		if len(items) == 0 {
			break
		}

		if wc, err = output.WriteData(items); err != nil {
			return err
		}

		if wc != len(items) {
			return fmt.Errorf("got items %d, but wrote %d", len(items), wc)
		}

		log.Info("Dump: query dump success = %d", wc)
	}

	return nil
}