From a574118649a9b052a1299175cff70991b01a317f Mon Sep 17 00:00:00 2001 From: loveuer Date: Thu, 2 Jan 2025 21:38:30 -0800 Subject: [PATCH] update: add flag(disable-ping) --- internal/cmd/cmd.go | 1 + internal/cmd/run.go | 2 +- internal/opt/opt.go | 7 ++- internal/xes/xes7.go | 25 +++++--- xes/es7/client.go | 134 ++++++++++++++++++++--------------------- xes/es7/client_test.go | 4 +- 6 files changed, 87 insertions(+), 86 deletions(-) diff --git a/internal/cmd/cmd.go b/internal/cmd/cmd.go index 48eba2f..06f5c1f 100644 --- a/internal/cmd/cmd.go +++ b/internal/cmd/cmd.go @@ -57,6 +57,7 @@ esgo2dump --input=http://127.0.0.1:9200/some_index --output=./data.json --query_ func init() { rootCommand.PersistentFlags().BoolVar(&opt.Cfg.Debug, "debug", false, "") rootCommand.PersistentFlags().BoolVar(&opt.Cfg.Dev, "dev", false, "") + rootCommand.PersistentFlags().BoolVar(&opt.Cfg.DisablePing, "disable-ping", false, "") rootCommand.PersistentFlags().BoolVarP(&opt.Cfg.Args.Version, "version", "v", false, "print esgo2dump version") rootCommand.Flags().IntVar(&opt.Cfg.Args.Timeout, "timeout", 30, "max timeout seconds per operation with limit") diff --git a/internal/cmd/run.go b/internal/cmd/run.go index e80ac0c..0a14d0a 100644 --- a/internal/cmd/run.go +++ b/internal/cmd/run.go @@ -268,7 +268,7 @@ func newIO(source string, ioType interfaces.IO, esv string) (interfaces.DumpIO, switch esv { case "7": - return xes.NewClient(iurl, ioType) + return xes.NewClient(source, ioType) case "6": return xes.NewClientV6(iurl, ioType) case "8": diff --git a/internal/opt/opt.go b/internal/opt/opt.go index 92d548e..926177e 100644 --- a/internal/opt/opt.go +++ b/internal/opt/opt.go @@ -15,9 +15,10 @@ type args struct { } type config struct { - Debug bool `json:"-"` - Dev bool `json:"-"` - Args args `json:"-"` + Debug bool `json:"-"` + Dev bool `json:"-"` + DisablePing bool `json:"-"` + Args args `json:"-"` } var Cfg = &config{} diff --git a/internal/xes/xes7.go b/internal/xes/xes7.go index 77210b3..13ae5ca 100644 --- a/internal/xes/xes7.go +++ b/internal/xes/xes7.go @@ -32,22 +32,27 @@ func (c *client) WriteData(ctx context.Context, docsCh <-chan []*model.ESSource) return es7.WriteData(ctx, c.client, c.index, docsCh, c) } -func NewClient(url *url.URL, iot interfaces.IO) (interfaces.DumpIO, error) { +func NewClient(uri string, iot interfaces.IO) (interfaces.DumpIO, error) { var ( - urlIndex = strings.TrimPrefix(url.Path, "/") - cli *elastic.Client - err error + cli *elastic.Client + err error + ins *url.URL + index string ) - if urlIndex == "" { - return nil, fmt.Errorf("please specify index name: (like => http://127.0.0.1:9200/my_index)") - } - - if cli, err = es7.NewClient(context.TODO(), url); err != nil { + if ins, err = url.Parse(uri); err != nil { return nil, err } - return &client{client: cli, iot: iot, index: urlIndex}, nil + if index = strings.TrimSpace(strings.TrimPrefix(ins.Path, "/")); index == "" { + return nil, fmt.Errorf("please specify index name: (like => http://127.0.0.1:9200/my_index)") + } + + if cli, err = es7.NewClient(context.TODO(), uri, es7.Config{DisablePing: opt.Cfg.DisablePing}); err != nil { + return nil, err + } + + return &client{client: cli, iot: iot, index: index}, nil } func (c *client) checkResponse(r *esapi.Response) error { diff --git a/xes/es7/client.go b/xes/es7/client.go index 154db9d..3c3f6e2 100644 --- a/xes/es7/client.go +++ b/xes/es7/client.go @@ -16,78 +16,74 @@ import ( "github.com/samber/lo" ) -func NewClient(ctx context.Context, url *url.URL) (*elastic.Client, error) { +type Config struct { + DisablePing bool +} + +// NewClient +// new esv7 client +// uri example: +// - http://127.0.0.1:9200 +// - https://:@node1.dev:9200,node2.dev:19200,node3.dev:29200 +func NewClient(ctx context.Context, uri string, configs ...Config) (*elastic.Client, error) { var ( - err error - urlUsername string - urlPassword string - client *elastic.Client - errCh = make(chan error) - cliCh = make(chan *elastic.Client) - endpoints = lo.Map( - strings.Split(url.Host, ","), - func(item string, index int) string { - return fmt.Sprintf("%s://%s", url.Scheme, item) - }, - ) + err error + username string + password string + client *elastic.Client + ins *url.URL ) - if url.User != nil { - urlUsername = url.User.Username() - if p, ok := url.User.Password(); ok { - urlPassword = p - } - } - - ncFunc := func(endpoints []string, username, password string) { - var ( - err error - cli *elastic.Client - infoResp *esapi.Response - ) - - if cli, err = elastic.NewClient( - elastic.Config{ - Addresses: endpoints, - Username: username, - Password: password, - CACert: nil, - RetryOnStatus: []int{429}, - MaxRetries: 3, - RetryBackoff: nil, - Transport: &http.Transport{ - TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, - DialContext: (&net.Dialer{Timeout: 10 * time.Second}).DialContext, - }, - }, - ); err != nil { - errCh <- err - return - } - - if infoResp, err = cli.Info(); err != nil { - errCh <- err - return - } - - if infoResp.StatusCode != 200 { - err = fmt.Errorf("info es7 status=%d", infoResp.StatusCode) - errCh <- err - return - } - - cliCh <- cli - } - - go ncFunc(endpoints, urlUsername, urlPassword) - timeout := tool.TimeoutCtx(ctx, 10) - - select { - case <-timeout.Done(): - return nil, fmt.Errorf("dial es=%v err=%v", endpoints, context.DeadlineExceeded) - case client = <-cliCh: - return client, nil - case err = <-errCh: + if ins, err = url.Parse(uri); err != nil { 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 { + return fmt.Sprintf("%s://%s", ins.Scheme, item) + }, + ) + + if ins.User != nil { + username = ins.User.Username() + password, _ = ins.User.Password() + } + + if client, err = elastic.NewClient( + elastic.Config{ + Addresses: endpoints, + Username: username, + Password: password, + CACert: nil, + RetryOnStatus: []int{429}, + MaxRetries: 3, + RetryBackoff: nil, + Transport: &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, + DialContext: (&net.Dialer{Timeout: 10 * time.Second}).DialContext, + }, + }, + ); err != nil { + return nil, err + } + + if !cfg.DisablePing { + var res *esapi.Response + if res, err = client.Ping(client.Ping.WithContext(tool.TimeoutCtx(ctx, 5))); err != nil { + return nil, err + } + + if res.StatusCode != 200 { + err = fmt.Errorf("ping es server response: %s", res.String()) + return nil, err + } + } + + return client, nil } diff --git a/xes/es7/client_test.go b/xes/es7/client_test.go index 405e97b..0433ce3 100644 --- a/xes/es7/client_test.go +++ b/xes/es7/client_test.go @@ -1,7 +1,6 @@ package es7 import ( - "net/url" "testing" "github.com/loveuer/esgo2dump/internal/tool" @@ -9,9 +8,8 @@ import ( func TestNewClient(t *testing.T) { uri := "http://es1.dev:9200,es2.dev:9200" - ins, _ := url.Parse(uri) - c, err := NewClient(tool.Timeout(5), ins) + c, err := NewClient(tool.Timeout(5), uri) if err != nil { t.Fatal(err.Error()) }