PBM-1422: Adjust Oplog and PITR to support cloned collection (#1046)

* Remove guard for cloning ns with PITR
* Expand oplog options with cloning namespace
* Add logic for getting UUID from namespace
* Add oplog entries cloning functionality
* Add guard for backup without source cloning collections
* Fix transaction handling for cloning namespace

Author: boris-ilijic

cmd/pbm/restore.go

@@ -798,10 +798,6 @@ func validateNSFromNSTo(o *restoreOpts) error {
 	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
 	}

cmd/pbm/restore_test.go

@@ -43,15 +43,6 @@ func TestCloningValidation(t *testing.T) {
 			},
 			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{

pbm/oplog/restore.go

@@ -81,6 +81,15 @@ var selectedNSSupportedCommands = map[string]struct{}{
 	"commitIndexBuild": {},
 }
 
+var cloningNSSupportedCommands = map[string]struct{}{
+	"createIndexes":    {},
+	"deleteIndex":      {},
+	"deleteIndexes":    {},
+	"dropIndex":        {},
+	"dropIndexes":      {},
+	"commitIndexBuild": {},
+}
+
 var dontPreserveUUID = []string{
 	"admin.system.users",
 	"admin.system.roles",
@@ -89,6 +98,18 @@ var dontPreserveUUID = []string{
 	"*.system.views",     // timeseries
 }
 
+var ErrNoCloningNamespace = errors.New("cloning namespace desn't exist")
+
+// cloneNS has all data related to cloning namespace within oplog
+type cloneNS struct {
+	snapshot.CloneNS
+	toUUID primitive.Binary
+}
+
+func (c *cloneNS) SetNSPair(nsPair snapshot.CloneNS) {
+	c.CloneNS = nsPair
+}
+
 // OplogRestore is the oplog applyer
 type OplogRestore struct {
 	dst               *mongo.Client
@@ -120,7 +141,8 @@ type OplogRestore struct {
 
 	unsafe bool
 
-	filter OpFilter
+	filter  OpFilter
+	cloneNS cloneNS
 }
 
 const saveLastDistTxns = 100
@@ -196,6 +218,7 @@ func (o *OplogRestore) Apply(src io.ReadCloser) (primitive.Timestamp, error) {
 	defer bsonSource.Close()
 
 	var lts primitive.Timestamp
+
 	for {
 		rawOplogEntry := bsonSource.LoadNext()
 		if rawOplogEntry == nil {
@@ -257,6 +280,27 @@ func (o *OplogRestore) SetIncludeNS(nss []string) {
 	o.includeNS = dbs
 }
 
+// SetCloneNS sets all needed data for cloning namespace:
+// collection names and target namespace UUID
+func (o *OplogRestore) SetCloneNS(ctx context.Context, ns snapshot.CloneNS) error {
+	if !ns.IsSpecified() {
+		return nil
+	}
+
+	o.cloneNS.SetNSPair(ns)
+
+	var err error
+	o.cloneNS.toUUID, err = getUUIDForNS(ctx, o.dst, o.cloneNS.ToNS)
+	if err != nil {
+		return errors.Wrap(err, "get to ns uuid")
+	}
+	if o.cloneNS.toUUID.IsZero() {
+		return ErrNoCloningNamespace
+	}
+
+	return nil
+}
+
 func isOpAllowed(oe *Record) bool {
 	coll, ok := strings.CutPrefix(oe.Namespace, "config.")
 	if !ok {
@@ -310,6 +354,50 @@ func (o *OplogRestore) isOpSelected(oe *Record) bool {
 	return false
 }
 
+// isOpForCloning returns whether op needs to be processed or not in case of cloning NS.
+// In case of non cloning use case, it's always true.
+func (o *OplogRestore) isOpForCloning(oe *db.Oplog) bool {
+	if !o.cloneNS.IsSpecified() {
+		return true
+	}
+
+	// i, u, d ops for cloning ns
+	if oe.Namespace == o.cloneNS.FromNS {
+		return true
+	}
+
+	// filter out all other i, u, d ops
+	if oe.Operation != "c" {
+		return false
+	}
+
+	db, coll, _ := strings.Cut(oe.Namespace, ".")
+	if coll != "$cmd" {
+		return false
+	}
+
+	cmd := oe.Object[0].Key
+	if cmd == "applyOps" {
+		return true // internal ops of applyOps are checked one by one later
+	}
+
+	cloneFromDB, cloneFromColl, _ := strings.Cut(o.cloneNS.FromNS, ".")
+	if db != cloneFromDB {
+		// it's command not relevant for db to clone from
+		return false
+	}
+
+	if _, ok := cloningNSSupportedCommands[cmd]; ok {
+		// check if command targets collection
+		collForCmd, _ := oe.Object[0].Value.(string)
+		if collForCmd == cloneFromColl {
+			return true
+		}
+	}
+
+	return false
+}
+
 func (o *OplogRestore) isOpExcluded(oe *Record) bool {
 	if o.excludeNS == nil {
 		return false
@@ -360,7 +448,8 @@ func (o *OplogRestore) handleOp(oe db.Oplog) error {
 		return nil
 	}
 
-	if o.isOpExcluded(&oe) || !isOpAllowed(&oe) || !o.isOpSelected(&oe) {
+	if o.isOpExcluded(&oe) || !isOpAllowed(&oe) ||
+		!o.isOpSelected(&oe) || !o.isOpForCloning(&oe) {
 		return nil
 	}
 
@@ -676,10 +765,21 @@ func (o *OplogRestore) HandleUncommittedTxn(
 	return partial, uncommitted, nil
 }
 
+func (o *OplogRestore) cloneEntry(op *db.Oplog) {
+	if !o.cloneNS.IsSpecified() {
+		return
+	}
+	if op.Namespace == o.cloneNS.FromNS {
+		*op.UI = o.cloneNS.toUUID
+		op.Namespace = o.cloneNS.ToNS
+	}
+}
+
 func (o *OplogRestore) handleNonTxnOp(op db.Oplog) error {
 	// have to handle it here one more time because before the op gets thru
 	// txnBuffer its namespace is `collection.$cmd` instead of the real one
-	if o.isOpExcluded(&op) || !isOpAllowed(&op) || !o.isOpSelected(&op) {
+	if o.isOpExcluded(&op) || !isOpAllowed(&op) ||
+		!o.isOpSelected(&op) || !o.isOpForCloning(&op) {
 		return nil
 	}
 
@@ -688,6 +788,8 @@ func (o *OplogRestore) handleNonTxnOp(op db.Oplog) error {
 		return errors.Wrap(err, "filtering UUIDs from oplog")
 	}
 
+	o.cloneEntry(&op)
+
 	dbName, collName, _ := strings.Cut(op.Namespace, ".")
 	if op.Operation == "c" {
 		if len(op.Object) == 0 {
@@ -1125,3 +1227,29 @@ func isTruthy(val interface{}) bool {
 func isFalsy(val interface{}) bool {
 	return !isTruthy(val)
 }
+
+// getUUIDForNS ruturns UUID of existing collection.
+// When ns doesn't exist, it returns zero value without an error.
+// In case of error, it returns zero value for UUID in addition to error.
+func getUUIDForNS(ctx context.Context, m *mongo.Client, ns string) (primitive.Binary, error) {
+	var uuid primitive.Binary
+
+	d, c, _ := strings.Cut(ns, ".")
+	cur, err := m.Database(d).ListCollections(ctx, bson.D{{"name", c}})
+	if err != nil {
+		return uuid, errors.Wrap(err, "list collections")
+	}
+	defer cur.Close(ctx)
+
+	for cur.Next(ctx) {
+		if subtype, data, ok := cur.Current.Lookup("info", "uuid").BinaryOK(); ok {
+			uuid = primitive.Binary{
+				Subtype: subtype,
+				Data:    data,
+			}
+			break
+		}
+	}
+
+	return uuid, errors.Wrap(cur.Err(), "list collections cursor")
+}

pbm/restore/logical.go

@@ -256,20 +256,19 @@ func (r *Restore) Snapshot(
 	oplogRanges := []oplogRange{
 		{chunks: chunks, storage: r.bcpStg},
 	}
-	oplogOption := &applyOplogOption{end: &bcp.LastWriteTS, nss: nss}
+	oplogOption := &applyOplogOption{
+		end:     &bcp.LastWriteTS,
+		nss:     nss,
+		cloudNS: cloneNS,
+	}
 	if r.nodeInfo.IsConfigSrv() && util.IsSelective(nss) {
 		oplogOption.nss = []string{"config.databases"}
 		oplogOption.filter = newConfigsvrOpFilter(nss)
 	}
-	if cloneNS.IsSpecified() {
-		// oplog doesn't need to be applied when cloning ns
-		// this restriction will be removed during PBM-1422
-		l.Debug("applying oplog is skipped when cloning collection")
-	} else {
-		err = r.applyOplog(ctx, oplogRanges, oplogOption)
-		if err != nil {
-			return err
-		}
+
+	err = r.applyOplog(ctx, oplogRanges, oplogOption)
+	if err != nil {
+		return err
 	}
 
 	err = r.restoreIndexes(ctx, oplogOption.nss, cloneNS)
@@ -424,7 +423,11 @@ func (r *Restore) PITR(
 		{chunks: bcpChunks, storage: r.bcpStg},
 		{chunks: chunks, storage: r.oplogStg},
 	}
-	oplogOption := applyOplogOption{end: &cmd.OplogTS, nss: nss}
+	oplogOption := applyOplogOption{
+		end:     &cmd.OplogTS,
+		nss:     nss,
+		cloudNS: cloneNS,
+	}
 	if r.nodeInfo.IsConfigSrv() && util.IsSelective(nss) {
 		oplogOption.nss = []string{"config.databases"}
 		oplogOption.filter = newConfigsvrOpFilter(nss)

pbm/restore/restore.go

@@ -20,6 +20,7 @@ import (
 	"github.com/percona/percona-backup-mongodb/pbm/log"
 	"github.com/percona/percona-backup-mongodb/pbm/oplog"
 	"github.com/percona/percona-backup-mongodb/pbm/restore/phys"
+	"github.com/percona/percona-backup-mongodb/pbm/snapshot"
 	"github.com/percona/percona-backup-mongodb/pbm/storage"
 	"github.com/percona/percona-backup-mongodb/pbm/topo"
 	"github.com/percona/percona-backup-mongodb/pbm/util"
@@ -286,11 +287,12 @@ func chunks(
 }
 
 type applyOplogOption struct {
-	start  *primitive.Timestamp
-	end    *primitive.Timestamp
-	nss    []string
-	unsafe bool
-	filter oplog.OpFilter
+	start   *primitive.Timestamp
+	end     *primitive.Timestamp
+	nss     []string
+	cloudNS snapshot.CloneNS
+	unsafe  bool
+	filter  oplog.OpFilter
 }
 
 type (
@@ -341,7 +343,14 @@ func applyOplog(
 		txnSyncErr chan error
 	)
 
-	oplogRestore, err := oplog.NewOplogRestore(node, ic, mgoV, options.unsafe, true, ctxn, txnSyncErr)
+	oplogRestore, err := oplog.NewOplogRestore(
+		node,
+		ic,
+		mgoV,
+		options.unsafe,
+		true,
+		ctxn,
+		txnSyncErr)
 	if err != nil {
 		return nil, errors.Wrap(err, "create oplog")
 	}
@@ -357,6 +366,13 @@ func applyOplog(
 	}
 	oplogRestore.SetTimeframe(startTS, endTS)
 	oplogRestore.SetIncludeNS(options.nss)
+	err = oplogRestore.SetCloneNS(ctx, options.cloudNS)
+	if errors.Is(err, oplog.ErrNoCloningNamespace) {
+		log.Info("cloning namespace doesn't exist so oplog will not be applied")
+		return partial, nil
+	} else if err != nil {
+		return nil, errors.Wrap(err, "set cloning ns")
+	}
 
 	var lts primitive.Timestamp
 	for _, oplogRange := range ranges {