bootstrap: add more tests for initialize-sql-file

session/bootstrap.go

@@ -2430,19 +2430,25 @@ func doDDLWorks(s Session) {
 
 // doBootstrapSQLFile executes SQL commands in a file as the last stage of bootstrap.
 // It is useful for setting the initial value of GLOBAL variables.
-func doBootstrapSQLFile(s Session) {
+func doBootstrapSQLFile(s Session) error {
 	sqlFile := config.GetGlobalConfig().InitializeSQLFile
 	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnBootstrap)
 	if sqlFile == "" {
-		return
+		return nil
 	}
 	logutil.BgLogger().Info("executing -initialize-sql-file", zap.String("file", sqlFile))
 	b, err := ioutil.ReadFile(sqlFile) //nolint:gosec
 	if err != nil {
+		if intest.InTest {
+			return err
+		}
 		logutil.BgLogger().Fatal("unable to read InitializeSQLFile", zap.Error(err))
 	}
 	stmts, err := s.Parse(ctx, string(b))
 	if err != nil {
+		if intest.InTest {
+			return err
+		}
 		logutil.BgLogger().Fatal("unable to parse InitializeSQLFile", zap.Error(err))
 	}
 	for _, stmt := range stmts {
@@ -2458,6 +2464,7 @@ func doBootstrapSQLFile(s Session) {
 			}
 		}
 	}
+	return nil
 }
 
 // doDMLWorks executes DML statements in bootstrap stage.

session/bootstrap_test.go

@@ -1081,6 +1081,28 @@ func TestUpgradeToVer85(t *testing.T) {
 }
 
 func TestInitializeSQLFile(t *testing.T) {
+	testEmptyInitSQLFile(t)
+	testInitSystemVariable(t)
+	testInitUsers(t)
+	testErrorHappenWhileInit(t)
+}
+
+func testEmptyInitSQLFile(t *testing.T) {
+	// An non-existent sql file would stop the bootstrap of the tidb cluster
+	store, err := mockstore.NewMockStore()
+	require.NoError(t, err)
+	config.GetGlobalConfig().InitializeSQLFile = "non-existent.sql"
+	defer func() {
+		config.GetGlobalConfig().InitializeSQLFile = ""
+	}()
+
+	dom, err := BootstrapSession(store)
+	require.Nil(t, dom)
+	require.NoError(t, err)
+	require.NoError(t, store.Close())
+}
+
+func testInitSystemVariable(t *testing.T) {
 	// We create an initialize-sql-file and then bootstrap the server with it.
 	// The observed behavior should be that tidb_enable_noop_variables is now
 	// disabled, and the feature works as expected.
@@ -1134,6 +1156,189 @@ func TestInitializeSQLFile(t *testing.T) {
 	require.NoError(t, r.Close())
 }
 
+func testInitUsers(t *testing.T) {
+	// Two sql files are set to 'initialize-sql-file' one after another,
+	// and only the first one are executed.
+	var err error
+	sqlFiles := make([]*os.File, 2)
+	for i, name := range []string{"1.sql", "2.sql"} {
+		sqlFiles[i], err = os.CreateTemp("", name)
+		require.NoError(t, err)
+	}
+	defer func() {
+		for _, sqlFile := range sqlFiles {
+			path := sqlFile.Name()
+			err = sqlFile.Close()
+			require.NoError(t, err)
+			err = os.Remove(path)
+			require.NoError(t, err)
+		}
+	}()
+	_, err = sqlFiles[0].WriteString(`
+CREATE USER cloud_admin;
+GRANT BACKUP_ADMIN, RESTORE_ADMIN ON *.* TO 'cloud_admin'@'%';
+GRANT DASHBOARD_CLIENT on *.* TO 'cloud_admin'@'%';
+GRANT SYSTEM_VARIABLES_ADMIN ON *.* TO 'cloud_admin'@'%';
+GRANT CONNECTION_ADMIN ON *.* TO 'cloud_admin'@'%';
+GRANT RESTRICTED_VARIABLES_ADMIN ON *.* TO 'cloud_admin'@'%';
+GRANT RESTRICTED_STATUS_ADMIN ON *.* TO 'cloud_admin'@'%';
+GRANT RESTRICTED_CONNECTION_ADMIN ON *.* TO 'cloud_admin'@'%';
+GRANT RESTRICTED_USER_ADMIN ON *.* TO 'cloud_admin'@'%';
+GRANT RESTRICTED_TABLES_ADMIN ON *.* TO 'cloud_admin'@'%';
+GRANT RESTRICTED_REPLICA_WRITER_ADMIN ON *.* TO 'cloud_admin'@'%';
+GRANT CREATE USER ON *.* TO 'cloud_admin'@'%';
+GRANT RELOAD ON *.* TO 'cloud_admin'@'%';
+GRANT PROCESS ON *.* TO 'cloud_admin'@'%';
+GRANT SELECT, INSERT, UPDATE, DELETE ON mysql.* TO 'cloud_admin'@'%';
+GRANT SELECT ON information_schema.* TO 'cloud_admin'@'%';
+GRANT SELECT ON performance_schema.* TO 'cloud_admin'@'%';
+GRANT SHOW DATABASES on *.* TO 'cloud_admin'@'%';
+GRANT REFERENCES ON *.* TO 'cloud_admin'@'%';
+GRANT SELECT ON *.* TO 'cloud_admin'@'%';
+GRANT INDEX ON *.* TO 'cloud_admin'@'%';
+GRANT INSERT ON *.* TO 'cloud_admin'@'%';
+GRANT UPDATE ON *.* TO 'cloud_admin'@'%';
+GRANT DELETE ON *.* TO 'cloud_admin'@'%';
+GRANT CREATE ON *.* TO 'cloud_admin'@'%';
+GRANT DROP ON *.* TO 'cloud_admin'@'%';
+GRANT ALTER ON *.* TO 'cloud_admin'@'%';
+GRANT CREATE VIEW ON *.* TO 'cloud_admin'@'%';
+GRANT SHUTDOWN, CONFIG ON *.* TO 'cloud_admin'@'%';
+REVOKE SHUTDOWN, CONFIG ON *.* FROM root;
+
+DROP USER root;
+`)
+	require.NoError(t, err)
+	_, err = sqlFiles[1].WriteString("drop user cloud_admin;")
+	require.NoError(t, err)
+
+	store, err := mockstore.NewMockStore()
+	require.NoError(t, err)
+	config.GetGlobalConfig().InitializeSQLFile = sqlFiles[0].Name()
+	defer func() {
+		require.NoError(t, store.Close())
+		config.GetGlobalConfig().InitializeSQLFile = ""
+	}()
+
+	// Bootstrap with the first sql file
+	dom, err := BootstrapSession(store)
+	require.NoError(t, err)
+	se := createSessionAndSetID(t, store)
+	ctx := context.Background()
+	// 'cloud_admin' has been created successfully
+	r, err := exec(se, `select user from mysql.user where user = 'cloud_admin'`)
+	require.NoError(t, err)
+	req := r.NewChunk(nil)
+	err = r.Next(ctx, req)
+	require.NoError(t, err)
+	require.Equal(t, 1, req.NumRows())
+	row := req.GetRow(0)
+	require.Equal(t, "cloud_admin", row.GetString(0))
+	require.NoError(t, r.Close())
+	// 'root' has been deleted successfully
+	r, err = exec(se, `select user from mysql.user where user = 'root'`)
+	require.NoError(t, err)
+	req = r.NewChunk(nil)
+	err = r.Next(ctx, req)
+	require.NoError(t, err)
+	require.Equal(t, 0, req.NumRows())
+	require.NoError(t, r.Close())
+	dom.Close()
+
+	runBootstrapSQLFile = false
+
+	// Bootstrap with the second sql file, which would not been executed.
+	config.GetGlobalConfig().InitializeSQLFile = sqlFiles[1].Name()
+	dom, err = BootstrapSession(store)
+	require.NoError(t, err)
+	se = createSessionAndSetID(t, store)
+	r, err = exec(se, `select user from mysql.user where user = 'cloud_admin'`)
+	require.NoError(t, err)
+	req = r.NewChunk(nil)
+	err = r.Next(ctx, req)
+	require.NoError(t, err)
+	require.Equal(t, 1, req.NumRows())
+	row = req.GetRow(0)
+	require.Equal(t, "cloud_admin", row.GetString(0))
+	require.NoError(t, r.Close())
+	dom.Close()
+}
+
+func testErrorHappenWhileInit(t *testing.T) {
+	// 1. parser error in sql file (1.sql) makes the bootstrap panic
+	// 2. other errors in sql file (2.sql) will be ignored
+	var err error
+	sqlFiles := make([]*os.File, 2)
+	for i, name := range []string{"1.sql", "2.sql"} {
+		sqlFiles[i], err = os.CreateTemp("", name)
+		require.NoError(t, err)
+	}
+	defer func() {
+		for _, sqlFile := range sqlFiles {
+			path := sqlFile.Name()
+			err = sqlFile.Close()
+			require.NoError(t, err)
+			err = os.Remove(path)
+			require.NoError(t, err)
+		}
+	}()
+	_, err = sqlFiles[0].WriteString("create table test.t (c in);")
+	require.NoError(t, err)
+	_, err = sqlFiles[1].WriteString(`
+create table test.t (c int);
+insert into test.t values ("abc"); -- invalid statement
+`)
+	require.NoError(t, err)
+
+	store, err := mockstore.NewMockStore()
+	require.NoError(t, err)
+	config.GetGlobalConfig().InitializeSQLFile = sqlFiles[0].Name()
+	defer func() {
+		config.GetGlobalConfig().InitializeSQLFile = ""
+	}()
+
+	// Bootstrap with the first sql file
+	dom, err := BootstrapSession(store)
+	require.Nil(t, dom)
+	require.NoError(t, err)
+	require.NoError(t, store.Close())
+
+	runBootstrapSQLFile = false
+
+	// Bootstrap with the second sql file, which would not been executed.
+	store, err = mockstore.NewMockStore()
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, store.Close())
+	}()
+	config.GetGlobalConfig().InitializeSQLFile = sqlFiles[1].Name()
+	dom, err = BootstrapSession(store)
+	require.NoError(t, err)
+	se := createSessionAndSetID(t, store)
+	ctx := context.Background()
+	_, err = exec(se, `use test;`)
+	require.NoError(t, err)
+	// Table t has been created.
+	r, err := exec(se, `show tables;`)
+	require.NoError(t, err)
+	req := r.NewChunk(nil)
+	err = r.Next(ctx, req)
+	require.NoError(t, err)
+	require.Equal(t, 1, req.NumRows())
+	row := req.GetRow(0)
+	require.Equal(t, "t", row.GetString(0))
+	require.NoError(t, r.Close())
+	// But data is failed to inserted since the error
+	r, err = exec(se, `select * from test.t`)
+	require.NoError(t, err)
+	req = r.NewChunk(nil)
+	err = r.Next(ctx, req)
+	require.NoError(t, err)
+	require.Equal(t, 0, req.NumRows())
+	require.NoError(t, r.Close())
+	dom.Close()
+}
+
 func TestTiDBEnablePagingVariable(t *testing.T) {
 	store, dom := createStoreAndBootstrap(t)
 	se := createSessionAndSetID(t, store)

session/session.go

@@ -3351,7 +3351,7 @@ func BootstrapSession(store kv.Storage) (*domain.Domain, error) {
 		}
 	}
 
-	//  Rebuild sysvar cache in a loop
+	// Rebuild sysvar cache in a loop
 	err = dom.LoadSysVarCacheLoop(ses[4])
 	if err != nil {
 		return nil, err
@@ -3416,12 +3416,15 @@ func BootstrapSession(store kv.Storage) (*domain.Domain, error) {
 	// setup historical stats worker
 	dom.SetupHistoricalStatsWorker(ses[8])
 	dom.StartHistoricalStatsWorker()
+	failToLoadOrParseSQLFile := false // only used for unit test
 	if runBootstrapSQLFile {
 		pm := &privileges.UserPrivileges{
 			Handle: dom.PrivilegeHandle(),
 		}
 		privilege.BindPrivilegeManager(ses[9], pm)
-		doBootstrapSQLFile(ses[9])
+		if err := doBootstrapSQLFile(ses[9]); err != nil {
+			failToLoadOrParseSQLFile = true
+		}
 	}
 	// A sub context for update table stats, and other contexts for concurrent stats loading.
 	cnt := 1 + concurrency
@@ -3489,6 +3492,12 @@ func BootstrapSession(store kv.Storage) (*domain.Domain, error) {
 		}
 	}
 
+	// This only happens in testing, since the failure of loading or parsing sql file
+	// would panic the bootstrapping.
+	if failToLoadOrParseSQLFile {
+		dom.Close()
+		return nil, err
+	}
 	return dom, err
 }