Improvement : Teams api to update and remove users (#18729)

openmetadata-service/src/main/java/org/openmetadata/service/exception/CatalogExceptionMessage.java

@@ -279,6 +279,11 @@ public final class CatalogExceptionMessage {
         "Team of type %s can't own entities. Only Team of type Group can own entities.", teamType);
   }
 
+  public static String invalidTeamUpdateUsers(TeamType teamType) {
+    return String.format(
+        "Team is of type %s. Users can be updated only in team of type Group.", teamType);
+  }
+
   public static String invalidOwnerType(String entityType) {
     return String.format(
         "Entity of type %s can't be the owner. Only Team of type Group or a User can own entities.",

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java

@@ -847,6 +847,13 @@ public interface CollectionDAO {
       bulkInsertTo(insertToRelationship);
     }
 
+    default void bulkRemoveToRelationship(
+        UUID fromId, List<UUID> toIds, String fromEntity, String toEntity, int relation) {
+
+      List<String> toIdsAsString = toIds.stream().map(UUID::toString).toList();
+      bulkRemoveTo(fromId, toIdsAsString, fromEntity, toEntity, relation);
+    }
+
     @ConnectionAwareSqlUpdate(
         value =
             "INSERT INTO entity_relationship(fromId, toId, fromEntity, toEntity, relation, json) "
@@ -882,6 +889,18 @@ public interface CollectionDAO {
                 propertyNames = {"fromId", "toId", "fromEntity", "toEntity", "relation"})
             List<EntityRelationshipObject> values);
 
+    @SqlUpdate(
+        value =
+            "DELETE FROM entity_relationship WHERE fromId = :fromId "
+                + "AND fromEntity = :fromEntity AND toId IN (<toIds>) "
+                + "AND toEntity = :toEntity AND relation = :relation")
+    void bulkRemoveTo(
+        @BindUUID("fromId") UUID fromId,
+        @BindList("toIds") List<String> toIds,
+        @Bind("fromEntity") String fromEntity,
+        @Bind("toEntity") String toEntity,
+        @Bind("relation") int relation);
+
     //
     // Find to operations
     //

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java

@@ -1835,6 +1835,14 @@ public abstract class EntityRepository<T extends EntityInterface> {
         .bulkInsertToRelationship(fromId, toId, fromEntity, toEntity, relationship.ordinal());
   }
 
+  @Transaction
+  public final void bulkRemoveToRelationship(
+      UUID fromId, List<UUID> toId, String fromEntity, String toEntity, Relationship relationship) {
+    daoCollection
+        .relationshipDAO()
+        .bulkRemoveToRelationship(fromId, toId, fromEntity, toEntity, relationship.ordinal());
+  }
+
   public final List<EntityReference> findBoth(
       UUID entity1, String entityType1, Relationship relationship, String entity2) {
     // Find bidirectional relationship

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TeamRepository.java

@@ -24,6 +24,7 @@ import static org.openmetadata.schema.api.teams.CreateTeam.TeamType.DEPARTMENT;
 import static org.openmetadata.schema.api.teams.CreateTeam.TeamType.DIVISION;
 import static org.openmetadata.schema.api.teams.CreateTeam.TeamType.GROUP;
 import static org.openmetadata.schema.api.teams.CreateTeam.TeamType.ORGANIZATION;
+import static org.openmetadata.schema.type.EventType.ENTITY_FIELDS_CHANGED;
 import static org.openmetadata.schema.type.Include.ALL;
 import static org.openmetadata.schema.type.Include.NON_DELETED;
 import static org.openmetadata.service.Entity.ADMIN_USER_NAME;
@@ -41,6 +42,7 @@ import static org.openmetadata.service.exception.CatalogExceptionMessage.UNEXPEC
 import static org.openmetadata.service.exception.CatalogExceptionMessage.invalidChild;
 import static org.openmetadata.service.exception.CatalogExceptionMessage.invalidParent;
 import static org.openmetadata.service.exception.CatalogExceptionMessage.invalidParentCount;
+import static org.openmetadata.service.util.EntityUtil.*;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -50,20 +52,26 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.UUID;
 import java.util.stream.Collectors;
+import javax.ws.rs.core.Response;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.csv.CSVPrinter;
 import org.apache.commons.csv.CSVRecord;
 import org.jdbi.v3.sqlobject.transaction.Transaction;
 import org.openmetadata.common.utils.CommonUtil;
 import org.openmetadata.csv.EntityCsv;
+import org.openmetadata.schema.api.teams.CreateTeam;
 import org.openmetadata.schema.api.teams.CreateTeam.TeamType;
 import org.openmetadata.schema.entity.teams.Team;
 import org.openmetadata.schema.entity.teams.TeamHierarchy;
+import org.openmetadata.schema.type.ChangeDescription;
+import org.openmetadata.schema.type.ChangeEvent;
 import org.openmetadata.schema.type.EntityReference;
+import org.openmetadata.schema.type.EventType;
 import org.openmetadata.schema.type.Include;
 import org.openmetadata.schema.type.Relationship;
 import org.openmetadata.schema.type.api.BulkAssets;
@@ -74,17 +82,20 @@ import org.openmetadata.schema.type.csv.CsvFile;
 import org.openmetadata.schema.type.csv.CsvHeader;
 import org.openmetadata.schema.type.csv.CsvImportResult;
 import org.openmetadata.service.Entity;
+import org.openmetadata.service.exception.CatalogExceptionMessage;
 import org.openmetadata.service.exception.EntityNotFoundException;
 import org.openmetadata.service.jdbi3.CollectionDAO.EntityRelationshipRecord;
 import org.openmetadata.service.resources.teams.TeamResource;
 import org.openmetadata.service.security.policyevaluator.SubjectContext;
 import org.openmetadata.service.util.EntityUtil;
 import org.openmetadata.service.util.EntityUtil.Fields;
+import org.openmetadata.service.util.RestUtil;
 import org.openmetadata.service.util.ResultList;
 
 @Slf4j
 public class TeamRepository extends EntityRepository<Team> {
   static final String PARENTS_FIELD = "parents";
+  static final String USERS_FIELD = "users";
   static final String TEAM_UPDATE_FIELDS =
       "profile,users,defaultRoles,parents,children,policies,teamType,email,domains";
   static final String TEAM_PATCH_FIELDS =
@@ -607,6 +618,111 @@ public class TeamRepository extends EntityRepository<Team> {
     }
   }
 
+  @Transaction
+  public RestUtil.PutResponse<Team> updateTeamUsers(
+      String updatedBy, UUID teamId, List<EntityReference> updatedUsers) {
+
+    if (updatedUsers == null) {
+      throw new IllegalArgumentException("Users list cannot be null");
+    }
+
+    Team team = Entity.getEntity(Entity.TEAM, teamId, USERS_FIELD, Include.NON_DELETED);
+    if (!team.getTeamType().equals(CreateTeam.TeamType.GROUP)) {
+      throw new IllegalArgumentException(
+          CatalogExceptionMessage.invalidTeamUpdateUsers(team.getTeamType()));
+    }
+
+    List<EntityReference> currentUsers = team.getUsers();
+
+    Set<UUID> oldUserIds =
+        currentUsers.stream().map(EntityReference::getId).collect(Collectors.toSet());
+    Set<UUID> updatedUserIds =
+        updatedUsers.stream().map(EntityReference::getId).collect(Collectors.toSet());
+    List<EntityReference> addedUsers =
+        updatedUsers.stream()
+            .filter(user -> !oldUserIds.contains(user.getId()))
+            .collect(Collectors.toList());
+
+    Optional.of(addedUsers).ifPresent(this::validateUsers);
+
+    List<UUID> addedUserIds =
+        updatedUsers.stream()
+            .map(EntityReference::getId)
+            .filter(id -> !oldUserIds.contains(id))
+            .collect(Collectors.toList());
+
+    List<UUID> removedUserIds =
+        currentUsers.stream()
+            .map(EntityReference::getId)
+            .filter(id -> !updatedUserIds.contains(id))
+            .collect(Collectors.toList());
+
+    Optional.of(addedUserIds)
+        .filter(ids -> !ids.isEmpty())
+        .ifPresent(
+            ids -> bulkAddToRelationship(teamId, ids, Entity.TEAM, Entity.USER, Relationship.HAS));
+
+    Optional.of(removedUserIds)
+        .filter(ids -> !ids.isEmpty())
+        .ifPresent(
+            ids ->
+                bulkRemoveToRelationship(teamId, ids, Entity.TEAM, Entity.USER, Relationship.HAS));
+
+    setFieldsInternal(team, new EntityUtil.Fields(allowedFields, USERS_FIELD));
+    ChangeDescription change = new ChangeDescription().withPreviousVersion(team.getVersion());
+    fieldAdded(change, USERS_FIELD, updatedUsers);
+
+    ChangeEvent changeEvent =
+        new ChangeEvent()
+            .withId(UUID.randomUUID())
+            .withEntity(team)
+            .withChangeDescription(change)
+            .withEventType(EventType.ENTITY_UPDATED)
+            .withEntityType(entityType)
+            .withEntityId(teamId)
+            .withEntityFullyQualifiedName(team.getFullyQualifiedName())
+            .withUserName(updatedBy)
+            .withTimestamp(System.currentTimeMillis())
+            .withCurrentVersion(team.getVersion())
+            .withPreviousVersion(change.getPreviousVersion());
+    team.setChangeDescription(change);
+
+    return new RestUtil.PutResponse<>(Response.Status.OK, changeEvent, ENTITY_FIELDS_CHANGED);
+  }
+
+  public final RestUtil.PutResponse<Team> deleteTeamUser(
+      String updatedBy, UUID teamId, UUID userId) {
+    Team team = find(teamId, NON_DELETED);
+    if (!team.getTeamType().equals(CreateTeam.TeamType.GROUP)) {
+      throw new IllegalArgumentException(
+          CatalogExceptionMessage.invalidTeamUpdateUsers(team.getTeamType()));
+    }
+
+    // Validate user
+    EntityReference user = Entity.getEntityReferenceById(Entity.USER, userId, NON_DELETED);
+
+    deleteRelationship(teamId, Entity.TEAM, userId, Entity.USER, Relationship.HAS);
+
+    ChangeDescription change = new ChangeDescription().withPreviousVersion(team.getVersion());
+    fieldDeleted(change, USERS_FIELD, List.of(user));
+
+    ChangeEvent changeEvent =
+        new ChangeEvent()
+            .withId(UUID.randomUUID())
+            .withEntity(team)
+            .withChangeDescription(change)
+            .withEventType(EventType.ENTITY_UPDATED)
+            .withEntityFullyQualifiedName(team.getFullyQualifiedName())
+            .withEntityType(entityType)
+            .withEntityId(teamId)
+            .withUserName(updatedBy)
+            .withTimestamp(System.currentTimeMillis())
+            .withCurrentVersion(team.getVersion())
+            .withPreviousVersion(change.getPreviousVersion());
+
+    return new RestUtil.PutResponse<>(Response.Status.OK, changeEvent, ENTITY_FIELDS_CHANGED);
+  }
+
   public void initOrganization() {
     organization = findByNameOrNull(ORGANIZATION_NAME, ALL);
     if (organization == null) {

openmetadata-service/src/main/java/org/openmetadata/service/resources/teams/TeamResource.java

@@ -57,6 +57,7 @@ import org.openmetadata.schema.entity.teams.Team;
 import org.openmetadata.schema.entity.teams.TeamHierarchy;
 import org.openmetadata.schema.type.ChangeEvent;
 import org.openmetadata.schema.type.EntityHistory;
+import org.openmetadata.schema.type.EntityReference;
 import org.openmetadata.schema.type.Include;
 import org.openmetadata.schema.type.MetadataOperation;
 import org.openmetadata.schema.type.api.BulkAssets;
@@ -71,6 +72,7 @@ import org.openmetadata.service.limits.Limits;
 import org.openmetadata.service.resources.Collection;
 import org.openmetadata.service.resources.EntityResource;
 import org.openmetadata.service.security.Authorizer;
+import org.openmetadata.service.security.policyevaluator.OperationContext;
 import org.openmetadata.service.util.CSVExportResponse;
 import org.openmetadata.service.util.EntityUtil;
 import org.openmetadata.service.util.JsonUtils;
@@ -666,6 +668,69 @@ public class TeamResource extends EntityResource<Team, TeamRepository> {
     return importCsvInternal(securityContext, name, csv, dryRun);
   }
 
+  @PUT
+  @Path("/{teamId}/users")
+  @Operation(
+      operationId = "updateTeamUsers",
+      summary = "Update team users",
+      description =
+          "Update the list of users for a team. Replaces existing users with the provided list.",
+      responses = {
+        @ApiResponse(
+            responseCode = "200",
+            description = "Updated team users",
+            content = @Content(mediaType = "application/json")),
+        @ApiResponse(responseCode = "404", description = "Team not found")
+      })
+  public Response updateTeamUsers(
+      @Context UriInfo uriInfo,
+      @Context SecurityContext securityContext,
+      @PathParam("teamId") UUID teamId,
+      List<EntityReference> users) {
+
+    OperationContext operationContext =
+        new OperationContext(entityType, MetadataOperation.EDIT_ALL);
+    authorizer.authorize(securityContext, operationContext, getResourceContextById(teamId));
+    return repository
+        .updateTeamUsers(securityContext.getUserPrincipal().getName(), teamId, users)
+        .toResponse();
+  }
+
+  @DELETE
+  @Path("/{teamId}/users/{userId}")
+  @Operation(
+      operationId = "deleteTeamUser",
+      summary = "Remove a user from a team",
+      description = "Remove the user identified by `userId` from the team identified by `teamId`.",
+      responses = {
+        @ApiResponse(
+            responseCode = "200",
+            description = "User removed from team",
+            content =
+                @Content(
+                    mediaType = "application/json",
+                    schema = @Schema(implementation = ChangeEvent.class))),
+        @ApiResponse(responseCode = "404", description = "Team or user not found")
+      })
+  public Response deleteTeamUser(
+      @Context UriInfo uriInfo,
+      @Context SecurityContext securityContext,
+      @Parameter(description = "Id of the team", schema = @Schema(type = "UUID"))
+          @PathParam("teamId")
+          UUID teamId,
+      @Parameter(description = "Id of the user being removed", schema = @Schema(type = "string"))
+          @PathParam("userId")
+          String userId) {
+
+    OperationContext operationContext =
+        new OperationContext(entityType, MetadataOperation.EDIT_ALL);
+    authorizer.authorize(securityContext, operationContext, getResourceContextById(teamId));
+    return repository
+        .deleteTeamUser(
+            securityContext.getUserPrincipal().getName(), teamId, UUID.fromString(userId))
+        .toResponse();
+  }
+
   @PUT
   @Path("/name/{name}/importAsync")
   @Consumes(MediaType.TEXT_PLAIN)

openmetadata-service/src/test/java/org/openmetadata/service/resources/teams/TeamResourceTest.java

@@ -72,6 +72,7 @@ import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import javax.ws.rs.client.WebTarget;
+import javax.ws.rs.core.Response;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.http.client.HttpResponseException;
 import org.junit.jupiter.api.Test;
@@ -93,7 +94,9 @@ import org.openmetadata.schema.entity.teams.TeamHierarchy;
 import org.openmetadata.schema.entity.teams.User;
 import org.openmetadata.schema.type.ApiStatus;
 import org.openmetadata.schema.type.ChangeDescription;
+import org.openmetadata.schema.type.ChangeEvent;
 import org.openmetadata.schema.type.EntityReference;
+import org.openmetadata.schema.type.EventType;
 import org.openmetadata.schema.type.ImageList;
 import org.openmetadata.schema.type.Include;
 import org.openmetadata.schema.type.MetadataOperation;
@@ -932,6 +935,86 @@ public class TeamResourceTest extends EntityResourceTest<Team, CreateTeam> {
     return entity;
   }
 
+  @Test
+  void put_addDeleteTeamUser_200(TestInfo test) throws IOException {
+    // Create a team of type GROUP
+    TeamResourceTest teamResourceTest = new TeamResourceTest();
+    Team team =
+        teamResourceTest.createEntity(teamResourceTest.createRequest(test, 1), ADMIN_AUTH_HEADERS);
+    UUID teamId = team.getId();
+
+    // Add user to the team
+    UserResourceTest userResourceTest = new UserResourceTest();
+    User user1 =
+        userResourceTest.createEntity(userResourceTest.createRequest(test, 1), ADMIN_AUTH_HEADERS);
+
+    User user2 =
+        userResourceTest.createEntity(userResourceTest.createRequest(test, 2), ADMIN_AUTH_HEADERS);
+
+    addAndCheckTeamUser(
+        teamId,
+        List.of(user1.getEntityReference(), user2.getEntityReference()),
+        OK,
+        2,
+        ADMIN_AUTH_HEADERS);
+
+    CreateTeam createDepartmentTeam =
+        createRequest(test)
+            .withDomains(List.of(DOMAIN.getFullyQualifiedName()))
+            .withTeamType(DEPARTMENT);
+    Team departmentTeam = createEntity(createDepartmentTeam, ADMIN_AUTH_HEADERS);
+
+    // Add user only for GROUP type team
+    assertResponse(
+        () ->
+            addAndCheckTeamUser(
+                departmentTeam.getId(),
+                List.of(user1.getEntityReference(), user2.getEntityReference()),
+                OK,
+                0,
+                ADMIN_AUTH_HEADERS),
+        BAD_REQUEST,
+        CatalogExceptionMessage.invalidTeamUpdateUsers(departmentTeam.getTeamType()));
+
+    deleteAndCheckTeamUser(teamId, user1.getId(), 1, ADMIN_AUTH_HEADERS);
+    deleteAndCheckTeamUser(teamId, user2.getId(), 0, ADMIN_AUTH_HEADERS);
+  }
+
+  private void addAndCheckTeamUser(
+      UUID teamId,
+      List<EntityReference> users,
+      Response.Status status,
+      int expectedUserCount,
+      Map<String, String> authHeaders)
+      throws IOException {
+    WebTarget target = getResource("teams/" + teamId + "/users");
+    ChangeEvent event = TestUtils.put(target, users, ChangeEvent.class, status, authHeaders);
+    Team team = getEntity(teamId, authHeaders);
+    validateEntityReferences(team.getUsers());
+    assertEquals(expectedUserCount, team.getUsers().size());
+    validateChangeEvents(
+        team,
+        event.getTimestamp(),
+        EventType.ENTITY_UPDATED,
+        event.getChangeDescription(),
+        authHeaders);
+  }
+
+  private void deleteAndCheckTeamUser(
+      UUID teamId, UUID userId, int expectedUserCount, Map<String, String> authHeaders)
+      throws IOException {
+    WebTarget target = getResource("teams/" + teamId + "/users/" + userId);
+    ChangeEvent change = TestUtils.delete(target, ChangeEvent.class, authHeaders);
+    Team team = getEntity(teamId, authHeaders);
+    assertEquals(expectedUserCount, team.getUsers().size());
+    validateChangeEvents(
+        team,
+        change.getTimestamp(),
+        EventType.ENTITY_UPDATED,
+        change.getChangeDescription(),
+        authHeaders);
+  }
+
   private static void validateTeam(
       Team team,
       String expectedDescription,