[CXP-1251] - Promote v0.2.6-test to main

pkg/connector/groups.go

@@ -170,7 +170,7 @@ func (g *groupBuilder) Entitlements(ctx context.Context, resource *v2.Resource,
 
 // Grants return all grants relevant to the group.
 // Databricks Groups have membership and role permissions grants (granting identity resource some permission to this specific group, e.g. group manager).
-func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) {
+func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, attr rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) {
 	l := ctxzap.Extract(ctx)
 
 	var rv []*v2.Grant
@@ -186,6 +186,11 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.S
 		workspaceId = parentId.Resource
 	}
 
+	bag, cursor, err := parseCursorToken(attr.PageToken.Token, &v2.ResourceId{ResourceType: groupResourceType.Id})
+	if err != nil {
+		return nil, nil, fmt.Errorf("databricks-connector: failed to parse page token: %w", err)
+	}
+
 	// membership grants
 	// Always fetch the group with members attribute to ensure we get the members
 	// regardless of authentication type (OAuth vs personal access token)
@@ -199,8 +204,19 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.S
 		annos.WithRateLimiting(rateLimitData)
 	}
 
-	l.Debug("grants: group resource", zap.String("group_id", groupId.Resource), zap.Int("member_count", len(group.Members)))
-	for _, member := range group.Members {
+	// The API returns every member at once and has no member pagination, so members are paged
+	// locally to keep each response under the Lambda payload limit. The API returns members in a
+	// different order on every call, so pageAfter sorts them for pages to line up across re-fetches.
+	members, nextCursor := pageAfter(group.Members, func(m databricks.Member) string { return m.Ref }, cursor)
+
+	l.Debug("grants: group resource",
+		zap.String("group_id", groupId.Resource),
+		zap.Int("member_count", len(group.Members)),
+		zap.Int("page_member_count", len(members)),
+		zap.String("cursor", cursor),
+	)
+
+	for _, member := range members {
 		// member.Ref contains the type and ID separated by "/", e.g., "Users/123" or "Groups/456"
 		pp := strings.Split(member.Ref, "/")
 		if len(pp) != 2 {
@@ -230,7 +246,16 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.S
 		rv = append(rv, grant.NewGrant(resource, groupMemberEntitlement, resourceId, grant.WithAnnotation(anns...)))
 	}
 
-	// role permissions grants
+	nextPage, err := bag.NextToken(nextCursor)
+	if err != nil {
+		return nil, nil, fmt.Errorf("databricks-connector: failed to create next page token: %w", err)
+	}
+
+	// role permissions grants are only emitted with the first page of members
+	if cursor != "" {
+		return rv, &rs.SyncOpResults{NextPageToken: nextPage, Annotations: annos}, nil
+	}
+
 	ruleSets, rateLimitDataRuleSets, err := g.client.ListRuleSets(ctx, workspaceId, GroupsType, groupId.Resource)
 	if err != nil {
 		return nil, nil, fmt.Errorf("databricks-connector: failed to list role rule sets for group %s: %w", resource.Id.Resource, err)
@@ -260,7 +285,7 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.S
 		}
 	}
 
-	return rv, &rs.SyncOpResults{Annotations: annos}, nil
+	return rv, &rs.SyncOpResults{NextPageToken: nextPage, Annotations: annos}, nil
 }
 
 func (g *groupBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) {

pkg/connector/groups_test.go

@@ -3,15 +3,19 @@ package connector
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
+	"math/rand/v2"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"slices"
+	"strings"
 	"testing"
 
 	"github.com/conductorone/baton-databricks/pkg/databricks"
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
+	"github.com/conductorone/baton-sdk/pkg/pagination"
 	ent "github.com/conductorone/baton-sdk/pkg/types/entitlement"
 	rs "github.com/conductorone/baton-sdk/pkg/types/resource"
 	"google.golang.org/grpc/codes"
@@ -373,3 +377,203 @@ func TestGroupRuleSetPrincipalLookup(t *testing.T) {
 		}
 	}
 }
+
+// shuffled returns a shuffled copy of items, mimicking Databricks returning members in a different
+// order on every call.
+func shuffled[T any](items []T) []T {
+	out := slices.Clone(items)
+	rand.Shuffle(len(out), func(i, j int) { out[i], out[j] = out[j], out[i] }) // #nosec G404 -- test data order, not security sensitive
+	return out
+}
+
+func testGroupMembers(count int) []databricks.Member {
+	members := make([]databricks.Member, 0, count)
+	for i := range count {
+		id := fmt.Sprintf("%05d", i)
+		members = append(members, databricks.Member{ID: id, Ref: "Users/" + id})
+	}
+	return members
+}
+
+// newTestGroupBuilder backs a group builder with a test server serving an account group with
+// *members as its members and one rule set granting a service principal the manager role.
+// Tests may change *members between pages. It counts the rule set requests so tests can assert
+// they are only made once.
+func newTestGroupBuilder(t *testing.T, members *[]databricks.Member, ruleSetCalls *int) *groupBuilder {
+	t.Helper()
+	// Every page re-fetches the group. Disable the HTTP cache so each page sees a fresh, reshuffled
+	// response, as it does when pages run in separate Lambda invocations.
+	t.Setenv("BATON_DISABLE_HTTP_CACHE", "true")
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		switch {
+		case strings.HasSuffix(r.URL.Path, "/scim/v2/Groups/g1"):
+			_ = json.NewEncoder(w).Encode(databricks.Group{ID: "g1", DisplayName: "all users", Members: shuffled(*members)})
+		case strings.HasSuffix(r.URL.Path, "/access-control/rule-sets"):
+			*ruleSetCalls++
+			_, _ = w.Write([]byte(`{"grant_rules":[{"role":"roles/group.manager","principals":["servicePrincipals/app-1"]}],"etag":"e1"}`))
+		case strings.HasSuffix(r.URL.Path, "/scim/v2/ServicePrincipals"):
+			_, _ = w.Write([]byte(`{"Resources":[{"id":"sp-1","applicationId":"app-1"}],"totalResults":1}`))
+		default:
+			t.Errorf("unexpected request: %s", r.URL.Path)
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	t.Cleanup(srv.Close)
+
+	target, err := url.Parse(srv.URL)
+	if err != nil {
+		t.Fatalf("parse server url: %v", err)
+	}
+
+	httpClient := &http.Client{Transport: &hostRewriter{target: target}}
+	c, err := databricks.NewClient(context.Background(), httpClient, testHostname, testAccountHost, testAccountId, "", &databricks.NoAuth{}, nil)
+	if err != nil {
+		t.Fatalf("NewClient: %v", err)
+	}
+	return newGroupBuilder(c)
+}
+
+func testAccountGroup(t *testing.T) *v2.Resource {
+	t.Helper()
+	parent := &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: testAccountId}
+	g, err := groupResource(context.Background(), &databricks.Group{ID: "g1", DisplayName: "all users"}, parent)
+	if err != nil {
+		t.Fatalf("groupResource: %v", err)
+	}
+	return g
+}
+
+// groupGrantPages pages through Grants for group until the page token is empty, calling
+// afterPage after each page. It returns how many times each member was granted, the manager
+// grant count and the page count, and checks no page exceeds GrantsPageSize member grants.
+func groupGrantPages(t *testing.T, g *groupBuilder, group *v2.Resource, afterPage func(page int)) (map[string]int, int, int) {
+	t.Helper()
+	seen := map[string]int{}
+	managerGrants := 0
+	pages := 0
+	token := ""
+	for {
+		grants, results, err := g.Grants(context.Background(), group, rs.SyncOpAttrs{PageToken: pagination.Token{Token: token}})
+		if err != nil {
+			t.Fatalf("Grants page %d: %v", pages, err)
+		}
+		pages++
+
+		memberGrants := 0
+		for _, gr := range grants {
+			switch {
+			case strings.HasSuffix(gr.Entitlement.Id, ":"+groupMemberEntitlement):
+				memberGrants++
+				seen[gr.Principal.Id.Resource]++
+			case strings.HasSuffix(gr.Entitlement.Id, ":"+groupManagerEntitlement):
+				managerGrants++
+			}
+		}
+		if memberGrants > int(GrantsPageSize) {
+			t.Errorf("page %d has %d member grants, want at most %d", pages, memberGrants, GrantsPageSize)
+		}
+
+		token = results.NextPageToken
+		if token == "" {
+			return seen, managerGrants, pages
+		}
+		if pages > 100 {
+			t.Fatal("pagination did not terminate")
+		}
+		if afterPage != nil {
+			afterPage(pages)
+		}
+	}
+}
+
+// A group's members come back in one API response, so Grants must split them across pages:
+// a single page for a group with thousands of members exceeds the 6MB Lambda response limit.
+func TestGroupGrantsPagesMembers(t *testing.T) {
+	memberCount := 2*int(GrantsPageSize) + 500
+	members := testGroupMembers(memberCount)
+	ruleSetCalls := 0
+	g := newTestGroupBuilder(t, &members, &ruleSetCalls)
+
+	seen, managerGrants, pages := groupGrantPages(t, g, testAccountGroup(t), nil)
+
+	if pages != 3 {
+		t.Errorf("got %d pages, want 3", pages)
+	}
+	if len(seen) != memberCount {
+		t.Errorf("got %d distinct member grants, want %d", len(seen), memberCount)
+	}
+	for id, n := range seen {
+		if n != 1 {
+			t.Errorf("member %s granted %d times, want 1", id, n)
+		}
+	}
+	if managerGrants != 1 || ruleSetCalls != 1 {
+		t.Errorf("got %d manager grants from %d rule set calls, want 1 from 1", managerGrants, ruleSetCalls)
+	}
+}
+
+// Members removed or added between page fetches must not shift the next page's start, or a
+// member near the boundary is skipped (C1 reads its grant as revoked) or granted twice.
+func TestGroupGrantsMembershipChurnBetweenPages(t *testing.T) {
+	memberCount := 2*int(GrantsPageSize) + 500
+	tests := []struct {
+		name   string
+		change func([]databricks.Member) []databricks.Member
+	}{
+		{
+			// With an offset cursor, every later member moves up one slot and the first member
+			// of page 2 is never granted.
+			name:   "member already emitted is removed",
+			change: func(m []databricks.Member) []databricks.Member { return slices.Clone(m[1:]) },
+		},
+		{
+			// With an offset cursor, every later member moves down one slot and the last member
+			// of page 1 is granted again on page 2.
+			name: "member sorting before the cursor is added",
+			change: func(m []databricks.Member) []databricks.Member {
+				return append(slices.Clone(m), databricks.Member{ID: "00000a", Ref: "Users/00000a"})
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			members := testGroupMembers(memberCount)
+			ruleSetCalls := 0
+			g := newTestGroupBuilder(t, &members, &ruleSetCalls)
+
+			seen, _, _ := groupGrantPages(t, g, testAccountGroup(t), func(page int) {
+				if page == 1 {
+					members = tt.change(members)
+				}
+			})
+
+			// Every original member is granted exactly once; the removed one on page 1, before
+			// its removal. A member added before the cursor is picked up by the next sync.
+			for _, m := range testGroupMembers(memberCount) {
+				if seen[m.ID] != 1 {
+					t.Errorf("member %s granted %d times, want 1", m.ID, seen[m.ID])
+				}
+			}
+		})
+	}
+}
+
+func TestGroupGrantsSmallGroupSinglePage(t *testing.T) {
+	members := testGroupMembers(3)
+	ruleSetCalls := 0
+	g := newTestGroupBuilder(t, &members, &ruleSetCalls)
+
+	grants, results, err := g.Grants(context.Background(), testAccountGroup(t), rs.SyncOpAttrs{})
+	if err != nil {
+		t.Fatalf("Grants: %v", err)
+	}
+	if results.NextPageToken != "" {
+		t.Errorf("got next page token %q, want none", results.NextPageToken)
+	}
+	if len(grants) != 4 {
+		t.Errorf("got %d grants, want 3 members + 1 manager", len(grants))
+	}
+}

pkg/connector/pagination.go

@@ -2,14 +2,20 @@ package connector
 
 import (
 	"fmt"
+	"slices"
 	"strconv"
+	"strings"
 
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 	"github.com/conductorone/baton-sdk/pkg/pagination"
 )
 
 const ResourcesPageSize uint = 50
 
+// GrantsPageSize caps grants built from a single unpaginated API response (e.g. group members).
+// Each grant serializes to ~550 bytes, so a page is ~1.1MB, well under the 6MB Lambda response limit.
+const GrantsPageSize uint = 2000
+
 func parsePageToken(i string, resourceID *v2.ResourceId) (*pagination.Bag, uint, error) {
 	b := &pagination.Bag{}
 	err := b.Unmarshal(i)
@@ -58,3 +64,47 @@ func prepareNextToken(page uint, pageTotal int, total uint) string {
 
 	return token
 }
+
+// parseCursorToken returns the pagination bag and the cursor stored in its current page token,
+// which is empty on the first page.
+func parseCursorToken(i string, resourceID *v2.ResourceId) (*pagination.Bag, string, error) {
+	b := &pagination.Bag{}
+	err := b.Unmarshal(i)
+	if err != nil {
+		return nil, "", err
+	}
+
+	if b.Current() == nil {
+		b.Push(pagination.PageState{
+			ResourceTypeID: resourceID.ResourceType,
+			ResourceID:     resourceID.Resource,
+		})
+	}
+
+	return b, b.PageToken(), nil
+}
+
+// pageAfter sorts items by key and returns up to GrantsPageSize of them that sort after cursor,
+// plus the cursor for the next page (empty when this is the last page). Keys must be unique.
+// Paging from the last emitted key rather than an offset keeps page boundaries stable when items
+// are added or removed between page fetches.
+func pageAfter[T any](items []T, key func(T) string, cursor string) ([]T, string) {
+	slices.SortFunc(items, func(a, b T) int { return strings.Compare(key(a), key(b)) })
+
+	start := 0
+	if cursor != "" {
+		var found bool
+		start, found = slices.BinarySearchFunc(items, cursor, func(item T, c string) int { return strings.Compare(key(item), c) })
+		if found {
+			start++
+		}
+	}
+
+	end := min(start+int(GrantsPageSize), len(items))
+	var next string
+	if end < len(items) {
+		next = key(items[end-1])
+	}
+
+	return items[start:end], next
+}