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