backupccl: Support RESTORE SYSTEM USERS from a backup

Support a new variant of RESTORE that recreates system users that don't exist in current cluster from a backup that contains system.users and also grant roles for these users. Example invocation: RESTORE SYSTEM USERS FROM 'nodelocal://foo/1';

Similar with full cluster restore, we firstly restore a temp system database which contains system.users and system.role_members into the restoring cluster and insert users and roles into the current system table from the temp system table.

Fixes: #45358

Release note (sql change): A special flavor of RESTORE, RESTORE SYSTEM USERS FROM ..., is added to support restoring system users from a backup. When executed, the statement recreates those users which are in a backup of system.users but do not currently exist (ignoring those who do) and re-grant roles for users if the backup contains system.role_members.

docs/generated/sql/bnf/restore.bnf

@@ -23,3 +23,15 @@ restore_stmt ::=
 	| 'RESTORE' ( 'TABLE' table_pattern ( ( ',' table_pattern ) )* | 'DATABASE' database_name ( ( ',' database_name ) )* ) 'FROM' subdirectory 'IN' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' )  'WITH' restore_options_list
 	| 'RESTORE' ( 'TABLE' table_pattern ( ( ',' table_pattern ) )* | 'DATABASE' database_name ( ( ',' database_name ) )* ) 'FROM' subdirectory 'IN' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' )  'WITH' 'OPTIONS' '(' restore_options_list ')'
 	| 'RESTORE' ( 'TABLE' table_pattern ( ( ',' table_pattern ) )* | 'DATABASE' database_name ( ( ',' database_name ) )* ) 'FROM' subdirectory 'IN' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' )  
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' ) 'AS' 'OF' 'SYSTEM' 'TIME' timestamp 'WITH' restore_options_list
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' ) 'AS' 'OF' 'SYSTEM' 'TIME' timestamp 'WITH' 'OPTIONS' '(' restore_options_list ')'
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' ) 'AS' 'OF' 'SYSTEM' 'TIME' timestamp 
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' )  'WITH' restore_options_list
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' )  'WITH' 'OPTIONS' '(' restore_options_list ')'
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' )  
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' subdirectory 'IN' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' ) 'AS' 'OF' 'SYSTEM' 'TIME' timestamp 'WITH' restore_options_list
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' subdirectory 'IN' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' ) 'AS' 'OF' 'SYSTEM' 'TIME' timestamp 'WITH' 'OPTIONS' '(' restore_options_list ')'
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' subdirectory 'IN' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' ) 'AS' 'OF' 'SYSTEM' 'TIME' timestamp 
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' subdirectory 'IN' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' )  'WITH' restore_options_list
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' subdirectory 'IN' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' )  'WITH' 'OPTIONS' '(' restore_options_list ')'
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' subdirectory 'IN' ( destination | '(' partitioned_backup_location ( ',' partitioned_backup_location )* ')' )  

docs/generated/sql/bnf/stmt_block.bnf

@@ -184,6 +184,8 @@ restore_stmt ::=
 	| 'RESTORE' 'FROM' string_or_placeholder 'IN' list_of_string_or_placeholder_opt_list opt_as_of_clause opt_with_restore_options
 	| 'RESTORE' targets 'FROM' list_of_string_or_placeholder_opt_list opt_as_of_clause opt_with_restore_options
 	| 'RESTORE' targets 'FROM' string_or_placeholder 'IN' list_of_string_or_placeholder_opt_list opt_as_of_clause opt_with_restore_options
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' list_of_string_or_placeholder_opt_list opt_as_of_clause opt_with_restore_options
+	| 'RESTORE' 'SYSTEM' 'USERS' 'FROM' string_or_placeholder 'IN' list_of_string_or_placeholder_opt_list opt_as_of_clause opt_with_restore_options
 	| 'RESTORE' targets 'FROM' 'REPLICATION' 'STREAM' 'FROM' string_or_placeholder_opt_list opt_as_of_clause
 
 resume_stmt ::=

pkg/ccl/backupccl/backup_test.go

@@ -8925,3 +8925,91 @@ DROP INDEX foo@bar;
 
 	sqlRunner.Exec(t, `BACKUP INTO LATEST IN 'nodelocal://0/foo' WITH revision_history`)
 }
+
+// TestBackupRestoreSystemUsers tests RESTORE SYSTEM USERS feature which allows user to
+// restore users from a backup into current cluster and regrant roles.
+func TestBackupRestoreSystemUsers(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	sqlDB, tempDir, cleanupFn := createEmptyCluster(t, singleNode)
+	_, _, sqlDBRestore, cleanupEmptyCluster := backupRestoreTestSetupEmpty(t, singleNode, tempDir, InitManualReplication, base.TestClusterArgs{})
+	defer cleanupFn()
+	defer cleanupEmptyCluster()
+
+	sqlDB.Exec(t, `CREATE USER app; CREATE USER test`)
+	sqlDB.Exec(t, `CREATE ROLE app_role; CREATE ROLE test_role`)
+	sqlDB.Exec(t, `GRANT app_role TO test_role;`) // 'test_role' is a member of 'app_role'
+	sqlDB.Exec(t, `GRANT admin, app_role TO app; GRANT test_role TO test`)
+	sqlDB.Exec(t, `CREATE DATABASE db; CREATE TABLE db.foo (ind INT)`)
+	sqlDB.Exec(t, `BACKUP TO $1`, LocalFoo+"/1")
+	sqlDB.Exec(t, `BACKUP DATABASE db TO $1`, LocalFoo+"/2")
+	sqlDB.Exec(t, `BACKUP TABLE system.users TO $1`, LocalFoo+"/3")
+
+	// User 'test' exists in both clusters but 'app' only exists in the backup
+	sqlDBRestore.Exec(t, `CREATE USER test`)
+	sqlDBRestore.Exec(t, `CREATE DATABASE db`)
+	// Create multiple databases to make max descriptor ID be larger than max descriptor ID
+	// in the backup to test if we correctly generate new descriptor IDs
+	sqlDBRestore.Exec(t, `CREATE DATABASE db1; CREATE DATABASE db2; CREATE DATABASE db3`)
+
+	t.Run("system users", func(t *testing.T) {
+		sqlDBRestore.Exec(t, "RESTORE SYSTEM USERS FROM $1", LocalFoo+"/1")
+
+		// Role 'app_role' and user 'app' will be added, and 'app' is granted with 'app_role'
+		// User test will remain untouched with no role granted
+		sqlDBRestore.CheckQueryResults(t, "SELECT * FROM system.users", [][]string{
+			{"admin", "", "true"},
+			{"app", "", "false"},
+			{"app_role", "", "true"},
+			{"root", "", "false"},
+			{"test", "", "false"},
+			{"test_role", "", "true"},
+		})
+		sqlDBRestore.CheckQueryResults(t, "SELECT * FROM system.role_members", [][]string{
+			{"admin", "app", "false"},
+			{"admin", "root", "true"},
+			{"app_role", "app", "false"},
+			{"app_role", "test_role", "false"},
+		})
+		sqlDBRestore.CheckQueryResults(t, "SHOW USERS", [][]string{
+			{"admin", "", "{}"},
+			{"app", "", "{admin,app_role}"},
+			{"app_role", "", "{}"},
+			{"root", "", "{admin}"},
+			{"test", "", "{}"},
+			{"test_role", "", "{app_role}"},
+		})
+	})
+
+	t.Run("restore-from-backup-with-no-system-users", func(t *testing.T) {
+		sqlDBRestore.ExpectErr(t, "cannot restore system users as no system.users table in the backup",
+			"RESTORE SYSTEM USERS FROM $1", LocalFoo+"/2")
+	})
+
+	_, _, sqlDBRestore1, cleanupEmptyCluster1 := backupRestoreTestSetupEmpty(t, singleNode, tempDir, InitManualReplication, base.TestClusterArgs{})
+	defer cleanupEmptyCluster1()
+	t.Run("restore-from-backup-with-no-system-role-members", func(t *testing.T) {
+		sqlDBRestore1.Exec(t, "RESTORE SYSTEM USERS FROM $1", LocalFoo+"/3")
+
+		sqlDBRestore1.CheckQueryResults(t, "SELECT * FROM system.users", [][]string{
+			{"admin", "", "true"},
+			{"app", "", "false"},
+			{"app_role", "", "true"},
+			{"root", "", "false"},
+			{"test", "", "false"},
+			{"test_role", "", "true"},
+		})
+		sqlDBRestore1.CheckQueryResults(t, "SELECT * FROM system.role_members", [][]string{
+			{"admin", "root", "true"},
+		})
+		sqlDBRestore1.CheckQueryResults(t, "SHOW USERS", [][]string{
+			{"admin", "", "{}"},
+			{"app", "", "{}"},
+			{"app_role", "", "{}"},
+			{"root", "", "{admin}"},
+			{"test", "", "{}"},
+			{"test_role", "", "{}"},
+		})
+	})
+}

pkg/ccl/backupccl/restore_job.go

@@ -1101,7 +1101,7 @@ func createImportingDescriptors(
 	}
 
 	tempSystemDBID := descpb.InvalidID
-	if details.DescriptorCoverage == tree.AllDescriptors {
+	if details.DescriptorCoverage == tree.AllDescriptors || details.RestoreSystemUsers {
 		tempSystemDBID = getTempSystemDBID(details)
 		databases = append(databases, dbdesc.NewInitial(tempSystemDBID, restoreTempSystemDB,
 			security.AdminRoleName()))
@@ -1740,6 +1740,15 @@ func (r *restoreResumer) doResume(ctx context.Context, execCtx interface{}) erro
 		// Reload the details as we may have updated the job.
 		details = r.job.Details().(jobspb.RestoreDetails)
 
+		if err := r.cleanupTempSystemTables(ctx, nil /* txn */); err != nil {
+			return err
+		}
+	} else if details.RestoreSystemUsers {
+		if err := r.restoreSystemUsers(ctx, p.ExecCfg().DB, mainData.systemTables); err != nil {
+			return err
+		}
+		details = r.job.Details().(jobspb.RestoreDetails)
+
 		if err := r.cleanupTempSystemTables(ctx, nil /* txn */); err != nil {
 			return err
 		}
@@ -2491,6 +2500,59 @@ type systemTableNameWithConfig struct {
 	config           systemBackupConfiguration
 }
 
+// Restore system.users from the backup into the restoring cluster. Only recreate users
+// which are in a backup of system.users but do not currently exist (ignoring those who do)
+// and re-grant roles for users if the backup has system.role_members.
+func (r *restoreResumer) restoreSystemUsers(
+	ctx context.Context, db *kv.DB, systemTables []catalog.TableDescriptor,
+) error {
+	executor := r.execCfg.InternalExecutor
+	return db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+		selectNonExistentUsers := "SELECT * FROM crdb_temp_system.users temp " +
+			"WHERE NOT EXISTS (SELECT * FROM system.users u WHERE temp.username = u.username)"
+		users, err := executor.QueryBuffered(ctx, "get-users",
+			txn, selectNonExistentUsers)
+		if err != nil {
+			return err
+		}
+
+		insertUser := `INSERT INTO system.users ("username", "hashedPassword", "isRole") VALUES ($1, $2, $3)`
+		newUsernames := make(map[string]bool)
+		for _, user := range users {
+			newUsernames[user[0].String()] = true
+			if _, err = executor.Exec(ctx, "insert-non-existent-users", txn, insertUser,
+				user[0], user[1], user[2]); err != nil {
+				return err
+			}
+		}
+
+		// We skip granting roles if the backup does not contain system.role_members.
+		if len(systemTables) == 1 {
+			return nil
+		}
+
+		selectNonExistentRoleMembers := "SELECT * FROM crdb_temp_system.role_members temp_rm WHERE " +
+			"NOT EXISTS (SELECT * FROM system.role_members rm WHERE temp_rm.role = rm.role AND temp_rm.member = rm.member)"
+		roleMembers, err := executor.QueryBuffered(ctx, "get-role-members",
+			txn, selectNonExistentRoleMembers)
+		if err != nil {
+			return err
+		}
+
+		insertRoleMember := `INSERT INTO system.role_members ("role", "member", "isAdmin") VALUES ($1, $2, $3)`
+		for _, roleMember := range roleMembers {
+			// Only grant roles to users that don't currently exist, i.e., new users we just added
+			if _, ok := newUsernames[roleMember[1].String()]; ok {
+				if _, err = executor.Exec(ctx, "insert-non-existent-role-members", txn, insertRoleMember,
+					roleMember[0], roleMember[1], roleMember[2]); err != nil {
+					return err
+				}
+			}
+		}
+		return nil
+	})
+}
+
 // restoreSystemTables atomically replaces the contents of the system tables
 // with the data from the restored system tables.
 func (r *restoreResumer) restoreSystemTables(

pkg/ccl/backupccl/restore_planning.go

@@ -320,6 +320,7 @@ func allocateDescriptorRewrites(
 	opts tree.RestoreOptions,
 	intoDB string,
 	newDBName string,
+	restoreSystemUsers bool,
 ) (DescRewriteMap, error) {
 	descriptorRewrites := make(DescRewriteMap)
 
@@ -432,29 +433,33 @@ func allocateDescriptorRewrites(
 
 	needsNewParentIDs := make(map[string][]descpb.ID)
 
-	// Increment the DescIDSequenceKey so that it is higher than the max desc ID
-	// in the backup. This generator keeps produced the next descriptor ID.
+	// Increment the DescIDSequenceKey so that it is higher than both the max desc ID
+	// in the backup and current max desc ID in the restoring cluster. This generator
+	// keeps produced the next descriptor ID.
 	var tempSysDBID descpb.ID
-	if descriptorCoverage == tree.AllDescriptors {
+	if descriptorCoverage == tree.AllDescriptors || restoreSystemUsers {
 		var err error
-		// Restore the key which generates descriptor IDs.
-		if err = p.ExecCfg().DB.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
-			b := txn.NewBatch()
-			// N.B. This key is usually mutated using the Inc command. That
-			// command warns that if the key was every Put directly, Inc will
-			// return an error. This is only to ensure that the type of the key
-			// doesn't change. Here we just need to be very careful that we only
-			// write int64 values.
-			// The generator's value should be set to the value of the next ID
-			// to generate.
-			b.Put(p.ExecCfg().Codec.DescIDSequenceKey(), maxDescIDInBackup+1)
-			return txn.Run(ctx, b)
-		}); err != nil {
-			return nil, err
-		}
-		tempSysDBID, err = catalogkv.GenerateUniqueDescID(ctx, p.ExecCfg().DB, p.ExecCfg().Codec)
-		if err != nil {
-			return nil, err
+		if descriptorCoverage == tree.AllDescriptors {
+			// Restore the key which generates descriptor IDs.
+			if err = p.ExecCfg().DB.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+				b := txn.NewBatch()
+					// write int64 values.
+					// The generator's value should be set to the value of the next ID
+					// to generate.
+					b.Put(p.ExecCfg().Codec.DescIDSequenceKey(), maxDescIDInBackup+1)
+				return txn.Run(ctx, b)
+			}); err != nil {
+				return nil, err
+			}
+			tempSysDBID, err = catalogkv.GenerateUniqueDescID(ctx, p.ExecCfg().DB, p.ExecCfg().Codec)
+			if err != nil {
+				return nil, err
+			}
+		} else if restoreSystemUsers {
+			tempSysDBID, err = catalogkv.GenerateUniqueDescID(ctx, p.ExecCfg().DB, p.ExecCfg().Codec)
+			if err != nil {
+				return nil, err
+			}
 		}
 		// Remap all of the descriptor belonging to system tables to the temp system
 		// DB.
@@ -833,7 +838,7 @@ func allocateDescriptorRewrites(
 	// backup should have the same ID as they do in the backup.
 	descriptorsToRemap := make([]catalog.Descriptor, 0, len(tablesByID))
 	for _, table := range tablesByID {
-		if descriptorCoverage == tree.AllDescriptors {
+		if descriptorCoverage == tree.AllDescriptors || restoreSystemUsers {
 			if table.ParentID == systemschema.SystemDB.GetID() {
 				// This is a system table that should be marked for descriptor creation.
 				descriptorsToRemap = append(descriptorsToRemap, table)
@@ -1420,6 +1425,7 @@ func restoreJobDescription(
 	kmsURIs []string,
 ) (string, error) {
 	r := &tree.Restore{
+		SystemUsers:        restore.SystemUsers,
 		DescriptorCoverage: restore.DescriptorCoverage,
 		AsOf:               restore.AsOf,
 		Targets:            restore.Targets,
@@ -1499,6 +1505,9 @@ func restorePlanHook(
 
 	var intoDBFn func() (string, error)
 	if restoreStmt.Options.IntoDB != nil {
+		if restoreStmt.SystemUsers {
+			return nil, nil, nil, false, errors.New("cannot set into_db option when only restoring system users")
+		}
 		intoDBFn, err = p.TypeAsString(ctx, restoreStmt.Options.IntoDB, "RESTORE")
 		if err != nil {
 			return nil, nil, nil, false, err
@@ -1520,6 +1529,9 @@ func restorePlanHook(
 				" database")
 			return nil, nil, nil, false, err
 		}
+		if restoreStmt.SystemUsers {
+			return nil, nil, nil, false, errors.New("cannot set new_db_name option when only restoring system users")
+		}
 		newDBNameFn, err = p.TypeAsString(ctx, restoreStmt.Options.NewDBName, "RESTORE")
 		if err != nil {
 			return nil, nil, nil, false, err
@@ -1855,7 +1867,7 @@ func doRestorePlan(
 	}
 
 	sqlDescs, restoreDBs, tenants, err := selectTargets(
-		ctx, p, mainBackupManifests, restoreStmt.Targets, restoreStmt.DescriptorCoverage, endTime,
+		ctx, p, mainBackupManifests, restoreStmt.Targets, restoreStmt.DescriptorCoverage, endTime, restoreStmt.SystemUsers,
 	)
 	if err != nil {
 		return errors.Wrap(err,
@@ -1971,7 +1983,8 @@ func doRestorePlan(
 		restoreStmt.DescriptorCoverage,
 		restoreStmt.Options,
 		intoDB,
-		newDBName)
+		newDBName,
+		restoreStmt.SystemUsers)
 	if err != nil {
 		return err
 	}
@@ -2052,6 +2065,7 @@ func doRestorePlan(
 			RevalidateIndexes:  revalidateIndexes,
 			DatabaseModifiers:  databaseModifiers,
 			DebugPauseOn:       debugPauseOn,
+			RestoreSystemUsers: restoreStmt.SystemUsers,
 		},
 		Progress: jobspb.RestoreProgress{},
 	}

pkg/ccl/backupccl/targets.go

@@ -326,9 +326,10 @@ func fullClusterTargetsBackup(
 	return fullClusterDescs, fullClusterDBIDs, nil
 }
 
-// selectTargets loads all descriptors from the selected backup manifest(s), and
-// filters the descriptors based on the targets specified in the restore. Post
-// filtering, the method returns:
+// selectTargets loads all descriptors from the selected backup manifest(s),
+// filters the descriptors based on the targets specified in the restore, and
+// calculates the max descriptor ID in the backup.
+// Post filtering, the method returns:
 //  - A list of all descriptors (table, type, database, schema) along with their
 //    parent databases.
 //  - A list of database descriptors IFF the user is restoring on the cluster or
@@ -341,13 +342,35 @@ func selectTargets(
 	targets tree.TargetList,
 	descriptorCoverage tree.DescriptorCoverage,
 	asOf hlc.Timestamp,
+	restoreSystemUsers bool,
 ) ([]catalog.Descriptor, []catalog.DatabaseDescriptor, []descpb.TenantInfoWithUsage, error) {
 	allDescs, lastBackupManifest := loadSQLDescsFromBackupsAtTime(backupManifests, asOf)
 
 	if descriptorCoverage == tree.AllDescriptors {
 		return fullClusterTargetsRestore(allDescs, lastBackupManifest)
 	}
 
+	if restoreSystemUsers {
+		systemTables := make([]catalog.Descriptor, 0)
+		var users catalog.Descriptor
+		for _, desc := range allDescs {
+			if desc.GetParentID() == systemschema.SystemDB.GetID() {
+				switch desc.GetName() {
+				case systemschema.UsersTable.GetName():
+					users = desc
+					systemTables = append(systemTables, desc)
+				case systemschema.RoleMembersTable.GetName():
+					systemTables = append(systemTables, desc)
+					// TODO(casper): should we handle role_options table?
+				}
+			}
+		}
+		if users == nil {
+			return nil, nil, nil, errors.Errorf("cannot restore system users as no system.users table in the backup")
+		}
+		return systemTables, nil, nil, nil
+	}
+
 	if targets.Tenant != (roachpb.TenantID{}) {
 		for _, tenant := range lastBackupManifest.GetTenants() {
 			// TODO(dt): for now it is zero-or-one but when that changes, we should