[cobra] vtgate and vttablet (#13943)

Signed-off-by: Andrew Mason <andrew@planetscale.com>

go/cmd/vtgate/cli/cli.go

@@ -0,0 +1,193 @@
+/*
+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 agreedto 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 (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	"vitess.io/vitess/go/acl"
+	"vitess.io/vitess/go/exit"
+	"vitess.io/vitess/go/vt/discovery"
+	"vitess.io/vitess/go/vt/servenv"
+	"vitess.io/vitess/go/vt/srvtopo"
+	"vitess.io/vitess/go/vt/topo"
+	"vitess.io/vitess/go/vt/topo/topoproto"
+	"vitess.io/vitess/go/vt/vterrors"
+	"vitess.io/vitess/go/vt/vtgate"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
+
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	"vitess.io/vitess/go/vt/proto/vtrpc"
+)
+
+var (
+	cell              string
+	tabletTypesToWait []topodatapb.TabletType
+	plannerName       string
+	resilientServer   *srvtopo.ResilientServer
+
+	Main = &cobra.Command{
+		Use:   "vtgate",
+		Short: "VTGate is a stateless proxy responsible for accepting requests from applications and routing them to the appropriate tablet server(s) for query execution. It speaks both the MySQL Protocol and a gRPC protocol.",
+		Long: `VTGate is a stateless proxy responsible for accepting requests from applications and routing them to the appropriate tablet server(s) for query execution. It speaks both the MySQL Protocol and a gRPC protocol.
+
+### Key Options
+` +
+			"\n* `--srv_topo_cache_ttl`: There may be instances where you will need to increase the cached TTL from the default of 1 second to a higher number:\n" +
+			`	* You may want to increase this option if you see that your topo leader goes down and keeps your queries waiting for a few seconds.`,
+		Example: `vtgate \
+	--topo_implementation etcd2 \
+	--topo_global_server_address localhost:2379 \
+	--topo_global_root /vitess/global \
+	--log_dir $VTDATAROOT/tmp \
+	--port 15001 \
+	--grpc_port 15991 \
+	--mysql_server_port 15306 \
+	--cell test \
+	--cells_to_watch test \
+	--tablet_types_to_wait PRIMARY,REPLICA \
+	--service_map 'grpc-vtgateservice' \
+	--pid_file $VTDATAROOT/tmp/vtgate.pid \
+	--mysql_auth_server_impl none`,
+		Args:    cobra.NoArgs,
+		Version: servenv.AppVersion.String(),
+		PreRunE: servenv.CobraPreRunE,
+		RunE:    run,
+	}
+)
+
+// CheckCellFlags will check validation of cell and cells_to_watch flag
+// it will help to avoid strange behaviors when vtgate runs but actually does not work
+func CheckCellFlags(ctx context.Context, serv srvtopo.Server, cell string, cellsToWatch string) error {
+	// topo check
+	var topoServer *topo.Server
+	if serv != nil {
+		var err error
+		topoServer, err = serv.GetTopoServer()
+		if err != nil {
+			return fmt.Errorf("Unable to create gateway: %w", err)
+		}
+	} else {
+		return fmt.Errorf("topo server cannot be nil")
+	}
+	cellsInTopo, err := topoServer.GetKnownCells(ctx)
+	if err != nil {
+		return err
+	}
+	if len(cellsInTopo) == 0 {
+		return vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "topo server should have at least one cell")
+	}
+
+	// cell valid check
+	if cell == "" {
+		return vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "cell flag must be set")
+	}
+	hasCell := false
+	for _, v := range cellsInTopo {
+		if v == cell {
+			hasCell = true
+			break
+		}
+	}
+	if !hasCell {
+		return vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "cell:[%v] does not exist in topo", cell)
+	}
+
+	// cells_to_watch valid check
+	cells := make([]string, 0, 1)
+	for _, c := range strings.Split(cellsToWatch, ",") {
+		if c == "" {
+			continue
+		}
+		// cell should contained in cellsInTopo
+		if exists := topo.InCellList(c, cellsInTopo); !exists {
+			return vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "cell: [%v] is not valid. Available cells: [%v]", c, strings.Join(cellsInTopo, ","))
+		}
+		cells = append(cells, c)
+	}
+	if len(cells) == 0 {
+		return vterrors.Errorf(vtrpc.Code_INVALID_ARGUMENT, "cells_to_watch flag cannot be empty")
+	}
+
+	return nil
+}
+
+func run(cmd *cobra.Command, args []string) error {
+	defer exit.Recover()
+
+	servenv.Init()
+	defer servenv.Close()
+
+	ts := topo.Open()
+	defer ts.Close()
+
+	resilientServer = srvtopo.NewResilientServer(context.Background(), ts, "ResilientSrvTopoServer")
+
+	tabletTypes := make([]topodatapb.TabletType, 0, 1)
+	for _, tt := range tabletTypesToWait {
+		if topoproto.IsServingType(tt) {
+			tabletTypes = append(tabletTypes, tt)
+		}
+	}
+
+	if len(tabletTypes) == 0 {
+		return fmt.Errorf("tablet_types_to_wait must contain at least one serving tablet type")
+	}
+
+	err := CheckCellFlags(context.Background(), resilientServer, cell, vtgate.CellsToWatch)
+	if err != nil {
+		return fmt.Errorf("cells_to_watch validation failed: %v", err)
+	}
+
+	plannerVersion, _ := plancontext.PlannerNameToVersion(plannerName)
+
+	// pass nil for HealthCheck and it will be created
+	vtg := vtgate.Init(context.Background(), nil, resilientServer, cell, tabletTypes, plannerVersion)
+
+	servenv.OnRun(func() {
+		// Flags are parsed now. Parse the template using the actual flag value and overwrite the current template.
+		discovery.ParseTabletURLTemplateFromFlag()
+		addStatusParts(vtg)
+	})
+	servenv.OnClose(func() {
+		_ = vtg.Gateway().Close(context.Background())
+	})
+	servenv.RunDefault()
+
+	return nil
+}
+
+func init() {
+	servenv.RegisterDefaultFlags()
+	servenv.RegisterFlags()
+	servenv.RegisterGRPCServerFlags()
+	servenv.RegisterGRPCServerAuthFlags()
+	servenv.RegisterServiceMapFlag()
+
+	servenv.MoveFlagsToCobraCommand(Main)
+
+	acl.RegisterFlags(Main.Flags())
+	Main.Flags().StringVar(&cell, "cell", cell, "cell to use")
+	Main.Flags().Var((*topoproto.TabletTypeListFlag)(&tabletTypesToWait), "tablet_types_to_wait", "Wait till connected for specified tablet types during Gateway initialization. Should be provided as a comma-separated set of tablet types.")
+	Main.Flags().StringVar(&plannerName, "planner-version", plannerName, "Sets the default planner to use when the session has not changed it. Valid values are: Gen4, Gen4Greedy, Gen4Left2Right")
+
+	Main.MarkFlagRequired("tablet_types_to_wait")
+}

go/cmd/vtgate/plugin_auth_clientcert.go → go/cmd/vtgate/cli/plugin_auth_clientcert.go

@@ -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 clientcert to register the client certificate implementation of AuthServer.
 

go/cmd/vtgate/plugin_auth_ldap.go → go/cmd/vtgate/cli/plugin_auth_ldap.go

@@ -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 ldapauthserver to register the LDAP implementation of AuthServer.
 

go/cmd/vtgate/plugin_auth_static.go → go/cmd/vtgate/cli/plugin_auth_static.go

@@ -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 staticauthserver to register the flat-file implementation of AuthServer.
 

go/cmd/vtgate/plugin_auth_vault.go → go/cmd/vtgate/cli/plugin_auth_vault.go

@@ -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 InitAuthServerVault to register the HashiCorp Vault implementation of AuthServer.
 

go/cmd/vtgate/plugin_consultopo.go → go/cmd/vtgate/cli/plugin_consultopo.go

@@ -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.
 

go/cmd/vtgate/plugin_etcd2topo.go → go/cmd/vtgate/cli/plugin_etcd2topo.go

@@ -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.
 

go/cmd/vttablet/plugin_grpctabletconn.go → go/cmd/vtgate/cli/plugin_grpctabletconn.go

@@ -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 tabletconn client
 

go/cmd/vtgate/plugin_grpcvtgateservice.go → go/cmd/vtgate/cli/plugin_grpcvtgateservice.go

@@ -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 vtgateservice server
 

go/cmd/vtgate/plugin_opentracing.go → go/cmd/vtgate/cli/plugin_opentracing.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package main
+package cli
 
 import (
 	"vitess.io/vitess/go/trace"

go/cmd/vtgate/plugin_opentsdb.go → go/cmd/vtgate/cli/plugin_opentsdb.go

@@ -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 opentsdb to register the opentsdb stats backend.
 

go/cmd/vtgate/plugin_prometheusbackend.go → go/cmd/vtgate/cli/plugin_prometheusbackend.go

@@ -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

go/cmd/vtgate/cli/plugin_statsd.go

@@ -0,0 +1,23 @@
+/*
+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 "vitess.io/vitess/go/stats/statsd"
+
+func init() {
+	statsd.Init("vtgate")
+}

go/cmd/vtgate/plugin_zk2topo.go → go/cmd/vtgate/cli/plugin_zk2topo.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package main
+package cli
 
 import (
 	// Imports and register the zk2 TopologyServer

go/cmd/vtgate/status.go → go/cmd/vtgate/cli/status.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package main
+package cli
 
 import (
 	"vitess.io/vitess/go/vt/discovery"

go/cmd/vtgate/docgen/main.go

@@ -0,0 +1,42 @@
+/*
+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"
+	"github.com/spf13/pflag"
+
+	"vitess.io/vitess/go/cmd/internal/docgen"
+	"vitess.io/vitess/go/cmd/vtgate/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)
+		},
+	}
+
+	// Here because we inadvertently transfer the required "tablet-types-to-wait"
+	// flag during vtgate/cli's init func.
+	pflag.CommandLine = cmd.Flags()
+
+	cmd.Flags().StringVarP(&dir, "dir", "d", "doc", "output directory to write documentation")
+	_ = cmd.Execute()
+}