deleting users + other dbs

Signed-off-by: Renan Rangel <rrangel@slack-corp.com>

go/vt/mysqlctl/mysqlshellbackupengine.go

@@ -26,15 +26,17 @@ import (
 	"os"
 	"os/exec"
 	"path"
+	"slices"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/spf13/pflag"
 
-	"vitess.io/vitess/go/constants/sidecar"
 	"vitess.io/vitess/go/mysql"
 	"vitess.io/vitess/go/mysql/capabilities"
+	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/logutil"
 	"vitess.io/vitess/go/vt/mysqlctl/backupstorage"
 	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
 	"vitess.io/vitess/go/vt/servenv"
@@ -49,7 +51,7 @@ var (
 	// flags passed to the Dump command, as a JSON string
 	mysqlShellDumpFlags = `{"threads": 2}`
 	// flags passed to the Load command, as a JSON string
-	mysqlShellLoadFlags = `{"threads": 4, "updateGtidSet": "replace", "skipBinlog": true, "progressFile": ""}`
+	mysqlShellLoadFlags = `{"threads": 4, "loadUsers": true, "excludeUsers": ["'root'@'localhost'"], "updateGtidSet": "replace", "skipBinlog": true, "progressFile": ""}`
 	// drain a tablet when taking a backup
 	mysqlShellBackupShouldDrain = false
 	// disable redo logging and double write buffer
@@ -59,6 +61,15 @@ var (
 	knownObjectStoreParams = []string{"s3BucketName", "osBucketName", "azureContainerName"}
 
 	MySQLShellPreCheckError = errors.New("MySQLShellPreCheckError")
+
+	// internal databases not backed up by MySQL Shell
+	internalDBs = []string{
+		"information_schema", "mysql", "ndbinfo", "performance_schema", "sys",
+	}
+	// reserved MySQL users https://dev.mysql.com/doc/refman/8.0/en/reserved-accounts.html
+	reservedUsers = []string{
+		"mysql.sys@localhost", "mysql.session@localhost", "mysql.infoschema@localhost",
+	}
 )
 
 // MySQLShellBackupManifest represents a backup.
@@ -248,38 +259,32 @@ func (be *MySQLShellBackupEngine) ExecuteRestore(ctx context.Context, params Res
 		return nil, vterrors.Wrap(err, "disable semi-sync failed")
 	}
 
-	// if we received a RestoreFromBackup API call instead of it being a command line argument,
-	// we need to first clean the host before we start the restore.
-	if params.DeleteBeforeRestore {
-		params.Logger.Infof("restoring on an existing tablet, so dropping database %q", params.DbName)
+	params.Logger.Infof("restoring on an existing tablet, so dropping database %q", params.DbName)
+
+	readonly, err := params.Mysqld.IsSuperReadOnly(ctx)
+	if err != nil {
+		return nil, vterrors.Wrap(err, fmt.Sprintf("checking if mysqld has super_read_only=enable: %v", err))
+	}
 
-		readonly, err := params.Mysqld.IsSuperReadOnly(ctx)
+	if readonly {
+		resetFunc, err := params.Mysqld.SetSuperReadOnly(ctx, false)
 		if err != nil {
-			return nil, vterrors.Wrap(err, fmt.Sprintf("checking if mysqld has super_read_only=enable: %v", err))
+			return nil, vterrors.Wrap(err, fmt.Sprintf("unable to disable super-read-only: %v", err))
 		}
 
-		if readonly {
-			resetFunc, err := params.Mysqld.SetSuperReadOnly(ctx, false)
+		defer func() {
+			err := resetFunc()
 			if err != nil {
-				return nil, vterrors.Wrap(err, fmt.Sprintf("unable to disable super-read-only: %v", err))
+				params.Logger.Errorf("Not able to set super_read_only to its original value after restore")
 			}
+		}()
+	}
 
-			defer func() {
-				err := resetFunc()
-				if err != nil {
-					params.Logger.Errorf("Not able to set super_read_only to its original value after restore")
-				}
-			}()
-		}
-
-		err = params.Mysqld.ExecuteSuperQueryList(ctx,
-			[]string{fmt.Sprintf("DROP DATABASE IF EXISTS `%s`", params.DbName),
-				fmt.Sprintf("DROP DATABASE IF EXISTS `%s`", sidecar.GetName())},
-		)
-		if err != nil {
-			return nil, vterrors.Wrap(err, fmt.Sprintf("dropping database %q failed", params.DbName))
-		}
-
+	err = cleanupMySQL(ctx, params.Mysqld, params.Logger)
+	if err != nil {
+		log.Errorf(err.Error())
+		// time.Sleep(time.Minute * 2)
+		return nil, vterrors.Wrap(err, "error cleaning MySQL")
 	}
 
 	// we need to get rid of all the current replication information on the host.
@@ -505,6 +510,64 @@ func releaseReadLock(ctx context.Context, reader io.Reader, params BackupParams,
 	}
 }
 
+func cleanupMySQL(ctx context.Context, mysql MysqlDaemon, logger logutil.Logger) error {
+	logger.Infof("Cleaning up MySQL ahead of a restore")
+	result, err := mysql.FetchSuperQuery(ctx, "SHOW DATABASES")
+	if err != nil {
+		return err
+	}
+
+	// drop all databases
+	for _, row := range result.Rows {
+		dbName := row[0].ToString()
+		if slices.Contains(internalDBs, dbName) {
+			continue // not dropping internal DBs
+		}
+
+		logger.Infof("Dropping DB %q", dbName)
+		err = mysql.ExecuteSuperQuery(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS `%s`", row[0].ToString()))
+		if err != nil {
+			return fmt.Errorf("error droppping database %q: %w", row[0].ToString(), err)
+		}
+	}
+
+	// get current user
+	var currentUser string
+	result, err = mysql.FetchSuperQuery(ctx, "SELECT user()")
+	if err != nil {
+		return fmt.Errorf("error fetching current user: %w", err)
+	}
+
+	for _, row := range result.Rows {
+		currentUser = row[0].ToString()
+	}
+
+	// drop all users except reserved ones
+	result, err = mysql.FetchSuperQuery(ctx, "SELECT user, host FROM mysql.user")
+	if err != nil {
+		return err
+	}
+
+	for _, row := range result.Rows {
+		user := fmt.Sprintf("%s@%s", row[0].ToString(), row[1].ToString())
+
+		if user == currentUser {
+			continue // we don't drop the current user
+		}
+		if slices.Contains(reservedUsers, user) {
+			continue // we skip reserved MySQL users
+		}
+
+		logger.Infof("Dropping User %q", user)
+		err = mysql.ExecuteSuperQuery(ctx, fmt.Sprintf("DROP USER '%s'@'%s'", row[0].ToString(), row[1].ToString()))
+		if err != nil {
+			return fmt.Errorf("error droppping user %q: %w", user, err)
+		}
+	}
+
+	return err
+}
+
 func init() {
 	BackupRestoreEngineMap[mysqlShellBackupEngineName] = &MySQLShellBackupEngine{}
 }

go/vt/mysqlctl/mysqlshellbackupengine_test.go

@@ -18,13 +18,16 @@ package mysqlctl
 
 import (
 	"context"
+	"fmt"
 	"path"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/mysql/fakesqldb"
+	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/vt/logutil"
 	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
 )
 
@@ -175,3 +178,93 @@ func TestShouldDrainForBackupMySQLShell(t *testing.T) {
 	assert.True(t, engine.ShouldDrainForBackup(nil))
 	assert.True(t, engine.ShouldDrainForBackup(&tabletmanagerdatapb.BackupRequest{}))
 }
+
+func TestCleanupMySQL(t *testing.T) {
+	type userRecord struct {
+		user, host string
+	}
+
+	tests := []struct {
+		name              string
+		existingDBs       []string
+		expectedDropDBs   []string
+		currentUser       string
+		existingUsers     []userRecord
+		expectedDropUsers []string
+	}{
+		{
+			name:            "testing only specific DBs",
+			existingDBs:     []string{"_vt", "vt_test"},
+			expectedDropDBs: []string{"_vt", "vt_test"},
+		},
+		{
+			name:            "testing with internal dbs",
+			existingDBs:     []string{"_vt", "mysql", "vt_test", "performance_schema"},
+			expectedDropDBs: []string{"_vt", "vt_test"},
+		},
+		{
+			name:            "with users",
+			existingDBs:     []string{"_vt", "mysql", "vt_test", "performance_schema"},
+			expectedDropDBs: []string{"_vt", "vt_test"},
+			existingUsers: []userRecord{
+				{"test", "localhost"},
+				{"app", "10.0.0.1"},
+			},
+			expectedDropUsers: []string{"'test'@'localhost'", "'app'@'10.0.0.1'"},
+		},
+		{
+			name:            "with reserved users",
+			existingDBs:     []string{"_vt", "mysql", "vt_test", "performance_schema"},
+			expectedDropDBs: []string{"_vt", "vt_test"},
+			existingUsers: []userRecord{
+				{"mysql.sys", "localhost"},
+				{"mysql.infoschema", "localhost"},
+				{"mysql.session", "localhost"},
+				{"test", "localhost"},
+				{"app", "10.0.0.1"},
+			},
+			expectedDropUsers: []string{"'test'@'localhost'", "'app'@'10.0.0.1'"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			fakedb := fakesqldb.New(t)
+			defer fakedb.Close()
+			mysql := NewFakeMysqlDaemon(fakedb)
+			defer mysql.Close()
+
+			databases := [][]sqltypes.Value{}
+			for _, db := range tt.existingDBs {
+				databases = append(databases, []sqltypes.Value{sqltypes.NewVarChar(db)})
+			}
+
+			users := [][]sqltypes.Value{}
+			for _, record := range tt.existingUsers {
+				users = append(users, []sqltypes.Value{sqltypes.NewVarChar(record.user), sqltypes.NewVarChar(record.host)})
+			}
+
+			mysql.FetchSuperQueryMap = map[string]*sqltypes.Result{
+				"SHOW DATABASES":                    {Rows: databases},
+				"SELECT user()":                     {Rows: [][]sqltypes.Value{{sqltypes.NewVarChar(tt.currentUser)}}},
+				"SELECT user, host FROM mysql.user": {Rows: users},
+			}
+
+			for _, drop := range tt.expectedDropDBs {
+				mysql.ExpectedExecuteSuperQueryList = append(mysql.ExpectedExecuteSuperQueryList,
+					fmt.Sprintf("DROP DATABASE IF EXISTS `%s`", drop),
+				)
+			}
+
+			for _, drop := range tt.expectedDropUsers {
+				mysql.ExpectedExecuteSuperQueryList = append(mysql.ExpectedExecuteSuperQueryList,
+					fmt.Sprintf("DROP USER %s", drop),
+				)
+			}
+
+			err := cleanupMySQL(context.Background(), mysql, logutil.NewMemoryLogger())
+			require.NoError(t, err)
+		})
+	}
+
+}