pr fixes

otterize · Dec 10, 2024 · 5d1b56a · 5d1b56a
1 parent 5daf132
commit 5d1b56a
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 29 deletions.
diff --git a/src/operator/controllers/intents_reconcilers/iam/iampolicyagents/azurepolicyagent/agent.go b/src/operator/controllers/intents_reconcilers/iam/iampolicyagents/azurepolicyagent/agent.go
@@ -95,7 +95,7 @@ func (a *Agent) ensureRoleAssignmentsForIntents(ctx context.Context, userAssigne
 	a.assignmentMutex.Lock()
 	defer a.assignmentMutex.Unlock()
 
-	existingRoleAssignments, err := a.ListRoleAssignments(ctx, userAssignedIdentity)
+	existingRoleAssignments, err := a.ListRoleAssignmentsAcrossSubscriptions(ctx, userAssignedIdentity)
 	if err != nil {
 		return errors.Wrap(err)
 	}
@@ -199,7 +199,7 @@ func (a *Agent) DeleteRolePolicyFromIntents(ctx context.Context, intents otteriz
 		return errors.Wrap(err)
 	}
 
-	existingRoleAssignments, err := a.ListRoleAssignments(ctx, userAssignedIdentity)
+	existingRoleAssignments, err := a.ListRoleAssignmentsAcrossSubscriptions(ctx, userAssignedIdentity)
 	if err != nil {
 		return errors.Wrap(err)
 	}

diff --git a/src/shared/azureagent/identities.go b/src/shared/azureagent/identities.go
@@ -108,7 +108,7 @@ func (a *Agent) DeleteUserAssignedIdentity(ctx context.Context, namespace string
 		return errors.Wrap(err)
 	}
 
-	roleAssignments, err := a.ListRoleAssignments(ctx, identity)
+	roleAssignments, err := a.ListRoleAssignmentsAcrossSubscriptions(ctx, identity)
 	if err != nil {
 		return errors.Wrap(err)
 	}

diff --git a/src/shared/azureagent/roleassignments.go b/src/shared/azureagent/roleassignments.go
@@ -59,48 +59,42 @@ func (a *Agent) DeleteRoleAssignment(ctx context.Context, roleAssignment armauth
 	return nil
 }
 
-func (a *Agent) ListRoleAssignments(ctx context.Context, userAssignedIdentity armmsi.Identity) ([]armauthorization.RoleAssignment, error) {
+func (a *Agent) ListRoleAssignmentsAcrossSubscriptions(ctx context.Context, userAssignedIdentity armmsi.Identity) ([]armauthorization.RoleAssignment, error) {
+	subscriptions, err := a.ListSubscriptions(ctx)
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
+
 	var roleAssignments []armauthorization.RoleAssignment
-	pager := a.roleAssignmentsClient.NewListForSubscriptionPager(nil)
-	for pager.More() {
-		page, err := pager.NextPage(ctx)
+	for _, sub := range subscriptions {
+		roleAssignmentsForSubscription, err := a.ListRoleAssignmentsForSubscription(ctx, *sub.SubscriptionID, userAssignedIdentity)
 		if err != nil {
 			return nil, errors.Wrap(err)
 		}
-		for _, roleAssignment := range page.Value {
-			if *roleAssignment.Properties.PrincipalID == *userAssignedIdentity.Properties.PrincipalID {
-				roleAssignments = append(roleAssignments, *roleAssignment)
-			}
-		}
+
+		roleAssignments = append(roleAssignments, roleAssignmentsForSubscription...)
 	}
 
 	return roleAssignments, nil
 }
 
-func (a *Agent) ListRoleAssignmentsAcrossSubscriptions(ctx context.Context, userAssignedIdentity armmsi.Identity) ([]armauthorization.RoleAssignment, error) {
-	subscriptions, err := a.ListSubscriptions(ctx)
+func (a *Agent) ListRoleAssignmentsForSubscription(ctx context.Context, subscriptionID string, userAssignedIdentity armmsi.Identity) ([]armauthorization.RoleAssignment, error) {
+	roleClient, err := a.GetRoleAssignmentClientForSubscription(subscriptionID)
 	if err != nil {
 		return nil, errors.Wrap(err)
 	}
 
 	var roleAssignments []armauthorization.RoleAssignment
-	for _, sub := range subscriptions {
-		roleClient, err := a.GetRoleAssignmentClientForSubscription(*sub.SubscriptionID)
+	pager := roleClient.NewListForSubscriptionPager(nil)
+	for pager.More() {
+		page, err := pager.NextPage(ctx)
 		if err != nil {
 			return nil, errors.Wrap(err)
 		}
 
-		pager := roleClient.NewListForSubscriptionPager(nil)
-		for pager.More() {
-			page, err := pager.NextPage(ctx)
-			if err != nil {
-				return nil, errors.Wrap(err)
-			}
-
-			for _, roleAssignment := range page.Value {
-				if *roleAssignment.Properties.PrincipalID == *userAssignedIdentity.Properties.PrincipalID {
-					roleAssignments = append(roleAssignments, *roleAssignment)
-				}
+		for _, roleAssignment := range page.Value {
+			if *roleAssignment.Properties.PrincipalID == *userAssignedIdentity.Properties.PrincipalID {
+				roleAssignments = append(roleAssignments, *roleAssignment)
 			}
 		}
 	}

diff --git a/src/shared/azureagent/subscriptions.go b/src/shared/azureagent/subscriptions.go
@@ -15,8 +15,8 @@ func (a *Agent) ListSubscriptions(ctx context.Context) ([]armsubscriptions.Subsc
 			return nil, errors.Wrap(err)
 		}
 
-		for _, roleAssignment := range page.Value {
-			subscriptions = append(subscriptions, *roleAssignment)
+		for _, subscription := range page.Value {
+			subscriptions = append(subscriptions, *subscription)
 		}
 	}