From 14205d61e28e8f35e9f500f7d8a610c62c92d137 Mon Sep 17 00:00:00 2001
From: Vaibhav Kamra <vkamra@alcion.ai>
Date: Tue, 13 Sep 2022 17:14:27 -0700
Subject: [PATCH] Refactor restore helper (#838)

## Description

Small refactor of the data collection restore helper so that we can plug-in the OneDrive helpers in a
follow-up PR

## Type of change

<!--- Please check the type of change your PR introduces: --->
- [ ] :sunflower: Feature
- [ ] :bug: Bugfix
- [ ] :world_map: Documentation
- [ ] :robot: Test
- [ ] :computer: CI/Deployment
- [x] :hamster: Trivial/Minor

## Test Plan

<!-- How will this be tested prior to merging.-->
- [ ] :muscle: Manual
- [ ] :zap: Unit test
- [x] :green_heart: E2E
---
 src/internal/connector/graph_connector.go | 27 ++++++++++++++++++-----
 src/internal/operations/restore.go        |  2 +-
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/internal/connector/graph_connector.go b/src/internal/connector/graph_connector.go
index cab564563..cb7437bc4 100644
--- a/src/internal/connector/graph_connector.go
+++ b/src/internal/connector/graph_connector.go
@@ -242,10 +242,9 @@ func (gc *GraphConnector) ExchangeDataCollection(
 	return collections, errs
 }
 
-// RestoreExchangeDataCollection: Utility function to connect to M365 backstore
-// and upload messages from DataCollection.
-// FullPath: tenantId, userId, <collectionCategory>, FolderId
-func (gc *GraphConnector) RestoreExchangeDataCollection(
+// RestoreDataCollections restores data from the specified collections
+// into M365
+func (gc *GraphConnector) RestoreDataCollections(
 	ctx context.Context,
 	dcs []data.Collection,
 ) error {
@@ -276,10 +275,23 @@ func (gc *GraphConnector) RestoreExchangeDataCollection(
 
 		category := directory.Category()
 		user := directory.ResourceOwner()
+		service := directory.Service()
+
+		// Check whether restoring data into the specified service is supported
+		switch service {
+		case path.ExchangeService:
+			// Supported
+		default:
+			return errors.Errorf("restore data from service %s not supported", service.String())
+		}
 
 		if _, ok := pathCounter[directory.String()]; !ok {
 			pathCounter[directory.String()] = true
-			folderID, errs = exchange.GetRestoreContainer(&gc.graphService, user, category)
+
+			switch service {
+			case path.ExchangeService:
+				folderID, errs = exchange.GetRestoreContainer(&gc.graphService, user, category)
+			}
 
 			if errs != nil {
 				return errs
@@ -305,7 +317,10 @@ func (gc *GraphConnector) RestoreExchangeDataCollection(
 					continue
 				}
 
-				err = exchange.RestoreExchangeObject(ctx, buf.Bytes(), category, policy, &gc.graphService, folderID, user)
+				switch service {
+				case path.ExchangeService:
+					err = exchange.RestoreExchangeObject(ctx, buf.Bytes(), category, policy, &gc.graphService, folderID, user)
+				}
 
 				if err != nil {
 					errs = support.WrapAndAppend(itemData.UUID(), err, errs)
diff --git a/src/internal/operations/restore.go b/src/internal/operations/restore.go
index 3501acd01..a04917ae2 100644
--- a/src/internal/operations/restore.go
+++ b/src/internal/operations/restore.go
@@ -139,7 +139,7 @@ func (op *RestoreOperation) Run(ctx context.Context) (err error) {
 		return err
 	}
 
-	err = gc.RestoreExchangeDataCollection(ctx, dcs)
+	err = gc.RestoreDataCollections(ctx, dcs)
 	if err != nil {
 		err = errors.Wrap(err, "restoring service data")
 		opStats.writeErr = err