-
Notifications
You must be signed in to change notification settings - Fork 4.1k
Expand file tree
/
Copy pathutil.go
More file actions
246 lines (205 loc) · 6.75 KB
/
util.go
File metadata and controls
246 lines (205 loc) · 6.75 KB
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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
// Copyright 2023 The Cockroach Authors.
//
// Use of this software is governed by the CockroachDB Software License
// included in the /LICENSE file.
package cli
import (
"fmt"
"os"
"strings"
"text/tabwriter"
"time"
rperrors "github.com/cockroachdb/cockroach/pkg/roachprod/errors"
"github.com/cockroachdb/cockroach/pkg/roachprod/install"
"github.com/cockroachdb/cockroach/pkg/roachprod/vm"
"github.com/cockroachdb/cockroach/pkg/roachprod/vm/gce"
"github.com/cockroachdb/errors"
"github.com/cockroachdb/errors/oserror"
"github.com/spf13/cobra"
)
func PromptYesNo(msg string, defaultYes bool) bool {
if defaultYes {
fmt.Printf("%s y[default]/n: ", msg)
} else {
fmt.Printf("%s y/n[default]: ", msg)
}
var answer string
_, _ = fmt.Scanln(&answer)
answer = strings.TrimSpace(answer)
isYes := answer == "y" || answer == "Y"
isEmpty := answer == ""
if defaultYes {
return isYes || isEmpty
}
return isYes
}
// SwapBinary attempts to swap the `old` file with the `new` file. Used to
// update a running roachprod binary.
// Note: there is special handling if `new` points to a file ending in `.bak`.
// In this case, it is assumed to be a `revert` operation, in which case we
// do *not* backup the old/current file.
func SwapBinary(old, new string) error {
destInfo, err := os.Stat(new)
if err != nil {
if oserror.IsNotExist(err) {
return errors.WithDetail(err, "binary does not exist: "+new)
}
return err
}
if destInfo.IsDir() {
return errors.Newf("binary path is a directory, not a file: %s", new)
}
oldInfo, err := os.Stat(old)
if err != nil {
return err
}
// Copy the current file permissions to the new binary and ensure it is executable.
err = os.Chmod(new, oldInfo.Mode())
if err != nil {
return err
}
// Backup only for upgrading, not when reverting which is assumed if the new binary ends in `.bak`.
if !strings.HasSuffix(new, ".bak") {
// Backup the current binary, so that it may be restored via `roachprod update --revert`.
err = os.Rename(old, old+".bak")
if err != nil {
return errors.WithDetail(err, "unable to backup current binary")
}
}
// Move the new binary into place.
return os.Rename(new, old)
}
// Computes the age of the current binary, relative to the given update time.
func TimeSinceUpdate(updateTime time.Time) (time.Duration, error) {
currentBinary, err := os.Executable()
if err != nil {
return -1, err
}
statInfo, err := os.Stat(currentBinary)
if err != nil {
return -1, err
}
return updateTime.Sub(statInfo.ModTime()), nil
}
// Provide `cobra.Command` functions with a standard return code handler.
// Exit codes come from rperrors.Error.ExitCode().
//
// If the wrapped error tree of an error does not contain an instance of
// rperrors.Error, the error will automatically be wrapped with
// rperrors.Unclassified.
func Wrap(f func(cmd *cobra.Command, args []string) error) func(cmd *cobra.Command, args []string) {
return func(cmd *cobra.Command, args []string) {
var err error
isSecure, err = isSecureCluster(cmd)
if err != nil {
cmd.Printf("Error: %v\n", err)
os.Exit(1)
}
err = f(cmd, args)
if err != nil {
roachprodError, ok := rperrors.AsError(err)
if !ok {
roachprodError = rperrors.Unclassified{Err: err}
err = roachprodError
}
cmd.Printf("Error: %+v\n", err)
os.Exit(roachprodError.ExitCode())
}
}
}
func isSecureCluster(cmd *cobra.Command) (install.ComplexSecureOption, error) {
hasSecureFlag := cmd.Flags().Changed("secure")
hasInsecureFlag := cmd.Flags().Changed("insecure")
switch {
case hasSecureFlag && hasInsecureFlag:
// Disallow passing both flags, even if they are consistent.
return install.ComplexSecureOption{}, fmt.Errorf("cannot pass both --secure and --insecure flags")
case hasSecureFlag:
return install.ComplexSecureOption{ForcedSecure: true}, nil
case hasInsecureFlag:
return install.ComplexSecureOption{ForcedInsecure: true}, nil
case insecureEnvSet:
// If COCKROACH_ROACHPROD_INSECURE env var was explicitly set, treat it
// as a forced setting that takes precedence over ephemeral project defaults.
if insecure {
return install.ComplexSecureOption{ForcedInsecure: true}, nil
}
return install.ComplexSecureOption{ForcedSecure: true}, nil
default:
return install.ComplexSecureOption{DefaultSecure: !insecure}, nil
}
}
func printPublicKeyTable(keys gce.AuthorizedKeys, includeSize bool) error {
// Align columns left and separate with at least two spaces.
tw := tabwriter.NewWriter(os.Stdout, 0, 8, 2, ' ', 0)
fmt.Fprintf(tw, "%s\t%s\n", "User", "Key")
for _, ak := range keys {
fmt.Fprintf(tw, "%s\t%s\n", ak.User, ak.Format(64 /* maxLen */))
}
err := tw.Flush()
if !includeSize {
return err
}
const maxProjectMetadataBytes = 262144 /* 256 KiB */
metadataLen := len(keys.AsProjectMetadata())
usage := int(float64(metadataLen*100) / float64(maxProjectMetadataBytes))
_, err = fmt.Printf("\nTOTAL: %d bytes (usage: %d%%)\n", metadataLen, usage)
return err
}
// addHelpAboutNodes adds help about nodes to each of the commands
func addHelpAboutNodes(cmd *cobra.Command) {
// Add help about specifying nodes
if cmd.Long == "" {
cmd.Long = cmd.Short
}
cmd.Long += fmt.Sprintf(`
Node specification
By default the operation is performed on all nodes in <cluster>. A subset of
nodes can be specified by appending :<nodes> to the cluster name. The syntax
of <nodes> is a comma separated list of specific node IDs or range of
IDs. For example:
roachprod %[1]s marc-test:1-3,8-9
will perform %[1]s on:
marc-test-1
marc-test-2
marc-test-3
marc-test-8
marc-test-9
`, cmd.Name())
}
// Before executing any command, validate and canonicalize args.
func ValidateAndConfigure(cmd *cobra.Command, args []string) {
// Skip validation for commands that are self-sufficient.
switch cmd.Name() {
case "help", "version", "list":
return
}
printErrAndExit := func(err error) {
if err != nil {
fmt.Fprintf(os.Stderr, "%s\n", err)
os.Exit(1)
}
}
// Validate architecture flag, if set.
if archOpt := cmd.Flags().Lookup("arch"); archOpt != nil && archOpt.Changed {
arch := vm.CPUArch(strings.ToLower(archOpt.Value.String()))
if arch != vm.ArchAMD64 && arch != vm.ArchARM64 && arch != vm.ArchFIPS && arch != vm.ArchS390x {
printErrAndExit(fmt.Errorf("unsupported architecture %q", arch))
}
if string(arch) != archOpt.Value.String() {
// Set the canonical value.
_ = cmd.Flags().Set("arch", string(arch))
}
}
// Validate cloud providers, if set.
providersSet := make(map[string]struct{})
for _, p := range createVMOpts.VMProviders {
if _, ok := vm.Providers[p]; !ok {
printErrAndExit(fmt.Errorf("unknown cloud provider %q", p))
}
if _, ok := providersSet[p]; ok {
printErrAndExit(fmt.Errorf("duplicate cloud provider specified %q", p))
}
providersSet[p] = struct{}{}
}
}