Terraform clean way to use 2 for_each

currently I have this code:

resource "azurerm_role_assignment" "health-reader-assignment" {
  for_each = var.env == "prod" ? local.prod_groups : local.non_prod_groups_by_env

  scope              = each.value
  role_definition_id = azurerm_role_definition.health-reader.role_definition_resource_id
  principal_id       = azuread_group.infrastructure-management[var.env].object_id
}

resource "azurerm_role_assignment" "rg-mgmt-assignment" {
  for_each = var.env == "prod" ? local.prod_groups : local.non_prod_groups_by_env

  scope              = each.value
  role_definition_id = azurerm_role_definition.rg-mgmt.role_definition_resource_id
  principal_id       = azuread_group.infrastructure-management[var.env].object_id
}

The only difference between them is the two separate roles. What is the cleanest way to merge this logic into a single role assignment? Should I merge the two together, or is it better to keep them separate? I would appreciate your input on the best practices.

I tried using some AI tools, but each one provided a different result with multiple for loops that I don't find clean or readable. I'm looking for the best practice for these types of scenarios. Should I place everything (scope + roles) into two separate locals and merge them into a map? I didn't like this approach because it feels unnecessary to create locals that are only used once.

Thank you!

From what you've shared it seems like the main part that might be beneficial to factor out is the policy decision of using a different local value for the production environment, so that you can change the rules for choosing groups in just one place in future:

locals {
  groups = (
    var.env == prod ?
    local.prod_groups :
    local.non_prod_groups_by_env
  )
}

resource "azurerm_role_assignment" "health-reader-assignment" {
  for_each = local.groups

  scope              = each.value
  role_definition_id = azurerm_role_definition.health-reader.role_definition_resource_id
  principal_id       = azuread_group.infrastructure-management[var.env].object_id
}

resource "azurerm_role_assignment" "rg-mgmt-assignment" {
  for_each = local.groups

  scope              = each.value
  role_definition_id = azurerm_role_definition.rg-mgmt.role_definition_resource_id
  principal_id       = azuread_group.infrastructure-management[var.env].object_id
}

I would personally not try to generalize this any further unless the set of roles were expected to change often, but if you do want to do that then the general approach would be to combine your collection of groups with your collection of roles to produce a collection that has one element for each combination elements from the two collections.

The typical way to achieve that is using the setproduct function, as described in Finding combinations for for_each.

Adapting the examples in that section to your situation would be something like this:

locals {
  groups = (
    var.env == prod ?
    local.prod_groups :
    local.non_prod_groups_by_env
  )
  roles = {
    health_reader = azurerm_role_definition.health_reader
    rg_mgmt       = azurerm_role_definition.rg_mgmt
  }
  group_roles = {
    for pair in setproduct(keys(local.groups), keys(local.roles)) :
    "${pair[0]}:${pair[1]}" => {
      group = local.groups[pair[0]]
      role  = local.roles[pair[1]]
    }
  }
}

resource "azurerm_role_assignment" "all" {
  # NOTE: You could inline the expression of local.group_roles
  # in here if you don't want to have it declared as a separate
  # local value. Some authors prefer to keep for_each expressions
  # relatively simple, but that's purely a style decision.
  for_each = local.group_roles

  scope              = each.value.group
  role_definition_id = each.value.role.role_definition_resource_id
  principal_id       = azuread_group.infrastructure_management[var.env].object_id
}

This formulation means that you can add or remove elements of local.roles in future and have Terraform figure out which role assignments it needs to create or destroy based on that, without needing to add any new resource blocks.

However, this comes at the cost of some configuration that's potentially harder to follow for someone who is not a Terraform language expert, so I suggest considering whether this additional complexity is worthwhile vs. the relative simplicity of the first solution.