PBM-1213: Selective Logical Restore with included Users&Roles (#934)

* Add users&roles flag for restore cmd
* Extract parse namespace into util
* Add validation for restoring user&roles for specific collection
* Extract user&roles validation for restore cmd
* Add tests for ns helper funcs
* Add logical selective restore with users and roles
* Extract method for restoring users&roles
* Add logic for config server selective restore
* Add tests for resolving ns and users&roles opt

Author: boris-ilijic

cmd/pbm/main.go

@@ -170,6 +170,8 @@ func main() {
 		StringVar(&restore.pitrBase)
 	restoreCmd.Flag("ns", `Namespaces to restore (e.g. "db1.*,db2.collection2"). If not set, restore all ("*.*")`).
 		StringVar(&restore.ns)
+	restoreCmd.Flag("with-users-and-roles", "Includes users and roles for selected database (--ns flag)").
+		BoolVar(&restore.usersAndRoles)
 	restoreCmd.Flag("wait", "Wait for the restore to finish.").
 		Short('w').
 		BoolVar(&restore.wait)

cmd/pbm/restore.go

@@ -29,15 +29,16 @@ import (
 )
 
 type restoreOpts struct {
-	bcp      string
-	pitr     string
-	pitrBase string
-	wait     bool
-	extern   bool
-	ns       string
-	rsMap    string
-	conf     string
-	ts       string
+	bcp           string
+	pitr          string
+	pitrBase      string
+	wait          bool
+	extern        bool
+	ns            string
+	usersAndRoles bool
+	rsMap         string
+	conf          string
+	ts            string
 }
 
 type restoreRet struct {
@@ -101,6 +102,9 @@ func runRestore(ctx context.Context, conn connect.Client, o *restoreOpts, outf o
 	if err != nil {
 		return nil, errors.Wrap(err, "parse --ns option")
 	}
+	if err := validateRestoreUsersAndRoles(o.usersAndRoles, nss); err != nil {
+		return nil, errors.Wrap(err, "parse --with-users-and-roles-option")
+	}
 
 	rsMap, err := parseRSNamesMapping(o.rsMap)
 	if err != nil {
@@ -315,11 +319,12 @@ func doRestore(
 	cmd := ctrl.Cmd{
 		Cmd: ctrl.CmdRestore,
 		Restore: &ctrl.RestoreCmd{
-			Name:       name,
-			BackupName: bcp,
-			Namespaces: nss,
-			RSMap:      rsMapping,
-			External:   o.extern,
+			Name:          name,
+			BackupName:    bcp,
+			Namespaces:    nss,
+			UsersAndRoles: o.usersAndRoles,
+			RSMap:         rsMapping,
+			External:      o.extern,
 		},
 	}
 	if o.pitr != "" {
@@ -668,3 +673,16 @@ func describeRestore(ctx context.Context, conn connect.Client, o descrRestoreOpt
 
 	return res, nil
 }
+
+func validateRestoreUsersAndRoles(usersAndRoles bool, nss []string) error {
+	if !util.IsSelective(nss) && usersAndRoles {
+		return errors.New("Including users and roles are only allowed for selected database " +
+			"(use --ns flag for selective backup)")
+	}
+	if len(nss) >= 1 && util.ContainsSpecifiedColl(nss) && usersAndRoles {
+		return errors.New("Including users and roles are not allowed for specific collection. " +
+			"Use --ns='db.*' to specify the whole database instead.")
+	}
+
+	return nil
+}

pbm/backup/logical.go

@@ -44,7 +44,7 @@ func (b *Backup) doLogical(
 		if inf.IsConfigSrv() {
 			db = "config"
 		} else {
-			db, coll = parseNS(bcp.Namespaces[0])
+			db, coll = util.ParseNS(bcp.Namespaces[0])
 		}
 	}
 
@@ -445,16 +445,3 @@ func getNamespacesSize(ctx context.Context, m *mongo.Client, db, coll string) (m
 	err = eg.Wait()
 	return rv, err
 }
-
-func parseNS(ns string) (string, string) {
-	db, coll, _ := strings.Cut(ns, ".")
-
-	if db == "*" {
-		db = ""
-	}
-	if coll == "*" {
-		coll = ""
-	}
-
-	return db, coll
-}

pbm/ctrl/cmd.go

@@ -126,10 +126,11 @@ func (b BackupCmd) String() string {
 }
 
 type RestoreCmd struct {
-	Name       string            `bson:"name"`
-	BackupName string            `bson:"backupName"`
-	Namespaces []string          `bson:"nss,omitempty"`
-	RSMap      map[string]string `bson:"rsMap,omitempty"`
+	Name          string            `bson:"name"`
+	BackupName    string            `bson:"backupName"`
+	Namespaces    []string          `bson:"nss,omitempty"`
+	UsersAndRoles bool              `bson:"usersAndRoles,omitempty"`
+	RSMap         map[string]string `bson:"rsMap,omitempty"`
 
 	OplogTS primitive.Timestamp `bson:"oplogTS,omitempty"`
 

pbm/restore/logical.go

@@ -64,6 +64,9 @@ type Restore struct {
 	indexCatalog *idx.IndexCatalog
 }
 
+// PBM restore from temp collections (pbmRUsers/pbmRRoles)should be used
+type restoreUsersAndRolesOption bool
+
 // New creates a new restore object
 func New(leadConn connect.Client, nodeConn *mongo.Client, brief topo.NodeBrief, rsMap map[string]string) *Restore {
 	if rsMap == nil {
@@ -100,6 +103,41 @@ func (r *Restore) exit(ctx context.Context, err error) {
 	r.Close()
 }
 
+// resolveNamespace resolves final namespace(s) based on the backup namespace,
+// restore namespace, and option whether we should restore users&roles
+func resolveNamespace(nssBackup, nssRestore []string, usingUsersAndRoles bool) []string {
+	if util.IsSelective(nssRestore) {
+		if usingUsersAndRoles {
+			var nss []string
+			nss = append(nss, nssRestore...)
+			nss = append(nss,
+				defs.DB+"."+defs.TmpUsersCollection,
+				defs.DB+"."+defs.TmpRolesCollection,
+			)
+			return nss
+		}
+
+		return nssRestore
+	}
+	if util.IsSelective(nssBackup) {
+		return nssBackup
+	}
+
+	return nssBackup
+}
+
+// shouldRestoreUsersAndRoles determines whether user&roles should be restored from the backup
+func shouldRestoreUsersAndRoles(nssBackup, nssRestore []string, usingUsersAndRoles bool) restoreUsersAndRolesOption {
+	if util.IsSelective(nssBackup) {
+		return false
+	}
+	if util.IsSelective(nssRestore) {
+		return restoreUsersAndRolesOption(usingUsersAndRoles)
+	}
+
+	return true
+}
+
 // Snapshot do the snapshot's (mongo dump) restore
 //
 //nolint:nonamedreturns
@@ -116,10 +154,8 @@ func (r *Restore) Snapshot(ctx context.Context, cmd *ctrl.RestoreCmd, opid ctrl.
 		return err
 	}
 
-	nss := cmd.Namespaces
-	if !util.IsSelective(nss) {
-		nss = bcp.Namespaces
-	}
+	nss := resolveNamespace(bcp.Namespaces, cmd.Namespaces, cmd.UsersAndRoles)
+	usersAndRolesOpt := shouldRestoreUsersAndRoles(bcp.Namespaces, cmd.Namespaces, cmd.UsersAndRoles)
 
 	err = setRestoreBackup(ctx, r.leadConn, r.name, cmd.BackupName, nss)
 	if err != nil {
@@ -150,7 +186,7 @@ func (r *Restore) Snapshot(ctx context.Context, cmd *ctrl.RestoreCmd, opid ctrl.
 		return err
 	}
 
-	err = r.RunSnapshot(ctx, dump, bcp, nss)
+	err = r.RunSnapshot(ctx, dump, bcp, nss, usersAndRolesOpt)
 	if err != nil {
 		return err
 	}
@@ -226,10 +262,8 @@ func (r *Restore) PITR(ctx context.Context, cmd *ctrl.RestoreCmd, opid ctrl.OPID
 			"Try to set an earlier snapshot. Or leave the snapshot empty so PBM will choose one.")
 	}
 
-	nss := cmd.Namespaces
-	if len(nss) == 0 {
-		nss = bcp.Namespaces
-	}
+	nss := resolveNamespace(bcp.Namespaces, cmd.Namespaces, cmd.UsersAndRoles)
+	usersAndRolesOpt := shouldRestoreUsersAndRoles(bcp.Namespaces, cmd.Namespaces, cmd.UsersAndRoles)
 
 	if r.nodeInfo.IsLeader() {
 		err = SetOplogTimestamps(ctx, r.leadConn, r.name, 0, int64(cmd.OplogTS.T))
@@ -281,7 +315,7 @@ func (r *Restore) PITR(ctx context.Context, cmd *ctrl.RestoreCmd, opid ctrl.OPID
 		return err
 	}
 
-	err = r.RunSnapshot(ctx, dump, bcp, nss)
+	err = r.RunSnapshot(ctx, dump, bcp, nss, usersAndRolesOpt)
 	if err != nil {
 		return err
 	}
@@ -647,7 +681,13 @@ func (r *Restore) toState(ctx context.Context, status defs.Status, wait *time.Du
 	return toState(ctx, r.leadConn, status, r.name, r.nodeInfo, r.reconcileStatus, wait)
 }
 
-func (r *Restore) RunSnapshot(ctx context.Context, dump string, bcp *backup.BackupMeta, nss []string) error {
+func (r *Restore) RunSnapshot(
+	ctx context.Context,
+	dump string,
+	bcp *backup.BackupMeta,
+	nss []string,
+	usersAndRolesOpt restoreUsersAndRolesOption,
+) error {
 	var rdr io.ReadCloser
 
 	var err error
@@ -673,7 +713,15 @@ func (r *Restore) RunSnapshot(ctx context.Context, dump string, bcp *backup.Back
 		mapRS := util.MakeReverseRSMapFunc(r.rsMap)
 		if r.nodeInfo.IsConfigSrv() && util.IsSelective(nss) {
 			// restore cluster specific configs only
-			return r.configsvrRestore(ctx, bcp, nss, mapRS)
+			if err := r.configsvrRestore(ctx, bcp, nss, mapRS); err != nil {
+				return err
+			}
+			if !usersAndRolesOpt {
+				return nil
+			}
+
+			// selective restore needs to process users and roles from the full backup,
+			// so we'll continue with selective restore
 		}
 
 		var cfg *config.Config
@@ -730,17 +778,23 @@ func (r *Restore) RunSnapshot(ctx context.Context, dump string, bcp *backup.Back
 		return errors.Wrap(err, "mongorestore")
 	}
 
-	if util.IsSelective(nss) {
-		return nil
+	if usersAndRolesOpt {
+		if err := r.restoreUsersAndRoles(ctx, nss); err != nil {
+			return errors.Wrap(err, "restoring users and roles")
+		}
 	}
 
+	return nil
+}
+
+func (r *Restore) restoreUsersAndRoles(ctx context.Context, nss []string) error {
 	r.log.Info("restoring users and roles")
 	cusr, err := topo.CurrentUser(ctx, r.nodeConn)
 	if err != nil {
 		return errors.Wrap(err, "get current user")
 	}
 
-	err = r.swapUsers(ctx, cusr)
+	err = r.swapUsers(ctx, cusr, nss)
 	if err != nil {
 		return errors.Wrap(err, "swap users 'n' roles")
 	}
@@ -1123,22 +1177,39 @@ func (r *Restore) Done(ctx context.Context) error {
 	return nil
 }
 
-func (r *Restore) swapUsers(ctx context.Context, exclude *topo.AuthInfo) error {
+func (r *Restore) swapUsers(ctx context.Context, exclude *topo.AuthInfo, nss []string) error {
+	dbs := []string{}
+	for _, ns := range nss {
+		// ns can be "*.*" or "admin.pbmRUsers" or "admin.pbmRRoles"
+		db, _ := util.ParseNS(ns)
+		if len(db) == 0 || strings.HasPrefix(db, defs.DB) {
+			continue
+		}
+		dbs = append(dbs, db)
+	}
+
 	rolesC := r.nodeConn.Database("admin").Collection("system.roles")
 
 	eroles := []string{}
 	for _, r := range exclude.UserRoles {
 		eroles = append(eroles, r.DB+"."+r.Role)
 	}
 
+	rolesFilter := bson.M{
+		"_id": bson.M{"$nin": eroles},
+	}
+	if len(dbs) > 0 {
+		rolesFilter["db"] = bson.M{"$in": dbs}
+	}
+
 	curr, err := r.nodeConn.Database(defs.DB).Collection(defs.TmpRolesCollection).
-		Find(ctx, bson.M{"_id": bson.M{"$nin": eroles}})
+		Find(ctx, rolesFilter)
 	if err != nil {
 		return errors.Wrap(err, "create cursor for tmpRoles")
 	}
 	defer curr.Close(ctx)
 
-	_, err = rolesC.DeleteMany(ctx, bson.M{"_id": bson.M{"$nin": eroles}})
+	_, err = rolesC.DeleteMany(ctx, rolesFilter)
 	if err != nil {
 		return errors.Wrap(err, "delete current roles")
 	}
@@ -1159,15 +1230,21 @@ func (r *Restore) swapUsers(ctx context.Context, exclude *topo.AuthInfo) error {
 	if len(exclude.Users) > 0 {
 		user = exclude.Users[0].DB + "." + exclude.Users[0].User
 	}
+	filterUsers := bson.M{
+		"_id": bson.M{"$ne": user},
+	}
+	if len(dbs) > 0 {
+		filterUsers["db"] = bson.M{"$in": dbs}
+	}
 	cur, err := r.nodeConn.Database(defs.DB).Collection(defs.TmpUsersCollection).
-		Find(ctx, bson.M{"_id": bson.M{"$ne": user}})
+		Find(ctx, filterUsers)
 	if err != nil {
 		return errors.Wrap(err, "create cursor for tmpUsers")
 	}
 	defer cur.Close(ctx)
 
 	usersC := r.nodeConn.Database("admin").Collection("system.users")
-	_, err = usersC.DeleteMany(ctx, bson.M{"_id": bson.M{"$ne": user}})
+	_, err = usersC.DeleteMany(ctx, filterUsers)
 	if err != nil {
 		return errors.Wrap(err, "delete current users")
 	}