PBM-1421: Enable mongorestore for cloning collection

cmd/pbm/main.go

@@ -246,6 +246,12 @@ func main() {
 		Int32Var(&restore.numParallelColls)
 	restoreCmd.Flag("ns", `Namespaces to restore (e.g. "db1.*,db2.collection2"). If not set, restore all ("*.*")`).
 		StringVar(&restore.ns)
+	restoreCmd.Flag("ns-from", "Allows collection cloning (creating from the backup with different name) "+
+		"and specifies source collection for cloning from.").
+		StringVar(&restore.nsFrom)
+	restoreCmd.Flag("ns-to", "Allows collection cloning (creating from the backup with different name) "+
+		"and specifies destination collection for cloning to.").
+		StringVar(&restore.nsTo)
 	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.").

cmd/pbm/restore.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/mongodb/mongo-tools/common/db"
+	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"gopkg.in/yaml.v2"
 
@@ -28,6 +29,15 @@ import (
 	"github.com/percona/percona-backup-mongodb/sdk"
 )
 
+var (
+	ErrNSFromMissing        = errors.New("--ns-from should be specified as the cloning source")
+	ErrNSToMissing          = errors.New("--ns-to should be specified as the cloning destination")
+	ErrSelAndCloning        = errors.New("cloning with selective restore is not possible (remove --ns option)")
+	ErrCloningWithUAndR     = errors.New("cloning with restoring users and rolles is not possible")
+	ErrCloningWithPITR      = errors.New("cloning with restore to the point-in-time is not possible")
+	ErrCloningWithWildCards = errors.New("cloning with wild-cards is not possible")
+)
+
 type restoreOpts struct {
 	bcp           string
 	pitr          string
@@ -36,6 +46,8 @@ type restoreOpts struct {
 	waitTime      time.Duration
 	extern        bool
 	ns            string
+	nsFrom        string
+	nsTo          string
 	usersAndRoles bool
 	rsMap         string
 	conf          string
@@ -116,6 +128,9 @@ func runRestore(
 	if err != nil {
 		return nil, errors.Wrap(err, "parse --ns option")
 	}
+	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 {
 		return nil, errors.Wrap(err, "parse --with-users-and-roles option")
 	}
@@ -139,7 +154,7 @@ func runRestore(
 	}
 	tdiff := time.Now().Unix() - int64(clusterTime.T)
 
-	m, err := doRestore(ctx, conn, o, numParallelColls, nss, rsMap, node, outf)
+	m, err := doRestore(ctx, conn, o, numParallelColls, nss, o.nsFrom, o.nsTo, rsMap, node, outf)
 	if err != nil {
 		return nil, err
 	}
@@ -283,6 +298,8 @@ func checkBackup(
 	conn connect.Client,
 	o *restoreOpts,
 	nss []string,
+	nsFrom string,
+	nsTo string,
 ) (string, defs.BackupType, error) {
 	if o.extern && o.bcp == "" {
 		return "", defs.ExternalBackup, nil
@@ -318,28 +335,65 @@ func checkBackup(
 	if len(nss) != 0 && bcp.Type != defs.LogicalBackup {
 		return "", "", errors.New("--ns flag is only allowed for logical restore")
 	}
+	if nsFrom != "" && nsTo != "" && bcp.Type != defs.LogicalBackup {
+		return "", "", errors.New("--ns-from and ns-to flags are only allowed for logical restore")
+	}
 	if bcp.Status != defs.StatusDone {
 		return "", "", errors.Errorf("backup '%s' didn't finish successfully", b)
 	}
 
 	return bcp.Name, bcp.Type, nil
 }
 
+// nsIsTaken returns error in case when specified namesapce is already in use (collection is created)
+// or when any other error ocurres within the checking process.
+func nsIsTaken(
+	ctx context.Context,
+	conn connect.Client,
+	ns string,
+) error {
+	ns = strings.TrimSpace(ns)
+	db, coll, ok := strings.Cut(ns, ".")
+	if !ok {
+		return errors.Wrap(ErrInvalidNamespace, ns)
+	}
+
+	collNames, err := conn.MongoClient().Database(db).ListCollectionNames(ctx, bson.D{{"name", coll}})
+	if err != nil {
+		return errors.Wrap(err, "list collection names for cloning target validation")
+	}
+
+	if len(collNames) > 0 {
+		return errors.New("cloning namespace (--ns-to) is already in use, specify another one that doesn't exist in database")
+	}
+
+	return nil
+}
+
 func doRestore(
 	ctx context.Context,
 	conn connect.Client,
 	o *restoreOpts,
 	numParallelColls *int32,
 	nss []string,
+	nsFrom string,
+	nsTo string,
 	rsMapping map[string]string,
 	node string,
 	outf outFormat,
 ) (*restore.RestoreMeta, error) {
-	bcp, bcpType, err := checkBackup(ctx, conn, o, nss)
+	bcp, bcpType, err := checkBackup(ctx, conn, o, nss, nsFrom, nsTo)
 	if err != nil {
 		return nil, err
 	}
 
+	// check if namespace exists when cloning collection
+	if nsFrom != "" && nsTo != "" {
+		if err := nsIsTaken(ctx, conn, nsTo); err != nil {
+			return nil, err
+		}
+	}
+
 	name := time.Now().UTC().Format(time.RFC3339Nano)
 
 	cmd := ctrl.Cmd{
@@ -349,6 +403,8 @@ func doRestore(
 			BackupName:       bcp,
 			NumParallelColls: numParallelColls,
 			Namespaces:       nss,
+			NamespaceFrom:    nsFrom,
+			NamespaceTo:      nsTo,
 			UsersAndRoles:    o.usersAndRoles,
 			RSMap:            rsMapping,
 			External:         o.extern,
@@ -715,3 +771,36 @@ func validateRestoreUsersAndRoles(usersAndRoles bool, nss []string) error {
 
 	return nil
 }
+
+func validateNSFromNSTo(o *restoreOpts) error {
+	if o.nsFrom == "" && o.nsTo == "" {
+		return nil
+	}
+	if o.nsFrom == "" && o.nsTo != "" {
+		return ErrNSFromMissing
+	}
+	if o.nsFrom != "" && o.nsTo == "" {
+		return ErrNSToMissing
+	}
+	if _, _, ok := strings.Cut(o.nsFrom, "."); !ok {
+		return errors.Wrap(ErrInvalidNamespace, o.nsFrom)
+	}
+	if _, _, ok := strings.Cut(o.nsTo, "."); !ok {
+		return errors.Wrap(ErrInvalidNamespace, o.nsTo)
+	}
+	if o.nsFrom != "" && o.nsTo != "" && o.ns != "" {
+		return ErrSelAndCloning
+	}
+	if o.nsFrom != "" && o.nsTo != "" && o.usersAndRoles {
+		return ErrCloningWithUAndR
+	}
+	if o.nsFrom != "" && o.nsTo != "" && o.pitr != "" {
+		// this check will be removed with: PBM-1422
+		return ErrCloningWithPITR
+	}
+	if strings.Contains(o.nsTo, "*") || strings.Contains(o.nsFrom, "*") {
+		return ErrCloningWithWildCards
+	}
+
+	return nil
+}

cmd/pbm/restore_test.go

@@ -0,0 +1,112 @@
+package main
+
+import (
+	"errors"
+	"testing"
+)
+
+func TestCloningValidation(t *testing.T) {
+	testCases := []struct {
+		desc    string
+		opts    restoreOpts
+		wantErr error
+	}{
+		{
+			desc: "ns-to options is missing when cloning",
+			opts: restoreOpts{
+				nsFrom: "d.c",
+			},
+			wantErr: ErrNSToMissing,
+		},
+		{
+			desc: "ns-from options is missing when cloning",
+			opts: restoreOpts{
+				nsTo: "d.c",
+			},
+			wantErr: ErrNSFromMissing,
+		},
+		{
+			desc: "cloning with selective restore is not allowed",
+			opts: restoreOpts{
+				nsFrom: "d.c1",
+				nsTo:   "d.c2",
+				ns:     "d.c",
+			},
+			wantErr: ErrSelAndCloning,
+		},
+		{
+			desc: "cloning with restoring users and roles are not allowed",
+			opts: restoreOpts{
+				nsFrom:        "d.c1",
+				nsTo:          "d.c2",
+				usersAndRoles: true,
+			},
+			wantErr: ErrCloningWithUAndR,
+		},
+		{
+			desc: "cloning with PITR is not allowed",
+			opts: restoreOpts{
+				nsFrom: "d.c1",
+				nsTo:   "d.c2",
+				pitr:   "2024-10-27T11:23:30",
+			},
+			wantErr: ErrCloningWithPITR,
+		},
+		{
+			desc: "cloning with wild cards within nsFrom",
+			opts: restoreOpts{
+				nsFrom: "d.*",
+				nsTo:   "d.c2",
+			},
+			wantErr: ErrCloningWithWildCards,
+		},
+		{
+			desc: "cloning with wild cards within nsTo",
+			opts: restoreOpts{
+				nsFrom: "d.c1",
+				nsTo:   "d.*",
+			},
+			wantErr: ErrCloningWithWildCards,
+		},
+		{
+			desc: "cloning with ns without dot within nsFrom",
+			opts: restoreOpts{
+				nsFrom: "c",
+				nsTo:   "c.d",
+			},
+			wantErr: ErrInvalidNamespace,
+		},
+		{
+			desc: "cloning with ns without dot within nsTo",
+			opts: restoreOpts{
+				nsFrom: "d.c",
+				nsTo:   "d",
+			},
+			wantErr: ErrInvalidNamespace,
+		},
+		{
+			desc: "no error without cloning options",
+			opts: restoreOpts{
+				nsFrom: "",
+				nsTo:   "",
+			},
+			wantErr: nil,
+		},
+		{
+			desc: "no error when cloning options are correct",
+			opts: restoreOpts{
+				nsFrom: "b.a",
+				nsTo:   "d.c",
+			},
+			wantErr: nil,
+		},
+	}
+	for _, tC := range testCases {
+		t.Run(tC.desc, func(t *testing.T) {
+			err := validateNSFromNSTo(&tC.opts)
+			if !errors.Is(err, tC.wantErr) {
+				t.Errorf("Invalid validation error: want=%v, got=%v", tC.wantErr, err)
+			}
+		})
+	}
+}

pbm/ctrl/cmd.go

@@ -152,6 +152,8 @@ type RestoreCmd struct {
 	Name          string            `bson:"name"`
 	BackupName    string            `bson:"backupName"`
 	Namespaces    []string          `bson:"nss,omitempty"`
+	NamespaceFrom string            `bson:"nsFrom,omitempty"`
+	NamespaceTo   string            `bson:"nsTo,omitempty"`
 	UsersAndRoles bool              `bson:"usersAndRoles,omitempty"`
 	RSMap         map[string]string `bson:"rsMap,omitempty"`