org.apache.atlas.model.instance.AtlasClassification Java Examples

/**
 * Map newly associated/defined classifications for the entity with given GUID
 * @param guid Entity guid
 * @param classifications new classifications added to the entity
 * @return Full text string ONLY for the added classifications
 * @throws AtlasBaseException
 */
public String getIndexTextForClassifications(String guid, List<AtlasClassification> classifications) throws AtlasBaseException {
    String                 ret     = null;
    AtlasEntityWithExtInfo entityWithExtInfo  = getAndCacheEntity(guid);

    if (entityWithExtInfo != null) {
        StringBuilder sb = new StringBuilder();

        if (CollectionUtils.isNotEmpty(classifications)) {
            for (AtlasClassification classification : classifications) {
                sb.append(classification.getTypeName()).append(FULL_TEXT_DELIMITER);

                mapAttributes(classification.getAttributes(), entityWithExtInfo, sb, new HashSet<String>());
            }
        }

        ret = sb.toString();
    }

    if (LOG.isDebugEnabled()) {
        LOG.debug("FullTextMapperV2.map({}): {}", guid, ret);
    }

    return ret;
}
 

protected AtlasEntity loadProcess(String name, String description, String user,
                                  List<AtlasEntity> inputTables, List<AtlasEntity> outputTables,
                                  String queryText, String queryPlan, String queryId, String queryGraph, String... traitNames) {
    AtlasEntity process = new AtlasEntity(HIVE_PROCESS_TYPE);
    process.setAttribute("name", name);
    process.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
    process.setAttribute("description", description);
    process.setAttribute("userName", user);
    process.setAttribute("startTime", System.currentTimeMillis());
    process.setAttribute("endTime", System.currentTimeMillis() + 10000);

    process.setAttribute("operationType", "load");
    process.setAttribute("inputs", getAtlasObjectIds(inputTables));
    process.setAttribute("outputs", getAtlasObjectIds(outputTables));

    process.setAttribute("queryText", queryText);
    process.setAttribute("queryPlan", queryPlan);
    process.setAttribute("queryId", queryId);
    process.setAttribute("queryGraph", queryGraph);

    process.setClassifications(Stream.of(traitNames).map(AtlasClassification::new).collect(Collectors.toList()));

    return process;
}
 

/**
 * Adds classification to the entity identified by its type and unique attributes.
 * @param typeName
 */
@POST
@Path("/uniqueAttribute/type/{typeName}/classifications")
public void addClassificationsByUniqueAttribute(@PathParam("typeName") String typeName, @Context HttpServletRequest servletRequest, List<AtlasClassification> classifications) throws AtlasBaseException {
    Servlets.validateQueryParamLength("typeName", typeName);

    AtlasPerfTracer perf = null;

    try {
        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.addClassificationsByUniqueAttribute(" + 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.addClassifications(guid, classifications);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

/**
 * Gets the list of classifications for a given entity represented by a guid.
 * @param guid globally unique identifier for the entity
 * @return a list of classifications for the given entity guid
 */
@GET
@Path("/guid/{guid}/classifications")
public AtlasClassification.AtlasClassifications getClassifications(@PathParam("guid") String guid) throws AtlasBaseException {
    Servlets.validateQueryParamLength("guid", guid);

    AtlasPerfTracer perf = null;

    try {
        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.getClassifications(" + guid + ")");
        }

        if (StringUtils.isEmpty(guid)) {
            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
        }

        return new AtlasClassification.AtlasClassifications(entitiesStore.getClassifications(guid));
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

@Override
public void apply(AtlasTransformableEntity transformableEntity) {
    AtlasEntity entity = transformableEntity.getEntity();

    if (entity.getClassifications() == null) {
        entity.setClassifications(new ArrayList<>());
    }

    boolean hasClassification = false;

    for (AtlasClassification c : entity.getClassifications()) {
        hasClassification = c.getTypeName().equals(classificationName);

        if (hasClassification) {
            break;
        }
    }

    if (!hasClassification) {
        entity.getClassifications().add(new AtlasClassification(classificationName));
    }
}
 

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, getListenerName(listener), "TraitUpdate");
        }
    }
}
 

/**
 * Updates classification on an entity identified by its type and unique attributes.
 * @param  typeName
 */
@PUT
@Path("/uniqueAttribute/type/{typeName}/classifications")
public void updateClassificationsByUniqueAttribute(@PathParam("typeName") String typeName, @Context HttpServletRequest servletRequest, List<AtlasClassification> classifications) throws AtlasBaseException {
    Servlets.validateQueryParamLength("typeName", typeName);

    AtlasPerfTracer perf = null;

    try {
        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.updateClassificationsByUniqueAttribute(" + 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.updateClassifications(guid, classifications);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

/**
 * 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")
public void updateClassifications(@PathParam("guid") final String guid, List<AtlasClassification> classifications) throws AtlasBaseException {
    Servlets.validateQueryParamLength("guid", guid);

    AtlasPerfTracer perf = null;

    try {
        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.updateClassifications(" + guid + ")");
        }

        if (StringUtils.isEmpty(guid)) {
            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
        }

        entitiesStore.updateClassifications(guid, classifications);

    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

@Override
public Object getNormalizedValueForUpdate(Object obj) {
    Object ret = null;

    if (obj != null) {
        if (isValidValueForUpdate(obj)) {
            if (obj instanceof AtlasClassification) {
                normalizeAttributeValuesForUpdate((AtlasClassification) obj);
                ret = obj;
            } else if (obj instanceof Map) {
                normalizeAttributeValuesForUpdate((Map) obj);
                ret = obj;
            }
        }
    }

    return ret;
}
 

private void notifyPropagatedEntities(Map<String, List<AtlasClassification>> entityPropagationMap, EntityAuditActionV2 action) throws AtlasBaseException {
    if (MapUtils.isEmpty(entityPropagationMap) || action == null) {
        return;
    }

    RequestContext context = RequestContext.get();

    for (String guid : entityPropagationMap.keySet()) {
        // if entity is deleted, don't send propagated classifications add/remove notifications.
        if (context.isDeletedEntity(guid)) {
            continue;
        }

        AtlasEntity entity = fullTextMapperV2.getAndCacheEntity(guid);

        if (entity == null) {
            continue;
        }

        if (action == PROPAGATED_CLASSIFICATION_ADD) {
            onClassificationAddedToEntity(entity, entityPropagationMap.get(guid));
        } else if (action == PROPAGATED_CLASSIFICATION_DELETE) {
            onClassificationDeletedFromEntity(entity, entityPropagationMap.get(guid));
        }
    }
}
 

/**
 * Gets the list of classifications for a given entity represented by a guid.
 * @param guid globally unique identifier for the entity
 * @return classification for the given entity guid
 */
@GET
@Path("/guid/{guid}/classification/{classificationName}")
@Produces(Servlets.JSON_MEDIA_TYPE)
public AtlasClassification getClassification(@PathParam("guid") String guid, @PathParam("classificationName") final String classificationName) throws AtlasBaseException {
    AtlasPerfTracer perf = null;

    try {
        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.getClassification(" + guid + "," + classificationName + ")");
        }

        if (StringUtils.isEmpty(guid)) {
            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
        }

        ensureClassificationType(classificationName);
        return entitiesStore.getClassification(guid, classificationName);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

@Test(dependsOnMethods = "testGetEntityWithAssociations")
public void testDeleteClassification() throws Exception {

    deleteClassification(dbEntity.getGuid(), TestUtilsV2.CLASSIFICATION);
    final AtlasClassification.AtlasClassifications retrievedClassifications = entityREST.getClassifications(dbEntity.getGuid());

    String expectedClsNames    = "";
    AtlasVertex vertex         = AtlasGraphUtilsV2.findByGuid(dbEntity.getGuid());
    String classificationNames = vertex.getProperty(Constants.CLASSIFICATION_NAMES_KEY, String.class);

    Assert.assertNotNull(classificationNames);
    Assert.assertEquals(classificationNames, expectedClsNames);

    Assert.assertNotNull(retrievedClassifications);
    Assert.assertEquals(retrievedClassifications.getList().size(), 0);
}
 

@Override
public Object apply(Object o) {
    if (!(o instanceof AtlasEntity)) {
        return o;
    }

    AtlasEntity entity = (AtlasEntity) o;
    if(!passThruFilters(entity)) {
        return o;
    }

    if(entity.getClassifications() == null) {
        entity.setClassifications(new ArrayList<AtlasClassification>());
    }

    for (AtlasClassification c : entity.getClassifications()) {
        if (c.getTypeName().equals(classificationName)) {
            return entity;
        }
    }

    entity.getClassifications().add(new AtlasClassification(classificationName));
    return entity;
}
 

@Override
public void onClassificationsDeleted(List<AtlasEntity> entities, List<AtlasClassification> classifications) throws AtlasBaseException {
    if (CollectionUtils.isNotEmpty(classifications) && CollectionUtils.isNotEmpty(entities)) {
        MetricRecorder           metric = RequestContext.get().startMetricRecord("entityAudit");
        List<EntityAuditEventV2> events = Collections.synchronizedList(new ArrayList<>());

        for (AtlasClassification classification : classifications) {
            for (AtlasEntity entity : entities) {
                if (StringUtils.equals(entity.getGuid(), classification.getEntityGuid())) {
                    events.add(createEvent(entity, CLASSIFICATION_DELETE, "Deleted classification: " + classification.getTypeName()));
                } else {
                    events.add(createEvent(entity, PROPAGATED_CLASSIFICATION_DELETE, "Deleted propagated classification: " + classification.getTypeName()));
                }
            }
        }

        auditRepository.putEventsV2(events);

        RequestContext.get().endMetricRecord(metric);
    }
}
 

@Override
public void onClassificationsAdded(AtlasEntity entity, List<AtlasClassification> classifications) throws AtlasBaseException {
    if (CollectionUtils.isNotEmpty(classifications)) {
        MetricRecorder metric = RequestContext.get().startMetricRecord("entityAudit");

        List<EntityAuditEventV2> events = new ArrayList<>();

        for (AtlasClassification classification : classifications) {
            if (entity.getGuid().equals(classification.getEntityGuid())) {
                events.add(createEvent(entity, CLASSIFICATION_ADD, "Added classification: " + AtlasType.toJson(classification)));
            } else {
                events.add(createEvent(entity, PROPAGATED_CLASSIFICATION_ADD, "Added propagated classification: " + AtlasType.toJson(classification)));
            }
        }

        auditRepository.putEventsV2(events);

        RequestContext.get().endMetricRecord(metric);
    }
}
 

@Test(dependsOnMethods = {"addClassification_removePropagationsTrue_DeleteCase"})
public void addClassification_removePropagationsFalse_DeleteCase() throws AtlasBaseException {
    AtlasEntity         employees = getEntity(EMPLOYEES_TABLE);
    AtlasClassification tag1      = new AtlasClassification("tag1");

    tag1.setEntityGuid(employees.getGuid());
    tag1.setPropagate(true);
    tag1.setRemovePropagationsOnEntityDelete(false);

    addClassification(employees, tag1);

    List<String> propagatedEntities = Arrays.asList(EMPLOYEES_PROCESS, US_EMPLOYEES_TABLE);

    assertClassificationExistInEntities(propagatedEntities, tag1);

    AtlasEntity employees_process = getEntity(EMPLOYEES_PROCESS);
    AtlasEntity us_employees      = getEntity(US_EMPLOYEES_TABLE);

    deletePropagatedClassificationExpectFail(employees_process, tag1);
    deletePropagatedClassificationExpectFail(us_employees, tag1);

    deleteEntity(EMPLOYEES_PROCESS);
    assertClassificationExistInEntity(EMPLOYEES_PROCESS, tag1);
    assertClassificationExistInEntity(US_EMPLOYEES_TABLE, tag1);
}
 

/**
 * Adds classifications to an existing entity represented by a guid.
 * @param guid globally unique identifier for the entity
 */
@POST
@Path("/guid/{guid}/classifications")
@Consumes({Servlets.JSON_MEDIA_TYPE, MediaType.APPLICATION_JSON})
@Produces(Servlets.JSON_MEDIA_TYPE)
public void addClassifications(@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.addClassifications(" + guid + ")");
        }

        if (StringUtils.isEmpty(guid)) {
            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
        }

        entitiesStore.addClassifications(guid, classifications);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

private void mapClassifications(AtlasVertex entityVertex, AtlasEntity entity) throws AtlasBaseException {
    List<AtlasEdge> edges = getAllClassificationEdges(entityVertex);

    if (CollectionUtils.isNotEmpty(edges)) {
        List<AtlasClassification> allClassifications = new ArrayList<>();

        for (AtlasEdge edge : edges) {
            AtlasVertex         classificationVertex = edge.getInVertex();
            AtlasClassification classification       = toAtlasClassification(classificationVertex);

            if (classification != null) {
                allClassifications.add(classification);
            }
        }

        entity.setClassifications(allClassifications);
    }
}
 

public void removeTagPropagation(AtlasVertex classificationVertex, List<AtlasVertex> entityVertices) throws AtlasBaseException {
    if (classificationVertex != null && CollectionUtils.isNotEmpty(entityVertices)) {
        String              classificationName = getClassificationName(classificationVertex);
        AtlasClassification classification     = entityRetriever.toAtlasClassification(classificationVertex);
        String              entityGuid         = getClassificationEntityGuid(classificationVertex);
        RequestContext      context            = RequestContext.get();

        for (AtlasVertex entityVertex : entityVertices) {
            AtlasEdge propagatedEdge = getPropagatedClassificationEdge(entityVertex, classificationName, entityGuid);

            if (propagatedEdge != null) {
                deletePropagatedEdge(propagatedEdge);

                // record remove propagation details to send notifications at the end
                context.recordRemovedPropagation(getGuid(entityVertex), classification);
            }
        }
    }
}
 

public void updateClassificationTextAndNames(AtlasVertex vertex) throws AtlasBaseException {
    if(CollectionUtils.isEmpty(vertex.getPropertyValues(Constants.TRAIT_NAMES_PROPERTY_KEY, String.class)) &&
            CollectionUtils.isEmpty(vertex.getPropertyValues(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, String.class))) {
        return;
    }

    String guid = graphHelper.getGuid(vertex);
    AtlasEntity entity = instanceConverter.getAndCacheEntity(guid, ENTITY_CHANGE_NOTIFY_IGNORE_RELATIONSHIP_ATTRIBUTES);
    List<String> classificationNames = new ArrayList<>();
    List<String> propagatedClassificationNames = new ArrayList<>();

    for (AtlasClassification classification : entity.getClassifications()) {
        if (isPropagatedClassification(classification, guid)) {
            propagatedClassificationNames.add(classification.getTypeName());
        } else {
            classificationNames.add(classification.getTypeName());
        }
    }

    vertex.setProperty(CLASSIFICATION_NAMES_KEY, getDelimitedClassificationNames(classificationNames));
    vertex.setProperty(PROPAGATED_CLASSIFICATION_NAMES_KEY, getDelimitedClassificationNames(propagatedClassificationNames));
    vertex.setProperty(CLASSIFICATION_TEXT_KEY, fullTextMapperV2.getClassificationTextForEntity(entity));
}
 

/**
 * 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);
    }
}
 

@Test(dependsOnMethods = {"addClassification_PropagateFalse"})
public void updateClassification_PropagateFalseToTrue() throws AtlasBaseException {
    AtlasEntity         hdfs_employees = getEntity(HDFS_PATH_EMPLOYEES);
    AtlasClassification tag2           = new AtlasClassification("tag2"); tag2.setEntityGuid(hdfs_employees.getGuid());

    //update tag2 propagate to 'true'
    tag2 = getClassification(hdfs_employees, tag2); tag2.setPropagate(true);

    updateClassifications(hdfs_employees, tag2);

    List<String> propagatedToEntities = Arrays.asList(EMPLOYEES1_PROCESS, EMPLOYEES2_PROCESS, EMPLOYEES1_TABLE,
                                                      EMPLOYEES2_TABLE, EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);

    assertClassificationExistInEntities(propagatedToEntities, tag2);

    deleteClassification(hdfs_employees, tag2);
}
 

private boolean hasPropagatedEntry(Map<String, List<AtlasClassification>> propagationsMap, String guid, AtlasClassification classification) {
    boolean ret = false;

    if (MapUtils.isNotEmpty(propagationsMap) && propagationsMap.containsKey(guid) && CollectionUtils.isNotEmpty(propagationsMap.get(guid))) {
        List<AtlasClassification> classifications    = propagationsMap.get(guid);
        String                    classificationName = classification.getTypeName();
        String                    entityGuid         = classification.getEntityGuid();

        for (AtlasClassification c : classifications) {
            if (StringUtils.equals(c.getTypeName(), classificationName) && StringUtils.equals(c.getEntityGuid(), entityGuid)) {
                ret = true;
                break;
            }
        }
    }

    return ret;
}
 

public 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);
}
 

@Override
public void onClassificationsDeleted(AtlasEntity entity, List<AtlasClassification> classifications) throws AtlasBaseException {
    if (CollectionUtils.isNotEmpty(classifications)) {
        MetricRecorder metric = RequestContext.get().startMetricRecord("entityAudit");

        List<EntityAuditEventV2> events = new ArrayList<>();

        for (AtlasClassification classification : classifications) {
            if (StringUtils.equals(entity.getGuid(), classification.getEntityGuid())) {
                events.add(createEvent(entity, CLASSIFICATION_DELETE, "Deleted classification: " + classification.getTypeName()));
            } else {
                events.add(createEvent(entity, PROPAGATED_CLASSIFICATION_DELETE, "Deleted propagated classification: " + classification.getTypeName()));
            }
        }

        auditRepository.putEventsV2(events);

        RequestContext.get().endMetricRecord(metric);
    }
}
 

public AtlasClassification toAtlasClassification(Struct classification) throws AtlasBaseException {
    AtlasFormatConverter    converter          = instanceFormatters.getConverter(TypeCategory.CLASSIFICATION);
    AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(classification.getTypeName());

    if (classificationType == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.CLASSIFICATION.name(), classification.getTypeName());
    }

    AtlasClassification  ret = (AtlasClassification)converter.fromV1ToV2(classification, classificationType, new AtlasFormatConverter.ConverterContext());

    return ret;
}
 

@Test(dependsOnMethods = "testCreate")
public void associateSameTagToMultipleEntities() throws AtlasBaseException {
    final String TAG_NAME            = "tagx";
    final String TAG_ATTRIBUTE_NAME  = "testAttribute";
    final String TAG_ATTRIBUTE_VALUE = "test-string";

    createTag(TAG_NAME, "string");
    List<AtlasClassification> addedClassifications = new ArrayList<>();
    addedClassifications.add(new AtlasClassification(TAG_NAME, TAG_ATTRIBUTE_NAME, TAG_ATTRIBUTE_VALUE));

    entityStore.addClassifications(dbEntityGuid, addedClassifications);
    entityStore.addClassifications(tblEntityGuid, addedClassifications);

    AtlasEntity dbEntityFromDb  = getEntityFromStore(dbEntityGuid);
    AtlasEntity tblEntityFromDb = getEntityFromStore(tblEntityGuid);

    Set<String> actualDBClassifications  = new HashSet<>(CollectionUtils.collect(dbEntityFromDb.getClassifications(), o -> ((AtlasClassification) o).getTypeName()));
    Set<String> actualTblClassifications = new HashSet<>(CollectionUtils.collect(tblEntityFromDb.getClassifications(), o -> ((AtlasClassification) o).getTypeName()));

    assertTrue(actualDBClassifications.contains(TAG_NAME));
    assertTrue(actualTblClassifications.contains(TAG_NAME));

    Set<String> actualDBAssociatedEntityGuid  = new HashSet<>(CollectionUtils.collect(dbEntityFromDb.getClassifications(), o -> ((AtlasClassification) o).getEntityGuid()));
    Set<String> actualTblAssociatedEntityGuid = new HashSet<>(CollectionUtils.collect(tblEntityFromDb.getClassifications(), o -> ((AtlasClassification) o).getEntityGuid()));

    assertTrue(actualDBAssociatedEntityGuid.contains(dbEntityGuid));
    assertTrue(actualTblAssociatedEntityGuid.contains(tblEntityGuid));

    entityStore.deleteClassification(dbEntityGuid, TAG_NAME);
    entityStore.deleteClassification(tblEntityGuid, TAG_NAME);
}
 

private AtlasEdge mapClassification(EntityOperation operation,  final EntityMutationContext context, AtlasClassification classification,
                                    AtlasEntityType entityType, AtlasVertex parentInstanceVertex, AtlasVertex traitInstanceVertex)
                                    throws AtlasBaseException {
    if (classification.getValidityPeriods() != null) {
        String strValidityPeriods = AtlasJson.toJson(classification.getValidityPeriods());

        AtlasGraphUtilsV2.setEncodedProperty(traitInstanceVertex, CLASSIFICATION_VALIDITY_PERIODS_KEY, strValidityPeriods);
    } else {
        // if 'null', don't update existing value in the classification
    }

    if (classification.isPropagate() != null) {
        AtlasGraphUtilsV2.setEncodedProperty(traitInstanceVertex, CLASSIFICATION_VERTEX_PROPAGATE_KEY, classification.isPropagate());
    }

    if (classification.getRemovePropagationsOnEntityDelete() != null) {
        AtlasGraphUtilsV2.setEncodedProperty(traitInstanceVertex, CLASSIFICATION_VERTEX_REMOVE_PROPAGATIONS_KEY, classification.getRemovePropagationsOnEntityDelete());
    }

    // map all the attributes to this newly created AtlasVertex
    mapAttributes(classification, traitInstanceVertex, operation, context);

    AtlasEdge ret = getClassificationEdge(parentInstanceVertex, traitInstanceVertex);

    if (ret == null) {
        ret = graphHelper.addClassificationEdge(parentInstanceVertex, traitInstanceVertex, false);
    }

    return ret;
}
 

@Test(dependsOnMethods = "testGetTraitNames")
public void testAddTraitWithAttribute() throws Exception {
    final String traitName = "PII_Trait" + randomString();
    AtlasClassificationDef piiTrait = AtlasTypeUtil
            .createTraitTypeDef(traitName, ImmutableSet.<String>of(),
                    AtlasTypeUtil.createRequiredAttrDef("type", "string"));
    AtlasTypesDef typesDef = new AtlasTypesDef();
    typesDef.getClassificationDefs().add(piiTrait);
    createType(typesDef);

    AtlasClassification traitInstance = new AtlasClassification(traitName);
    traitInstance.setAttribute("type", "SSN");

    final String guid = createHiveTable().getGuid();
    atlasClientV2.addClassifications(guid, ImmutableList.of(traitInstance));

    // verify the response
    AtlasEntity withAssociationByGuid = atlasClientV2.getEntityByGuid(guid).getEntity();
    assertNotNull(withAssociationByGuid);
    assertFalse(withAssociationByGuid.getClassifications().isEmpty());

    boolean found = false;
    for (AtlasClassification atlasClassification : withAssociationByGuid.getClassifications()) {
        String attribute = (String)atlasClassification.getAttribute("type");
        if (attribute != null && attribute.equals("SSN")) {
            found = true;
            break;
        }
    }
    assertTrue(found);
}
 

private List<AtlasClassification> toAtlasClassifications(String[] classificationNames) {
    List<AtlasClassification> ret             = new ArrayList<>();
    List<String>              classifications = asList(classificationNames);

    if (CollectionUtils.isNotEmpty(classifications)) {
        for (String classificationName : classifications) {
            ret.add(new AtlasClassification(classificationName));
        }
    }

    return ret;
}