Repository: incubator-atlas Updated Branches: refs/heads/master 2615b308e -> 4feee3bf6
ATLAS-1751: Implement REST endpoint to support update of classification attribute Project: http://git-wip-us.apache.org/repos/asf/incubator-atlas/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-atlas/commit/4feee3bf Tree: http://git-wip-us.apache.org/repos/asf/incubator-atlas/tree/4feee3bf Diff: http://git-wip-us.apache.org/repos/asf/incubator-atlas/diff/4feee3bf Branch: refs/heads/master Commit: 4feee3bf61a032f742c0c4738e6ffdb423cdd14d Parents: 2615b30 Author: Sarath Subramanian <[email protected]> Authored: Wed Apr 26 16:04:09 2017 -0700 Committer: Sarath Subramanian <[email protected]> Committed: Wed Apr 26 16:04:09 2017 -0700 ---------------------------------------------------------------------- .../java/org/apache/atlas/EntityAuditEvent.java | 4 +- dashboardv2/public/js/utils/Enums.js | 1 + .../test/java/org/apache/atlas/TestUtilsV2.java | 8 ++- .../notification/entity/EntityNotification.java | 3 +- .../repository/audit/EntityAuditListener.java | 15 +++++ .../store/graph/AtlasEntityStore.java | 5 ++ .../graph/v1/AtlasEntityChangeNotifier.java | 20 ++++++ .../store/graph/v1/AtlasEntityStoreV1.java | 64 +++++++++++++++++++- .../store/graph/v1/EntityGraphMapper.java | 29 +++++++++ .../service/DefaultMetadataServiceTest.java | 5 ++ .../atlas/listener/EntityChangeListener.java | 10 +++ .../NotificationEntityChangeListener.java | 5 ++ .../org/apache/atlas/web/rest/EntityREST.java | 27 +++++++++ .../atlas/web/adapters/TestEntityREST.java | 56 +++++++++++++++++ 14 files changed, 247 insertions(+), 5 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/client/src/main/java/org/apache/atlas/EntityAuditEvent.java ---------------------------------------------------------------------- diff --git a/client/src/main/java/org/apache/atlas/EntityAuditEvent.java b/client/src/main/java/org/apache/atlas/EntityAuditEvent.java index e541531..904674d 100644 --- a/client/src/main/java/org/apache/atlas/EntityAuditEvent.java +++ b/client/src/main/java/org/apache/atlas/EntityAuditEvent.java @@ -28,8 +28,8 @@ import java.util.Objects; */ public class EntityAuditEvent { public enum EntityAuditAction { - ENTITY_CREATE, ENTITY_UPDATE, ENTITY_DELETE, TAG_ADD, TAG_DELETE, - ENTITY_IMPORT_CREATE, ENTITY_IMPORT_UPDATE, ENTITY_IMPORT_DELETE + ENTITY_CREATE, ENTITY_UPDATE, ENTITY_DELETE, TAG_ADD, TAG_DELETE, TAG_UPDATE, + ENTITY_IMPORT_CREATE, ENTITY_IMPORT_UPDATE, ENTITY_IMPORT_DELETE, } private String entityId; http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/dashboardv2/public/js/utils/Enums.js ---------------------------------------------------------------------- diff --git a/dashboardv2/public/js/utils/Enums.js b/dashboardv2/public/js/utils/Enums.js index a44490a..90510bc 100644 --- a/dashboardv2/public/js/utils/Enums.js +++ b/dashboardv2/public/js/utils/Enums.js @@ -27,6 +27,7 @@ define(['require'], function(require) { ENTITY_DELETE: "Entity Deleted", TAG_ADD: "Tag Added", TAG_DELETE: "Tag Deleted", + TAG_UPDATE: "Tag Updated", ENTITY_IMPORT_CREATE: "Entity Created by import", ENTITY_IMPORT_UPDATE: "Entity Updated by import", ENTITY_IMPORT_DELETE: "Entity Deleted by import" http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/intg/src/test/java/org/apache/atlas/TestUtilsV2.java ---------------------------------------------------------------------- diff --git a/intg/src/test/java/org/apache/atlas/TestUtilsV2.java b/intg/src/test/java/org/apache/atlas/TestUtilsV2.java index 7b1f2ad..2a6ea92 100755 --- a/intg/src/test/java/org/apache/atlas/TestUtilsV2.java +++ b/intg/src/test/java/org/apache/atlas/TestUtilsV2.java @@ -516,6 +516,7 @@ public final class TestUtilsV2 { public static final String TABLE_NAME = "bar"; public static final String CLASSIFICATION = "classification"; public static final String PII = "PII"; + public static final String PHI = "PHI"; public static final String SUPER_TYPE_NAME = "Base"; public static final String STORAGE_DESC_TYPE = "hive_storagedesc"; public static final String PARTITION_STRUCT_TYPE = "partition_struct_type"; @@ -787,9 +788,14 @@ public final class TestUtilsV2 { AtlasTypeUtil.createTraitTypeDef("fetl" + CLASSIFICATION, "fetl" + CLASSIFICATION + _description, ImmutableSet.of(CLASSIFICATION), AtlasTypeUtil.createRequiredAttrDef("tag", "string")); + AtlasClassificationDef phiTypeDefinition = AtlasTypeUtil.createTraitTypeDef(PHI, PHI + _description, ImmutableSet.<String>of(), + AtlasTypeUtil.createRequiredAttrDef("stringAttr", "string"), + AtlasTypeUtil.createRequiredAttrDef("booleanAttr", "boolean"), + AtlasTypeUtil.createRequiredAttrDef("integerAttr", "int")); + return AtlasTypeUtil.getTypesDef(ImmutableList.of(enumTypeDefinition), ImmutableList.of(structTypeDefinition, partitionDefinition), - ImmutableList.of(classificationTypeDefinition, fetlClassificationTypeDefinition, piiTypeDefinition), + ImmutableList.of(classificationTypeDefinition, fetlClassificationTypeDefinition, piiTypeDefinition, phiTypeDefinition), ImmutableList.of(superTypeDefinition, databaseTypeDefinition, columnsDefinition, tableTypeDefinition, storageDescClsDef, partClsDef, processClsType)); } http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/notification/src/main/java/org/apache/atlas/notification/entity/EntityNotification.java ---------------------------------------------------------------------- diff --git a/notification/src/main/java/org/apache/atlas/notification/entity/EntityNotification.java b/notification/src/main/java/org/apache/atlas/notification/entity/EntityNotification.java index 82a1100..379e815 100644 --- a/notification/src/main/java/org/apache/atlas/notification/entity/EntityNotification.java +++ b/notification/src/main/java/org/apache/atlas/notification/entity/EntityNotification.java @@ -35,7 +35,8 @@ public interface EntityNotification { ENTITY_UPDATE, ENTITY_DELETE, TRAIT_ADD, - TRAIT_DELETE + TRAIT_DELETE, + TRAIT_UPDATE } http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/repository/src/main/java/org/apache/atlas/repository/audit/EntityAuditListener.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/audit/EntityAuditListener.java b/repository/src/main/java/org/apache/atlas/repository/audit/EntityAuditListener.java index 9d0f802..eab86c4 100644 --- a/repository/src/main/java/org/apache/atlas/repository/audit/EntityAuditListener.java +++ b/repository/src/main/java/org/apache/atlas/repository/audit/EntityAuditListener.java @@ -100,6 +100,18 @@ public class EntityAuditListener implements EntityChangeListener { } @Override + public void onTraitsUpdated(ITypedReferenceableInstance entity, Collection<? extends IStruct> traits) throws AtlasException { + if (traits != null) { + for (IStruct trait : traits) { + EntityAuditEvent event = createEvent(entity, EntityAuditAction.TAG_UPDATE, + "Updated trait: " + InstanceSerialization.toJson(trait, true)); + + auditRepository.putEvents(event); + } + } + } + + @Override public void onEntitiesDeleted(Collection<ITypedReferenceableInstance> entities, boolean isImport) throws AtlasException { List<EntityAuditEvent> events = new ArrayList<>(); for (ITypedReferenceableInstance entity : entities) { @@ -279,6 +291,9 @@ public class EntityAuditListener implements EntityChangeListener { case TAG_DELETE: ret = "Deleted trait: "; break; + case TAG_UPDATE: + ret = "Updated trait: "; + break; case ENTITY_IMPORT_CREATE: ret = "Created by import: "; break; http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasEntityStore.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasEntityStore.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasEntityStore.java index c256ae2..e873e91 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasEntityStore.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasEntityStore.java @@ -127,6 +127,11 @@ public interface AtlasEntityStore { */ void addClassifications(String guid, List<AtlasClassification> classification) throws AtlasBaseException; + /** + * Update classification(s) + */ + void updateClassifications(String guid, List<AtlasClassification> classifications) throws AtlasBaseException; + @GraphTransaction void addClassification(List<String> guids, AtlasClassification classification) throws AtlasBaseException; http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityChangeNotifier.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityChangeNotifier.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityChangeNotifier.java index f3d9ca7..d9bc924 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityChangeNotifier.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityChangeNotifier.java @@ -125,6 +125,26 @@ public class AtlasEntityChangeNotifier { } } + public void onClassificationUpdatedToEntity(String entityId, List<AtlasClassification> classifications) throws AtlasBaseException { + // Since the classification attributes are updated in the graph, we need to recursively remap the entityText + doFullTextMapping(entityId); + + ITypedReferenceableInstance entity = toITypedReferenceable(entityId); + List<ITypedStruct> traits = toITypedStructs(classifications); + + if (entity == null || CollectionUtils.isEmpty(traits)) { + return; + } + + for (EntityChangeListener listener : entityChangeListeners) { + try { + listener.onTraitsUpdated(entity, traits); + } catch (AtlasException e) { + throw new AtlasBaseException(AtlasErrorCode.NOTIFICATION_FAILED, e); + } + } + } + private void notifyListeners(List<AtlasEntityHeader> entityHeaders, EntityOperation operation, boolean isImport) throws AtlasBaseException { if (CollectionUtils.isEmpty(entityHeaders)) { return; http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java index 27f6928..70a904b 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java @@ -441,6 +441,51 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore { @Override @GraphTransaction + public void updateClassifications(String guid, List<AtlasClassification> newClassifications) throws AtlasBaseException { + if (LOG.isDebugEnabled()) { + LOG.debug("Updating classifications={} for entity={}", newClassifications, guid); + } + + if (StringUtils.isEmpty(guid)) { + throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "Guid not specified"); + } + + if (CollectionUtils.isEmpty(newClassifications)) { + throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "classifications(s) not specified"); + } + + EntityGraphMapper graphMapper = new EntityGraphMapper(deleteHandler, typeRegistry); + List<AtlasClassification> updatedClassifications = new ArrayList<>(); + + for (AtlasClassification newClassification : newClassifications) { + String classificationName = newClassification.getTypeName(); + AtlasClassification oldClassification = getClassification(guid, classificationName); + + if (oldClassification == null) { + throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATION_NOT_FOUND, classificationName); + } + + validateAndNormalizeForUpdate(newClassification); + + Map<String, Object> newAttrs = newClassification.getAttributes(); + + if (MapUtils.isNotEmpty(newAttrs)) { + for (String attrName : newAttrs.keySet()) { + oldClassification.setAttribute(attrName, newAttrs.get(attrName)); + } + } + + graphMapper.updateClassification(new EntityMutationContext(), guid, oldClassification); + + updatedClassifications.add(oldClassification); + } + + // notify listeners on update to classifications + entityChangeNotifier.onClassificationUpdatedToEntity(guid, updatedClassifications); + } + + @Override + @GraphTransaction public void addClassification(final List<String> guids, final AtlasClassification classification) throws AtlasBaseException { if (CollectionUtils.isEmpty(guids)) { throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "Guid(s) not specified"); @@ -514,7 +559,6 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore { return graphRetriever.getClassification(guid, classificationName); } - private EntityMutationContext preCreateOrUpdate(EntityStream entityStream, EntityGraphMapper entityGraphMapper, boolean isPartialUpdate) throws AtlasBaseException { EntityGraphDiscovery graphDiscoverer = new AtlasEntityGraphDiscoveryV1(typeRegistry, entityStream); EntityGraphDiscoveryContext discoveryContext = graphDiscoverer.discoverEntities(); @@ -607,6 +651,24 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore { type.getNormalizedValue(classification); } + private void validateAndNormalizeForUpdate(AtlasClassification classification) throws AtlasBaseException { + AtlasClassificationType type = typeRegistry.getClassificationTypeByName(classification.getTypeName()); + + if (type == null) { + throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATION_NOT_FOUND, classification.getTypeName()); + } + + List<String> messages = new ArrayList<>(); + + type.validateValueForUpdate(classification, classification.getTypeName(), messages); + + if (!messages.isEmpty()) { + throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, messages); + } + + type.getNormalizedValueForUpdate(classification); + } + /** * Validate if classification is not already associated with the entities * @param guid unique entity id http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java index 9d11aa5..0fe748e 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java @@ -941,6 +941,35 @@ public class EntityGraphMapper { } } + public void updateClassification(final EntityMutationContext context, String guid, AtlasClassification classification) + throws AtlasBaseException { + + AtlasVertex instanceVertex = AtlasGraphUtilsV1.findByGuid(guid); + + if (instanceVertex == null) { + throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid); + } + + String entityTypeName = AtlasGraphUtilsV1.getTypeName(instanceVertex); + + final AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entityTypeName); + + if (LOG.isDebugEnabled()) { + LOG.debug("Updating classification {} for entity {}", classification, guid); + } + + // get the classification vertex from entity + String relationshipLabel = GraphHelper.getTraitLabel(entityTypeName, classification.getTypeName()); + AtlasEdge classificationEdge = graphHelper.getEdgeForLabel(instanceVertex, relationshipLabel); + AtlasVertex classificationVertex = classificationEdge.getInVertex(); + + if (LOG.isDebugEnabled()) { + LOG.debug("updating vertex {} for trait {}", string(classificationVertex), classification.getTypeName()); + } + + mapClassification(EntityOperation.UPDATE, context, classification, entityType, instanceVertex, classificationVertex); + } + private AtlasEdge mapClassification(EntityOperation operation, final EntityMutationContext context, AtlasClassification classification, AtlasEntityType entityType, AtlasVertex parentInstanceVertex, AtlasVertex traitInstanceVertex) throws AtlasBaseException { http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/repository/src/test/java/org/apache/atlas/service/DefaultMetadataServiceTest.java ---------------------------------------------------------------------- diff --git a/repository/src/test/java/org/apache/atlas/service/DefaultMetadataServiceTest.java b/repository/src/test/java/org/apache/atlas/service/DefaultMetadataServiceTest.java index c6d7e9d..2b72f2a 100644 --- a/repository/src/test/java/org/apache/atlas/service/DefaultMetadataServiceTest.java +++ b/repository/src/test/java/org/apache/atlas/service/DefaultMetadataServiceTest.java @@ -1268,6 +1268,11 @@ public class DefaultMetadataServiceTest { } @Override + public void onTraitsUpdated(ITypedReferenceableInstance entity, Collection<? extends IStruct> traits) + throws AtlasException { + } + + @Override public void onEntitiesDeleted(Collection<ITypedReferenceableInstance> entities, boolean isImport) throws AtlasException { deletedEntities.clear(); http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/server-api/src/main/java/org/apache/atlas/listener/EntityChangeListener.java ---------------------------------------------------------------------- diff --git a/server-api/src/main/java/org/apache/atlas/listener/EntityChangeListener.java b/server-api/src/main/java/org/apache/atlas/listener/EntityChangeListener.java index 346c8a2..e05a775 100644 --- a/server-api/src/main/java/org/apache/atlas/listener/EntityChangeListener.java +++ b/server-api/src/main/java/org/apache/atlas/listener/EntityChangeListener.java @@ -67,6 +67,16 @@ public interface EntityChangeListener { * @throws AtlasException if the listener notification fails */ void onTraitsDeleted(ITypedReferenceableInstance entity, Collection<String> traitNames) throws AtlasException; + + /** + * This is upon updating a trait from a typed instance. + * + * @param entity the entity + * @param traits trait that needs to be added to entity + * + * @throws AtlasException if the listener notification fails + */ + void onTraitsUpdated(ITypedReferenceableInstance entity, Collection<? extends IStruct> traits) throws AtlasException; /** * This is upon deleting entities from the repository. http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/webapp/src/main/java/org/apache/atlas/notification/NotificationEntityChangeListener.java ---------------------------------------------------------------------- diff --git a/webapp/src/main/java/org/apache/atlas/notification/NotificationEntityChangeListener.java b/webapp/src/main/java/org/apache/atlas/notification/NotificationEntityChangeListener.java index 565eea7..f3af37d 100644 --- a/webapp/src/main/java/org/apache/atlas/notification/NotificationEntityChangeListener.java +++ b/webapp/src/main/java/org/apache/atlas/notification/NotificationEntityChangeListener.java @@ -97,6 +97,11 @@ public class NotificationEntityChangeListener implements EntityChangeListener { } @Override + public void onTraitsUpdated(ITypedReferenceableInstance entity, Collection<? extends IStruct> traits) throws AtlasException { + notifyOfEntityEvent(Collections.singleton(entity), EntityNotification.OperationType.TRAIT_UPDATE); + } + + @Override public void onEntitiesDeleted(Collection<ITypedReferenceableInstance> entities, boolean isImport) throws AtlasException { notifyOfEntityEvent(entities, EntityNotification.OperationType.ENTITY_DELETE); } http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java ---------------------------------------------------------------------- diff --git a/webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java b/webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java index 362fb67..88222db 100644 --- a/webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java +++ b/webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java @@ -353,6 +353,33 @@ public class EntityREST { } /** + * Updates classifications to an existing entity represented by a guid. + * @param guid globally unique identifier for the entity + * @return classification for the given entity guid + */ + @PUT + @Path("/guid/{guid}/classifications") + @Produces(Servlets.JSON_MEDIA_TYPE) + public void updateClassification(@PathParam("guid") final String guid, List<AtlasClassification> classifications) throws AtlasBaseException { + AtlasPerfTracer perf = null; + + try { + if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) { + perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.updateClassification(" + guid + ")"); + } + + if (StringUtils.isEmpty(guid)) { + throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid); + } + + entitiesStore.updateClassifications(guid, classifications); + + } finally { + AtlasPerfTracer.log(perf); + } + } + + /** * Deletes a given classification from an existing entity represented by a guid. * @param guid globally unique identifier for the entity * @param classificationName name of the classifcation http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/4feee3bf/webapp/src/test/java/org/apache/atlas/web/adapters/TestEntityREST.java ---------------------------------------------------------------------- diff --git a/webapp/src/test/java/org/apache/atlas/web/adapters/TestEntityREST.java b/webapp/src/test/java/org/apache/atlas/web/adapters/TestEntityREST.java index 3161a0d..f079d63 100644 --- a/webapp/src/test/java/org/apache/atlas/web/adapters/TestEntityREST.java +++ b/webapp/src/test/java/org/apache/atlas/web/adapters/TestEntityREST.java @@ -22,6 +22,7 @@ import org.apache.atlas.RequestContext; import org.apache.atlas.RequestContextV1; import org.apache.atlas.TestUtilsV2; import org.apache.atlas.model.instance.AtlasClassification; +import org.apache.atlas.model.instance.AtlasClassification.AtlasClassifications; import org.apache.atlas.model.instance.AtlasEntity; import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo; import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo; @@ -44,6 +45,7 @@ import org.testng.annotations.Test; import javax.inject.Inject; import javax.servlet.http.HttpServletRequest; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -60,6 +62,8 @@ public class TestEntityREST { private AtlasEntity dbEntity; private AtlasClassification testClassification; + + private AtlasClassification phiClassification; @BeforeClass public void setUp() throws Exception { @@ -123,7 +127,59 @@ public class TestEntityREST { Assert.assertNotNull(retrievedClassification); Assert.assertEquals(retrievedClassification, testClassification); + } + + @Test(dependsOnMethods = "testGetEntityById") + public void testAddAndUpdateClassificationWithAttributes() throws Exception { + phiClassification = new AtlasClassification(TestUtilsV2.PHI, new HashMap<String, Object>() {{ + put("stringAttr", "sample_string"); + put("booleanAttr", true); + put("integerAttr", 100); + }}); + + testClassification = new AtlasClassification(TestUtilsV2.CLASSIFICATION, new HashMap<String, Object>() {{ + put("tag", "tagName"); + }}); + + entityREST.addClassifications(dbEntity.getGuid(), new ArrayList<>(Arrays.asList(phiClassification))); + + final AtlasClassifications retrievedClassifications = entityREST.getClassifications(dbEntity.getGuid()); + Assert.assertNotNull(retrievedClassifications); + + final List<AtlasClassification> retrievedClassificationsList = retrievedClassifications.getList(); + Assert.assertNotNull(retrievedClassificationsList); + + final AtlasClassification retrievedClassification = entityREST.getClassification(dbEntity.getGuid(), TestUtilsV2.PHI); + Assert.assertNotNull(retrievedClassification); + Assert.assertEquals(retrievedClassification, phiClassification); + + for (String attrName : retrievedClassification.getAttributes().keySet()) { + Assert.assertEquals(retrievedClassification.getAttribute(attrName), phiClassification.getAttribute(attrName)); + } + + // update multiple tags attributes + phiClassification = new AtlasClassification(TestUtilsV2.PHI, new HashMap<String, Object>() {{ + put("stringAttr", "sample_string_v2"); + put("integerAttr", 200); + }}); + + testClassification = new AtlasClassification(TestUtilsV2.CLASSIFICATION, new HashMap<String, Object>() {{ + put("tag", "tagName_updated"); + }}); + + entityREST.updateClassification(dbEntity.getGuid(), new ArrayList<>(Arrays.asList(phiClassification, testClassification))); + + AtlasClassification updatedClassification = entityREST.getClassification(dbEntity.getGuid(), TestUtilsV2.PHI); + Assert.assertNotNull(updatedClassification); + Assert.assertEquals(updatedClassification.getAttribute("stringAttr"), "sample_string_v2"); + Assert.assertEquals(updatedClassification.getAttribute("integerAttr"), 200); + Assert.assertEquals(updatedClassification.getAttribute("booleanAttr"), true); + + updatedClassification = entityREST.getClassification(dbEntity.getGuid(), TestUtilsV2.CLASSIFICATION); + Assert.assertNotNull(updatedClassification); + Assert.assertEquals(updatedClassification.getAttribute("tag"), testClassification.getAttribute("tag")); + entityREST.deleteClassification(dbEntity.getGuid(), TestUtilsV2.PHI); } @Test(dependsOnMethods = "testAddAndGetClassification")
