Fix #577: Users API should support put op (#578)

* Fix #577: Users API should support put op

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/UserRepository.java

@@ -33,6 +33,7 @@ import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.util.EntityUtil;
 import org.openmetadata.catalog.util.EntityUtil.Fields;
 import org.openmetadata.catalog.util.JsonUtils;
+import org.openmetadata.catalog.util.RestUtil;
 import org.openmetadata.common.utils.CipherText;
 import org.skife.jdbi.v2.sqlobject.Bind;
 import org.skife.jdbi.v2.sqlobject.CreateSqlObject;
@@ -43,6 +44,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import javax.json.JsonPatch;
+import javax.ws.rs.core.Response;
 import java.io.IOException;
 import java.security.GeneralSecurityException;
 import java.text.ParseException;
@@ -169,12 +171,7 @@ public abstract class UserRepository {
 
   @Transaction
   public User create(User user, List<UUID> teamIds) throws IOException {
-    List<Team> teams = validateTeams(teamIds);
+    return createInternal(user, teamIds);
-    userDAO().insert(JsonUtils.pojoToJson(user));
-    assignTeams(user, teams);
-    List<EntityReference> entityRefs = toEntityReference(teams);
-    user.setTeams(entityRefs.isEmpty() ? null : entityRefs);
-    return user;
   }
 
   @Transaction
@@ -189,6 +186,28 @@ public abstract class UserRepository {
     relationshipDAO().deleteFrom(id, FOLLOWS.ordinal());
   }
 
+  @Transaction
+  public RestUtil.PutResponse<User> createOrUpdate(User updatedUser) throws
+          IOException, ParseException {
+    User storedUser = JsonUtils.readValue(userDAO().findByName(updatedUser.getName()), User.class);
+    List<UUID> teamIds = new ArrayList<>();
+    if (updatedUser.getTeams() != null) {
+      for (EntityReference team : updatedUser.getTeams()) {
+        teamIds.add(team.getId());
+      }
+    }
+    if (storedUser == null) {
+      return new RestUtil.PutResponse<>(Response.Status.CREATED, createInternal(updatedUser, teamIds));
+    }
+    updatedUser.setId(storedUser.getId());
+    userDAO().update(updatedUser.getId().toString(), JsonUtils.pojoToJson(updatedUser));
+    List<Team> teams = validateTeams(teamIds);
+    if (!teams.isEmpty()) {
+      assignTeams(updatedUser, teams);
+    }
+    return new RestUtil.PutResponse<>(Response.Status.OK, updatedUser);
+  }
+
   @Transaction
   public User patch(String id, JsonPatch patch) throws IOException {
     User original = setFields(validateUser(id), USER_PATCH_FIELDS); // Query 1 - find user by Id
@@ -198,6 +217,11 @@ public abstract class UserRepository {
     return updated;
   }
 
+  @Transaction
+  public EntityReference getOwnerReference(User user) throws IOException {
+    return EntityUtil.getEntityReference(user);
+  }
+
   private void patch(User original, User updated) throws IOException {
     String userId = original.getId().toString();
     if (!updated.getId().equals(original.getId())) {
@@ -262,6 +286,15 @@ public abstract class UserRepository {
     return EntityUtil.validate(userId, userDAO().findById(userId), User.class);
   }
 
+  private User createInternal(User user, List<UUID> teamIds) throws IOException {
+    List<Team> teams = validateTeams(teamIds);
+    userDAO().insert(JsonUtils.pojoToJson(user));
+    assignTeams(user, teams);
+    List<EntityReference> entityRefs = toEntityReference(teams);
+    user.setTeams(entityRefs.isEmpty() ? null : entityRefs);
+    return user;
+  }
+
 
   private List<Team> validateTeams(List<UUID> teamIds) throws IOException {
     if (teamIds == null) {

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/teams/UserResource.java

@@ -50,6 +50,7 @@ import javax.ws.rs.DELETE;
 import javax.ws.rs.DefaultValue;
 import javax.ws.rs.GET;
 import javax.ws.rs.POST;
+import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
 import javax.ws.rs.Produces;
@@ -237,6 +238,30 @@ public class UserResource {
     return Response.created(user.getHref()).entity(user).build();
   }
 
+  @PUT
+  @Operation(summary = "Create or Update a user", tags = "users",
+          description = "Create or Update a user.",
+          responses = {
+                  @ApiResponse(responseCode = "200", description = "The user ",
+                          content = @Content(mediaType = "application/json",
+                                  schema = @Schema(implementation = User.class))),
+                  @ApiResponse(responseCode = "400", description = "Bad request")
+          })
+  public Response createOrUpdateUser(@Context UriInfo uriInfo,
+                                     @Context SecurityContext securityContext,
+                                     @Valid CreateUser create) throws IOException, ParseException {
+    if (create.getIsAdmin() != null && create.getIsAdmin()) {
+      SecurityUtil.checkAdminOrBotRole(authorizer, securityContext);
+    }
+    User user = new User().withId(UUID.randomUUID()).withName(create.getName()).withEmail(create.getEmail())
+            .withDisplayName(create.getDisplayName()).withIsBot(create.getIsBot()).withIsAdmin(create.getIsAdmin())
+            .withProfile(create.getProfile()).withTimezone(create.getTimezone());
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOwnerReference(user));
+    RestUtil.PutResponse<User> response = dao.createOrUpdate(user);
+    user = addHref(uriInfo, response.getEntity());
+    return Response.status(response.getStatus()).entity(user).build();
+  }
+
   @PATCH
   @Valid
   @Path("/{id}")

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/teams/UserResourceTest.java

@@ -40,6 +40,7 @@ import org.slf4j.LoggerFactory;
 
 import javax.json.JsonPatch;
 import javax.ws.rs.client.WebTarget;
+import javax.ws.rs.core.Response;
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -54,6 +55,7 @@ import static javax.ws.rs.core.Response.Status.CONFLICT;
 import static javax.ws.rs.core.Response.Status.CREATED;
 import static javax.ws.rs.core.Response.Status.FORBIDDEN;
 import static javax.ws.rs.core.Response.Status.NOT_FOUND;
+import static javax.ws.rs.core.Response.Status.OK;
 import static javax.ws.rs.core.Response.Status.UNAUTHORIZED;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -156,6 +158,18 @@ public class UserResourceTest extends CatalogApplicationTest {
     createAndCheckUser(create, adminAuthHeaders());
   }
 
+  @Test
+  public void put_validUser_200_ok(TestInfo test) throws HttpResponseException {
+    // Create user with different optional fields
+    CreateUser create = create(test, 1);
+    User user = createOrUpdateAndCheckUser(create, CREATED, adminAuthHeaders());
+    CreateUser update = new CreateUser().withName(user.getName()).withEmail("test1@email.com")
+            .withDisplayName("displayName1");
+    createOrUpdateAndCheckUser(update, OK, adminAuthHeaders());
+  }
+
+
+
   @Test
   public void post_validAdminUser_Non_Admin_401(TestInfo test) {
     CreateUser create = create(test, 6)
@@ -485,6 +499,12 @@ public class UserResourceTest extends CatalogApplicationTest {
     assertResponse(exception, NOT_FOUND, entityNotFound("User", TestUtils.NON_EXISTENT_ENTITY));
   }
 
+  public static User putUser(CreateUser user, Response.Status expectedStatus, Map<String, String> authHeaders)
+          throws HttpResponseException {
+    WebTarget target = CatalogApplicationTest.getResource("users");
+    return TestUtils.put(target, user, User.class, expectedStatus, authHeaders);
+  }
+
   private User patchUser(UUID userId, String originalJson, User updated, Map<String, String> headers)
           throws JsonProcessingException, HttpResponseException {
     String updatedJson = JsonUtils.pojoToJson(updated);
@@ -540,6 +560,24 @@ public class UserResourceTest extends CatalogApplicationTest {
     return user;
   }
 
+  public static User createOrUpdateAndCheckUser(CreateUser create, Response.Status expectedStatus,
+                                                Map<String, String> authHeaders)
+          throws HttpResponseException {
+    final User user = putUser(create, expectedStatus, authHeaders);
+    List<Team> expectedTeams = new ArrayList<>();
+    for (UUID teamId : Optional.ofNullable(create.getTeams()).orElse(Collections.emptyList())) {
+      expectedTeams.add(new Team().withId(teamId));
+    }
+    validateUser(user, create.getName(), create.getDisplayName(), expectedTeams, create.getProfile(),
+            create.getTimezone(), create.getIsBot(), create.getIsAdmin());
+
+    // GET the newly created user and validate
+    User getUser = getUser(user.getId(), "profile,teams", authHeaders);
+    validateUser(getUser, create.getName(), create.getDisplayName(), expectedTeams, create.getProfile(),
+            create.getTimezone(), create.getIsBot(), create.getIsAdmin());
+    return user;
+  }
+
   public static CreateUser create(TestInfo test, int index) {
     return new CreateUser().withName(getUserName(test) + index).withEmail(getUserName(test) + "@open-metadata.org");
   }

ingestion/setup.py

@@ -99,7 +99,7 @@ plugins: Dict[str, Set[str]] = {
     "data-profiler": {"openmetadata-data-profiler"},
     "snowflake": {"snowflake-sqlalchemy<=1.2.4"},
     "snowflake-usage": {"snowflake-sqlalchemy<=1.2.4"},
-    "sample-data": {"faker~=8.1.1","pandas~=1.3.1"},
+    "sample-data": {"faker~=8.1.1", "pandas~=1.3.1"},
     "superset": {},
     "tableau": {"tableau-api-lib==0.1.22"},
     "vertica": {"sqlalchemy-vertica[vertica-python]>=0.0.5"}

ingestion/src/metadata/ingestion/sink/metadata_rest_users.py

@@ -86,7 +86,7 @@ class MetadataRestUsersSink(Sink):
                                      email=record.email,
                                      teams=teams)
         try:
-            self.client.post(self.api_users, data=metadata_user.to_json())
+            self.client.put(self.api_users, data=metadata_user.to_json())
             self.status.records_written(record.github_username)
             logger.info("Sink: {}".format(record.github_username))
         except APIError: