From 6384cfa15b868e88d4d3f57395ca1185e3da92ed Mon Sep 17 00:00:00 2001 From: agustin-conductor Date: Tue, 27 Oct 2026 10:24:18 -0300 Subject: [PATCH] [CXP-1251] - Paginate group member grants to stay under the Lambda 6MB response limit (#81) * Paginate group and workspace member grants to stay under Lambda payload limit Databricks returns every group member in a single Get group response and has no member pagination. Grants() emitted one grant per member in one response, so a group with ~11k+ members (e.g. an all-users account group) exceeded the 6MB Lambda response limit and failed with HTTP 413 on every retry. Page members locally (GrantsPageSize = 2000, ~1.1MB per page) using the page token as a member offset, and emit role rule-set grants on the first page only. Databricks returns members in a different order on every call, so they are sorted before slicing so pages line up across re-fetches in separate Lambda invocations. Apply the same paging to workspace assignment grants. Co-Authored-By: Claude Opus 5.5 * Page member grants by last emitted key instead of offset An offset page token into a list that is re-fetched for every page drifts when members are added or removed between page invocations: a removal skips the first member of the next page (C1 reads that grant as revoked) and an addition grants one member twice. Store the key of the last emitted item (member Ref for groups, principal ID for workspaces) and start the next page after it, so membership churn no longer shifts page boundaries. Add churn tests for both paths, a workspace multi-page test, and mark the test-only shuffle as not security sensitive for gosec. Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Claude Opus 5.5 --- pkg/connector/groups.go | 35 +++++- pkg/connector/groups_test.go | 204 +++++++++++++++++++++++++++++++ pkg/connector/pagination.go | 50 ++++++++ pkg/connector/workspaces.go | 36 +++++- pkg/connector/workspaces_test.go | 90 ++++++++++++++ 5 files changed, 406 insertions(+), 9 deletions(-) create mode 100644 pkg/connector/workspaces_test.go diff --git a/pkg/connector/groups.go b/pkg/connector/groups.go index a394e5a6..b5979b73 100644 --- a/pkg/connector/groups.go +++ b/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) { diff --git a/pkg/connector/groups_test.go b/pkg/connector/groups_test.go index b9478870..9882f7d6 100644 --- a/pkg/connector/groups_test.go +++ b/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)) + } +} diff --git a/pkg/connector/pagination.go b/pkg/connector/pagination.go index ca0c111d..f9157b61 100644 --- a/pkg/connector/pagination.go +++ b/pkg/connector/pagination.go @@ -2,7 +2,9 @@ package connector import ( "fmt" + "slices" "strconv" + "strings" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/pagination" @@ -10,6 +12,10 @@ import ( 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 +} diff --git a/pkg/connector/workspaces.go b/pkg/connector/workspaces.go index 49ec7cea..20990ca6 100644 --- a/pkg/connector/workspaces.go +++ b/pkg/connector/workspaces.go @@ -23,6 +23,13 @@ import ( const workspaceMemberEntitlement = "member" +func principalID(a databricks.WorkspaceAssignment) int { + if a.Principal == nil { + return 0 + } + return a.Principal.ID +} + type workspaceBuilder struct { client *databricks.Client resourceType *v2.ResourceType @@ -166,7 +173,7 @@ func (w *workspaceBuilder) Entitlements(_ context.Context, resource *v2.Resource // Grants returns slice of grants representing workspace members. // To get workspace members, we can only use the account API. -func (w *workspaceBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) { +func (w *workspaceBuilder) Grants(ctx context.Context, resource *v2.Resource, attr rs.SyncOpAttrs) ([]*v2.Grant, *rs.SyncOpResults, error) { l := ctxzap.Extract(ctx) if !w.client.IsAccountAPIAvailable() { @@ -180,6 +187,11 @@ func (w *workspaceBuilder) Grants(ctx context.Context, resource *v2.Resource, _ return nil, nil, fmt.Errorf("databricks-connector: failed to get workspace ID") } + bag, cursor, err := parseCursorToken(attr.PageToken.Token, &v2.ResourceId{ResourceType: workspaceResourceType.Id}) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to parse page token: %w", err) + } + workspace := strconv.Itoa(int(workspaceId)) assignments, rateLimitDesc, err := w.client.ListWorkspaceMembers(ctx, workspace) annos := annotations.Annotations{} @@ -212,9 +224,20 @@ func (w *workspaceBuilder) Grants(ctx context.Context, resource *v2.Resource, _ return nil, &rs.SyncOpResults{Annotations: annos}, fmt.Errorf("databricks-connector: failed to list workspace members: %w", err) } + // The API returns every assignment at once, so they are paged locally to keep each + // response under the Lambda payload limit. + assignmentsPage, nextCursor := pageAfter(assignments, func(a databricks.WorkspaceAssignment) string { + return strconv.Itoa(principalID(a)) + }, cursor) + var rv []*v2.Grant - l.Debug("grants: workspace resource", zap.String("workspace_id", workspace), zap.Int("assignments_count", len(assignments))) - for _, assignment := range assignments { + l.Debug("grants: workspace resource", + zap.String("workspace_id", workspace), + zap.Int("assignments_count", len(assignments)), + zap.Int("page_assignments_count", len(assignmentsPage)), + zap.String("cursor", cursor), + ) + for _, assignment := range assignmentsPage { resourceType, err := prepareResourceType(assignment.Principal) if err != nil { return nil, nil, fmt.Errorf("databricks-connector: failed to prepare resource type: %w", err) @@ -238,7 +261,12 @@ func (w *workspaceBuilder) Grants(ctx context.Context, resource *v2.Resource, _ rv = append(rv, grant.NewGrant(resource, workspaceMemberEntitlement, resourceID, grant.WithAnnotation(annotations...))) } - return rv, nil, nil + nextPage, err := bag.NextToken(nextCursor) + if err != nil { + return nil, nil, fmt.Errorf("databricks-connector: failed to create next page token: %w", err) + } + + return rv, &rs.SyncOpResults{NextPageToken: nextPage, Annotations: annos}, nil } // workspacePrincipalID returns the numeric principal ID the permission assignments API expects. diff --git a/pkg/connector/workspaces_test.go b/pkg/connector/workspaces_test.go new file mode 100644 index 00000000..3d3bc54c --- /dev/null +++ b/pkg/connector/workspaces_test.go @@ -0,0 +1,90 @@ +package connector + +import ( + "context" + "slices" + "strconv" + "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" + rs "github.com/conductorone/baton-sdk/pkg/types/resource" +) + +func testUserAssignments(count int) []map[string]any { + assignments := make([]map[string]any, 0, count) + for i := range count { + assignments = append(assignments, map[string]any{ + "principal": map[string]any{"principal_id": 1000 + i, "user_name": "user" + strconv.Itoa(i) + "@example.com"}, + "permissions": []string{"USER"}, + }) + } + return assignments +} + +func testWorkspace(t *testing.T) *v2.Resource { + t.Helper() + r, err := workspaceResource(context.Background(), + &databricks.Workspace{ID: testProvisioningWorkspaceID, Name: "ws", DeploymentName: "ws-deploy"}, + &v2.ResourceId{ResourceType: accountResourceType.Id, Resource: testAccountId}, + ) + if err != nil { + t.Fatalf("workspaceResource: %v", err) + } + return r +} + +// Workspace assignments come back in one API response, so Grants must split them across pages +// without skipping or repeating a principal when assignments change between page fetches. +func TestWorkspaceGrantsPagesAssignments(t *testing.T) { + // Every page re-fetches the assignments; disable the HTTP cache so each page sees the change. + t.Setenv("BATON_DISABLE_HTTP_CACHE", "true") + + count := 2*int(GrantsPageSize) + 500 + m := &workspaceAssignmentsServer{assignments: shuffled(testUserAssignments(count))} + w := newWorkspaceBuilderWithMock(t, m) + workspace := testWorkspace(t) + + seen := map[string]int{} + pages := 0 + token := "" + for { + grants, results, err := w.Grants(context.Background(), workspace, rs.SyncOpAttrs{PageToken: pagination.Token{Token: token}}) + if err != nil { + t.Fatalf("Grants page %d: %v", pages, err) + } + pages++ + if len(grants) > int(GrantsPageSize) { + t.Errorf("page %d has %d grants, want at most %d", pages, len(grants), GrantsPageSize) + } + for _, gr := range grants { + seen[gr.Principal.Id.Resource]++ + } + + token = results.NextPageToken + if token == "" { + break + } + if pages > 100 { + t.Fatal("pagination did not terminate") + } + if pages == 1 { + // Remove the first assignment, already emitted on page 1, and reshuffle. + m.mu.Lock() + all := testUserAssignments(count) + m.assignments = shuffled(slices.Clone(all[1:])) + m.mu.Unlock() + } + } + + if pages != 3 { + t.Errorf("got %d pages, want 3", pages) + } + for i := range count { + id := strconv.Itoa(1000 + i) + if seen[id] != 1 { + t.Errorf("principal %s granted %d times, want 1", id, seen[id]) + } + } +}