Create a ContainerResolver that returns items in priority order (#4607)

Wrapper for existing graph.ContainerResolver structs that returns items
in priority order in the Items() function. Also has logic to filter out
containers if desired

This can help with things like preview backups where we only want to
backup a subset of items because it allows us to ensure we'll see the
containers we care about most first

---

#### 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

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

#### Test Plan

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

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
 }

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)
+}

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
 // ---------------------------------------------------------------------------

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
 }