JPA method to update roles for users In my Spring MVC app I have User and Role entities.

User:

public class User {
    private Long id;
    private Set&lt;Role&gt; roles;

    // other fields, getters and setters
}


Role:

public class Role {
    private Long id;

    // other fields, getters and setters
}


I have a service method with the following signature:

@Transactional
void setRolesByUserId(Map&lt;Long, Set&lt;Long&gt;&gt; rolesByUserId)


The key of the Map is a User ID and the corresponding value is a Set of Role IDs that needs to be assigned to the user.

My implementation:

public void setRolesByUserId(Map&lt;Long, Set&lt;Long&gt;&gt; rolesByUserId) {
    // 1.
    Map&lt;Long, Role&gt; roles = roleRepository.findAll()
        .stream().collect(Collectors.toMap(Role::getId, Function.identity()));

    // 2.
    userRepository.findAllById(rolesByUserId.keySet())
            // 3.
            .forEach(user -&gt; user.setRoles(rolesByUserId.get(user.getId())
                    .stream().map(roles::get).collect(Collectors.toSet())));
}


So basically the flow is the following:


fetch all the Role entities from the DB into a Map&lt;Long, Role&gt; (where the key is the Role ID and the value is the Role entity itself)
fetch the User entities from the DB that needs to be modified
for each User, iterate through the Role IDs to be assigned, fetch the corresponding Role entities from the Map, collect them into a Set and assign it to the User


Can this method be optimized?
Well, you created a good code so far. Your lamda is OK. But there are some things about naming which could be improved


About the API setRolesByUserId, it is OK to accept userId as long. But about role, it could be better to change the Set to Set, and i think you're about to update for List of user, hence the name should change to updateUserWithRoleSet.
IMHO, It is totally OK to update one user at once => updateUserWithRole(long userId, Set&lt;Role&gt; roles&gt;