-
Notifications
You must be signed in to change notification settings - Fork 0
/
cobrautils.go
123 lines (104 loc) · 2.94 KB
/
cobrautils.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package cobrautils
import (
"context"
"io"
"os"
"strings"
"sync"
"github.com/cockroachdb/errors"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
)
var (
stdoutCaptureMu sync.Mutex
)
// InitRootCmd sets up the root command for the shell and calls
// some of the same functions that would run behind the scenes for
// `rootCmd.Execute()`.
//
// These are needed to do stuff like add the internal autocomplete command
func InitRootCmd(rootCmd *cobra.Command) {
addDefaultShellCommands(rootCmd)
rootCmd.SilenceErrors = true
rootCmd.CompletionOptions.HiddenDefaultCmd = true
rootCmd.Hidden = true
for _, subCmd := range rootCmd.Commands() {
subCmd.SilenceUsage = true
}
rootCmd.InitDefaultHelpCmd()
rootCmd.InitDefaultCompletionCmd()
}
// ExecuteCmd executes a command and captures the output
func ExecuteCmd(ctx context.Context, rootCmd *cobra.Command, line string, in io.Reader, stdout io.Writer, stderr io.Writer) error {
// Capture stdout and stderr and then restore them when we leave here
stdoutCaptureMu.Lock()
originalStdOut := os.Stdout
originalStdErr := os.Stderr
defer func() {
os.Stdout = originalStdOut
os.Stderr = originalStdErr
stdoutCaptureMu.Unlock()
}()
// Set up our stdout and stderr pipes
stdoutR, stdoutW, _ := os.Pipe()
stderrR, stderrW, _ := os.Pipe()
os.Stdout = stdoutW
os.Stderr = stderrW
// Copy the output to the correct places
var wg sync.WaitGroup
wg.Add(2)
go func() {
_, _ = io.Copy(stdout, stdoutR)
wg.Done()
}()
go func() {
_, _ = io.Copy(stderr, stderrR)
wg.Done()
}()
// Reset the internal state of the command
args := strings.Split(line, " ")
if cmd, _, err := rootCmd.Find(args); err == nil {
// Reset the context to nil
cmd.SetContext(nil)
// Reset flag values between runs due to a limitation in Cobra
cmd.Flags().VisitAll(func(flag *pflag.Flag) {
if val, ok := flag.Value.(pflag.SliceValue); ok {
_ = val.Replace([]string{})
} else {
_ = flag.Value.Set(flag.DefValue)
}
})
cmd.InitDefaultHelpFlag()
cmd.InitDefaultVersionFlag()
}
// Setup the cobra command to output to the write places
rootCmd.SetIn(in)
rootCmd.SetOut(stdoutW)
rootCmd.SetErr(stderrW)
rootCmd.SetContext(ctx)
rootCmd.SetArgs(args)
// Finally execute it!
err := rootCmd.Execute()
// Close the pipes for all the IO copies to finish
_ = stderrW.Close()
_ = stdoutW.Close()
wg.Wait()
return errors.WithStack(err)
}
// addDefaultShellCommands adds the default shell commands to the root command
func addDefaultShellCommands(rootCmd *cobra.Command) {
rootCmd.AddCommand(&cobra.Command{
Use: "exit",
Short: "Exit the shell",
Aliases: []string{"quit"},
Args: cobra.NoArgs,
RunE: func(cmd *cobra.Command, args []string) error {
// no-op which is needed otherwise Cobra won't list the help for it
// when the user runs `help exit`
//
// However the shell will intercept this command and exit the shell
// before running this function
return nil
},
})
}