sql,backupccl: set system.role_members ID columns to NOT NULL

This patch sets the newly added `role_id` and `member_id` columns in
`system.role_members` to be NOT NULL. It also changes the `RESTORE`
logic to be able to handle the case when restoring from a backup where
the `system.role_members` table did not have the columns.

Release note: None

pkg/ccl/backupccl/backup_test.go

@@ -9611,10 +9611,10 @@ func TestBackupRestoreSystemUsers(t *testing.T) {
 			{"test_role", "NULL", "true", "103"},
 		})
 		sqlDBRestore.CheckQueryResults(t, "SELECT * FROM system.role_members", [][]string{
-			{"admin", "app", "false", "NULL", "NULL"},
+			{"admin", "app", "false", "2", "101"},
 			{"admin", "root", "true", "2", "1"},
-			{"app_role", "app", "false", "NULL", "NULL"},
-			{"app_role", "test_role", "false", "NULL", "NULL"},
+			{"app_role", "app", "false", "102", "101"},
+			{"app_role", "test_role", "false", "102", "103"},
 		})
 		sqlDBRestore.CheckQueryResults(t, "SHOW USERS", [][]string{
 			{"admin", "", "{}"},

pkg/ccl/backupccl/restore_job.go

@@ -2835,12 +2835,23 @@ func (r *restoreResumer) restoreSystemUsers(
 				return err
 			}
 
-			insertRoleMember := `INSERT INTO system.role_members ("role", "member", "isAdmin") VALUES ($1, $2, $3)`
+			roleMembersHasIDColumns := r.execCfg.Settings.Version.IsActive(ctx, clusterversion.V23_1RoleMembersTableHasIDColumns)
+			insertRoleMember := `
+INSERT INTO system.role_members ("role", "member", "isAdmin", role_id, member_id)
+VALUES ($1, $2, $3, (SELECT user_id FROM system.users WHERE username = $1), (SELECT user_id FROM system.users WHERE username = $2))`
+			if !roleMembersHasIDColumns {
+				insertRoleMember = `INSERT INTO system.role_members ("role", "member", "isAdmin") VALUES ($1, $2, $3)`
+			}
+
 			for _, roleMember := range roleMembers {
+				member := tree.MustBeDString(roleMember[1])
 				// 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 {
+				if _, ok := newUsernames[member.String()]; ok {
+					role := tree.MustBeDString(roleMember[0])
+					isAdmin := tree.MustBeDBool(roleMember[2])
+					if _, err := executor.Exec(ctx, "insert-non-existent-role-members", txn,
+						insertRoleMember, role, member, isAdmin,
+					); err != nil {
 						return err
 					}
 				}

pkg/ccl/backupccl/restore_old_versions_test.go

@@ -75,6 +75,7 @@ func TestRestoreOldVersions(t *testing.T) {
 		multiRegionDirs             = testdataBase + "/multi-region"
 		publicSchemaDirs            = testdataBase + "/public-schema-remap"
 		systemUsersDirs             = testdataBase + "/system-users-restore"
+		systemRoleMembersDirs       = testdataBase + "/system-role-members-restore"
 	)
 
 	t.Run("table-restore", func(t *testing.T) {
@@ -300,6 +301,17 @@ ORDER BY object_type, object_name`, [][]string{
 			t.Run(dir.Name(), fullClusterRestoreUsersWithoutIDs(exportDir))
 		}
 	})
+
+	t.Run("full-cluster-restore-system-role-members-without-ids", func(t *testing.T) {
+		dirs, err := os.ReadDir(systemRoleMembersDirs)
+		require.NoError(t, err)
+		for _, dir := range dirs {
+			require.True(t, dir.IsDir())
+			exportDir, err := filepath.Abs(filepath.Join(systemRoleMembersDirs, dir.Name()))
+			require.NoError(t, err)
+			t.Run(dir.Name(), fullClusterRestoreSystemRoleMembersWithoutIDs(exportDir))
+		}
+	})
 }
 
 func restoreOldVersionTestWithInterleave(exportDir string) func(t *testing.T) {
@@ -1160,3 +1172,30 @@ func restoreSystemUsersWithoutIDs(exportDir string) func(t *testing.T) {
 
 	}
 }
+
+func fullClusterRestoreSystemRoleMembersWithoutIDs(exportDir string) func(t *testing.T) {
+	return func(t *testing.T) {
+		const numAccounts = 1000
+		_, _, tmpDir, cleanupFn := backupRestoreTestSetup(t, multiNode, numAccounts, InitManualReplication)
+		defer cleanupFn()
+
+		_, sqlDB, cleanup := backupRestoreTestSetupEmpty(t, singleNode, tmpDir,
+			InitManualReplication, base.TestClusterArgs{
+				ServerArgs: base.TestServerArgs{
+					Knobs: base.TestingKnobs{
+						JobsTestingKnobs: jobs.NewTestingKnobsWithShortIntervals(),
+					},
+				}})
+		defer cleanup()
+		err := os.Symlink(exportDir, filepath.Join(tmpDir, "foo"))
+		require.NoError(t, err)
+
+		sqlDB.Exec(t, fmt.Sprintf("RESTORE FROM '%s'", localFoo))
+
+		sqlDB.CheckQueryResults(t, "SELECT * FROM system.role_members", [][]string{
+			{"admin", "root", "true", "2", "1"},
+			{"testrole", "testuser1", "false", "100", "101"},
+			{"testrole", "testuser2", "true", "100", "102"},
+		})
+	}
+}

pkg/ccl/backupccl/system_schema.go

@@ -169,7 +169,7 @@ func usersRestoreFunc(
 	execCfg *sql.ExecutorConfig,
 	txn *kv.Txn,
 	systemTableName, tempTableName string,
-) error {
+) (retErr error) {
 	if !execCfg.Settings.Version.IsActive(ctx, clusterversion.V22_2RoleOptionsTableHasIDColumn) {
 		return defaultSystemTableRestoreFunc(
 			ctx, execCfg, txn, systemTableName, tempTableName,
@@ -203,6 +203,9 @@ func usersRestoreFunc(
 	if err != nil {
 		return err
 	}
+	defer func() {
+		retErr = errors.CombineErrors(retErr, it.Close())
+	}()
 
 	for {
 		ok, err := it.Next(ctx)
@@ -239,12 +242,68 @@ func usersRestoreFunc(
 	return nil
 }
 
-func roleOptionsRestoreFunc(
+func roleMembersRestoreFunc(
 	ctx context.Context,
 	execCfg *sql.ExecutorConfig,
 	txn *kv.Txn,
 	systemTableName, tempTableName string,
 ) error {
+	if !execCfg.Settings.Version.IsActive(ctx, clusterversion.V23_1RoleMembersTableHasIDColumns) {
+		return defaultSystemTableRestoreFunc(ctx, execCfg, txn, systemTableName, tempTableName)
+	}
+
+	executor := execCfg.InternalExecutor
+
+	// It's enough to just check if role_id exists since member_id was added at
+	// the same time.
+	hasIDColumns, err := tableHasColumnName(ctx, txn, executor, tempTableName, "role_id")
+	if err != nil {
+		return err
+	}
+	if hasIDColumns {
+		return defaultSystemTableRestoreFunc(ctx, execCfg, txn, systemTableName, tempTableName)
+	}
+
+	deleteQuery := fmt.Sprintf("DELETE FROM system.%s WHERE true", systemTableName)
+	log.Eventf(ctx, "clearing data from system table %s with query %q", systemTableName, deleteQuery)
+
+	_, err = executor.Exec(ctx, systemTableName+"-data-deletion", txn, deleteQuery)
+	if err != nil {
+		return errors.Wrapf(err, "deleting data from system.%s", systemTableName)
+	}
+
+	roleMembers, err := executor.QueryBufferedEx(ctx, systemTableName+"-query-all-rows",
+		txn, sessiondata.NodeUserSessionDataOverride,
+		fmt.Sprintf(`SELECT * FROM %s`, tempTableName),
+	)
+	if err != nil {
+		return err
+	}
+
+	restoreQuery := fmt.Sprintf(`
+INSERT INTO system.%s ("role", "member", "isAdmin", role_id, member_id)
+VALUES ($1, $2, $3, (SELECT user_id FROM system.users WHERE username = $1), (SELECT user_id FROM system.users WHERE username = $2))`, systemTableName)
+	for _, roleMember := range roleMembers {
+		role := tree.MustBeDString(roleMember[0])
+		member := tree.MustBeDString(roleMember[1])
+		isAdmin := tree.MustBeDBool(roleMember[2])
+		if _, err := executor.ExecEx(ctx, systemTableName+"-data-insert",
+			txn, sessiondata.NodeUserSessionDataOverride,
+			restoreQuery, role, member, isAdmin,
+		); err != nil {
+			return errors.Wrapf(err, "inserting data to system.%s", systemTableName)
+		}
+	}
+
+	return nil
+}
+
+func roleOptionsRestoreFunc(
+	ctx context.Context,
+	execCfg *sql.ExecutorConfig,
+	txn *kv.Txn,
+	systemTableName, tempTableName string,
+) (retErr error) {
 	executor := execCfg.InternalExecutor
 	hasIDColumn, err := tableHasColumnName(ctx, txn, executor, tempTableName, "user_id")
 	if err != nil {
@@ -272,6 +331,9 @@ func roleOptionsRestoreFunc(
 	if err != nil {
 		return err
 	}
+	defer func() {
+		retErr = errors.CombineErrors(retErr, it.Close())
+	}()
 
 	for {
 		ok, err := it.Next(ctx)
@@ -402,6 +464,8 @@ var systemTableBackupConfiguration = map[string]systemBackupConfiguration{
 	},
 	systemschema.RoleMembersTable.GetName(): {
 		shouldIncludeInClusterBackup: optInToClusterBackup, // No desc ID columns.
+		customRestoreFunc:            roleMembersRestoreFunc,
+		restoreInOrder:               1, // Restore after system.users.
 	},
 	systemschema.RoleOptionsTable.GetName(): {
 		shouldIncludeInClusterBackup: optInToClusterBackup, // No desc ID columns.

pkg/ccl/backupccl/testdata/restore_old_versions/create-system-role-members.sql

@@ -0,0 +1,18 @@
+-- The below SQL is used to create the data that is then exported with BACKUP
+-- for use in the TestRestoreOldVersions test. This should be run on a v22.2
+-- cluster and used to test that after a restore on a v23.1 cluster, the ID
+-- columns in the system.role_members table are backfilled.
+
+CREATE DATABASE test;
+
+SET database = test;
+
+CREATE ROLE testrole;
+
+CREATE USER testuser1;
+
+CREATE USER testuser2;
+
+GRANT testrole TO testuser1;
+
+GRANT testrole TO testuser2 WITH ADMIN OPTION;

pkg/ccl/backupccl/testdata/restore_old_versions/system-role-members-restore/v22.2.1/BACKUP-LOCK-824672390336479233

@@ -0,0 +1 @@
+lock

pkg/ccl/backupccl/testdata/restore_old_versions/system-role-members-restore/v22.2.1/BACKUP_MANIFEST-CHECKSUM

@@ -0,0 +1 @@
+�*4S

pkg/ccl/backupccl/testdata/restore_old_versions/system-role-members-restore/v22.2.1/progress/BACKUP-CHECKPOINT-13cdcfcdcdd2cecdd2cdcddfcdcfc5cdcfc5cbced1c6c9cfcac7cadfd4cfcfcfcfdfaaabbcdf92c2d4cdcbd1cdcbc8c9c7cbccc8c9fffe-CHECKSUM

@@ -0,0 +1 @@
+ΰ"

pkg/ccl/backupccl/testdata/restore_old_versions/system-role-members-restore/v22.2.1/progress/BACKUP-CHECKPOINT-13cdcfcdcdd2cecdd2cdcddfcdcfc5cdcfc5cbced1c8cec9cec9c7dfd4cfcfcfcfdfaaabbcdf92c2d4cdcbd1cfcfcccdc9c7cfcbccfffe-CHECKSUM

@@ -0,0 +1 @@
+D6,

pkg/sql/catalog/systemschema/system.go

@@ -294,8 +294,8 @@ CREATE TABLE system.role_members (
   "role"    STRING NOT NULL,
   "member"  STRING NOT NULL,
   "isAdmin" BOOL NOT NULL,
-  role_id   OID,
-  member_id OID,
+  role_id   OID NOT NULL,
+  member_id OID NOT NULL,
   CONSTRAINT "primary" PRIMARY KEY ("role", "member"),
   INDEX ("role"),
   INDEX ("member"),
@@ -1655,8 +1655,8 @@ var (
 				{Name: "role", ID: 1, Type: types.String},
 				{Name: "member", ID: 2, Type: types.String},
 				{Name: "isAdmin", ID: 3, Type: types.Bool},
-				{Name: "role_id", ID: 4, Type: types.Oid, Nullable: true},
-				{Name: "member_id", ID: 5, Type: types.Oid, Nullable: true},
+				{Name: "role_id", ID: 4, Type: types.Oid},
+				{Name: "member_id", ID: 5, Type: types.Oid},
 			},
 			[]descpb.ColumnFamilyDescriptor{
 				{