Make vttablet wait for vt_dba user to be granted privileges (#14565)

Signed-off-by: Manan Gupta <[email protected]>

go/cmd/vttablet/cli/cli.go

@@ -102,6 +102,10 @@ vttablet \
 	}
 )
 
+const (
+	dbaGrantWaitTime = 10 * time.Second
+)
+
 func run(cmd *cobra.Command, args []string) error {
 	servenv.Init()
 
@@ -244,6 +248,11 @@ func createTabletServer(ctx context.Context, config *tabletenv.TabletConfig, ts
 	} else if enforceTableACLConfig {
 		return nil, fmt.Errorf("table acl config has to be specified with table-acl-config flag because enforce-tableacl-config is set.")
 	}
+
+	err := tabletserver.WaitForDBAGrants(config, dbaGrantWaitTime)
+	if err != nil {
+		return nil, err
+	}
 	// creates and registers the query service
 	qsc := tabletserver.NewTabletServer(ctx, "", config, ts, tabletAlias)
 	servenv.OnRun(func() {

go/vt/vttablet/tabletserver/tabletserver.go

@@ -33,16 +33,16 @@ import (
 	"syscall"
 	"time"
 
+	"vitess.io/vitess/go/acl"
 	"vitess.io/vitess/go/mysql/sqlerror"
 	"vitess.io/vitess/go/pools/smartconnpool"
-
-	"vitess.io/vitess/go/acl"
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/stats"
 	"vitess.io/vitess/go/tb"
 	"vitess.io/vitess/go/trace"
 	"vitess.io/vitess/go/vt/callerid"
 	"vitess.io/vitess/go/vt/dbconfigs"
+	"vitess.io/vitess/go/vt/dbconnpool"
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/logutil"
 	"vitess.io/vitess/go/vt/mysqlctl"
@@ -232,6 +232,36 @@ func NewTabletServer(ctx context.Context, name string, config *tabletenv.TabletC
 	return tsv
 }
 
+// WaitForDBAGrants waits for DBA user to have the required privileges to function properly.
+func WaitForDBAGrants(config *tabletenv.TabletConfig, waitTime time.Duration) error {
+	if waitTime == 0 {
+		return nil
+	}
+	timer := time.NewTimer(waitTime)
+	ctx, cancel := context.WithTimeout(context.Background(), waitTime)
+	defer cancel()
+	for {
+		conn, err := dbconnpool.NewDBConnection(ctx, config.DB.DbaConnector())
+		if err == nil {
+			res, fetchErr := conn.ExecuteFetch("SHOW GRANTS", 1000, false)
+			if fetchErr == nil && res != nil && len(res.Rows) > 0 && len(res.Rows[0]) > 0 {
+				privileges := res.Rows[0][0].ToString()
+				// In MySQL 8.0, all the privileges are listed out explicitly, so we can search for SUPER in the output.
+				// In MySQL 5.7, all the privileges are not listed explicitly, instead ALL PRIVILEGES is written, so we search for that too.
+				if strings.Contains(privileges, "SUPER") || strings.Contains(privileges, "ALL PRIVILEGES") {
+					return nil
+				}
+			}
+		}
+		select {
+		case <-timer.C:
+			return fmt.Errorf("waited %v for dba user to have the required permissions", waitTime)
+		default:
+			time.Sleep(100 * time.Millisecond)
+		}
+	}
+}
+
 func (tsv *TabletServer) loadQueryTimeout() time.Duration {
 	return time.Duration(tsv.QueryTimeout.Load())
 }

go/vt/vttablet/tabletserver/tabletserver_test.go

@@ -31,7 +31,10 @@ import (
 	"time"
 
 	"vitess.io/vitess/go/mysql/sqlerror"
+	"vitess.io/vitess/go/vt/dbconfigs"
+	vttestpb "vitess.io/vitess/go/vt/proto/vttest"
 	"vitess.io/vitess/go/vt/sidecardb"
+	"vitess.io/vitess/go/vt/vttest"
 
 	"vitess.io/vitess/go/vt/callerid"
 
@@ -2626,3 +2629,154 @@ func addTabletServerSupportedQueries(db *fakesqldb.DB) {
 		}},
 	})
 }
+
+func TestWaitForDBAGrants(t *testing.T) {
+	tests := []struct {
+		name      string
+		waitTime  time.Duration
+		errWanted string
+		setupFunc func(t *testing.T) (*tabletenv.TabletConfig, func())
+	}{
+		{
+			name:      "Success without any wait",
+			waitTime:  1 * time.Second,
+			errWanted: "",
+			setupFunc: func(t *testing.T) (*tabletenv.TabletConfig, func()) {
+				// Create a new mysql instance, and the dba user with required grants.
+				// Since all the grants already exist, this should pass without any waiting to be needed.
+				testUser := "vt_test_dba"
+				cluster, err := startMySQLAndCreateUser(t, testUser)
+				require.NoError(t, err)
+				grantAllPrivilegesToUser(t, cluster.MySQLConnParams(), testUser)
+				tc := &tabletenv.TabletConfig{
+					DB: &dbconfigs.DBConfigs{},
+				}
+				connParams := cluster.MySQLConnParams()
+				connParams.Uname = testUser
+				tc.DB.SetDbParams(connParams, mysql.ConnParams{}, mysql.ConnParams{})
+				return tc, func() {
+					cluster.TearDown()
+				}
+			},
+		},
+		{
+			name:      "Success with wait",
+			waitTime:  1 * time.Second,
+			errWanted: "",
+			setupFunc: func(t *testing.T) (*tabletenv.TabletConfig, func()) {
+				// Create a new mysql instance, but delay granting the privileges to the dba user.
+				// This makes the waitForDBAGrants function retry the grant check.
+				testUser := "vt_test_dba"
+				cluster, err := startMySQLAndCreateUser(t, testUser)
+				require.NoError(t, err)
+
+				go func() {
+					time.Sleep(500 * time.Millisecond)
+					grantAllPrivilegesToUser(t, cluster.MySQLConnParams(), testUser)
+				}()
+
+				tc := &tabletenv.TabletConfig{
+					DB: &dbconfigs.DBConfigs{},
+				}
+				connParams := cluster.MySQLConnParams()
+				connParams.Uname = testUser
+				tc.DB.SetDbParams(connParams, mysql.ConnParams{}, mysql.ConnParams{})
+				return tc, func() {
+					cluster.TearDown()
+				}
+			},
+		}, {
+			name:      "Failure due to timeout",
+			waitTime:  300 * time.Millisecond,
+			errWanted: "waited 300ms for dba user to have the required permissions",
+			setupFunc: func(t *testing.T) (*tabletenv.TabletConfig, func()) {
+				// Create a new mysql but don't give the grants to the vt_dba user at all.
+				// This should cause a timeout after waiting, since the privileges are never granted.
+				testUser := "vt_test_dba"
+				cluster, err := startMySQLAndCreateUser(t, testUser)
+				require.NoError(t, err)
+
+				tc := &tabletenv.TabletConfig{
+					DB: &dbconfigs.DBConfigs{},
+				}
+				connParams := cluster.MySQLConnParams()
+				connParams.Uname = testUser
+				tc.DB.SetDbParams(connParams, mysql.ConnParams{}, mysql.ConnParams{})
+				return tc, func() {
+					cluster.TearDown()
+				}
+			},
+		}, {
+			name:      "Empty timeout",
+			waitTime:  0,
+			errWanted: "",
+			setupFunc: func(t *testing.T) (*tabletenv.TabletConfig, func()) {
+				return nil, func() {}
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config, cleanup := tt.setupFunc(t)
+			defer cleanup()
+			err := WaitForDBAGrants(config, tt.waitTime)
+			if tt.errWanted == "" {
+				require.NoError(t, err)
+			} else {
+				require.EqualError(t, err, tt.errWanted)
+			}
+		})
+	}
+}
+
+// startMySQLAndCreateUser starts a MySQL instance and creates the given user
+func startMySQLAndCreateUser(t *testing.T, testUser string) (vttest.LocalCluster, error) {
+	// Launch MySQL.
+	// We need a Keyspace in the topology, so the DbName is set.
+	// We need a Shard too, so the database 'vttest' is created.
+	cfg := vttest.Config{
+		Topology: &vttestpb.VTTestTopology{
+			Keyspaces: []*vttestpb.Keyspace{
+				{
+					Name: "vttest",
+					Shards: []*vttestpb.Shard{
+						{
+							Name:           "0",
+							DbNameOverride: "vttest",
+						},
+					},
+				},
+			},
+		},
+		OnlyMySQL: true,
+		Charset:   "utf8mb4",
+	}
+	cluster := vttest.LocalCluster{
+		Config: cfg,
+	}
+	err := cluster.Setup()
+	if err != nil {
+		return cluster, nil
+	}
+
+	connParams := cluster.MySQLConnParams()
+	conn, err := mysql.Connect(context.Background(), &connParams)
+	require.NoError(t, err)
+	_, err = conn.ExecuteFetch(fmt.Sprintf(`CREATE USER '%v'@'localhost';`, testUser), 1000, false)
+	conn.Close()
+
+	return cluster, err
+}
+
+// grantAllPrivilegesToUser grants all the privileges to the user specified.
+func grantAllPrivilegesToUser(t *testing.T, connParams mysql.ConnParams, testUser string) {
+	conn, err := mysql.Connect(context.Background(), &connParams)
+	require.NoError(t, err)
+	_, err = conn.ExecuteFetch(fmt.Sprintf(`GRANT ALL ON *.* TO '%v'@'localhost';`, testUser), 1000, false)
+	require.NoError(t, err)
+	_, err = conn.ExecuteFetch(fmt.Sprintf(`GRANT GRANT OPTION ON *.* TO '%v'@'localhost';`, testUser), 1000, false)
+	require.NoError(t, err)
+	_, err = conn.ExecuteFetch("FLUSH PRIVILEGES;", 1000, false)
+	require.NoError(t, err)
+	conn.Close()
+}