diff --git a/src/internal/m365/collection/exchange/backup_test.go b/src/internal/m365/collection/exchange/backup_test.go index eaa81a435..c032cb063 100644 --- a/src/internal/m365/collection/exchange/backup_test.go +++ b/src/internal/m365/collection/exchange/backup_test.go @@ -121,6 +121,16 @@ func newMockResolver(items ...mockContainer) mockResolver { return mockResolver{items: is} } +func (m mockResolver) ItemByID(id string) graph.CachedContainer { + for _, c := range m.items { + if ptr.Val(c.GetId()) == id { + return c + } + } + + return nil +} + func (m mockResolver) Items() []graph.CachedContainer { return m.items } diff --git a/src/internal/m365/collection/exchange/container_resolver.go b/src/internal/m365/collection/exchange/container_resolver.go index 40b80dd70..730850b4e 100644 --- a/src/internal/m365/collection/exchange/container_resolver.go +++ b/src/internal/m365/collection/exchange/container_resolver.go @@ -4,6 +4,7 @@ import ( "context" "github.com/alcionai/clues" + "golang.org/x/exp/slices" "github.com/alcionai/corso/src/internal/common/ptr" "github.com/alcionai/corso/src/pkg/fault" @@ -353,6 +354,10 @@ func (cr *containerResolver) addFolder(cf graph.CachedContainer) error { return nil } +func (cr *containerResolver) ItemByID(id string) graph.CachedContainer { + return cr.cache[id] +} + func (cr *containerResolver) Items() []graph.CachedContainer { res := make([]graph.CachedContainer, 0, len(cr.cache)) @@ -411,3 +416,137 @@ func (cr *containerResolver) populatePaths( return lastErr } + +// --------------------------------------------------------------------------- +// rankedContainerResolver +// --------------------------------------------------------------------------- + +type rankedContainerResolver struct { + graph.ContainerResolver + // resolvedInclude is the ordered list of resolved container IDs to add to the + // start of the Items result set. + resolvedInclude []string + // resolvedExclude is the set of items that shouldn't be included in the + // result of Items or ItemByID. Uses actual container IDs instead of + // well-known names. + resolvedExclude map[string]struct{} +} + +// newRankedContainerResolver creates a wrapper around base that returns results +// from Items in priority order. Priority is defined by includeRankedIDs. All +// items that don't appear in includeRankedIDs are considered to have equal +// priority but lower priority than those in includeRankedIDs. +// +// includeRankedIDs is the set of containers to place at the start of the result +// of Items in the order they should appear. IDs can either be actual +// container IDs or well-known container IDs like "inbox". +// +// excludeIDs is the set of IDs that shouldn't be in the results returned by +// Items. IDs can either be actual container IDs or well-known container IDs +// like "inbox". +// +// The include set takes priority over the exclude set, so container IDs +// appearing in both will be considered included and be returned by calls like +// Items and ItemByID. +func newRankedContainerResolver( + ctx context.Context, + base graph.ContainerResolver, + getter containerGetter, + userID string, + includeRankedIDs []string, + excludeIDs []string, +) (*rankedContainerResolver, error) { + if base == nil { + return nil, clues.New("nil base ContainerResolver") + } + + cr := &rankedContainerResolver{ + resolvedInclude: make([]string, 0, len(includeRankedIDs)), + resolvedExclude: make(map[string]struct{}, len(excludeIDs)), + ContainerResolver: base, + } + + // For both includes and excludes we need to get the container IDs from graph. + // This is required because the user could hand us one of the "well-known" + // IDs, which we don't use in the underlying container resolver. Resolving + // these here will allow us to match by ID later on. + for _, id := range includeRankedIDs { + ictx := clues.Add(ctx, "container_id", id) + + c, err := getter.GetContainerByID(ctx, userID, id) + if err != nil { + return nil, clues.Wrap(err, "getting ranked container").WithClues(ictx) + } + + gotID := ptr.Val(c.GetId()) + if len(gotID) == 0 { + return nil, clues.New("ranked include container missing ID"). + WithClues(ictx) + } + + cr.resolvedInclude = append(cr.resolvedInclude, gotID) + } + + for _, id := range excludeIDs { + ictx := clues.Add(ctx, "container_id", id) + + c, err := getter.GetContainerByID(ctx, userID, id) + if err != nil { + return nil, clues.Wrap(err, "getting exclude container").WithClues(ictx) + } + + gotID := ptr.Val(c.GetId()) + if len(gotID) == 0 { + return nil, clues.New("exclude container missing ID"). + WithClues(ictx) + } + + cr.resolvedExclude[gotID] = struct{}{} + } + + return cr, nil +} + +func (cr *rankedContainerResolver) Items() []graph.CachedContainer { + found := cr.ContainerResolver.Items() + res := make([]graph.CachedContainer, 0, len(found)) + + // Add the ranked items first. + // + // TODO(ashmrtn): If we need to handle a large number of ranked items we + // should think about making a map of the ranked items for fast lookups later + // in the function. + for _, include := range cr.resolvedInclude { + if c := cr.ContainerResolver.ItemByID(include); c != nil { + res = append(res, c) + } + } + + // Add the remaining, filtering out any of the ones we need to exclude or that + // we already added because they were ranked. + for _, c := range found { + if _, ok := cr.resolvedExclude[ptr.Val(c.GetId())]; ok { + continue + } + + if slices.Contains(cr.resolvedInclude, ptr.Val(c.GetId())) { + continue + } + + res = append(res, c) + } + + return res +} + +func (cr *rankedContainerResolver) ItemByID(id string) graph.CachedContainer { + // Includes take priority over excludes so check those too. + _, exclude := cr.resolvedExclude[id] + includeIdx := slices.Index(cr.resolvedInclude, id) + + if exclude && includeIdx == -1 { + return nil + } + + return cr.ContainerResolver.ItemByID(id) +} diff --git a/src/internal/m365/collection/exchange/container_resolver_test.go b/src/internal/m365/collection/exchange/container_resolver_test.go index 9ec895945..e0cbe9fe4 100644 --- a/src/internal/m365/collection/exchange/container_resolver_test.go +++ b/src/internal/m365/collection/exchange/container_resolver_test.go @@ -11,6 +11,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" + "golang.org/x/exp/maps" + "golang.org/x/exp/slices" "github.com/alcionai/corso/src/internal/common/ptr" "github.com/alcionai/corso/src/internal/tester" @@ -53,6 +55,346 @@ func strPtr(s string) *string { return &s } +var _ graph.ContainerResolver = mockContainerResolver{} + +type mockContainerResolver struct { + containersByID map[string]graph.CachedContainer +} + +func (mc mockContainerResolver) IDToPath( + ctx context.Context, + id string, +) (*path.Builder, *path.Builder, error) { + return nil, nil, clues.New("not implemented") +} + +func (mc mockContainerResolver) Populate( + ctx context.Context, + errs *fault.Bus, + baseFolderID string, + baseContainerPath ...string, +) error { + return clues.New("not implemented") +} + +func (mc mockContainerResolver) PathInCache(p string) (string, bool) { + return "", false +} + +func (mc mockContainerResolver) LocationInCache(p string) (string, bool) { + return "", false +} + +func (mc mockContainerResolver) AddToCache( + ctx context.Context, + c graph.Container, +) error { + return clues.New("not implemented") +} + +func (mc mockContainerResolver) ItemByID(id string) graph.CachedContainer { + return mc.containersByID[id] +} + +func (mc mockContainerResolver) Items() []graph.CachedContainer { + return maps.Values(mc.containersByID) +} + +var _ containerGetter = mockContainerGetter{} + +type containerGetterRes struct { + c graph.Container + err error +} + +type mockContainerGetter struct { + itemsByID map[string]containerGetterRes +} + +func (mcg mockContainerGetter) GetContainerByID( + ctx context.Context, + userID string, + containerID string, +) (graph.Container, error) { + res := mcg.itemsByID[containerID] + return res.c, res.err +} + +// --------------------------------------------------------------------------- +// rankedContainerResolver unit tests +// --------------------------------------------------------------------------- + +type RankedContainerResolverUnitSuite struct { + tester.Suite +} + +func TestRankedContainerResolverUnitSuite(t *testing.T) { + suite.Run(t, &RankedContainerResolverUnitSuite{ + Suite: tester.NewUnitSuite(t), + }) +} + +func (suite *RankedContainerResolverUnitSuite) TestItemByID() { + // Containers available to operate on directly in tests. + const ( + id1 = "id1" + idNotInBase = "idNotInBase" + ) + + mcg := mockContainerGetter{ + itemsByID: map[string]containerGetterRes{ + id1: {c: mockContainer{id: ptr.To(id1)}}, + idNotInBase: {c: mockContainer{id: ptr.To(idNotInBase)}}, + }, + } + + // Configure base containers we want. + mcr := &mockContainerResolver{ + containersByID: map[string]graph.CachedContainer{ + id1: &mockCachedContainer{id: id1}, + }, + } + + table := []struct { + name string + includes []string + excludes []string + itemID string + expectFound bool + }{ + { + name: "NeitherIncludedNorExcluded", + itemID: id1, + expectFound: true, + }, + { + name: "NeitherIncludedOrExcluded NotInBase", + itemID: idNotInBase, + expectFound: false, + }, + { + name: "Included", + includes: []string{id1}, + itemID: id1, + expectFound: true, + }, + { + name: "Excluded", + excludes: []string{id1}, + itemID: id1, + expectFound: false, + }, + { + name: "IncludedAndExcluded", + includes: []string{id1}, + excludes: []string{id1}, + itemID: id1, + expectFound: true, + }, + { + name: "IncludedAndExcluded NotInBaseResolver", + includes: []string{idNotInBase}, + excludes: []string{idNotInBase}, + itemID: idNotInBase, + expectFound: false, + }, + } + + for _, test := range table { + suite.Run(test.name, func() { + t := suite.T() + + ctx, flush := tester.NewContext(t) + defer flush() + + rcr, err := newRankedContainerResolver( + ctx, + mcr, + mcg, + "userID", + test.includes, + test.excludes) + require.NoError(t, err, clues.ToCore(err)) + + item := rcr.ItemByID(test.itemID) + if test.expectFound { + require.NotNil(t, item) + assert.Equal(t, test.itemID, ptr.Val(item.GetId()), "returned item ID") + } else { + assert.Nil(t, item, "unexpected item returned: %+v", item) + } + }) + } +} + +func (suite *RankedContainerResolverUnitSuite) TestItems() { + // Containers available to operate on directly in tests. + const ( + id1 = "id1" + id2 = "id2" + id3 = "id3" + idErr = "idErr" + idEmpty = "idEmpty" + idNotInBase = "idNotInBase" + ) + + mcg := mockContainerGetter{ + itemsByID: map[string]containerGetterRes{ + id1: {c: mockContainer{id: ptr.To(id1)}}, + id2: {c: mockContainer{id: ptr.To(id2)}}, + id3: {c: mockContainer{id: ptr.To(id3)}}, + idErr: {err: assert.AnError}, + idEmpty: {c: mockContainer{}}, + idNotInBase: {c: mockContainer{id: ptr.To(idNotInBase)}}, + }, + } + + // Configure base containers we want. + mcr := &mockContainerResolver{ + containersByID: map[string]graph.CachedContainer{ + id1: &mockCachedContainer{id: id1}, + id2: &mockCachedContainer{id: id2}, + id3: &mockCachedContainer{id: id3}, + }, + } + + // Add a bunch more containers so we're more likely to get a random order from + // the map. + var otherContainerIDs []string + + for i := 0; i < 100; i++ { + id := fmt.Sprintf("extra-id%d", i) + mcr.containersByID[id] = &mockCachedContainer{id: id} + otherContainerIDs = append(otherContainerIDs, id) + } + + table := []struct { + name string + includes []string + excludes []string + // expectPrefix is the prefix of container IDs that should be in the result. + expectPrefix []string + // expectExtraUnordered allows specifying additional IDs in the set of IDs + // available to work with that will appear in the unordered set. For + // example, if only id1 was part of includes and excludes was empty then id2 + // and id3 should appear somewhere in the output but don't have a particluar + // order requirement. + expectExtraUnordered []string + expectErr assert.ErrorAssertionFunc + }{ + { + name: "ReturnsRankedItems", + includes: []string{id2, id1, id3}, + expectPrefix: []string{id2, id1, id3}, + expectErr: assert.NoError, + }, + { + name: "ReturnsRankedItems SomeUnordered", + includes: []string{id2, id1}, + expectPrefix: []string{id2, id1}, + expectExtraUnordered: []string{id3}, + expectErr: assert.NoError, + }, + { + name: "ReturnsRankedItems SomeExcluded", + includes: []string{id2}, + excludes: []string{id3, id1}, + expectPrefix: []string{id2}, + expectErr: assert.NoError, + }, + { + name: "ReturnsRankedItems SomeIncludesNotInBase", + includes: []string{id2, id1, idNotInBase}, + expectPrefix: []string{id2, id1}, + expectExtraUnordered: []string{id3}, + expectErr: assert.NoError, + }, + { + name: "ReturnsRankedItems IncludedAndExcluded", + includes: []string{id2}, + excludes: []string{id2}, + expectPrefix: []string{id2}, + expectExtraUnordered: []string{id1, id3}, + expectErr: assert.NoError, + }, + { + name: "ReturnsRankedItems IncludedAndExcluded NotInBase", + includes: []string{idNotInBase}, + excludes: []string{idNotInBase}, + expectPrefix: []string{}, + expectExtraUnordered: []string{id1, id2, id3}, + expectErr: assert.NoError, + }, + { + name: "ReturnsRankedItems SomeExcludesNotInBase", + includes: []string{id2}, + excludes: []string{id1, idNotInBase, id3}, + expectPrefix: []string{id2}, + expectErr: assert.NoError, + }, + { + name: "FailsOnIncludeError", + includes: []string{id2, id1, idErr}, + expectErr: assert.Error, + }, + { + name: "FailsOnExcludeError", + excludes: []string{id2, id1, idErr}, + expectErr: assert.Error, + }, + { + name: "FailsOnIncludeContainerWithEmptyID", + includes: []string{id2, id1, idEmpty}, + expectErr: assert.Error, + }, + { + name: "FailsOnExcludeContainerWithEmptyID", + excludes: []string{id2, id1, idEmpty}, + expectErr: assert.Error, + }, + } + + for _, test := range table { + suite.Run(test.name, func() { + t := suite.T() + + ctx, flush := tester.NewContext(t) + defer flush() + + rcr, err := newRankedContainerResolver( + ctx, + mcr, + mcg, + "userID", + test.includes, + test.excludes) + test.expectErr(t, err, clues.ToCore(err)) + + if err != nil { + return + } + + items := rcr.Items() + resIDs := make([]string, 0, len(items)) + + for _, item := range items { + resIDs = append(resIDs, ptr.Val(item.GetId())) + } + + assert.Equal( + t, + test.expectPrefix, + resIDs[:len(test.expectPrefix)], + "ordered prefix of result") + assert.ElementsMatch( + t, + append(slices.Clone(test.expectExtraUnordered), otherContainerIDs...), + resIDs[len(test.expectPrefix):], + "unordered remainder of result") + }) + } +} + // --------------------------------------------------------------------------- // unit suite // --------------------------------------------------------------------------- diff --git a/src/pkg/services/m365/api/graph/cache_container.go b/src/pkg/services/m365/api/graph/cache_container.go index ad3866501..41ab4a4d2 100644 --- a/src/pkg/services/m365/api/graph/cache_container.go +++ b/src/pkg/services/m365/api/graph/cache_container.go @@ -62,6 +62,11 @@ type ContainerResolver interface { AddToCache(ctx context.Context, m365Container Container) error + // ItemByID returns the container with the given ID if it's in the container + // resolver. If the item isn't in the resolver then it returns nil. Assumes + // resolved IDs are used not well-known names for containers. + ItemByID(id string) CachedContainer + // Items returns the containers in the cache. Items() []CachedContainer }