From 2b247c1b30bb496f0ca445d4bdadc47ab754758d Mon Sep 17 00:00:00 2001
From: Hossain Mahmud <hmsayem@appscode.com>
Date: Wed, 4 Oct 2023 12:32:19 +0600
Subject: [PATCH] Add support for restoring specific snapshot (#1927)

/cherry-pick

Signed-off-by: hmsayem <hmsayem@appscode.com>
---
 pkg/restore.go | 29 +++++++++++++++++++++++------
 pkg/utils.go   |  9 +++++++++
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/pkg/restore.go b/pkg/restore.go
index b48133dfc..2851ef32e 100644
--- a/pkg/restore.go
+++ b/pkg/restore.go
@@ -232,12 +232,14 @@ func (opt *mongoOptions) restoreMongoDB(targetRef api_v1beta1.TargetRef) (*resti
 	// So, for stand-alone MongoDB and MongoDB ReplicaSet, we don't have to do anything.
 	// We only need to update totalHosts field for sharded MongoDB
 
+	restoreSession, err := opt.stashClient.StashV1beta1().RestoreSessions(opt.namespace).Get(context.TODO(), opt.restoreSessionName, metav1.GetOptions{})
+	if err != nil {
+		return nil, err
+	}
+	opt.totalHosts = 1
 	// For sharded MongoDB, parameter.ConfigServer will not be empty
 	if parameters.ConfigServer != "" {
-		restoreSession, err := opt.stashClient.StashV1beta1().RestoreSessions(opt.namespace).Get(context.TODO(), opt.restoreSessionName, metav1.GetOptions{})
-		if err != nil {
-			return nil, err
-		}
+		opt.totalHosts = len(parameters.ReplicaSets) + 1 // for each shard there will be one key in parameters.ReplicaSet
 		_, err = stash_cs_util.UpdateRestoreSessionStatus(
 			context.TODO(),
 			opt.stashClient.StashV1beta1(),
@@ -315,9 +317,8 @@ func (opt *mongoOptions) restoreMongoDB(targetRef api_v1beta1.TargetRef) (*resti
 			Host:       hostKey,
 			SourceHost: hostKey,
 			FileName:   opt.defaultDumpOptions.FileName,
-			Snapshot:   opt.defaultDumpOptions.Snapshot,
+			Snapshot:   opt.getSnapshotForHost(hostKey, restoreSession.Spec.Target.Rules),
 		}
-
 		// setup pipe command
 		restoreCmd := restic.Command{
 			Name: MongoRestoreCMD,
@@ -423,3 +424,19 @@ func dropTempReshardCollection(configsvrDSN string) error {
 
 	return sh.Command(MongoCMD, args...).Command("/usr/bin/tail", "-1").Run()
 }
+
+func (opt *mongoOptions) getSnapshotForHost(hostname string, rules []api_v1beta1.Rule) string {
+	var hostSnapshot string
+	for _, rule := range rules {
+		if len(rule.TargetHosts) == 0 || containsString(rule.TargetHosts, hostname) {
+			hostSnapshot = rule.Snapshots[0]
+			// if rule has empty targetHost then check further rules to see if any other rule with non-empty targetHost matches
+			if len(rule.TargetHosts) == 0 {
+				continue
+			} else {
+				return hostSnapshot
+			}
+		}
+	}
+	return hostSnapshot
+}
diff --git a/pkg/utils.go b/pkg/utils.go
index c28de5cf1..347b81234 100644
--- a/pkg/utils.go
+++ b/pkg/utils.go
@@ -94,3 +94,12 @@ func containsArg(args []string, checklist sets.String) bool {
 	}
 	return false
 }
+
+func containsString(a []string, e string) bool {
+	for _, s := range a {
+		if s == e {
+			return true
+		}
+	}
+	return false
+}