Add initialize_with_vt_dba_tcp flag to enable TCP/IP connection access to the underlying MySQL instance (#15354)

Signed-off-by: Yohei Yoshimuta <yoheimuta@gmail.com>

Author: yoheimuta

go/cmd/vttestserver/cli/main.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
+	"path"
 	"strconv"
 	"strings"
 	"syscall"
@@ -47,13 +48,14 @@ type topoFlags struct {
 }
 
 var (
-	basePort  int
-	config    vttest.Config
-	doSeed    bool
-	mycnf     string
-	protoTopo string
-	seed      vttest.SeedConfig
-	topo      topoFlags
+	basePort        int
+	config          vttest.Config
+	doSeed          bool
+	mycnf           string
+	protoTopo       string
+	seed            vttest.SeedConfig
+	topo            topoFlags
+	doCreateTCPUser bool
 )
 
 func (t *topoFlags) buildTopology() (*vttestpb.VTTestTopology, error) {
@@ -220,25 +222,44 @@ func New() (cmd *cobra.Command) {
 	cmd.Flags().StringVar(&config.ExternalTopoGlobalRoot, "external_topo_global_root", "", "the path of the global topology data in the global topology server for vtcombo process")
 
 	cmd.Flags().DurationVar(&config.VtgateTabletRefreshInterval, "tablet_refresh_interval", 10*time.Second, "Interval at which vtgate refreshes tablet information from topology server.")
+
+	cmd.Flags().BoolVar(&doCreateTCPUser, "initialize-with-vt-dba-tcp", false, "If this flag is enabled, MySQL will be initialized with an additional user named vt_dba_tcp, who will have access via TCP/IP connection.")
 	acl.RegisterFlags(cmd.Flags())
 
 	return cmd
 }
 
-func newEnv() (env vttest.Environment, err error) {
-	if basePort != 0 {
+func newEnv() (env *vttest.LocalTestEnv, err error) {
+	if basePort == 0 {
+		env, err = vttest.NewLocalTestEnv(0)
+	} else {
 		if config.DataDir == "" {
 			env, err = vttest.NewLocalTestEnv(basePort)
-			if err != nil {
-				return
-			}
 		} else {
 			env, err = vttest.NewLocalTestEnvWithDirectory(basePort, config.DataDir)
-			if err != nil {
-				return
-			}
 		}
 	}
+	if err != nil {
+		return
+	}
+
+	if doCreateTCPUser {
+		// The original initFile does not have any users who can access through TCP/IP connection.
+		// Here we update the init file to create the user.
+		mysqlInitFile := env.InitDBFile
+		createUserCmd := `
+			# Admin user for TCP/IP connection with all privileges.
+			CREATE USER 'vt_dba_tcp'@'%';
+			GRANT ALL ON *.* TO 'vt_dba_tcp'@'%';
+			GRANT GRANT OPTION ON *.* TO 'vt_dba_tcp'@'%';
+		`
+		newInitFile := path.Join(env.Directory(), "init_db_with_vt_dba_tcp.sql")
+		err = vttest.WriteInitDBFile(mysqlInitFile, createUserCmd, newInitFile)
+		if err != nil {
+			return
+		}
+		env.InitDBFile = newInitFile
+	}
 
 	if protoTopo == "" {
 		config.Topology, err = topo.buildTopology()

go/cmd/vttestserver/cli/main_test.go

@@ -186,6 +186,39 @@ func TestForeignKeysAndDDLModes(t *testing.T) {
 	assert.NoError(t, err)
 }
 
+// TestCreateDbaTCPUser tests that the vt_dba_tcp user is created and can connect through TCP/IP connection
+// when --initialize-with-vt-dba-tcp is set to true.
+func TestCreateDbaTCPUser(t *testing.T) {
+	conf := config
+	defer resetConfig(conf)
+
+	clusterInstance, err := startCluster("--initialize-with-vt-dba-tcp=true")
+	assert.NoError(t, err)
+	defer clusterInstance.TearDown()
+
+	defer func() {
+		if t.Failed() {
+			cluster.PrintFiles(t, clusterInstance.Env.Directory(), "init_db_with_vt_dba_tcp.sql")
+		}
+	}()
+
+	// Ensure that the vt_dba_tcp user was created and can connect through TCP/IP connection.
+	ctx := context.Background()
+	vtParams := mysql.ConnParams{
+		Host:  "127.0.0.1",
+		Uname: "vt_dba_tcp",
+		Port:  clusterInstance.Env.PortForProtocol("mysql", ""),
+	}
+	conn, err := mysql.Connect(ctx, &vtParams)
+	assert.NoError(t, err)
+	defer conn.Close()
+
+	// Ensure that the existing vt_dba user remains unaffected, meaning it cannot connect through TCP/IP connection.
+	vtParams.Uname = "vt_dba"
+	_, err = mysql.Connect(ctx, &vtParams)
+	assert.Error(t, err)
+}
+
 func TestCanGetKeyspaces(t *testing.T) {
 	conf := config
 	defer resetConfig(conf)

go/flags/endtoend/vttestserver.txt

@@ -72,6 +72,7 @@ Flags:
       --grpc_server_keepalive_time duration                              After a duration of this time, if the server doesn't see any activity, it pings the client to see if the transport is still alive. (default 10s)
       --grpc_server_keepalive_timeout duration                           After having pinged for keepalive check, the server waits for a duration of Timeout and if no activity is seen even after that the connection is closed. (default 10s)
   -h, --help                                                             help for vttestserver
+      --initialize-with-vt-dba-tcp                                       If this flag is enabled, MySQL will be initialized with an additional user named vt_dba_tcp, who will have access via TCP/IP connection.
       --initialize_with_random_data                                      If this flag is each table-shard will be initialized with random data. See also the 'rng_seed' and 'min_shard_size' and 'max_shard_size' flags.
       --keep_logs duration                                               keep logs for this long (using ctime) (zero to keep forever)
       --keep_logs_by_mtime duration                                      keep logs for this long (using mtime) (zero to keep forever)

go/vt/vttest/environment.go

@@ -99,6 +99,7 @@ type LocalTestEnv struct {
 	BasePort        int
 	TmpPath         string
 	DefaultMyCnf    []string
+	InitDBFile      string
 	Env             []string
 	EnableToxiproxy bool
 }
@@ -133,7 +134,7 @@ func (env *LocalTestEnv) BinaryPath(binary string) string {
 func (env *LocalTestEnv) MySQLManager(mycnf []string, snapshot string) (MySQLManager, error) {
 	mysqlctl := &Mysqlctl{
 		Binary:    env.BinaryPath("mysqlctl"),
-		InitFile:  path.Join(os.Getenv("VTROOT"), "config/init_db.sql"),
+		InitFile:  env.InitDBFile,
 		Directory: env.TmpPath,
 		Port:      env.PortForProtocol("mysql", ""),
 		MyCnf:     append(env.DefaultMyCnf, mycnf...),
@@ -281,6 +282,7 @@ func NewLocalTestEnvWithDirectory(basePort int, directory string) (*LocalTestEnv
 		BasePort:     basePort,
 		TmpPath:      directory,
 		DefaultMyCnf: mycnf,
+		InitDBFile:   path.Join(os.Getenv("VTROOT"), "config/init_db.sql"),
 		Env: []string{
 			fmt.Sprintf("VTDATAROOT=%s", directory),
 			"VTTEST=endtoend",

go/vt/vttest/toxiproxyctl.go

@@ -63,21 +63,16 @@ func NewToxiproxyctl(binary string, apiPort, mysqlPort int, mysqlctl *Mysqlctl,
 
 	// The original initFile does not have any users who can access through TCP/IP connection.
 	// Here we update the init file to create the user.
-	initDb, _ := os.ReadFile(mysqlctl.InitFile)
 	createUserCmd := fmt.Sprintf(`
 		# Admin user for TCP/IP connection with all privileges.
 		CREATE USER '%s'@'127.0.0.1';
 		GRANT ALL ON *.* TO '%s'@'127.0.0.1';
 		GRANT GRANT OPTION ON *.* TO '%s'@'127.0.0.1';
 	`, dbaUser, dbaUser, dbaUser)
-	sql, err := getInitDBSQL(string(initDb), createUserCmd)
-	if err != nil {
-		return nil, vterrors.Wrap(err, "failed to get a modified init db sql")
-	}
 	newInitFile := path.Join(mysqlctl.Directory, "init_db_toxiproxyctl.sql")
-	err = os.WriteFile(newInitFile, []byte(sql), 0600)
+	err := WriteInitDBFile(mysqlctl.InitFile, createUserCmd, newInitFile)
 	if err != nil {
-		return nil, vterrors.Wrap(err, "failed to write a modified init db file")
+		return nil, vterrors.Wrap(err, "failed to get a modified init db sql")
 	}
 	mysqlctl.InitFile = newInitFile
 
@@ -235,6 +230,20 @@ func (ctl *Toxiproxyctl) RemoveTimeoutToxic() error {
 	return ctl.proxy.RemoveToxic("my-timeout")
 }
 
+// WriteInitDBFile is a helper function that writes a modified init_db.sql file with custom SQL statements.
+func WriteInitDBFile(initFile, customSQL, newInitFile string) error {
+	initDb, _ := os.ReadFile(initFile)
+	sql, err := getInitDBSQL(string(initDb), customSQL)
+	if err != nil {
+		return vterrors.Wrap(err, "failed to get a modified init db sql")
+	}
+	err = os.WriteFile(newInitFile, []byte(sql), 0600)
+	if err != nil {
+		return vterrors.Wrap(err, "failed to write a modified init db file")
+	}
+	return nil
+}
+
 // getInitDBSQL is a helper function that retrieves the modified contents of the init_db.sql file with custom SQL statements.
 // We avoid using vitess.io/vitess/go/test/endtoend/utils.GetInitDBSQL as importing this package adds unnecessary flags to vttestserver.
 func getInitDBSQL(initDBSQL string, customSQL string) (string, error) {