Support users and roles for selective backup

Author: jcechace

cmd/pbm/backup.go

@@ -38,6 +38,7 @@ type backupOpts struct {
 	wait             bool
 	waitTime         time.Duration
 	externList       bool
+	usersAndRoles    bool
 
 	numParallelColls int32
 }
@@ -108,6 +109,9 @@ func runBackup(
 	if len(nss) != 0 && b.typ != string(defs.LogicalBackup) {
 		return nil, errors.New("--ns flag is only allowed for logical backup")
 	}
+	if err := util.ValidateUsersAndRolesOpt(b.usersAndRoles, nss); err != nil {
+		return nil, errors.Wrap(err, "parse --with-users-and-roles option")
+	}
 
 	if err := topo.CheckTopoForBackup(ctx, conn, defs.BackupType(b.typ)); err != nil {
 		return nil, errors.Wrap(err, "backup pre-check")
@@ -145,6 +149,7 @@ func runBackup(
 			IncrBase:         b.base,
 			Name:             b.name,
 			Namespaces:       nss,
+			UsersAndRoles:    b.usersAndRoles,
 			Compression:      compression,
 			CompressionLevel: level,
 			NumParallelColls: numParallelColls,
@@ -343,6 +348,7 @@ type bcpDesc struct {
 	LastWriteTime      string          `json:"last_write_time" yaml:"last_write_time"`
 	LastTransitionTime string          `json:"last_transition_time" yaml:"last_transition_time"`
 	Namespaces         []string        `json:"namespaces,omitempty" yaml:"namespaces,omitempty"`
+	SelUserAndRoles    bool            `json:"sel_user_and_roles" yaml:"sel_user_and_roles"`
 	MongoVersion       string          `json:"mongodb_version" yaml:"mongodb_version"`
 	FCV                string          `json:"fcv" yaml:"fcv"`
 	PBMVersion         string          `json:"pbm_version" yaml:"pbm_version"`
@@ -434,6 +440,7 @@ func describeBackup(
 		OPID:               bcp.OPID,
 		Type:               bcp.Type,
 		Namespaces:         bcp.Namespaces,
+		SelUserAndRoles:    bcp.SelUsersAndRoles,
 		MongoVersion:       bcp.MongoVersion,
 		FCV:                bcp.FCV,
 		PBMVersion:         bcp.PBMVersion,

cmd/pbm/main.go

@@ -295,6 +295,10 @@ func (app *pbmApp) buildBackupCmd() *cobra.Command {
 	backupCmd.Flags().StringVar(
 		&backupOptions.ns, "ns", "", `Namespaces to backup (e.g. "db.*", "db.collection"). If not set, backup all ("*.*")`,
 	)
+	backupCmd.Flags().BoolVar(
+		&backupOptions.usersAndRoles, "with-users-and-roles", false,
+		"Includes users and roles for selected database (--ns flag)",
+	)
 	backupCmd.Flags().BoolVarP(
 		&backupOptions.wait, "wait", "w", false, "Wait for the backup to finish",
 	)

cmd/pbm/restore.go

@@ -138,7 +138,7 @@ func runRestore(
 	if err := validateNSFromNSTo(o); err != nil {
 		return nil, errors.Wrap(err, "parse --ns-from and --ns-to options")
 	}
-	if err := validateRestoreUsersAndRoles(o.usersAndRoles, nss); err != nil {
+	if err := util.ValidateUsersAndRolesOpt(o.usersAndRoles, nss); err != nil {
 		return nil, errors.Wrap(err, "parse --with-users-and-roles option")
 	}
 	if err := validateFallbackOpts(o); err != nil {
@@ -805,19 +805,6 @@ func describeRestore(
 	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
-}
-
 func validateNSFromNSTo(o *restoreOpts) error {
 	if o.nsFrom == "" && o.nsTo == "" {
 		return nil

pbm/backup/backup.go

@@ -110,11 +110,12 @@ func (b *Backup) Init(
 	}
 
 	meta := &BackupMeta{
-		Type:        b.typ,
-		OPID:        opid.String(),
-		Name:        bcp.Name,
-		Namespaces:  bcp.Namespaces,
-		Compression: bcp.Compression,
+		Type:             b.typ,
+		OPID:             opid.String(),
+		Name:             bcp.Name,
+		Namespaces:       bcp.Namespaces,
+		SelUsersAndRoles: bcp.UsersAndRoles,
+		Compression:      bcp.Compression,
 		Store: Storage{
 			Name:        b.config.Name,
 			IsProfile:   b.config.IsProfile,

pbm/backup/logical.go

@@ -112,11 +112,11 @@ func (b *Backup) doLogical(
 	// ensure slicer is stopped in any case (done, error or canceled)
 	defer stopOplogSlicer() //nolint:errcheck
 
-	if !util.IsSelective(bcp.Namespaces) {
+	if !util.IsSelective(bcp.Namespaces) || bcp.UsersAndRoles {
 		// Save users and roles to the tmp collections so the restore would copy that data
 		// to the system collections. Have to do this because of issues with the restore and preserverUUID.
 		// see: https://jira.percona.com/browse/PBM-636 and comments
-		lw, err := copyUsersNRolles(ctx, b.brief.URI)
+		lw, err := copyUsersNRolles(ctx, b.brief.URI, bcp.Namespaces)
 		if err != nil {
 			return errors.Wrap(err, "copy users and roles for the restore")
 		}
@@ -152,6 +152,12 @@ func (b *Backup) doLogical(
 	nsFilter := archive.DefaultNSFilter
 	docFilter := archive.DefaultDocFilter
 	if util.IsSelective(bcp.Namespaces) {
+		if bcp.UsersAndRoles {
+			bcp.Namespaces = append(bcp.Namespaces,
+				defs.DB+"."+defs.TmpUsersCollection,
+				defs.DB+"."+defs.TmpRolesCollection,
+			)
+		}
 		if inf.IsConfigSrv() {
 			chunkSelector, err := createBackupChunkSelector(ctx, b.leadConn, bcp.Namespaces)
 			if err != nil {
@@ -286,7 +292,7 @@ func waitForWrite(ctx context.Context, m *mongo.Client, ts primitive.Timestamp)
 }
 
 //nolint:nonamedreturns
-func copyUsersNRolles(ctx context.Context, uri string) (lastWrite primitive.Timestamp, err error) {
+func copyUsersNRolles(ctx context.Context, uri string, nss []string) (lastWrite primitive.Timestamp, err error) {
 	cn, err := connect.MongoConnect(ctx, uri)
 	if err != nil {
 		return lastWrite, errors.Wrap(err, "connect to primary")
@@ -298,18 +304,19 @@ func copyUsersNRolles(ctx context.Context, uri string) (lastWrite primitive.Time
 		return lastWrite, errors.Wrap(err, "drop tmp collections before copy")
 	}
 
+	filter := util.MakeDBMatchFilter(nss)
 	err = copyColl(ctx,
 		cn.Database("admin").Collection("system.roles"),
 		cn.Database(defs.DB).Collection(defs.TmpRolesCollection),
-		bson.M{},
+		filter,
 	)
 	if err != nil {
 		return lastWrite, errors.Wrap(err, "copy admin.system.roles")
 	}
 	err = copyColl(ctx,
 		cn.Database("admin").Collection("system.users"),
 		cn.Database(defs.DB).Collection(defs.TmpUsersCollection),
-		bson.M{},
+		filter,
 	)
 	if err != nil {
 		return lastWrite, errors.Wrap(err, "copy admin.system.users")

pbm/backup/types.go

@@ -40,6 +40,7 @@ type BackupMeta struct {
 	ShardRemap map[string]string `bson:"shardRemap,omitempty" json:"shardRemap,omitempty"`
 
 	Namespaces       []string                 `bson:"nss,omitempty" json:"nss,omitempty"`
+	SelUsersAndRoles bool                     `bson:"sel_users_and_roles,omitempty" json:"sel_users_and_roles,omitempty"`
 	Replsets         []BackupReplset          `bson:"replsets" json:"replsets"`
 	Compression      compress.CompressionType `bson:"compression" json:"compression"`
 	Store            Storage                  `bson:"store" json:"store"`

pbm/ctrl/cmd.go

@@ -137,6 +137,7 @@ type BackupCmd struct {
 	NumParallelColls *int32                   `bson:"numParallelColls,omitempty"`
 	Filelist         bool                     `bson:"filelist,omitempty"`
 	Profile          string                   `bson:"profile,omitempty"`
+	UsersAndRoles    bool                     `bson:"usersAndRoles,omitempty"`
 }
 
 func (b BackupCmd) String() string {

pbm/restore/logical.go

@@ -146,12 +146,16 @@ func (r *Restore) exit(ctx context.Context, err error) {
 
 // resolveNamespace resolves final namespace(s) based on the backup namespace,
 // restore namespace, cloning options and option whether we should restore users&roles
-func resolveNamespace(nssBackup, nssRestore []string, cloneNS snapshot.CloneNS, usingUsersAndRoles bool) []string {
+func resolveNamespace(
+	nssBackup, nssRestore []string,
+	cloneNS snapshot.CloneNS,
+	usingUsersAndRolesBackup, usingUsersAndRolesRestore bool,
+) []string {
 	if cloneNS.IsSpecified() {
 		return []string{cloneNS.FromNS}
 	}
 	if util.IsSelective(nssRestore) {
-		if usingUsersAndRoles {
+		if usingUsersAndRolesRestore {
 			var nss []string
 			nss = append(nss, nssRestore...)
 			nss = append(nss,
@@ -164,23 +168,34 @@ func resolveNamespace(nssBackup, nssRestore []string, cloneNS snapshot.CloneNS,
 		return nssRestore
 	}
 
+	// Full restore from selective backup should include users & roles
+	if util.IsSelective(nssBackup) && usingUsersAndRolesBackup {
+		var nss []string
+		nss = append(nss, nssBackup...)
+		nss = append(nss,
+			defs.DB+"."+defs.TmpUsersCollection,
+			defs.DB+"."+defs.TmpRolesCollection,
+		)
+		return nss
+	}
+
 	return nssBackup
 }
 
 // shouldRestoreUsersAndRoles determines whether user&roles should be restored from the backup
 func shouldRestoreUsersAndRoles(
 	nssBackup, nssRestore []string,
 	cloneNS snapshot.CloneNS,
-	usingUsersAndRoles bool,
+	usingUsersAndRolesBackup, usingUsersAndRolesRestore bool,
 ) restoreUsersAndRolesOption {
 	if cloneNS.IsSpecified() {
 		return false
 	}
-	if util.IsSelective(nssBackup) {
+	if util.IsSelective(nssBackup) && !usingUsersAndRolesBackup {
 		return false
 	}
 	if util.IsSelective(nssRestore) {
-		return restoreUsersAndRolesOption(usingUsersAndRoles)
+		return restoreUsersAndRolesOption(usingUsersAndRolesRestore)
 	}
 
 	return true
@@ -218,11 +233,13 @@ func (r *Restore) Snapshot(
 		bcp.Namespaces,
 		cmd.Namespaces,
 		cloneNS,
+		bcp.SelUsersAndRoles,
 		cmd.UsersAndRoles)
 	usersAndRolesOpt := shouldRestoreUsersAndRoles(
 		bcp.Namespaces,
 		cmd.Namespaces,
 		cloneNS,
+		bcp.SelUsersAndRoles,
 		cmd.UsersAndRoles)
 
 	err = setRestoreBackup(ctx, r.leadConn, r.name, cmd.BackupName, nss)
@@ -362,11 +379,13 @@ func (r *Restore) PITR(
 		bcp.Namespaces,
 		cmd.Namespaces,
 		cloneNS,
+		bcp.SelUsersAndRoles,
 		cmd.UsersAndRoles)
 	usersAndRolesOpt := shouldRestoreUsersAndRoles(
 		bcp.Namespaces,
 		cmd.Namespaces,
 		cloneNS,
+		bcp.SelUsersAndRoles,
 		cmd.UsersAndRoles)
 
 	if r.nodeInfo.IsLeader() {