Add mongo-rs-members parameter to binary mongo restore

cmd/mongo/binary_backup_fetch.go

@@ -16,10 +16,13 @@ const (
 	binaryBackupFetchCommandName = "binary-backup-fetch"
 	MinimalConfigPathFlag        = "minimal-mongod-config-path"
 	MinimalConfigPathDescription = "Path to mongod config with minimal working configuration"
+	RsMembersFlag                = "mongo-rs-members"
+	RsMembersDescription         = "Comma separated host:port records from wished rs members (like rs.initiate())"
 )
 
 var (
 	minimalConfigPath = ""
+	rsMembers         = ""
 )
 
 var binaryBackupFetchCmd = &cobra.Command{
@@ -34,15 +37,17 @@ var binaryBackupFetchCmd = &cobra.Command{
 		defer func() { _ = signalHandler.Close() }()
 
 		backupName := args[0]
-		mongodbConfigPath := args[1]
+		mongodConfigPath := args[1]
 		mongodVersion := args[2]
 
-		err := mongo.HandleBinaryFetchPush(ctx, mongodbConfigPath, minimalConfigPath, backupName, mongodVersion)
+		err := mongo.HandleBinaryFetchPush(ctx, mongodConfigPath, minimalConfigPath, backupName, mongodVersion,
+			rsMembers)
 		tracelog.ErrorLogger.FatalOnError(err)
 	},
 }
 
 func init() {
 	binaryBackupFetchCmd.Flags().StringVar(&minimalConfigPath, MinimalConfigPathFlag, "", MinimalConfigPathDescription)
+	binaryBackupFetchCmd.Flags().StringVar(&rsMembers, RsMembersFlag, "", RsMembersDescription)
 	cmd.AddCommand(binaryBackupFetchCmd)
 }

internal/databases/mongo/binary/mongod.go

@@ -128,7 +128,7 @@ func (mongodService *MongodService) GetBackupCursorExtended(backupCursorMeta *Ba
 	})
 }
 
-func (mongodService *MongodService) FixSystemDataAfterRestore() error {
+func (mongodService *MongodService) FixSystemDataAfterRestore(rsMembers string) error {
 	ctx := mongodService.Context
 	localDatabase := mongodService.MongoClient.Database("local")
 
@@ -137,9 +137,28 @@ func (mongodService *MongodService) FixSystemDataAfterRestore() error {
 		return errors.Wrap(err, "unable to fix data in local.replset.election")
 	}
 
-	err = replaceData(ctx, localDatabase.Collection("system.replset"), false, nil)
-	if err != nil {
-		return errors.Wrap(err, "unable to fix data in local.system.replset")
+	if rsMembers == "" {
+		err = replaceData(ctx, localDatabase.Collection("system.replset"), false, nil)
+		if err != nil {
+			return errors.Wrap(err, "unable to fix data in local.system.replset")
+		}
+	} else {
+		var rsConfig bson.M
+		err = localDatabase.Collection("system.replset").FindOne(ctx, bson.D{}).Decode(&rsConfig)
+		if err != nil {
+			return errors.Wrap(err, "unable to read rs config from system.replset")
+		}
+		rsConfig["members"] = makeBsonRsMembers(rsMembers)
+
+		updateResult, err := localDatabase.Collection("system.replset").
+			UpdateOne(ctx, bson.D{{Key: "_id", Value: rsConfig["_id"]}}, rsConfig)
+		if err != nil {
+			return errors.Wrap(err, "unable to update rs config in system.replset")
+		}
+		if updateResult.MatchedCount != 1 {
+			return errors.Errorf("MatchedCount = %v during update rs config in system.replset",
+				updateResult.MatchedCount)
+		}
 	}
 
 	adminDatabase := mongodService.MongoClient.Database(adminDB)
@@ -191,6 +210,18 @@ func replaceData(ctx context.Context, collection *mongo.Collection, drop bool, i
 	return nil
 }
 
+func makeBsonRsMembers(rsMembers string) bson.A {
+	bsonMembers := bson.A{}
+
+	if rsMembers != "" {
+		for id, member := range strings.Split(rsMembers, ",") {
+			bsonMembers = append(bsonMembers, bson.M{"_id": id, "host": member})
+		}
+	}
+
+	return bsonMembers
+}
+
 func (mongodService *MongodService) Shutdown() error {
 	err := mongodService.MongoClient.Database(adminDB).RunCommand(context.Background(),
 		bson.D{{Key: "shutdown", Value: 1}},

internal/databases/mongo/binary/mongod_test.go

@@ -0,0 +1,20 @@
+package binary
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.mongodb.org/mongo-driver/bson"
+)
+
+func TestMakeBsonRsMembers(t *testing.T) {
+	assert.Equal(t, bson.A{}, makeBsonRsMembers(""))
+	assert.Equal(t, bson.A{bson.M{"_id": 0, "host": "localhost:1234"}}, makeBsonRsMembers("localhost:1234"))
+	assert.Equal(t,
+		bson.A{
+			bson.M{"_id": 0, "host": "localhost:1234"},
+			bson.M{"_id": 1, "host": "localhost:5678"},
+			bson.M{"_id": 2, "host": "remotehost:9876"},
+		},
+		makeBsonRsMembers("localhost:1234,localhost:5678,remotehost:9876"))
+}

internal/databases/mongo/binary/restore.go

@@ -27,7 +27,7 @@ func CreateRestoreService(ctx context.Context, localStorage *LocalStorage, uploa
 	}, nil
 }
 
-func (restoreService *RestoreService) DoRestore(backupName, restoreMongodVersion string) error {
+func (restoreService *RestoreService) DoRestore(backupName, restoreMongodVersion, rsMembers string) error {
 	sentinel, err := common.DownloadSentinel(restoreService.Uploader.Folder(), backupName)
 	if err != nil {
 		return err
@@ -54,7 +54,7 @@ func (restoreService *RestoreService) DoRestore(backupName, restoreMongodVersion
 		return err
 	}
 
-	if err = restoreService.fixSystemData(); err != nil {
+	if err = restoreService.fixSystemData(rsMembers); err != nil {
 		return err
 	}
 	if err = restoreService.recoverFromOplogAsStandalone(); err != nil {
@@ -69,7 +69,7 @@ func (restoreService *RestoreService) downloadFromTarArchives(backupName string)
 	return downloader.Download(backupName, restoreService.LocalStorage.MongodDBPath)
 }
 
-func (restoreService *RestoreService) fixSystemData() error {
+func (restoreService *RestoreService) fixSystemData(rsMembers string) error {
 	mongodProcess, err := StartMongodWithDisableLogicalSessionCacheRefresh(restoreService.minimalConfigPath)
 	if err != nil {
 		return errors.Wrap(err, "unable to start mongod in special mode")
@@ -80,7 +80,7 @@ func (restoreService *RestoreService) fixSystemData() error {
 		return errors.Wrap(err, "unable to create mongod service")
 	}
 
-	err = mongodService.FixSystemDataAfterRestore()
+	err = mongodService.FixSystemDataAfterRestore(rsMembers)
 	if err != nil {
 		return err
 	}

internal/databases/mongo/binary_backup_fetch_handler.go

@@ -8,8 +8,8 @@ import (
 	"github.com/wal-g/wal-g/utility"
 )
 
-func HandleBinaryFetchPush(ctx context.Context, mongodConfigPath, minimalConfigPath, backupName,
-	restoreMongodVersion string,
+func HandleBinaryFetchPush(ctx context.Context, mongodConfigPath, minimalConfigPath, backupName, restoreMongodVersion,
+	rsMembers string,
 ) error {
 	config, err := binary.CreateMongodConfig(mongodConfigPath)
 	if err != nil {
@@ -36,5 +36,5 @@ func HandleBinaryFetchPush(ctx context.Context, mongodConfigPath, minimalConfigP
 		return err
 	}
 
-	return restoreService.DoRestore(backupName, restoreMongodVersion)
+	return restoreService.DoRestore(backupName, restoreMongodVersion, rsMembers)
 }

tests_func/features/mongodb/binary_backup.feature

@@ -51,6 +51,12 @@ Feature: MongoDB binary backups
     And we restore binary mongo-backup #1 to mongodb02
     Then we got same mongodb data at mongodb01 mongodb02
 
+    # Restore initialized
+    Given mongodb02 has no data
+    And mongodb initialized on mongodb02
+    When we restore binary mongo-backup #1 to mongodb02 as initialized
+    Then we got same mongodb data at mongodb01 mongodb02
+
     # Fifth backup was done successfully
     Given mongodb01 has test mongodb data test5
     When we create binary mongo-backup on mongodb01

tests_func/helpers/walg.go

@@ -163,8 +163,12 @@ func (w *WalgUtil) FetchBackupByNum(backupNum int) error {
 	return err
 }
 
-func (w *WalgUtil) FetchBinaryBackup(backup, mongodConfigPath, mongodbVersion string) error {
-	_, err := w.runCmd("binary-backup-fetch", backup, mongodConfigPath, mongodbVersion)
+func (w *WalgUtil) FetchBinaryBackup(backup, mongodConfigPath, mongodbVersion, rsMembers string) error {
+	cli := []string{"binary-backup-fetch", backup, mongodConfigPath, mongodbVersion}
+	if rsMembers != "" {
+		cli = append(cli, "--mongo-rs-members", rsMembers)
+	}
+	_, err := w.runCmd(cli...)
 	return err
 }
 

tests_func/mongo_binary_backup_steps.go

@@ -8,7 +8,9 @@ import (
 
 func SetupMongodbBinaryBackupSteps(ctx *godog.ScenarioContext, tctx *TestContext) {
 	ctx.Step(`^we create binary mongo-backup on ([^\s]*)$`, tctx.createMongoBinaryBackup)
-	ctx.Step(`^we restore binary mongo-backup #(\d+) to ([^\s]*)`, tctx.restoreMongoBinaryBackup)
+	ctx.Step(`^we restore binary mongo-backup #(\d+) to ([^\s]*)`, tctx.restoreMongoBinaryBackupAsNonInitialized)
+	ctx.Step(`^we restore binary mongo-backup #(\d+) to ([^\s]*) as initialized`,
+		tctx.restoreMongoBinaryBackupAsInitialized)
 }
 
 func (tctx *TestContext) createMongoBinaryBackup(container string) error {
@@ -29,7 +31,15 @@ func (tctx *TestContext) createMongoBinaryBackup(container string) error {
 	return nil
 }
 
-func (tctx *TestContext) restoreMongoBinaryBackup(backupNumber int, container string) error {
+func (tctx *TestContext) restoreMongoBinaryBackupAsNonInitialized(backupNumber int, container string) error {
+	return tctx.restoreMongoBinaryBackup(backupNumber, container, false)
+}
+
+func (tctx *TestContext) restoreMongoBinaryBackupAsInitialized(backupNumber int, container string) error {
+	return tctx.restoreMongoBinaryBackup(backupNumber, container, true)
+}
+
+func (tctx *TestContext) restoreMongoBinaryBackup(backupNumber int, container string, initialized bool) error {
 	walg := WalgUtilFromTestContext(tctx, container)
 
 	backup, err := walg.GetBackupByNumber(backupNumber)
@@ -57,7 +67,11 @@ func (tctx *TestContext) restoreMongoBinaryBackup(backupNumber int, container st
 		return err
 	}
 
-	err = walg.FetchBinaryBackup(backup, configPath, mongodbVersion)
+	rsMembers := ""
+	if initialized {
+		rsMembers = container
+	}
+	err = walg.FetchBinaryBackup(backup, configPath, mongodbVersion, rsMembers)
 	if err != nil {
 		return err
 	}
@@ -70,5 +84,11 @@ func (tctx *TestContext) restoreMongoBinaryBackup(backupNumber int, container st
 		return err
 	}
 
-	return tctx.initiateReplSet(container)
+	if !initialized {
+		if err := tctx.initiateReplSet(container); err != nil {
+			return err
+		}
+	}
+
+	return nil
 }