From 8634b1b1aca135b417aed861b1aff45d3a8ca236 Mon Sep 17 00:00:00 2001 From: Vaibhav Kamra Date: Wed, 31 Aug 2022 11:39:18 -0700 Subject: [PATCH] Helpers to create OneDrive Collections (#692) ## Description This adds the logic that materializes collections for a specified user The collection paths are currently derived from OneDrive metadata but a future PR will introduce `paths` pkg support to create these. ## Type of change Please check the type of change your PR introduces: - [x] :sunflower: Feature - [ ] :bug: Bugfix - [ ] :world_map: Documentation - [ ] :robot: Test - [ ] :hamster: Trivial/Minor ## Issue(s) #388 ## Test Plan - [ ] :muscle: Manual - [x] :zap: Unit test - [ ] :green_heart: E2E --- .../connector/onedrive/collections.go | 117 +++++++++++++++ .../connector/onedrive/collections_test.go | 135 ++++++++++++++++++ 2 files changed, 252 insertions(+) create mode 100644 src/internal/connector/onedrive/collections.go create mode 100644 src/internal/connector/onedrive/collections_test.go diff --git a/src/internal/connector/onedrive/collections.go b/src/internal/connector/onedrive/collections.go new file mode 100644 index 000000000..367a01869 --- /dev/null +++ b/src/internal/connector/onedrive/collections.go @@ -0,0 +1,117 @@ +package onedrive + +import ( + "context" + "path" + + "github.com/microsoftgraph/msgraph-sdk-go/models" + "github.com/pkg/errors" + + "github.com/alcionai/corso/internal/connector/graph" + "github.com/alcionai/corso/internal/connector/support" + "github.com/alcionai/corso/internal/data" +) + +// Collections is used to retrieve OneDrive data for a +// specified user +type Collections struct { + user string + // collectionMap allows lookup of the data.Collection + // for a OneDrive folder + collectionMap map[string]data.Collection + service graph.Service + statusCh chan<- *support.ConnectorOperationStatus + + // Track stats from drive enumeration + numItems int + numDirs int + numFiles int + numPackages int +} + +func NewCollections( + user string, + service graph.Service, + statusCh chan<- *support.ConnectorOperationStatus, +) *Collections { + return &Collections{ + user: user, + collectionMap: map[string]data.Collection{}, + service: service, + statusCh: statusCh, + } +} + +// Retrieves OneDrive data as set of `data.Collections` +func (c *Collections) Get(ctx context.Context) ([]data.Collection, error) { + // Enumerate drives for the specified user + drives, err := drives(ctx, c.service, c.user) + if err != nil { + return nil, err + } + + // Update the collection map with items from each drive + for _, d := range drives { + err = collectItems(ctx, c.service, *d.GetId(), c.updateCollections) + if err != nil { + return nil, err + } + } + + collections := make([]data.Collection, 0, len(c.collectionMap)) + for _, c := range c.collectionMap { + collections = append(collections, c) + } + + return collections, nil +} + +// updateCollections initializes and adds the provided OneDrive items to Collections +// A new collection is created for every OneDrive folder (or package) +func (c *Collections) updateCollections(ctx context.Context, driveID string, items []models.DriveItemable) error { + for _, item := range items { + err := c.stats(item) + if err != nil { + return err + } + if item.GetParentReference() == nil || item.GetParentReference().GetPath() == nil { + return errors.Errorf("item does not have a parent reference. item name : %s", *item.GetName()) + } + // Create a collection for the parent of this item + collectionPath := *item.GetParentReference().GetPath() + if _, found := c.collectionMap[collectionPath]; !found { + c.collectionMap[collectionPath] = NewCollection(collectionPath, driveID, c.service, c.statusCh) + } + switch { + case item.GetFolder() != nil, item.GetPackage() != nil: + // For folders and packages we also create a collection to represent those + // TODO: This is where we might create a "special file" to represent these in the backup repository + // e.g. a ".folderMetadataFile" + itemPath := path.Join(*item.GetParentReference().GetPath(), *item.GetName()) + if _, found := c.collectionMap[itemPath]; !found { + c.collectionMap[itemPath] = NewCollection(itemPath, driveID, c.service, c.statusCh) + } + case item.GetFile() != nil: + collection := c.collectionMap[collectionPath].(*Collection) + collection.Add(*item.GetId()) + default: + return errors.Errorf("item type not supported. item name : %s", *item.GetName()) + } + } + return nil +} + +func (c *Collections) stats(item models.DriveItemable) error { + switch { + case item.GetFolder() != nil: + c.numDirs++ + case item.GetPackage() != nil: + c.numPackages++ + case item.GetFile() != nil: + c.numFiles++ + default: + return errors.Errorf("item type not supported. item name : %s", *item.GetName()) + } + c.numItems++ + return nil +} diff --git a/src/internal/connector/onedrive/collections_test.go b/src/internal/connector/onedrive/collections_test.go new file mode 100644 index 000000000..9cf02f6a8 --- /dev/null +++ b/src/internal/connector/onedrive/collections_test.go @@ -0,0 +1,135 @@ +package onedrive + +import ( + "context" + "testing" + + msgraphsdk "github.com/microsoftgraph/msgraph-sdk-go" + "github.com/microsoftgraph/msgraph-sdk-go/models" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/suite" +) + +type OneDriveCollectionsSuite struct { + suite.Suite +} + +func TestOneDriveCollectionsSuite(t *testing.T) { + suite.Run(t, new(OneDriveCollectionsSuite)) +} + +func (suite *OneDriveCollectionsSuite) TestUpdateCollections() { + tests := []struct { + testCase string + items []models.DriveItemable + expect assert.ErrorAssertionFunc + expectedCollectionPaths []string + expectedItemCount int + expectedFolderCount int + expectedPackageCount int + expectedFileCount int + }{ + { + testCase: "Invalid item", + items: []models.DriveItemable{ + driveItem("item", "/root", false, false, false), + }, + expect: assert.Error, + }, + { + testCase: "Single File", + items: []models.DriveItemable{ + driveItem("file", "/root", true, false, false), + }, + expect: assert.NoError, + expectedCollectionPaths: []string{"/root"}, + expectedItemCount: 1, + expectedFileCount: 1, + }, + { + testCase: "Single Folder", + items: []models.DriveItemable{ + driveItem("folder", "/root", false, true, false), + }, + expect: assert.NoError, + expectedCollectionPaths: []string{"/root", "/root/folder"}, + expectedItemCount: 1, + expectedFolderCount: 1, + }, + { + testCase: "Single Package", + items: []models.DriveItemable{ + driveItem("package", "/root", false, false, true), + }, + expect: assert.NoError, + expectedCollectionPaths: []string{"/root", "/root/package"}, + expectedItemCount: 1, + expectedPackageCount: 1, + }, + { + testCase: "1 root file, 1 folder, 1 package, 2 files, 3 collections", + items: []models.DriveItemable{ + driveItem("fileInRoot", "/root", true, false, false), + driveItem("folder", "/root", false, true, false), + driveItem("package", "/root", false, false, true), + driveItem("fileInFolder", "/root/folder", true, false, false), + driveItem("fileInPackage", "/root/package", true, false, false), + }, + expect: assert.NoError, + expectedCollectionPaths: []string{"/root", "/root/folder", "/root/package"}, + expectedItemCount: 5, + expectedFileCount: 3, + expectedFolderCount: 1, + expectedPackageCount: 1, + }, + } + for _, tt := range tests { + suite.T().Run(tt.testCase, func(t *testing.T) { + c := NewCollections("user", &MockGraphService{}, nil) + err := c.updateCollections(context.Background(), "driveID", tt.items) + tt.expect(t, err) + assert.Equal(t, len(tt.expectedCollectionPaths), len(c.collectionMap)) + assert.Equal(t, tt.expectedItemCount, c.numItems) + assert.Equal(t, tt.expectedFileCount, c.numFiles) + assert.Equal(t, tt.expectedFolderCount, c.numDirs) + assert.Equal(t, tt.expectedPackageCount, c.numPackages) + for _, collPath := range tt.expectedCollectionPaths { + assert.Contains(t, c.collectionMap, collPath) + } + }) + } +} + +func driveItem(name string, path string, isFile, isFolder, isPackage bool) models.DriveItemable { + item := models.NewDriveItem() + item.SetName(&name) + item.SetId(&name) + + parentReference := models.NewItemReference() + parentReference.SetPath(&path) + item.SetParentReference(parentReference) + + switch { + case isFile: + item.SetFile(models.NewFile()) + case isFolder: + item.SetFolder(models.NewFolder()) + case isPackage: + item.SetPackage(models.NewPackage_escaped()) + } + return item +} + +type MockGraphService struct{} + +func (ms *MockGraphService) Client() *msgraphsdk.GraphServiceClient { + return nil +} + +func (ms *MockGraphService) Adapter() *msgraphsdk.GraphRequestAdapter { + return nil +} + +func (ms *MockGraphService) ErrPolicy() bool { + return false +}