Skip to content

Commit

Permalink
migrate vtorc to use cobra commands (#13917)
Browse files Browse the repository at this point in the history
Co-authored-by: Florent Poinsard <35779988+frouioui@users.noreply.github.com>
  • Loading branch information
Andrew Mason and frouioui authored Sep 6, 2023
1 parent 04ad6ec commit 86c1d23
Show file tree
Hide file tree
Showing 12 changed files with 172 additions and 196 deletions.
118 changes: 118 additions & 0 deletions go/cmd/vtorc/cli/cli.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
/*
Copyright 2023 The Vitess Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cli

import (
"flag"
"fmt"

"github.com/spf13/cobra"

"vitess.io/vitess/go/acl"
_flag "vitess.io/vitess/go/internal/flag"
"vitess.io/vitess/go/viperutil"
viperdebug "vitess.io/vitess/go/viperutil/debug"
"vitess.io/vitess/go/vt/log"
"vitess.io/vitess/go/vt/servenv"
"vitess.io/vitess/go/vt/vtorc/config"
"vitess.io/vitess/go/vt/vtorc/inst"
"vitess.io/vitess/go/vt/vtorc/logic"
"vitess.io/vitess/go/vt/vtorc/server"
)

var (
configFile string
Main = &cobra.Command{
Use: "vtorc",
Short: "", // TODO
Args: cobra.NoArgs,
Version: servenv.AppVersion.String(),
PreRunE: func(cmd *cobra.Command, args []string) error {
_flag.TrickGlog()

watchCancel, err := viperutil.LoadConfig()
if err != nil {
return fmt.Errorf("%s: failed to read in config: %s", cmd.Name(), err)
}

servenv.OnTerm(watchCancel)
servenv.HTTPHandleFunc("/debug/config", viperdebug.HandlerFunc)
return nil
},
Run: run,
}
)

func run(cmd *cobra.Command, args []string) {
servenv.Init()
config.UpdateConfigValuesFromFlags()
inst.RegisterStats()

log.Info("starting vtorc")
if len(configFile) > 0 {
config.ForceRead(configFile)
} else {
config.Read("/etc/vtorc.conf.json", "conf/vtorc.conf.json", "vtorc.conf.json")
}
if config.Config.AuditToSyslog {
inst.EnableAuditSyslog()
}
config.MarkConfigurationLoaded()

// Log final config values to debug if something goes wrong.
config.LogConfigValues()
server.StartVTOrcDiscovery()

server.RegisterVTOrcAPIEndpoints()
servenv.OnRun(func() {
addStatusParts()
})

// For backward compatability, we require that VTOrc functions even when the --port flag is not provided.
// In this case, it should function like before but without the servenv pages.
// Therefore, currently we don't check for the --port flag to be necessary, but release 16+ that check
// can be added to always have the serenv page running in VTOrc.
servenv.RunDefault()
}

// addStatusParts adds UI parts to the /debug/status page of VTOrc
func addStatusParts() {
servenv.AddStatusPart("Recent Recoveries", logic.TopologyRecoveriesTemplate, func() any {
recoveries, _ := logic.ReadRecentRecoveries(false, 0)
return recoveries
})
}

func init() {
servenv.RegisterDefaultFlags()
servenv.RegisterFlags()

Main.Flags().AddFlagSet(servenv.GetFlagSetFor("vtorc"))

// glog flags, no better way to do this
_flag.PreventGlogVFlagFromClobberingVersionFlagShorthand(Main.Flags())
Main.Flags().AddGoFlag(flag.Lookup("logtostderr"))
Main.Flags().AddGoFlag(flag.Lookup("alsologtostderr"))
Main.Flags().AddGoFlag(flag.Lookup("stderrthreshold"))
Main.Flags().AddGoFlag(flag.Lookup("log_dir"))

logic.RegisterFlags(Main.Flags())
server.RegisterFlags(Main.Flags())
config.RegisterFlags(Main.Flags())
acl.RegisterFlags(Main.Flags())
Main.Flags().StringVar(&configFile, "config", "", "config file name")
}
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
limitations under the License.
*/

package main
package cli

// This plugin imports consultopo to register the consul implementation of TopoServer.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
limitations under the License.
*/

package main
package cli

// This plugin imports etcd2topo to register the etcd2 implementation of TopoServer.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
limitations under the License.
*/

package main
package cli

// Imports and register the gRPC tabletmanager client

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
limitations under the License.
*/

package main
package cli

// This plugin imports Prometheus to allow for instrumentation
// with the Prometheus client library
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
limitations under the License.
*/

package main
package cli

// Imports and register the zk2 TopologyServer

Expand Down
37 changes: 37 additions & 0 deletions go/cmd/vtorc/docgen/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
/*
Copyright 2023 The Vitess Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
"github.com/spf13/cobra"

"vitess.io/vitess/go/cmd/internal/docgen"
"vitess.io/vitess/go/cmd/vtorc/cli"
)

func main() {
var dir string
cmd := cobra.Command{
Use: "docgen [-d <dir>]",
RunE: func(cmd *cobra.Command, args []string) error {
return docgen.GenerateMarkdownTree(cli.Main, dir)
},
}

cmd.Flags().StringVarP(&dir, "dir", "d", "doc", "output directory to write documentation")
_ = cmd.Execute()
}
110 changes: 4 additions & 106 deletions go/cmd/vtorc/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,120 +17,18 @@
package main

import (
"strings"

_ "github.com/go-sql-driver/mysql"
"github.com/spf13/pflag"
_ "modernc.org/sqlite"

"vitess.io/vitess/go/acl"
"vitess.io/vitess/go/cmd/vtorc/cli"
"vitess.io/vitess/go/vt/log"
"vitess.io/vitess/go/vt/servenv"
"vitess.io/vitess/go/vt/vtorc/config"
"vitess.io/vitess/go/vt/vtorc/inst"
"vitess.io/vitess/go/vt/vtorc/logic"
"vitess.io/vitess/go/vt/vtorc/server"
)

// transformArgsForPflag turns a slice of raw args passed on the command line,
// possibly incompatible with pflag (because the user is expecting stdlib flag
// parsing behavior) and transforms them into the arguments that should have
// been passed to conform to pflag parsing behavior.
//
// the primary function is to catch any cases where the user specified a longopt
// with only a single hyphen (e.g. `-myflag`) and correct it to be
// double-hyphenated.
//
// note that this transformation does _not_ actually validate the arguments; for
// example if the user specifies `--myflag`, but the FlagSet has no such flag
// defined, that will still appear in the returned result and will (correctly)
// cause a parse error later on in `main`, at which point the CLI usage will
// be printed.
//
// note also that this transformation is incomplete. pflag allows interspersing
// of flag and positional arguments, whereas stdlib flag does not. however, for
// vtorc specifically, with the exception of `vtorc help <topic>`, the CLI only
// consumes flag arguments (in other words, there are no supported subcommands),
// so this is a non-issue, and is not implemented here in order to make this
// function a bit simpler.
func transformArgsForPflag(fs *pflag.FlagSet, args []string) (result []string) {
for i, arg := range args {
switch {
case arg == "--":
// pflag stops parsing at `--`, so we're done transforming the CLI
// arguments. Just append everything remaining and be done.
result = append(result, args[i:]...)
return result
case strings.HasPrefix(arg, "--"):
// Long-hand flag. Append it and continue.
result = append(result, arg)
case strings.HasPrefix(arg, "-"):
// Most complex case. This is either:
// 1. A legacy long-hand flag that needs a double-dash (e.g. `-myflag` => `--myflag`).
// 2. One _or more_ pflag shortopts all shoved together (think `rm -rf` as `rm -r -f`).
//
// In the latter case, we don't need to do any transformations, but
// in the former, we do.
name := strings.SplitN(arg[1:], "=", 2)[0] // discard any potential value (`-myflag` and `-myflag=10` both have the name of `myflag`)
if fs.Lookup(name) != nil || name == "help" {
// Case 1: We have a long opt with this name, so we need to
// prepend an additional hyphen.
result = append(result, "-"+arg)
} else {
// Case 2: No transformation needed.
result = append(result, arg)
}
default:
// Just a flag argument. Nothing to transform.
result = append(result, arg)
}
}

return result
}

// main is the application's entry point. It will spawn an HTTP interface.
func main() {
servenv.RegisterDefaultFlags()
servenv.RegisterFlags()

var configFile string
servenv.OnParseFor("vtorc", func(fs *pflag.FlagSet) {
logic.RegisterFlags(fs)
server.RegisterFlags(fs)
config.RegisterFlags(fs)
acl.RegisterFlags(fs)
// TODO: viperutil.BindFlags()

fs.StringVar(&configFile, "config", "", "config file name")
})
servenv.ParseFlags("vtorc")
servenv.Init()
config.UpdateConfigValuesFromFlags()
inst.RegisterStats()

log.Info("starting vtorc")
if len(configFile) > 0 {
config.ForceRead(configFile)
} else {
config.Read("/etc/vtorc.conf.json", "conf/vtorc.conf.json", "vtorc.conf.json")
}
if config.Config.AuditToSyslog {
inst.EnableAuditSyslog()
if err := cli.Main.Execute(); err != nil {
log.Exit(err)
}
config.MarkConfigurationLoaded()

// Log final config values to debug if something goes wrong.
config.LogConfigValues()
server.StartVTOrcDiscovery()

server.RegisterVTOrcAPIEndpoints()
servenv.OnRun(func() {
addStatusParts()
})

// For backward compatability, we require that VTOrc functions even when the --port flag is not provided.
// In this case, it should function like before but without the servenv pages.
// Therefore, currently we don't check for the --port flag to be necessary, but release 16+ that check
// can be added to always have the serenv page running in VTOrc.
servenv.RunDefault()
}
48 changes: 0 additions & 48 deletions go/cmd/vtorc/main_test.go

This file was deleted.

Loading

0 comments on commit 86c1d23

Please sign in to comment.