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

private void notifyV1Listeners(List<AtlasEntityHeader> entityHeaders, EntityOperation operation, boolean isImport) throws AtlasBaseException {
    if (operation != EntityOperation.PURGE && instanceConverter != null) {
        List<Referenceable> typedRefInsts = toReferenceables(entityHeaders, operation);

        for (EntityChangeListener listener : entityChangeListeners) {
            try {
                switch (operation) {
                    case CREATE:
                        listener.onEntitiesAdded(typedRefInsts, isImport);
                        break;
                    case UPDATE:
                    case PARTIAL_UPDATE:
                        listener.onEntitiesUpdated(typedRefInsts, isImport);
                        break;
                    case DELETE:
                        listener.onEntitiesDeleted(typedRefInsts, isImport);
                        break;
                }
            } catch (AtlasException e) {
                throw new AtlasBaseException(AtlasErrorCode.NOTIFICATION_FAILED, e, getListenerName(listener), operation.toString());
            }
        }
    }
}
 

@Test
public void testGetEntityHeaderByUniqueAttributes() throws Exception {
    createTestEntity();

    String[] attrVal = {String.valueOf(dbEntity.getAttribute("name"))};

    Map<String, String[]> paramMap = new HashMap<>();
    paramMap.put("attr:name", attrVal);

    HttpServletRequest mockRequest = Mockito.mock(HttpServletRequest.class);
    Mockito.when(mockRequest.getParameterMap()).thenReturn(paramMap);

    AtlasEntityHeader response = entityREST.getEntityHeaderByUniqueAttributes(dbEntity.getTypeName(), mockRequest);

    Assert.assertNotNull(response);
    Assert.assertEquals(dbEntity.getAttribute("name"), response.getAttribute("name"));
    Assert.assertEquals(dbEntity.getAttribute("description"), response.getAttribute("description"));
}
 

@Override
@GraphTransaction
public AtlasEntityWithExtInfo getByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes, boolean isMinExtInfo, boolean ignoreRelationships) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> getByUniqueAttribute({}, {})", entityType.getTypeName(), uniqAttributes);
    }

    AtlasVertex entityVertex = AtlasGraphUtilsV2.getVertexByUniqueAttributes(graph, entityType, uniqAttributes);

    EntityGraphRetriever entityRetriever = new EntityGraphRetriever(graph, typeRegistry, ignoreRelationships);

    AtlasEntityWithExtInfo ret = entityRetriever.toAtlasEntityWithExtInfo(entityVertex, isMinExtInfo);

    if (ret == null) {
        throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, entityType.getTypeName(),
                uniqAttributes.toString());
    }

    AtlasAuthorizationUtils.verifyAccess(new AtlasEntityAccessRequest(typeRegistry, AtlasPrivilege.ENTITY_READ, new AtlasEntityHeader(ret.getEntity())), "read entity: typeName=", entityType.getTypeName(), ", uniqueAttributes=", uniqAttributes);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== getByUniqueAttribute({}, {}): {}", entityType.getTypeName(), uniqAttributes, ret);
    }

    return ret;
}
 

@BeforeClass
public void setUp() throws Exception {
    new GraphBackedSearchIndexer(typeRegistry);

    // create employee relationship types
    AtlasTypesDef employeeTypes = getDepartmentEmployeeTypes();
    typeDefStore.createTypesDef(employeeTypes);

    AtlasEntitiesWithExtInfo employeeInstances = getDepartmentEmployeeInstances();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(employeeInstances), false);

    for (AtlasEntityHeader entityHeader : response.getCreatedEntities()) {
        employeeNameIdMap.put((String) entityHeader.getAttribute(NAME), getAtlasObjectId(entityHeader));
    }

    init();
    AtlasTypesDef testTypes = getInverseReferenceTestTypes();
    typeDefStore.createTypesDef(testTypes);
}
 

private void notifyListeners(List<AtlasEntityHeader> entityHeaders, EntityOperation operation, boolean isImport) throws AtlasBaseException {
    if (CollectionUtils.isEmpty(entityHeaders)) {
        return;
    }

    List<ITypedReferenceableInstance> typedRefInsts = toITypedReferenceable(entityHeaders);

    for (EntityChangeListener listener : entityChangeListeners) {
        try {
            switch (operation) {
                case CREATE:
                    listener.onEntitiesAdded(typedRefInsts, isImport);
                    break;
                case UPDATE:
                case PARTIAL_UPDATE:
                    listener.onEntitiesUpdated(typedRefInsts, isImport);
                    break;
                case DELETE:
                    listener.onEntitiesDeleted(typedRefInsts, isImport);
                    break;
            }
        } catch (AtlasException e) {
            throw new AtlasBaseException(AtlasErrorCode.NOTIFICATION_FAILED, e, getListenerName(listener), operation.toString());
        }
    }
}
 

private List<AtlasAuditEntry> toAtlasAuditEntries(AtlasSearchResult result) {
    List<AtlasAuditEntry> ret = new ArrayList<>();

    if(CollectionUtils.isNotEmpty(result.getEntities())) {
        for (AtlasEntityHeader entityHeader : result.getEntities()) {
            AtlasAuditEntry entry = AtlasAuditEntryDTO.from(entityHeader.getGuid(),
                    entityHeader.getAttributes());
            if (entry == null) {
                continue;
            }

            ret.add(entry);
        }
    }

    return ret;
}
 

@VisibleForTesting
AtlasEntityWithExtInfo createEntityInAtlas(AtlasEntityWithExtInfo entity) throws Exception {
    AtlasEntityWithExtInfo  ret      = null;
    EntityMutationResponse  response = atlasClientV2.createEntity(entity);
    List<AtlasEntityHeader> entities = response.getCreatedEntities();

    if (CollectionUtils.isNotEmpty(entities)) {
        AtlasEntityWithExtInfo getByGuidResponse = atlasClientV2.getEntityByGuid(entities.get(0).getGuid());

        ret = getByGuidResponse;

        LOG.info("Created {} entity: name={}, guid={}", ret.getEntity().getTypeName(), ret.getEntity().getAttribute(ATTRIBUTE_QUALIFIED_NAME), ret.getEntity().getGuid());
    }

    return ret;
}
 

@Override
@GraphTransaction
public AtlasEntityHeader getHeaderById(final String guid) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> getHeaderById({})", guid);
    }

    EntityGraphRetriever entityRetriever = new EntityGraphRetriever(graph, typeRegistry);

    AtlasEntityHeader ret = entityRetriever.toAtlasEntityHeader(guid);

    if (ret == null) {
        throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
    }

    AtlasAuthorizationUtils.verifyAccess(new AtlasEntityAccessRequest(typeRegistry, AtlasPrivilege.ENTITY_READ, ret), "read entity: guid=", guid);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== getHeaderById({}): {}", guid, ret);
    }

    return ret;
}
 

@Override
@GraphTransaction
public AtlasEntityWithExtInfo getById(final String guid, final boolean isMinExtInfo, boolean ignoreRelationships) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> getById({}, {})", guid, isMinExtInfo);
    }

    EntityGraphRetriever entityRetriever = new EntityGraphRetriever(graph, typeRegistry, ignoreRelationships);

    AtlasEntityWithExtInfo ret = entityRetriever.toAtlasEntityWithExtInfo(guid, isMinExtInfo);

    if (ret == null) {
        throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
    }

    AtlasAuthorizationUtils.verifyAccess(new AtlasEntityAccessRequest(typeRegistry, AtlasPrivilege.ENTITY_READ, new AtlasEntityHeader(ret.getEntity())), "read entity: guid=", guid);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== getById({}, {}): {}", guid, isMinExtInfo, ret);
    }

    return ret;
}
 

@Test
public void testDefaultValueForPrimitiveTypes() throws Exception  {
    init();

    EntityMutationResponse  response                   = entityStore.createOrUpdate(new AtlasEntityStream(primitiveEntity), false);
    List<AtlasEntityHeader> entitiesCreatedResponse    = response.getEntitiesByOperation(EntityOperation.CREATE);
    List<AtlasEntityHeader> entitiesCreatedwithdefault = response.getMutatedEntities().get(EntityOperation.CREATE);
    AtlasEntity             entityCreated              = getEntityFromStore(entitiesCreatedResponse.get(0));

    Map     attributesMap = entityCreated.getAttributes();
    String  description   = (String) attributesMap.get("description");
    String  check         = (String) attributesMap.get("check");
    String  sourceCode    = (String) attributesMap.get("sourcecode");
    float   diskUsage     = (float) attributesMap.get("diskUsage");
    boolean isstoreUse    = (boolean) attributesMap.get("isstoreUse");
    int     cost          = (int) attributesMap.get("Cost");

    assertEquals(description,"test");
    assertEquals(check,"check");

    //defaultValue
    assertEquals(diskUsage,70.5f);
    assertEquals(isstoreUse,true);
    assertEquals(sourceCode,"Hello World");
    assertEquals(cost,30);
}
 

public void onEntitiesMutated(EntityMutationResponse entityMutationResponse, boolean isImport) throws AtlasBaseException {
    if (CollectionUtils.isEmpty(entityChangeListeners) || instanceConverter == null) {
        return;
    }

    List<AtlasEntityHeader> createdEntities          = entityMutationResponse.getCreatedEntities();
    List<AtlasEntityHeader> updatedEntities          = entityMutationResponse.getUpdatedEntities();
    List<AtlasEntityHeader> partiallyUpdatedEntities = entityMutationResponse.getPartialUpdatedEntities();
    List<AtlasEntityHeader> deletedEntities          = entityMutationResponse.getDeletedEntities();

    // complete full text mapping before calling toITypedReferenceable(), from notifyListners(), to
    // include all vertex updates in the current graph-transaction
    doFullTextMapping(createdEntities);
    doFullTextMapping(updatedEntities);
    doFullTextMapping(partiallyUpdatedEntities);

    notifyListeners(createdEntities, EntityOperation.CREATE, isImport);
    notifyListeners(updatedEntities, EntityOperation.UPDATE, isImport);
    notifyListeners(partiallyUpdatedEntities, EntityOperation.PARTIAL_UPDATE, isImport);
    notifyListeners(deletedEntities, EntityOperation.DELETE, isImport);
}
 

@BeforeClass
public void setUp() throws Exception {
    new GraphBackedSearchIndexer(typeRegistry);

    // create employee relationship types
    AtlasTypesDef employeeTypes = getDepartmentEmployeeTypes();
    typeDefStore.createTypesDef(employeeTypes);

    AtlasEntitiesWithExtInfo employeeInstances = getDepartmentEmployeeInstances();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(employeeInstances), false);

    for (AtlasEntityHeader entityHeader : response.getCreatedEntities()) {
        employeeNameIdMap.put((String) entityHeader.getAttribute(NAME), getAtlasObjectId(entityHeader));
    }

    init();
    AtlasTypesDef typesDef = getInverseReferenceTestTypes();

    AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(typesDef, typeRegistry);

    if (!typesToCreate.isEmpty()) {
        typeDefStore.createTypesDef(typesToCreate);
    }
}
 

/**
 * Get entity header given its GUID.
 * @param guid GUID for the entity
 * @return AtlasEntity
 * @throws AtlasBaseException
 */
@GET
@Path("/guid/{guid}/header")
public AtlasEntityHeader getHeaderById(@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.getHeaderById(" + guid + ")");
        }

        return entitiesStore.getHeaderById(guid);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

@Override
AtlasEntityHeader getByUniqueAttributes(AtlasEntityType entityType, String qualifiedName, Map<String, Object> attrValues) {
    try {
        if (StringUtils.isEmpty(entityFileName)) {
            return null;
        }

        return getEntityHeaderFromFile(entityFileName);
    } catch (IOException e) {
        fail(entityFileName + " could not be loaded.");
        return null;
    }
}
 

public Set<String> getClassificationNames(AtlasEntityHeader entity) {
    final Set<String> ret;

    if (entity == null || entity.getClassifications() == null) {
        ret = Collections.emptySet();
    } else {
        ret = new HashSet<>();

        for (AtlasClassification classify : entity.getClassifications()) {
            ret.add(classify.getTypeName());
        }
    }

    return ret;
}
 

@Test
public void searchByTagAndTagFilters() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setClassification(DIMENSIONAL_CLASSIFICATION);
    FilterCriteria filterCriteria = getSingleFilterCondition("attr1", Operator.EQ, "Test");
    params.setTagFilters(filterCriteria);

    List<AtlasEntityHeader> entityHeaders = discoveryService.searchWithParameters(params).getEntities();

    assertEquals(entityHeaders.size(), 1);
    assertEquals(entityHeaders.get(0).getGuid(), dimensionalTagGuid);

}
 

public void searchByTagAndGraphSysFilters() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setClassification(DIMENSION_CLASSIFICATION);
    FilterCriteria filterCriteria = getSingleFilterCondition("__entityStatus", Operator.EQ, "DELETED");
    params.setTagFilters(filterCriteria);
    params.setExcludeDeletedEntities(false);

    List<AtlasEntityHeader> entityHeaders = discoveryService.searchWithParameters(params).getEntities();

    assertEquals(entityHeaders.size(), 1);
    assertEquals(entityHeaders.get(0).getGuid(), dimensionTagDeleteGuid);

}
 

private List<ITypedReferenceableInstance> toITypedReferenceable(List<AtlasEntityHeader> entityHeaders) throws AtlasBaseException {
    List<ITypedReferenceableInstance> ret = new ArrayList<>(entityHeaders.size());

    for (AtlasEntityHeader entityHeader : entityHeaders) {
        ret.add(instanceConverter.getITypedReferenceable(entityHeader.getGuid()));
    }

    return ret;
}
 

private String getUniqueName(AtlasEntityHeader header) {
    String uniqueAttributeName = ATTR_NAME_QUALIFIED_NAME;
    if (!header.getAttributes().containsKey(ATTR_NAME_QUALIFIED_NAME)) {
        uniqueAttributeName = getUniqueAttribute(header.getTypeName());
    }

    Object attrValue = header.getAttribute(uniqueAttributeName);
    if (attrValue == null) {
        LOG.warn("Unique Attribute Value: empty: {}", AtlasJson.toJson(header));
        return StringUtils.EMPTY;
    }

    return attrValue.toString();
}
 

private void lineage() throws AtlasServiceException {
    System.out.println("\nSample Lineage Info: ");

    AtlasLineageInfo               lineageInfo   = atlasClientV2.getLineageInfo(getTableId(SALES_FACT_DAILY_MV_TABLE), LineageDirection.BOTH, 0);
    Set<LineageRelation>           relations     = lineageInfo.getRelations();
    Map<String, AtlasEntityHeader> guidEntityMap = lineageInfo.getGuidEntityMap();

    for (LineageRelation relation : relations) {
        AtlasEntityHeader fromEntity = guidEntityMap.get(relation.getFromEntityId());
        AtlasEntityHeader toEntity   = guidEntityMap.get(relation.getToEntityId());

        System.out.println(fromEntity.getDisplayText() + "(" + fromEntity.getTypeName() + ") -> " +
                           toEntity.getDisplayText()   + "(" + toEntity.getTypeName() + ")");
    }
}
 

@Test(dependsOnMethods = "testDeleteClassificationByUniqueAttribute")
public void  testDeleteEntityById() throws Exception {

    EntityMutationResponse response = entityREST.deleteByGuid(dbEntity.getGuid());
    List<AtlasEntityHeader> entitiesMutated = response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE);
    Assert.assertNotNull(entitiesMutated);
    Assert.assertEquals(entitiesMutated.get(0).getGuid(), dbEntity.getGuid());
}
 

private int getEntityCount() throws AtlasBaseException {
    SearchParameters params = new SearchParameters();
    params.setTypeName(ALL_ENTITY_TYPES);

    List<AtlasEntityHeader> entityHeaders = discoveryService.searchWithParameters(params).getEntities();
    return entityHeaders.size();
}
 

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

    EntityMutationResponse response = entityREST.deleteByGuid(dbEntity.getGuid());
    List<AtlasEntityHeader> entitiesMutated = response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE);
    Assert.assertNotNull(entitiesMutated);
    Assert.assertEquals(entitiesMutated.get(0).getGuid(), dbEntity.getGuid());
}
 

@Test
public void testDefaultValueForPrimitiveTypes() throws Exception  {

    init();

    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(primitiveEntity), false);
    List<AtlasEntityHeader> entitiesCreatedResponse = response.getEntitiesByOperation(EntityOperation.CREATE);
    final Map<EntityOperation, List<AtlasEntityHeader>> entitiesMutated = response.getMutatedEntities();
    List<AtlasEntityHeader> entitiesCreatedwithdefault = entitiesMutated.get(EntityOperation.CREATE);

    AtlasEntity entityCreated   = getEntityFromStore(entitiesCreatedResponse.get(0));


    Map attributesMap = entityCreated.getAttributes();
    String description = (String) attributesMap.get("description");
    String check = (String) attributesMap.get("check");
    String   sourceCode =  (String) attributesMap.get("sourcecode");
    float   diskUsage =  (float) attributesMap.get("diskUsage");
    boolean   isstoreUse =  (boolean) attributesMap.get("isstoreUse");
    int cost = (int)attributesMap.get("Cost");


    assertEquals(description,"test");
    assertEquals(check,"check");

    //defaultValue
    assertEquals(diskUsage,70.5f);
    assertEquals(isstoreUse,true);
    assertEquals(sourceCode,"Hello World");
    assertEquals(cost,30);
}
 

private AtlasEntityHeaders removeEntityGuids(AtlasEntityHeaders headers) {
    Map<String, AtlasEntityHeader> uniqueNameEntityHeaderMap = new HashMap<>();

    for (AtlasEntityHeader header : headers.getGuidHeaderMap().values()) {
        String uniqueName = getUniqueName(header);
        if (StringUtils.isEmpty(uniqueName)) {
            displayCrLf("UniqueName is empty.  Ignoring: " + header.getGuid());
            LOG.warn("UniqueName is empty. Ignoring: {}", AtlasJson.toJson(header));
            continue;
        }

        displayCrLf("Processing: " + uniqueName);
        if (header.getStatus() == DELETED) {
            continue;
        }

        String key = String.format("%s:%s", header.getTypeName(), uniqueName);
        header.setGuid(null);
        boolean keyFound = uniqueNameEntityHeaderMap.containsKey(key);
        if (!keyFound) {
            uniqueNameEntityHeaderMap.put(key, header);
        }

        updateClassificationsForHeader(header, uniqueNameEntityHeaderMap.get(key), keyFound);
        displayCrLf("Processing: " + uniqueName);
    }

    displayCrLf("Processed: " + uniqueNameEntityHeaderMap.size());
    headers.setGuidHeaderMap(uniqueNameEntityHeaderMap);
    return headers;
}
 

public void removeEntity(AtlasEntityHeader entity) {
    List<AtlasEntityHeader> entities = this.entities;

    if (CollectionUtils.isNotEmpty(entities)) {
        Iterator<AtlasEntityHeader> iter = entities.iterator();
        while (iter.hasNext()) {
            AtlasEntityHeader currEntity = iter.next();
            if (StringUtils.equals(currEntity.getGuid(), entity.getGuid())) {
                iter.remove();
            }
        }
    }
}
 

private void purgeDeletedEntities(List<AtlasEntityHeader> entities) {
    if (entities != null) {
        for (ListIterator<AtlasEntityHeader> iter = entities.listIterator(); iter.hasNext(); ) {
            AtlasEntityHeader entity = iter.next();

            if (entity.getStatus() == AtlasEntity.Status.DELETED) {
                if (LOG.isDebugEnabled()) {
                    LOG.debug("purgeDeletedEntities(guid={}, status={}): REMOVED", entity.getGuid(), entity.getStatus());
                }

                iter.remove();
            }
        }
    }
}
 

@Override
protected void assertProcessForTestDeleteReference(final AtlasEntityHeader processInstance) throws Exception {
    //assert that outputs is empty
    ITypedReferenceableInstance newProcess =
        metadataService.getEntityDefinition(processInstance.getGuid());
    assertNull(newProcess.get(AtlasClient.PROCESS_ATTRIBUTE_OUTPUTS));
}
 

protected List<String> extractGuids(final List<AtlasEntityHeader> updatedEntities) {
    List<String> guids = new ArrayList<>();
    for (AtlasEntityHeader header : updatedEntities ) {
        guids.add(header.getGuid());
    }
    return guids;
}
 

private void doFullTextMapping(List<AtlasEntityHeader> entityHeaders) {
    if(AtlasRepositoryConfiguration.isFreeTextSearchEnabled() || !AtlasRepositoryConfiguration.isFullTextSearchEnabled()) {
        return;
    }

    if (CollectionUtils.isEmpty(entityHeaders)) {
        return;
    }

    MetricRecorder metric = RequestContext.get().startMetricRecord("fullTextMapping");

    for (AtlasEntityHeader entityHeader : entityHeaders) {
        if(GraphHelper.isInternalType(entityHeader.getTypeName())) {
            continue;
        }

        String      guid   = entityHeader.getGuid();
        AtlasVertex vertex = AtlasGraphUtilsV2.findByGuid(guid);

        if(vertex == null) {
            continue;
        }

        try {
            String fullText = fullTextMapperV2.getIndexTextForEntity(guid);

            AtlasGraphUtilsV2.setEncodedProperty(vertex, ENTITY_TEXT_PROPERTY_KEY, fullText);
        } catch (AtlasBaseException e) {
            LOG.error("FullText mapping failed for Vertex[ guid = {} ]", guid, e);
        }
    }

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