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 }