mirror of
https://github.com/neilotoole/sq.git
synced 2024-12-25 17:25:36 +03:00
a419e0b693
* debugging * generic retry functionality; some additions to stringz.go * fiddling with magefile_sakila * added test lock function testh.Lock * bunch of linting, goimports, what not
197 lines
5.1 KiB
Go
197 lines
5.1 KiB
Go
package cli
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/spf13/cobra"
|
|
|
|
"github.com/neilotoole/lg"
|
|
|
|
"github.com/neilotoole/sq/cli/output"
|
|
"github.com/neilotoole/sq/libsq/core/errz"
|
|
"github.com/neilotoole/sq/libsq/driver"
|
|
"github.com/neilotoole/sq/libsq/source"
|
|
)
|
|
|
|
func newPingCmd() *cobra.Command {
|
|
cmd := &cobra.Command{
|
|
Use: "ping [@HANDLE [@HANDLE_N]]",
|
|
RunE: execPing,
|
|
ValidArgsFunction: completeHandle(0),
|
|
|
|
Short: "Ping data sources",
|
|
Long: `Ping data sources to check connection health. If no arguments provided, the
|
|
active data source is pinged. Provide the handles of one or more sources
|
|
to ping those sources, or --all to ping all sources.
|
|
|
|
The exit code is 1 if ping fails for any of the sources.`,
|
|
Example: ` # ping active data source
|
|
$ sq ping
|
|
|
|
# ping all data sources
|
|
$ sq ping --all
|
|
|
|
# ping @my1 and @pg1
|
|
$ sq ping @my1 @pg1
|
|
|
|
# ping @my1 with 2s timeout
|
|
$ sq ping @my1 --timeout=2s
|
|
|
|
# output in TSV format
|
|
$ sq ping --tsv @my1`,
|
|
}
|
|
|
|
cmd.Flags().BoolP(flagTable, flagTableShort, false, flagTableUsage)
|
|
cmd.Flags().BoolP(flagCSV, flagCSVShort, false, flagCSVUsage)
|
|
cmd.Flags().BoolP(flagTSV, flagTSVShort, false, flagTSVUsage)
|
|
cmd.Flags().Duration(flagPingTimeout, time.Second*10, flagPingTimeoutUsage)
|
|
cmd.Flags().BoolP(flagPingAll, flagPingAllShort, false, flagPingAllUsage)
|
|
|
|
return cmd
|
|
}
|
|
|
|
func execPing(cmd *cobra.Command, args []string) error {
|
|
rc := RunContextFrom(cmd.Context())
|
|
cfg := rc.Config
|
|
var srcs []*source.Source
|
|
|
|
// args can be:
|
|
// [empty] : ping active source
|
|
// @handle1 @handleN: ping multiple sources
|
|
|
|
var pingAll bool
|
|
if cmd.Flags().Changed(flagPingAll) {
|
|
pingAll, _ = cmd.Flags().GetBool(flagPingAll)
|
|
}
|
|
|
|
switch {
|
|
case pingAll:
|
|
srcs = cfg.Sources.Items()
|
|
case len(args) == 0:
|
|
src := cfg.Sources.Active()
|
|
if src == nil {
|
|
return errz.New(msgNoActiveSrc)
|
|
}
|
|
srcs = []*source.Source{src}
|
|
default:
|
|
for _, arg := range args {
|
|
err := source.VerifyLegalHandle(arg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
src, err := cfg.Sources.Get(arg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
srcs = append(srcs, src)
|
|
}
|
|
}
|
|
|
|
timeout := cfg.Defaults.PingTimeout
|
|
if cmdFlagChanged(cmd, flagPingTimeout) {
|
|
timeout, _ = cmd.Flags().GetDuration(flagPingTimeout)
|
|
}
|
|
|
|
rc.Log.Debugf("Using ping timeout value: %s", timeout)
|
|
|
|
return pingSources(cmd.Context(), rc.Log, rc.registry, srcs, rc.writers.pingw, timeout)
|
|
}
|
|
|
|
// pingSources pings each of the sources in srcs, and prints results
|
|
// to w. If any error occurs pinging any of srcs, that error is printed
|
|
// inline as part of the ping results, and an errNoMsg is returned.
|
|
//
|
|
// NOTE: This ping code has an ancient lineage, in that it was written
|
|
//
|
|
// originally laid down before context.Context was a thing. Thus,
|
|
// the entire thing could probably be rewritten for simplicity.
|
|
func pingSources(ctx context.Context, log lg.Log, dp driver.Provider, srcs []*source.Source, w output.PingWriter, timeout time.Duration) error {
|
|
w.Open(srcs)
|
|
defer log.WarnIfFuncError(w.Close)
|
|
|
|
resultCh := make(chan pingResult, len(srcs))
|
|
|
|
// pingErrExists is set to true if there was an error for
|
|
// any of the pings. This later determines if an error
|
|
// is returned from this func.
|
|
var pingErrExists bool
|
|
|
|
for _, src := range srcs {
|
|
go pingSource(ctx, dp, src, timeout, resultCh)
|
|
}
|
|
|
|
// This func doesn't check for context.Canceled itself; instead
|
|
// it checks if any of the goroutines return that value on
|
|
// resultCh.
|
|
for i := 0; i < len(srcs); i++ {
|
|
result := <-resultCh
|
|
|
|
switch {
|
|
case result.err == context.Canceled:
|
|
// If any one of the goroutines have received context.Canceled,
|
|
// then we'll bubble that up and ignore the remaining goroutines.
|
|
return context.Canceled
|
|
|
|
case result.err == context.DeadlineExceeded:
|
|
// If timeout occurred, set the duration to timeout.
|
|
result.duration = timeout
|
|
pingErrExists = true
|
|
|
|
case result.err != nil:
|
|
pingErrExists = true
|
|
}
|
|
|
|
w.Result(result.src, result.duration, result.err)
|
|
}
|
|
|
|
// If there's at least one error, we return the
|
|
// sentinel errNoMsg so that sq can os.Exit(1) without printing
|
|
// an additional error message (as the error message will already have
|
|
// been printed by PingWriter).
|
|
if pingErrExists {
|
|
return errNoMsg
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// pingSource pings an individual driver.Source. It always returns a
|
|
// result on resultCh, even when ctx is done.
|
|
func pingSource(ctx context.Context, dp driver.Provider, src *source.Source, timeout time.Duration, resultCh chan<- pingResult) {
|
|
drvr, err := dp.DriverFor(src.Type)
|
|
if err != nil {
|
|
resultCh <- pingResult{src: src, err: err}
|
|
return
|
|
}
|
|
|
|
if timeout > 0 {
|
|
var cancelFn context.CancelFunc
|
|
ctx, cancelFn = context.WithTimeout(ctx, timeout)
|
|
defer cancelFn()
|
|
}
|
|
|
|
doneCh := make(chan pingResult)
|
|
start := time.Now()
|
|
|
|
go func() {
|
|
err = drvr.Ping(ctx, src)
|
|
doneCh <- pingResult{src: src, duration: time.Since(start), err: err}
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
resultCh <- pingResult{src: src, err: ctx.Err()}
|
|
case result := <-doneCh:
|
|
resultCh <- result
|
|
}
|
|
}
|
|
|
|
type pingResult struct {
|
|
src *source.Source
|
|
duration time.Duration
|
|
err error
|
|
}
|