pass ratelimitData through client

pkg/connector/account.go

@@ -85,22 +85,28 @@ func (a *accountBuilder) List(ctx context.Context, parentResourceID *v2.Resource
 }
 
 // Entitlements returns slice of entitlements for marketplace admins under account.
-func (a *accountBuilder) Entitlements(_ context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {
+func (a *accountBuilder) Entitlements(
+	_ context.Context,
+	resource *v2.Resource,
+	_ *pagination.Token,
+) (
+	[]*v2.Entitlement,
+	string,
+	annotations.Annotations,
+	error,
+) {
 	if !a.client.IsAccountAPIAvailable() {
 		return nil, "", nil, nil
 	}
-
-	var rv []*v2.Entitlement
-
-	permissionOptions := []ent.EntitlementOption{
-		ent.WithGrantableTo(userResourceType, groupResourceType, servicePrincipalResourceType),
-		ent.WithDisplayName(fmt.Sprintf("%s %s role", resource.DisplayName, MarketplaceAdminRole)),
-		ent.WithDescription(fmt.Sprintf("%s %s role in Databricks", resource.DisplayName, MarketplaceAdminRole)),
-	}
-
-	rv = append(rv, ent.NewPermissionEntitlement(resource, MarketplaceAdminRole, permissionOptions...))
-
-	return rv, "", nil, nil
+	return []*v2.Entitlement{
+		ent.NewPermissionEntitlement(
+			resource,
+			MarketplaceAdminRole,
+			ent.WithGrantableTo(userResourceType, groupResourceType, servicePrincipalResourceType),
+			ent.WithDisplayName(fmt.Sprintf("%s %s role", resource.DisplayName, MarketplaceAdminRole)),
+			ent.WithDescription(fmt.Sprintf("%s %s role in Databricks", resource.DisplayName, MarketplaceAdminRole)),
+		),
+	}, "", nil, nil
 }
 
 // Grants returns grants for marketplace admins under account.

pkg/connector/connector.go

@@ -52,7 +52,7 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err
 
 	// Check if we can list users from Account API (unless we are using token auth specific to a single workspace).
 	if !d.client.IsTokenAuth() {
-		_, err := d.client.ListRoles(ctx, "", "")
+		_, _, err := d.client.ListRoles(ctx, "", "")
 		if err == nil {
 			isAccAPIAvailable = true
 		}
@@ -63,7 +63,7 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err
 		for _, workspace := range d.workspaces {
 			d.client.SetWorkspaceConfig(workspace)
 
-			_, err := d.client.ListRoles(ctx, "", "")
+			_, _, err := d.client.ListRoles(ctx, "", "")
 			if err != nil && !isAccAPIAvailable {
 				return nil, fmt.Errorf("databricks-connector: failed to validate credentials for workspace %s: %w", workspace, err)
 			}
@@ -74,15 +74,15 @@ func (d *Databricks) Validate(ctx context.Context) (annotations.Annotations, err
 
 	// Validate that credentials are valid for every workspace.
 	if len(d.workspaces) == 0 {
-		workspaces, err := d.client.ListWorkspaces(ctx)
+		workspaces, _, err := d.client.ListWorkspaces(ctx)
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err)
 		}
 
 		for _, workspace := range workspaces {
 			d.client.SetWorkspaceConfig(workspace.Host)
 
-			_, err := d.client.ListRoles(ctx, "", "")
+			_, _, err := d.client.ListRoles(ctx, "", "")
 			if err != nil && !isAccAPIAvailable {
 				return nil, fmt.Errorf("databricks-connector: failed to validate credentials for workspace %s: %w", workspace.Host, err)
 			}

pkg/connector/groups.go

@@ -93,7 +93,7 @@ func (g *groupBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to parse page token: %w", err)
 	}
 
-	groups, total, err := g.client.ListGroups(
+	groups, total, _, err := g.client.ListGroups(
 		ctx,
 		databricks.NewPaginationVars(page, ResourcesPageSize),
 		databricks.NewGroupAttrVars(),
@@ -152,7 +152,7 @@ func (g *groupBuilder) Entitlements(_ context.Context, resource *v2.Resource, _
 
 	// role permissions entitlements
 	// get all assignable roles for this specific group resource
-	roles, err := g.client.ListRoles(context.Background(), GroupsType, groupId.Resource)
+	roles, _, err := g.client.ListRoles(context.Background(), GroupsType, groupId.Resource)
 	if err != nil {
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to list roles for group %s: %w", groupId.Resource, err)
 	}
@@ -229,7 +229,7 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, pToken
 	}
 
 	// role permissions grants
-	ruleSets, err := g.client.ListRuleSets(ctx, GroupsType, groupId.Resource)
+	ruleSets, _, err := g.client.ListRuleSets(ctx, 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)
 	}
@@ -291,7 +291,7 @@ func (g *groupBuilder) Grant(ctx context.Context, principal *v2.Resource, entitl
 
 	// If the entitlement is a member entitlement
 	if entitlement.Slug == groupMemberEntitlement {
-		group, err := g.client.GetGroup(ctx, groupId.Resource)
+		group, _, err := g.client.GetGroup(ctx, groupId.Resource)
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
 		}
@@ -313,7 +313,7 @@ func (g *groupBuilder) Grant(ctx context.Context, principal *v2.Resource, entitl
 			ID: principal.Id.Resource,
 		})
 
-		err = g.client.UpdateGroup(ctx, group)
+		_, err = g.client.UpdateGroup(ctx, group)
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to update group %s: %w", groupId.Resource, err)
 		}
@@ -322,7 +322,7 @@ func (g *groupBuilder) Grant(ctx context.Context, principal *v2.Resource, entitl
 	}
 
 	// If the entitlement is a role permission entitlement
-	ruleSets, err := g.client.ListRuleSets(ctx, GroupsType, groupId.Resource)
+	ruleSets, _, err := g.client.ListRuleSets(ctx, GroupsType, groupId.Resource)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to list rule sets for group %s (%s): %w", principal.Id.Resource, groupId.Resource, err)
 	}
@@ -361,7 +361,7 @@ func (g *groupBuilder) Grant(ctx context.Context, principal *v2.Resource, entitl
 		})
 	}
 
-	err = g.client.UpdateRuleSets(ctx, GroupsType, groupId.Resource, ruleSets)
+	_, err = g.client.UpdateRuleSets(ctx, GroupsType, groupId.Resource, ruleSets)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to update rule sets for group %s (%s): %w", principal.Id.Resource, groupId.Resource, err)
 	}
@@ -403,7 +403,7 @@ func (g *groupBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations
 	}
 
 	if entitlement.Slug == groupMemberEntitlement {
-		group, err := g.client.GetGroup(ctx, groupId.Resource)
+		group, _, err := g.client.GetGroup(ctx, groupId.Resource)
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
 		}
@@ -415,12 +415,12 @@ func (g *groupBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations
 			}
 		}
 
-		err = g.client.UpdateGroup(ctx, group)
+		_, err = g.client.UpdateGroup(ctx, group)
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to update group %s: %w", groupId.Resource, err)
 		}
 	} else {
-		ruleSets, err := g.client.ListRuleSets(ctx, GroupsType, groupId.Resource)
+		ruleSets, _, err := g.client.ListRuleSets(ctx, GroupsType, groupId.Resource)
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to list rule sets for group %s (%s): %w", principal.Id.Resource, groupId.Resource, err)
 		}
@@ -466,7 +466,7 @@ func (g *groupBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations
 			return nil, nil
 		}
 
-		err = g.client.UpdateRuleSets(ctx, GroupsType, groupId.Resource, ruleSets)
+		_, err = g.client.UpdateRuleSets(ctx, GroupsType, groupId.Resource, ruleSets)
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to update rule sets for group %s (%s): %w", principal.Id.Resource, groupId.Resource, err)
 		}

pkg/connector/helpers.go

@@ -77,21 +77,21 @@ func prepareResourceID(ctx context.Context, c *databricks.Client, principal stri
 
 	switch principalType {
 	case UsersType:
-		userID, err := c.FindUserID(ctx, principal)
+		userID, _, err := c.FindUserID(ctx, principal)
 		if err != nil {
 			return nil, fmt.Errorf("failed to find user %s: %w", principal, err)
 		}
 
 		resourceId = &v2.ResourceId{ResourceType: userResourceType.Id, Resource: userID}
 	case GroupsType:
-		groupID, err := c.FindGroupID(ctx, principal)
+		groupID, _, err := c.FindGroupID(ctx, principal)
 		if err != nil {
 			return nil, fmt.Errorf("failed to find group %s: %w", principal, err)
 		}
 
 		resourceId = &v2.ResourceId{ResourceType: groupResourceType.Id, Resource: groupID}
 	case ServicePrincipalsType:
-		servicePrincipalID, err := c.FindServicePrincipalID(ctx, principal)
+		servicePrincipalID, _, err := c.FindServicePrincipalID(ctx, principal)
 		if err != nil {
 			return nil, fmt.Errorf("failed to find service principal %s: %w", principal, err)
 		}
@@ -126,21 +126,21 @@ func preparePrincipalID(ctx context.Context, c *databricks.Client, principalType
 
 	switch principalType {
 	case userResourceType.Id:
-		username, err := c.FindUsername(ctx, principalID)
+		username, _, err := c.FindUsername(ctx, principalID)
 		if err != nil {
 			return "", fmt.Errorf("failed to find user %s: %w", principalID, err)
 		}
 
 		result = fmt.Sprintf("%s/%s", UsersType, username)
 	case groupResourceType.Id:
-		displayName, err := c.FindGroupDisplayName(ctx, principalID)
+		displayName, _, err := c.FindGroupDisplayName(ctx, principalID)
 		if err != nil {
 			return "", fmt.Errorf("failed to find group %s: %w", principalID, err)
 		}
 
 		result = fmt.Sprintf("%s/%s", GroupsType, displayName)
 	case servicePrincipalResourceType.Id:
-		appID, err := c.FindServicePrincipalAppID(ctx, principalID)
+		appID, _, err := c.FindServicePrincipalAppID(ctx, principalID)
 		if err != nil {
 			return "", fmt.Errorf("failed to find service principal %s: %w", principalID, err)
 		}

pkg/connector/roles.go

@@ -109,18 +109,25 @@ func (r *roleBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId,
 }
 
 // Entitlements returns membership entitlements for a given role.
-func (r *roleBuilder) Entitlements(_ context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {
-	var rv []*v2.Entitlement
-
-	entitlementOptions := []ent.EntitlementOption{
-		ent.WithGrantableTo(userResourceType, groupResourceType, servicePrincipalResourceType),
-		ent.WithDisplayName(fmt.Sprintf("%s role", resource.DisplayName)),
-		ent.WithDescription(fmt.Sprintf("%s Databricks role", resource.DisplayName)),
-	}
-
-	rv = append(rv, ent.NewAssignmentEntitlement(resource, RoleMemberEntitlement, entitlementOptions...))
-
-	return rv, "", nil, nil
+func (r *roleBuilder) Entitlements(
+	_ context.Context,
+	resource *v2.Resource,
+	_ *pagination.Token,
+) (
+	[]*v2.Entitlement,
+	string,
+	annotations.Annotations,
+	error,
+) {
+	return []*v2.Entitlement{
+		ent.NewAssignmentEntitlement(
+			resource,
+			RoleMemberEntitlement,
+			ent.WithGrantableTo(userResourceType, groupResourceType, servicePrincipalResourceType),
+			ent.WithDisplayName(fmt.Sprintf("%s role", resource.DisplayName)),
+			ent.WithDescription(fmt.Sprintf("%s Databricks role", resource.DisplayName)),
+		),
+	}, "", nil, nil
 }
 
 // Grants returns all the grants for a given role.

pkg/connector/service-principals.go

@@ -77,7 +77,7 @@ func (s *servicePrincipalBuilder) List(ctx context.Context, parentResourceID *v2
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to parse page token: %w", err)
 	}
 
-	servicePrincipals, total, err := s.client.ListServicePrincipals(
+	servicePrincipals, total, _, err := s.client.ListServicePrincipals(
 		ctx,
 		databricks.NewPaginationVars(page, ResourcesPageSize),
 		databricks.NewServicePrincipalAttrVars(),
@@ -134,7 +134,7 @@ func (s *servicePrincipalBuilder) Entitlements(_ context.Context, resource *v2.R
 
 	// role permissions entitlements
 	// get all assignable roles for this specific service principal resource
-	roles, err := s.client.ListRoles(context.Background(), ServicePrincipalsType, applicationId)
+	roles, _, err := s.client.ListRoles(context.Background(), ServicePrincipalsType, applicationId)
 	if err != nil {
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to list roles for service principal %s (%s): %w", resource.Id.Resource, applicationId, err)
 	}
@@ -176,7 +176,7 @@ func (s *servicePrincipalBuilder) Grants(ctx context.Context, resource *v2.Resou
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to get application_id from service principal profile")
 	}
 
-	ruleSets, err := s.client.ListRuleSets(ctx, ServicePrincipalsType, applicationId)
+	ruleSets, _, err := s.client.ListRuleSets(ctx, ServicePrincipalsType, applicationId)
 	if err != nil {
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to list rule sets for service principal %s (%s): %w", resource.Id.Resource, applicationId, err)
 	}
@@ -240,7 +240,7 @@ func (s *servicePrincipalBuilder) Grant(ctx context.Context, principal *v2.Resou
 		return nil, fmt.Errorf("databricks-connector: failed to get application_id from service principal profile")
 	}
 
-	ruleSets, err := s.client.ListRuleSets(ctx, ServicePrincipalsType, applicationId)
+	ruleSets, _, err := s.client.ListRuleSets(ctx, ServicePrincipalsType, applicationId)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to list rule sets for service principal %s (%s): %w", principal.Id.Resource, applicationId, err)
 	}
@@ -279,7 +279,7 @@ func (s *servicePrincipalBuilder) Grant(ctx context.Context, principal *v2.Resou
 		})
 	}
 
-	err = s.client.UpdateRuleSets(ctx, ServicePrincipalsType, applicationId, ruleSets)
+	_, err = s.client.UpdateRuleSets(ctx, ServicePrincipalsType, applicationId, ruleSets)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to update rule sets for service principal %s (%s): %w", principal.Id.Resource, applicationId, err)
 	}
@@ -324,7 +324,7 @@ func (s *servicePrincipalBuilder) Revoke(ctx context.Context, grant *v2.Grant) (
 		return nil, fmt.Errorf("databricks-connector: failed to get application_id from service principal profile")
 	}
 
-	ruleSets, err := s.client.ListRuleSets(ctx, ServicePrincipalsType, applicationId)
+	ruleSets, _, err := s.client.ListRuleSets(ctx, ServicePrincipalsType, applicationId)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to list rule sets for service principal %s (%s): %w", principal.Id.Resource, applicationId, err)
 	}
@@ -371,7 +371,7 @@ func (s *servicePrincipalBuilder) Revoke(ctx context.Context, grant *v2.Grant) (
 		return nil, nil
 	}
 
-	err = s.client.UpdateRuleSets(ctx, ServicePrincipalsType, applicationId, ruleSets)
+	_, err = s.client.UpdateRuleSets(ctx, ServicePrincipalsType, applicationId, ruleSets)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to update rule sets for service principal %s (%s): %w", principal.Id.Resource, applicationId, err)
 	}

pkg/connector/users.go

@@ -131,12 +131,30 @@ func (u *userBuilder) List(ctx context.Context, parentResourceID *v2.ResourceId,
 }
 
 // Entitlements always returns an empty slice for users.
-func (u *userBuilder) Entitlements(_ context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {
+func (u *userBuilder) Entitlements(
+	_ context.Context,
+	_ *v2.Resource,
+	_ *pagination.Token,
+) (
+	[]*v2.Entitlement,
+	string,
+	annotations.Annotations,
+	error,
+) {
 	return nil, "", nil, nil
 }
 
 // Grants always returns an empty slice for users since they don't have any entitlements.
-func (u *userBuilder) Grants(ctx context.Context, resource *v2.Resource, pToken *pagination.Token) ([]*v2.Grant, string, annotations.Annotations, error) {
+func (u *userBuilder) Grants(
+	_ context.Context,
+	_ *v2.Resource,
+	_ *pagination.Token,
+) (
+	[]*v2.Grant,
+	string,
+	annotations.Annotations,
+	error,
+) {
 	return nil, "", nil, nil
 }
 

pkg/connector/workspaces.go

@@ -83,7 +83,7 @@ func (w *workspaceBuilder) List(ctx context.Context, parentResourceID *v2.Resour
 
 	var rv []*v2.Resource
 	if w.client.IsAccountAPIAvailable() {
-		workspaces, err := w.client.ListWorkspaces(ctx)
+		workspaces, _, err := w.client.ListWorkspaces(ctx)
 		if err != nil {
 			return nil, "", nil, fmt.Errorf("databricks-connector: failed to list workspaces: %w", err)
 		}
@@ -162,7 +162,7 @@ func (w *workspaceBuilder) Grants(ctx context.Context, resource *v2.Resource, pT
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to get workspace ID: %w", err)
 	}
 
-	assignments, err := w.client.ListWorkspaceMembers(ctx, int(workspaceID))
+	assignments, _, err := w.client.ListWorkspaceMembers(ctx, int(workspaceID))
 	if err != nil {
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to list workspace members: %w", err)
 	}
@@ -218,7 +218,7 @@ func (w *workspaceBuilder) Grant(ctx context.Context, principal *v2.Resource, en
 		return nil, fmt.Errorf("databricks-connector: failed to get workspace ID: %w", err)
 	}
 
-	err = w.client.CreateOrUpdateWorkspaceMember(ctx, workspaceID, principal.Id.Resource)
+	_, err = w.client.CreateOrUpdateWorkspaceMember(ctx, workspaceID, principal.Id.Resource)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to create or update workspace member: %w", err)
 	}
@@ -252,7 +252,7 @@ func (w *workspaceBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotat
 		return nil, fmt.Errorf("databricks-connector: failed to get workspace ID: %w", err)
 	}
 
-	err = w.client.RemoveWorkspaceMember(ctx, workspaceID, principal.Id.Resource)
+	_, err = w.client.RemoveWorkspaceMember(ctx, workspaceID, principal.Id.Resource)
 	if err != nil {
 		return nil, fmt.Errorf("databricks-connector: failed to create or update workspace member: %w", err)
 	}