ATLAS-3483 Add REST support to add/delete labels

repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasEntityStore.java

@@ -234,7 +234,18 @@ public interface AtlasEntityStore {
     String setClassifications(AtlasEntityHeaders entityHeaders);
 
     /**
-     * Set Labels
+     * Set labels to given entity, if labels is null/empty, existing labels will all be removed.
      */
     void setLabels(String guid, Set<String> labels) throws AtlasBaseException;
+
+    /**
+     * Remove given labels, if labels is null/empty, no labels will be removed. If any labels in
+     * labels set are non-existing labels, they will be ignored, only existing labels will be removed.
+     */
+    void removeLabels(String guid, Set<String> labels) throws AtlasBaseException;
+
+    /**
+     * Add given labels to the given entity, if labels is null/empty, no labels will be added.
+     */
+    void addLabels(String guid, Set<String> labels) throws AtlasBaseException;
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2.java

@@ -767,6 +767,58 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
         }
     }
 
+    @Override
+    @GraphTransaction
+    public void removeLabels(String guid, Set<String> labels) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> removeLabels()");
+        }
+
+        if (StringUtils.isEmpty(guid)) {
+            throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "guid is null/empty");
+        }
+
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+
+        if (entityVertex == null) {
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
+        }
+
+        validateLabels(labels);
+
+        entityGraphMapper.removeLabels(entityVertex, labels);
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== removeLabels()");
+        }
+    }
+
+    @Override
+    @GraphTransaction
+    public void addLabels(String guid, Set<String> labels) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> addLabels()");
+        }
+
+        if (StringUtils.isEmpty(guid)) {
+            throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "guid is null/empty");
+        }
+
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+
+        if (entityVertex == null) {
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
+        }
+
+        validateLabels(labels);
+
+        entityGraphMapper.addLabels(entityVertex, labels);
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== addLabels()");
+        }
+    }
+
     private EntityMutationResponse createOrUpdate(EntityStream entityStream, boolean isPartialUpdate, boolean replaceClassifications) throws AtlasBaseException {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> createOrUpdate()");

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasGraphUtilsV2.java

@@ -62,6 +62,7 @@ import static org.apache.atlas.repository.Constants.CLASSIFICATION_NAMES_KEY;
 import static org.apache.atlas.repository.Constants.ENTITY_TYPE_PROPERTY_KEY;
 import static org.apache.atlas.repository.Constants.INDEX_SEARCH_VERTEX_PREFIX_DEFAULT;
 import static org.apache.atlas.repository.Constants.INDEX_SEARCH_VERTEX_PREFIX_PROPERTY;
+import static org.apache.atlas.repository.Constants.LABELS_PROPERTY_KEY;
 import static org.apache.atlas.repository.Constants.PROPAGATED_CLASSIFICATION_NAMES_KEY;
 import static org.apache.atlas.repository.Constants.STATE_PROPERTY_KEY;
 import static org.apache.atlas.repository.Constants.TYPE_NAME_PROPERTY_KEY;
@@ -649,7 +650,7 @@ public class AtlasGraphUtilsV2 {
         List<String> classificationNames = null;
         String classificationNamesString =  entityVertex.getProperty(propertyKey, String.class);
         if (StringUtils.isNotEmpty(classificationNamesString)) {
-            classificationNames = Arrays.asList(classificationNamesString.split("\\|"));
+            classificationNames = Arrays.asList(StringUtils.split(classificationNamesString, "\\|"));
         }
         return classificationNames;
     }

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java

@@ -339,7 +339,37 @@ public class EntityGraphMapper {
         }
     }
 
-    private String getLabelString(Set<String> labels) {
+    public void addLabels(AtlasVertex vertex, Set<String> labels) {
+        if (CollectionUtils.isNotEmpty(labels)) {
+            final Set<String> existingLabels = GraphHelper.getLabels(vertex);
+            final Set<String> updatedLabels;
+
+            if (CollectionUtils.isEmpty(existingLabels)) {
+                updatedLabels = labels;
+            } else {
+                updatedLabels = existingLabels;
+                updatedLabels.addAll(labels);
+            }
+
+            setLabels(vertex, updatedLabels);
+        }
+    }
+
+    public void removeLabels(AtlasVertex vertex, Set<String> labels) {
+        if (CollectionUtils.isNotEmpty(labels)) {
+            final Set<String> existingLabels = GraphHelper.getLabels(vertex);
+            Set<String> updatedLabels = null;
+
+            if (CollectionUtils.isNotEmpty(existingLabels)) {
+                updatedLabels = existingLabels;
+                updatedLabels.removeAll(labels);
+            }
+
+            setLabels(vertex, updatedLabels);
+        }
+    }
+
+    private String getLabelString(Collection<String> labels) {
         String ret = null;
 
         if (!labels.isEmpty()) {

repository/src/test/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2Test.java

@@ -1185,4 +1185,48 @@ public class AtlasEntityStoreV2Test extends AtlasEntityTestBase {
             assertEquals(ex.getAtlasErrorCode(), INVALID_LABEL_CHARACTERS);
         }
     }
+
+    @Test (dependsOnMethods = "invalidLabelCharactersToEntity")
+    public void addMoreLabelsToEntity() throws AtlasBaseException {
+        Set<String> labels = new HashSet<>();
+        labels.add("label_1_add");
+        labels.add("label_2_add");
+        labels.add("label_3_add");
+
+        entityStore.addLabels(tblEntityGuid, labels);
+        AtlasEntity tblEntity = getEntityFromStore(tblEntityGuid);
+        Assert.assertTrue(tblEntity.getLabels().containsAll(labels));
+
+        tblEntity.setAttribute("description", "tbl for labels");
+        AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tblEntity);
+        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), true);
+        validateMutationResponse(response, EntityOperation.PARTIAL_UPDATE, 1);
+        tblEntity = getEntityFromStore(response.getFirstEntityPartialUpdated());
+        Assert.assertEquals(tblEntity.getLabels(), labels );
+    }
+
+    @Test (dependsOnMethods = "addMoreLabelsToEntity")
+    public void deleteLabelsToEntity() throws AtlasBaseException {
+        Set<String> labels = new HashSet<>();
+        labels.add("label_1_add");
+        labels.add("label_2_add");
+
+        entityStore.removeLabels(tblEntityGuid, labels);
+        AtlasEntity tblEntity = getEntityFromStore(tblEntityGuid);
+        Assert.assertNotNull(tblEntity.getLabels());
+        Assert.assertEquals(tblEntity.getLabels().size(), 1);
+
+        labels.clear();
+        labels.add("label_4_add");
+        entityStore.removeLabels(tblEntityGuid, labels);
+        tblEntity = getEntityFromStore(tblEntityGuid);
+        Assert.assertNotNull(tblEntity.getLabels());
+        Assert.assertEquals(tblEntity.getLabels().size(), 1);
+
+        labels.clear();
+        labels.add("label_3_add");
+        entityStore.removeLabels(tblEntityGuid, labels);
+        tblEntity = getEntityFromStore(tblEntityGuid);
+        Assert.assertNull(tblEntity.getLabels());
+    }
 }
\ No newline at end of file

webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java

@@ -838,6 +838,132 @@ public class EntityREST {
         }
     }
 
+    /**
+     * delete given labels to a given entity
+     * @param guid - Unique entity identifier
+     * @throws AtlasBaseException
+     */
+    @DELETE
+    @Path("/guid/{guid}/labels")
+    @Produces(Servlets.JSON_MEDIA_TYPE)
+    @Consumes(Servlets.JSON_MEDIA_TYPE)
+    public void removeLabels(@PathParam("guid") final String guid, Set<String> labels) throws AtlasBaseException {
+        AtlasPerfTracer perf = null;
+
+        try {
+            if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
+                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.deleteLabels()");
+            }
+
+            entitiesStore.removeLabels(guid, labels);
+        } finally {
+            AtlasPerfTracer.log(perf);
+        }
+    }
+
+    /**
+     * add given labels to a given entity
+     * @param guid - Unique entity identifier
+     * @throws AtlasBaseException
+     */
+    @PUT
+    @Path("/guid/{guid}/labels")
+    @Produces(Servlets.JSON_MEDIA_TYPE)
+    @Consumes(Servlets.JSON_MEDIA_TYPE)
+    public void addLabels(@PathParam("guid") final String guid, Set<String> labels) throws AtlasBaseException {
+        AtlasPerfTracer perf = null;
+
+        try {
+            if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
+                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.addLabels()");
+            }
+
+            entitiesStore.addLabels(guid, labels);
+        } finally {
+            AtlasPerfTracer.log(perf);
+        }
+    }
+
+    @POST
+    @Path("/uniqueAttribute/type/{typeName}/labels")
+    public void setLabels(@PathParam("typeName") String typeName, Set<String> labels,
+                          @Context HttpServletRequest servletRequest) throws AtlasBaseException {
+
+        Servlets.validateQueryParamLength("typeName", typeName);
+        AtlasPerfTracer perf = null;
+
+        try {
+            if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
+                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.setLabels(" + typeName + ")");
+            }
+
+            AtlasEntityType     entityType = ensureEntityType(typeName);
+            Map<String, Object> attributes = getAttributes(servletRequest);
+            String              guid       = entitiesStore.getGuidByUniqueAttributes(entityType, attributes);
+
+            if (guid == null) {
+                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, typeName, attributes.toString());
+            }
+
+            entitiesStore.setLabels(guid, labels);
+        } finally {
+            AtlasPerfTracer.log(perf);
+        }
+    }
+
+    @PUT
+    @Path("/uniqueAttribute/type/{typeName}/labels")
+    public void addLabels(@PathParam("typeName") String typeName, Set<String> labels,
+                          @Context HttpServletRequest servletRequest) throws AtlasBaseException {
+        Servlets.validateQueryParamLength("typeName", typeName);
+        AtlasPerfTracer perf = null;
+
+        try {
+            if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
+                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.addLabels(" + typeName + ")");
+            }
+
+            AtlasEntityType     entityType = ensureEntityType(typeName);
+            Map<String, Object> attributes = getAttributes(servletRequest);
+            String              guid       = entitiesStore.getGuidByUniqueAttributes(entityType, attributes);
+
+            if (guid == null) {
+                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, typeName, attributes.toString());
+            }
+
+            entitiesStore.addLabels(guid, labels);
+        } finally {
+            AtlasPerfTracer.log(perf);
+        }
+    }
+
+    @DELETE
+    @Path("/uniqueAttribute/type/{typeName}/labels")
+    public void removeLabels(@PathParam("typeName") String typeName, Set<String> labels,
+                             @Context HttpServletRequest servletRequest) throws AtlasBaseException {
+
+        Servlets.validateQueryParamLength("typeName", typeName);
+        AtlasPerfTracer perf = null;
+
+        try {
+            if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
+                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.removeLabels(" + typeName + ")");
+            }
+
+            AtlasEntityType     entityType = ensureEntityType(typeName);
+            Map<String, Object> attributes = getAttributes(servletRequest);
+            String              guid       = entitiesStore.getGuidByUniqueAttributes(entityType, attributes);
+
+            if (guid == null) {
+                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, typeName, attributes.toString());
+            }
+
+            entitiesStore.removeLabels(guid, labels);
+        } finally {
+            AtlasPerfTracer.log(perf);
+        }
+    }
+
     private AtlasEntityType ensureEntityType(String typeName) throws AtlasBaseException {
         AtlasEntityType ret = typeRegistry.getEntityTypeByName(typeName);