Return stats directly from restore function (#4713)

Rework restore return status so that later PRs will have a smaller diff

Instead of returning a status and then waiting on the message just return the restore stats directly. The status getter was only setup to wait for one item anyway and was setup to wait after the entire restore operation already completed (at end of
m365/restore.go:ConsumeRestoreCollections())

---

#### Does this PR need a docs update or release note?

- [ ] ✅ Yes, it's included
- [ ] 🕐 Yes, but in a later PR
- [x] ⛔ No

#### Type of change

- [ ] 🌻 Feature
- [ ] 🐛 Bugfix
- [ ] 🗺️ Documentation
- [ ] 🤖 Supportability/Tests
- [ ] 💻 CI/Deployment
- [x] 🧹 Tech Debt/Cleanup

#### Issue(s)

* #4254

#### Test Plan

- [ ] 💪 Manual
- [x] ⚡ Unit test
- [x] 💚 E2E

src/cmd/factory/impl/common.go

@@ -111,7 +111,14 @@ func generateAndRestoreItems(
 		Selector:          sel,
 	}
 
-	return ctrl.ConsumeRestoreCollections(ctx, rcc, dataColls, errs, ctr)
+	deets, _, err := ctrl.ConsumeRestoreCollections(
+		ctx,
+		rcc,
+		dataColls,
+		errs,
+		ctr)
+
+	return deets, clues.Wrap(err, "restoring items").OrNil()
 }
 
 // ------------------------------------------------------------------------------------------
@@ -455,5 +462,12 @@ func generateAndRestoreDriveItems(
 		Selector:          sel,
 	}
 
-	return ctrl.ConsumeRestoreCollections(ctx, rcc, collections, errs, ctr)
+	deets, _, err := ctrl.ConsumeRestoreCollections(
+		ctx,
+		rcc,
+		collections,
+		errs,
+		ctr)
+
+	return deets, clues.Wrap(err, "restoring items").OrNil()
 }

src/internal/m365/controller_test.go

@@ -520,7 +520,7 @@ func (suite *ControllerIntegrationSuite) TestEmptyCollections() {
 				Selector:          test.sel,
 			}
 
-			deets, err := suite.ctrl.ConsumeRestoreCollections(
+			deets, _, err := suite.ctrl.ConsumeRestoreCollections(
 				ctx,
 				rcc,
 				test.col,
@@ -562,7 +562,7 @@ func runRestore(
 		Selector:          restoreSel,
 	}
 
-	deets, err := restoreCtrl.ConsumeRestoreCollections(
+	deets, status, err := restoreCtrl.ConsumeRestoreCollections(
 		ctx,
 		rcc,
 		collections,
@@ -571,7 +571,6 @@ func runRestore(
 	require.NoError(t, err, clues.ToCore(err))
 	assert.NotNil(t, deets)
 
-	status := restoreCtrl.Wait()
 	runTime := time.Since(start)
 
 	assert.Equal(t, numRestoreItems, status.Objects, "restored status.Objects")
@@ -1195,7 +1194,7 @@ func (suite *ControllerIntegrationSuite) TestMultiFolderBackupDifferentNames() {
 					Selector:          restoreSel,
 				}
 
-				deets, err := restoreCtrl.ConsumeRestoreCollections(
+				deets, status, err := restoreCtrl.ConsumeRestoreCollections(
 					ctx,
 					rcc,
 					collections,
@@ -1204,7 +1203,6 @@ func (suite *ControllerIntegrationSuite) TestMultiFolderBackupDifferentNames() {
 				require.NoError(t, err, clues.ToCore(err))
 				require.NotNil(t, deets)
 
-				status := restoreCtrl.Wait()
 				// Always just 1 because it's just 1 collection.
 				assert.Equal(t, totalItems, status.Objects, "status.Objects")
 				assert.Equal(t, totalItems, status.Successes, "status.Successes")

src/internal/m365/mock/connector.go

@@ -76,8 +76,8 @@ func (ctrl Controller) ConsumeRestoreCollections(
 	_ []data.RestoreCollection,
 	_ *fault.Bus,
 	_ *count.Bus,
-) (*details.Details, error) {
+) (*details.Details, *data.CollectionStats, error) {
-	return ctrl.Deets, ctrl.Err
+	return ctrl.Deets, &ctrl.Stats, ctrl.Err
 }
 
 func (ctrl Controller) CacheItemInfo(dii details.ItemInfo) {}

src/internal/m365/restore.go

@@ -12,7 +12,6 @@ import (
 	"github.com/alcionai/corso/src/internal/m365/service/groups"
 	"github.com/alcionai/corso/src/internal/m365/service/onedrive"
 	"github.com/alcionai/corso/src/internal/m365/service/sharepoint"
-	"github.com/alcionai/corso/src/internal/m365/support"
 	"github.com/alcionai/corso/src/internal/operations/inject"
 	"github.com/alcionai/corso/src/pkg/backup/details"
 	"github.com/alcionai/corso/src/pkg/count"
@@ -23,14 +22,13 @@ import (
 
 // ConsumeRestoreCollections restores data from the specified collections
 // into M365 using the GraphAPI.
-// SideEffect: status is updated at the completion of operation
 func (ctrl *Controller) ConsumeRestoreCollections(
 	ctx context.Context,
 	rcc inject.RestoreConsumerConfig,
 	dcs []data.RestoreCollection,
 	errs *fault.Bus,
 	ctr *count.Bus,
-) (*details.Details, error) {
+) (*details.Details, *data.CollectionStats, error) {
 	ctx, end := diagnostics.Span(ctx, "m365:restore")
 	defer end()
 
@@ -38,19 +36,19 @@ func (ctrl *Controller) ConsumeRestoreCollections(
 	ctx = clues.Add(ctx, "restore_config", rcc.RestoreConfig)
 
 	if len(dcs) == 0 {
-		return nil, clues.New("no data collections to restore")
+		return nil, nil, clues.New("no data collections to restore")
 	}
 
 	var (
 		service = rcc.Selector.PathService()
-		status  *support.ControllerOperationStatus
+		stats   *data.CollectionStats
 		deets   *details.Details
 		err     error
 	)
 
 	switch service {
 	case path.ExchangeService:
-		deets, status, err = exchange.ConsumeRestoreCollections(
+		deets, stats, err = exchange.ConsumeRestoreCollections(
 			ctx,
 			ctrl.AC,
 			rcc,
@@ -58,7 +56,7 @@ func (ctrl *Controller) ConsumeRestoreCollections(
 			errs,
 			ctr)
 	case path.OneDriveService:
-		deets, status, err = onedrive.ConsumeRestoreCollections(
+		deets, stats, err = onedrive.ConsumeRestoreCollections(
 			ctx,
 			drive.NewUserDriveRestoreHandler(ctrl.AC),
 			rcc,
@@ -67,7 +65,7 @@ func (ctrl *Controller) ConsumeRestoreCollections(
 			errs,
 			ctr)
 	case path.SharePointService:
-		deets, status, err = sharepoint.ConsumeRestoreCollections(
+		deets, stats, err = sharepoint.ConsumeRestoreCollections(
 			ctx,
 			rcc,
 			ctrl.AC,
@@ -76,7 +74,7 @@ func (ctrl *Controller) ConsumeRestoreCollections(
 			errs,
 			ctr)
 	case path.GroupsService:
-		deets, status, err = groups.ConsumeRestoreCollections(
+		deets, stats, err = groups.ConsumeRestoreCollections(
 			ctx,
 			rcc,
 			ctrl.AC,
@@ -89,8 +87,5 @@ func (ctrl *Controller) ConsumeRestoreCollections(
 		err = clues.Wrap(clues.New(service.String()), "service not supported")
 	}
 
-	ctrl.incrementAwaitingMessages()
+	return deets, stats, err
-	ctrl.UpdateStatus(status)
-
-	return deets, err
 }

src/internal/m365/service/exchange/restore.go

@@ -26,9 +26,9 @@ func ConsumeRestoreCollections(
 	dcs []data.RestoreCollection,
 	errs *fault.Bus,
 	ctr *count.Bus,
-) (*details.Details, *support.ControllerOperationStatus, error) {
+) (*details.Details, *data.CollectionStats, error) {
 	if len(dcs) == 0 {
-		return nil, support.CreateStatus(ctx, support.Restore, 0, support.CollectionMetrics{}, ""), nil
+		return nil, nil, clues.New("no data collections to restore")
 	}
 
 	var (
@@ -119,5 +119,5 @@ func ConsumeRestoreCollections(
 		metrics,
 		rcc.RestoreConfig.Location)
 
-	return deets.Details(), status, el.Failure()
+	return deets.Details(), status.ToCollectionStats(), el.Failure()
 }

src/internal/m365/service/groups/restore.go

@@ -34,7 +34,7 @@ func ConsumeRestoreCollections(
 	dcs []data.RestoreCollection,
 	errs *fault.Bus,
 	ctr *count.Bus,
-) (*details.Details, *support.ControllerOperationStatus, error) {
+) (*details.Details, *data.CollectionStats, error) {
 	var (
 		deets             = &details.Builder{}
 		restoreMetrics    support.CollectionMetrics
@@ -136,7 +136,7 @@ func ConsumeRestoreCollections(
 		restoreMetrics,
 		rcc.RestoreConfig.Location)
 
-	return deets.Details(), status, el.Failure()
+	return deets.Details(), status.ToCollectionStats(), el.Failure()
 }
 
 func getSiteName(

src/internal/m365/service/onedrive/restore.go

@@ -28,7 +28,7 @@ func ConsumeRestoreCollections(
 	dcs []data.RestoreCollection,
 	errs *fault.Bus,
 	ctr *count.Bus,
-) (*details.Details, *support.ControllerOperationStatus, error) {
+) (*details.Details, *data.CollectionStats, error) {
 	var (
 		deets             = &details.Builder{}
 		restoreMetrics    support.CollectionMetrics
@@ -91,7 +91,7 @@ func ConsumeRestoreCollections(
 		restoreMetrics,
 		rcc.RestoreConfig.Location)
 
-	return deets.Details(), status, el.Failure()
+	return deets.Details(), status.ToCollectionStats(), el.Failure()
 }
 
 // Augment restore path to add extra files(meta) needed for restore as

src/internal/m365/service/sharepoint/restore.go

@@ -30,7 +30,7 @@ func ConsumeRestoreCollections(
 	dcs []data.RestoreCollection,
 	errs *fault.Bus,
 	ctr *count.Bus,
-) (*details.Details, *support.ControllerOperationStatus, error) {
+) (*details.Details, *data.CollectionStats, error) {
 	var (
 		deets          = &details.Builder{}
 		lrh            = drive.NewSiteRestoreHandler(ac, rcc.Selector.PathService())
@@ -118,5 +118,5 @@ func ConsumeRestoreCollections(
 		restoreMetrics,
 		rcc.RestoreConfig.Location)
 
-	return deets.Details(), status, el.Failure()
+	return deets.Details(), status.ToCollectionStats(), el.Failure()
 }

src/internal/m365/support/status.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 
 	"github.com/dustin/go-humanize"
+
+	"github.com/alcionai/corso/src/internal/data"
 )
 
 // ControllerOperationStatus is a data type used to describe the state of
@@ -105,3 +107,18 @@ func (cos *ControllerOperationStatus) String() string {
 		operationStatement,
 		cos.details)
 }
+
+// ToCollectionStats turns the called ControllerOperationStatus into a
+// *data.CollectionStats instance
+//
+// TODO(ashmrtn): Remove this when we rework status handling. It's really here
+// to avoid repeated code in each service handler.
+func (cos ControllerOperationStatus) ToCollectionStats() *data.CollectionStats {
+	return &data.CollectionStats{
+		Folders:   cos.Folders,
+		Objects:   cos.Metrics.Objects,
+		Successes: cos.Metrics.Successes,
+		Bytes:     cos.Metrics.Bytes,
+		Details:   cos.String(),
+	}
+}

src/internal/operations/inject/inject.go

@@ -61,7 +61,7 @@ type (
 			dcs []data.RestoreCollection,
 			errs *fault.Bus,
 			ctr *count.Bus,
-		) (*details.Details, error)
+		) (*details.Details, *data.CollectionStats, error)
 
 		IsServiceEnableder
 

src/internal/operations/restore.go

@@ -304,7 +304,7 @@ func (op *RestoreOperation) do(
 	opStats.resourceCount = 1
 	opStats.cs = dcs
 
-	deets, err = consumeRestoreCollections(
+	deets, colStats, err := consumeRestoreCollections(
 		ctx,
 		op.rc,
 		bup.Version,
@@ -319,7 +319,7 @@ func (op *RestoreOperation) do(
 		return nil, clues.Stack(err)
 	}
 
-	opStats.ctrl = op.rc.Wait()
+	opStats.ctrl = colStats
 
 	logger.Ctx(ctx).Debug(opStats.ctrl)
 
@@ -392,7 +392,7 @@ func consumeRestoreCollections(
 	dcs []data.RestoreCollection,
 	errs *fault.Bus,
 	ctr *count.Bus,
-) (*details.Details, error) {
+) (*details.Details, *data.CollectionStats, error) {
 	progressBar := observe.MessageWithCompletion(ctx, observe.ProgressCfg{}, "Restoring data")
 	defer close(progressBar)
 
@@ -404,12 +404,9 @@ func consumeRestoreCollections(
 		Selector:          sel,
 	}
 
-	deets, err := rc.ConsumeRestoreCollections(ctx, rcc, dcs, errs, ctr)
+	deets, status, err := rc.ConsumeRestoreCollections(ctx, rcc, dcs, errs, ctr)
-	if err != nil {
-		return nil, clues.Wrap(err, "restoring collections")
-	}
 
-	return deets, nil
+	return deets, status, clues.Wrap(err, "restoring collections").OrNil()
 }
 
 // formatDetailsForRestoration reduces the provided detail entries according to the

src/internal/operations/test/helper_test.go

@@ -590,7 +590,7 @@ func generateContainerOfItems(
 		Selector:          sel,
 	}
 
-	deets, err := ctrl.ConsumeRestoreCollections(
+	deets, _, err := ctrl.ConsumeRestoreCollections(
 		ctx,
 		rcc,
 		dataColls,
@@ -598,10 +598,6 @@ func generateContainerOfItems(
 		count.New())
 	require.NoError(t, err, clues.ToCore(err))
 
-	// have to wait here, both to ensure the process
-	// finishes, and also to clean up the status
-	ctrl.Wait()
-
 	return deets
 }