vitessio · ajm188 · Sep 14, 2023 · Sep 7, 2023 · Sep 7, 2023 · Sep 8, 2023
@@ -0,0 +1,188 @@
+/*
+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.",
+ 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")
+}
@@ -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.
 

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

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

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

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

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

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

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

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

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

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

@@ -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")
+}
@@ -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

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

@@ -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/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)
+ },
+ }
+
+ cmd.Flags().StringVarP(&dir, "dir", "d", "doc", "output directory to write documentation")
+ _ = cmd.Execute()
+}