Skip to content

Commit 7b93c52

Browse files
vecerajveceraj
committed
update speed-test to cobra
1 parent 2e3b32c commit 7b93c52

File tree

1 file changed

+146
-48
lines changed

1 file changed

+146
-48
lines changed

‎cmd/pbm-speed-test/main.go

Lines changed: 146 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -7,76 +7,160 @@ import (
77
"os"
88
"time"
99

10-
"github.com/alecthomas/kingpin"
10+
"github.com/spf13/cobra"
11+
"github.com/spf13/viper"
1112
"go.mongodb.org/mongo-driver/mongo"
1213

1314
"github.com/percona/percona-backup-mongodb/pbm/compress"
1415
"github.com/percona/percona-backup-mongodb/pbm/connect"
16+
"github.com/percona/percona-backup-mongodb/pbm/errors"
1517
"github.com/percona/percona-backup-mongodb/pbm/log"
1618
"github.com/percona/percona-backup-mongodb/pbm/storage/blackhole"
1719
"github.com/percona/percona-backup-mongodb/pbm/util"
1820
"github.com/percona/percona-backup-mongodb/pbm/version"
1921
)
2022

23+
type rootOpts struct {
24+
mURL string
25+
sampleColF string
26+
sampleSizeF float64
27+
compression string
28+
compressLevelArg []int
29+
}
30+
31+
var validCompressions = []string{
32+
string(compress.CompressionTypeNone),
33+
string(compress.CompressionTypeGZIP),
34+
string(compress.CompressionTypeSNAPPY),
35+
string(compress.CompressionTypeLZ4),
36+
string(compress.CompressionTypeS2),
37+
string(compress.CompressionTypePGZIP),
38+
string(compress.CompressionTypeZstandard),
39+
}
40+
2141
func main() {
22-
var (
23-
tCmd = kingpin.New("pbm-speed-test", "Percona Backup for MongoDB compression and upload speed test")
24-
mURL = tCmd.Flag("mongodb-uri", "MongoDB connection string").Envar("PBM_MONGODB_URI").String()
25-
sampleColF = tCmd.Flag("sample-collection", "Set collection as the data source").Short('c').String()
26-
sampleSizeF = tCmd.Flag("size-gb", "Set data size in GB. Default 1").Short('s').Float64()
27-
28-
compressLevelArg []int
29-
compressLevel *int
30-
31-
compressType = tCmd.Flag("compression", "Compression type <none>/<gzip>/<snappy>/<lz4>/<s2>/<pgzip>/<zstd>").
32-
Default(string(compress.CompressionTypeS2)).
33-
Enum(string(compress.CompressionTypeNone), string(compress.CompressionTypeGZIP),
34-
string(compress.CompressionTypeSNAPPY), string(compress.CompressionTypeLZ4),
35-
string(compress.CompressionTypeS2), string(compress.CompressionTypePGZIP),
36-
string(compress.CompressionTypeZstandard),
37-
)
42+
rootOptions := rootOpts{}
43+
rootCmd := &cobra.Command{
44+
Use: "pbm-speed-test",
45+
Short: "Percona Backup for MongoDB compression and upload speed test",
46+
SilenceUsage: true,
47+
}
3848

39-
compressionCmd = tCmd.Command("compression", "Run compression test")
40-
storageCmd = tCmd.Command("storage", "Run storage test")
49+
rootCmd.PersistentFlags().StringVar(&rootOptions.mURL, "mongodb-uri", "", "MongoDB connection string")
50+
_ = viper.BindPFlag("mongodb-uri", rootCmd.PersistentFlags().Lookup("mongodb-uri"))
51+
_ = viper.BindEnv("mongodb-uri", "PBM_MONGODB_URI")
4152

42-
versionCmd = tCmd.Command("version", "PBM version info")
43-
versionShort = versionCmd.Flag("short", "Only version info").Default("false").Bool()
44-
versionCommit = versionCmd.Flag("commit", "Only git commit info").Default("false").Bool()
45-
versionFormat = versionCmd.Flag("format", "Output format <json or \"\">").Default("").String()
53+
rootCmd.PersistentFlags().StringVarP(
54+
&rootOptions.sampleColF, "sample-collection", "c", "",
55+
"Set collection as the data source",
56+
)
57+
rootCmd.PersistentFlags().Float64VarP(
58+
&rootOptions.sampleSizeF, "size-gb", "s", 1,
59+
"Set data size in GB. Default 1",
60+
)
61+
rootCmd.PersistentFlags().StringVar(
62+
&rootOptions.compression, "compression", "",
63+
"Compression type <none>/<gzip>/<snappy>/<lz4>/<s2>/<pgzip>/<zstd>",
64+
)
65+
rootCmd.PersistentFlags().IntSliceVar(
66+
&rootOptions.compressLevelArg, "compression-level", nil, "Compression level (specific to the compression type)",
4667
)
4768

48-
tCmd.Flag("compression-level", "Compression level (specific to the compression type)").IntsVar(&compressLevelArg)
69+
compressionCmd := &cobra.Command{
70+
Use: "compression",
71+
Short: "Run compression test",
72+
RunE: func(cmd *cobra.Command, args []string) error {
73+
if rootOptions.mURL == "" {
74+
rootOptions.mURL = viper.GetString("mongodb-uri")
75+
}
76+
77+
if err := validateEnum("compression", rootOptions.compression, validCompressions); err != nil {
78+
return err
79+
}
80+
81+
var compressLevel *int
82+
83+
if len(rootOptions.compressLevelArg) > 0 {
84+
compressLevel = &rootOptions.compressLevelArg[0]
85+
}
86+
87+
cmd.Print("Test started ")
88+
testCompression(
89+
rootOptions.mURL,
90+
compress.CompressionType(rootOptions.compression),
91+
compressLevel,
92+
rootOptions.sampleSizeF,
93+
rootOptions.sampleColF,
94+
)
4995

50-
if len(compressLevelArg) > 0 {
51-
compressLevel = &compressLevelArg[0]
96+
return nil
97+
},
5298
}
5399

54-
cmd, err := tCmd.DefaultEnvars().Parse(os.Args[1:])
55-
if err != nil && cmd != versionCmd.FullCommand() {
56-
stdlog.Println("Error: Parse command line parameters:", err)
57-
return
100+
rootCmd.AddCommand(compressionCmd)
101+
102+
storageCmd := &cobra.Command{
103+
Use: "storage",
104+
Short: "Run storage test",
105+
RunE: func(cmd *cobra.Command, args []string) error {
106+
if rootOptions.mURL == "" {
107+
rootOptions.mURL = viper.GetString("mongodb-uri")
108+
}
109+
110+
if err := validateEnum("compression", rootOptions.compression, validCompressions); err != nil {
111+
return err
112+
}
113+
114+
var compressLevel *int
115+
116+
if len(rootOptions.compressLevelArg) > 0 {
117+
compressLevel = &rootOptions.compressLevelArg[0]
118+
}
119+
120+
cmd.Print("Test started ")
121+
testStorage(
122+
rootOptions.mURL,
123+
compress.CompressionType(rootOptions.compression),
124+
compressLevel,
125+
rootOptions.sampleSizeF,
126+
rootOptions.sampleColF,
127+
)
128+
129+
return nil
130+
},
58131
}
59132

60-
if *sampleSizeF == 0 {
61-
*sampleSizeF = 1
133+
rootCmd.AddCommand(storageCmd)
134+
135+
var (
136+
versionShort bool
137+
versionCommit bool
138+
versionFormat string
139+
)
140+
141+
versionCmd := &cobra.Command{
142+
Use: "version",
143+
Short: "PBM version info",
144+
Run: func(cmd *cobra.Command, args []string) {
145+
switch {
146+
case versionShort:
147+
cmd.Println(version.Current().Short())
148+
case versionCommit:
149+
cmd.Println(version.Current().GitCommit)
150+
default:
151+
cmd.Println(version.Current().All(versionFormat))
152+
}
153+
},
62154
}
63155

64-
switch cmd {
65-
case compressionCmd.FullCommand():
66-
fmt.Print("Test started ")
67-
testCompression(*mURL, compress.CompressionType(*compressType), compressLevel, *sampleSizeF, *sampleColF)
68-
case storageCmd.FullCommand():
69-
fmt.Print("Test started ")
70-
testStorage(*mURL, compress.CompressionType(*compressType), compressLevel, *sampleSizeF, *sampleColF)
71-
case versionCmd.FullCommand():
72-
switch {
73-
case *versionCommit:
74-
fmt.Println(version.Current().GitCommit)
75-
case *versionShort:
76-
fmt.Println(version.Current().Short())
77-
default:
78-
fmt.Println(version.Current().All(*versionFormat))
79-
}
156+
versionCmd.Flags().BoolVar(&versionShort, "short", false, "Only version info")
157+
versionCmd.Flags().BoolVar(&versionCommit, "commit", false, "Only git commit info")
158+
versionCmd.Flags().StringVar(&versionFormat, "format", "", "Output format <json or \"\">")
159+
160+
rootCmd.AddCommand(versionCmd)
161+
162+
if err := rootCmd.Execute(); err != nil {
163+
os.Exit(1)
80164
}
81165
}
82166

@@ -148,3 +232,17 @@ func printw(done <-chan struct{}) {
148232
}
149233
}
150234
}
235+
236+
func validateEnum(fieldName, value string, valid []string) error {
237+
if value == "" {
238+
return nil
239+
}
240+
241+
for _, validItem := range valid {
242+
if value == validItem {
243+
return nil
244+
}
245+
}
246+
247+
return errors.New(fmt.Sprintf("invalid %s value: %q (must be one of %v)", fieldName, value, valid))
248+
}

0 commit comments

Comments
 (0)