Java Code Examples for org.apache.atlas.model.instance.AtlasEntity#getAttributes()

public void prepareForPostUpdate() {
    if (postUpdateEntities != null) {
        ListIterator<AtlasEntity> iter = postUpdateEntities.listIterator();

        while (iter.hasNext()) {
            AtlasEntity entity       = iter.next();
            String      assignedGuid = getAssignedGuid(entity.getGuid());

            // no need to perform partial-update for entities that are created/deleted while processing this message
            if (createdEntities.contains(assignedGuid) || deletedEntities.contains(assignedGuid)) {
                iter.remove();
            } else {
                entity.setGuid(assignedGuid);

                if (entity.getAttributes() != null) {
                    setAssignedGuids(entity.getAttributes().values());
                }

                if (entity.getRelationshipAttributes() != null) {
                    setAssignedGuids(entity.getRelationshipAttributes().values());
                }
            }
        }
    }
}
 

@Test
public void testProcessIsAdded() throws AtlasServiceException {
    Map<String, String> attributes         = Collections.singletonMap(REFERENCEABLE_ATTRIBUTE_NAME, LOAD_SALES_DAILY_PROCESS + CLUSTER_SUFFIX);
    AtlasEntity         loadProcess        = atlasClientV2.getEntityByAttribute(LOAD_PROCESS_TYPE, attributes).getEntity();
    Map                 loadProcessAttribs = loadProcess.getAttributes();

    assertEquals(LOAD_SALES_DAILY_PROCESS, loadProcessAttribs.get(NAME));
    assertEquals("hive query for daily summary", loadProcessAttribs.get("description"));

    List inputs = (List) loadProcessAttribs.get("inputs");
    List outputs = (List) loadProcessAttribs.get("outputs");

    assertEquals(2, inputs.size());

    String salesFactTableId   = getTableId(SALES_FACT_TABLE);
    String timeDimTableId     = getTableId(TIME_DIM_TABLE);
    String salesFactDailyMVId = getTableId(SALES_FACT_DAILY_MV_TABLE);

    assertEquals(salesFactTableId, ((Map) inputs.get(0)).get("guid"));
    assertEquals(timeDimTableId, ((Map) inputs.get(1)).get("guid"));
    assertEquals(salesFactDailyMVId, ((Map) outputs.get(0)).get("guid"));
}
 

private Map<String, Object> pruneEntityAttributesForAudit(AtlasEntity entity) {
    Map<String, Object> ret               = null;
    Map<String, Object> entityAttributes  = entity.getAttributes();
    List<String>        excludeAttributes = auditRepository.getAuditExcludeAttributes(entity.getTypeName());
    AtlasEntityType     entityType        = typeRegistry.getEntityTypeByName(entity.getTypeName());

    if (CollectionUtils.isNotEmpty(excludeAttributes) && MapUtils.isNotEmpty(entityAttributes) && entityType != null) {
        for (AtlasAttribute attribute : entityType.getAllAttributes().values()) {
            String attrName  = attribute.getName();
            Object attrValue = entityAttributes.get(attrName);

            if (excludeAttributes.contains(attrName)) {
                if (ret == null) {
                    ret = new HashMap<>();
                }

                ret.put(attrName, attrValue);
                entityAttributes.remove(attrName);
            }
        }
    }

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

@Test
public void testProcessIsAdded() throws AtlasServiceException, JSONException {
    Map<String, String> attributes = new HashMap<>();
    attributes.put(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, QuickStartV2.LOAD_SALES_DAILY_PROCESS);
    AtlasEntity loadProcess = atlasClientV2.getEntityByAttribute(QuickStartV2.LOAD_PROCESS_TYPE, attributes).getEntity();

    Map loadProcessAttribs = loadProcess.getAttributes();
    assertEquals(QuickStartV2.LOAD_SALES_DAILY_PROCESS, loadProcessAttribs.get(AtlasClient.NAME));
    assertEquals("hive query for daily summary", loadProcessAttribs.get("description"));

    List inputs = (List) loadProcessAttribs.get("inputs");
    List outputs = (List) loadProcessAttribs.get("outputs");
    assertEquals(2, inputs.size());

    String salesFactTableId = getTableId(QuickStartV2.SALES_FACT_TABLE);
    String timeDimTableId = getTableId(QuickStartV2.TIME_DIM_TABLE);
    String salesFactDailyMVId = getTableId(QuickStartV2.SALES_FACT_DAILY_MV_TABLE);

    assertEquals(salesFactTableId, ((Map) inputs.get(0)).get("guid"));
    assertEquals(timeDimTableId, ((Map) inputs.get(1)).get("guid"));
    assertEquals(salesFactDailyMVId, ((Map) outputs.get(0)).get("guid"));
}
 

@Test
public void testDBIsAdded() throws Exception {
    AtlasEntity         db           = getDB(SALES_DB);
    Map<String, Object> dbAttributes = db.getAttributes();

    assertEquals(SALES_DB, dbAttributes.get("name"));
    assertEquals("sales database", dbAttributes.get("description"));
}
 

@Test
public void testViewIsAdded() throws AtlasServiceException {
    Map<String, String> attributes                 = Collections.singletonMap(REFERENCEABLE_ATTRIBUTE_NAME, PRODUCT_DIM_VIEW + CLUSTER_SUFFIX);
    AtlasEntity         view                       = atlasClientV2.getEntityByAttribute(VIEW_TYPE, attributes).getEntity();
    Map<String, Object> viewAttributes             = view.getAttributes();
    Map<String, Object> viewRelationshipAttributes = view.getRelationshipAttributes();

    assertEquals(PRODUCT_DIM_VIEW, viewAttributes.get(NAME));

    String productDimId   = getTable(PRODUCT_DIM_TABLE).getGuid();
    List   inputTables    = (List) viewRelationshipAttributes.get("inputTables");
    Map    inputTablesMap = (Map) inputTables.get(0);

    assertEquals(productDimId, inputTablesMap.get("guid"));
}
 

@Test
public void testDBIsAdded() throws Exception {
    AtlasEntity db = getDB(QuickStartV2.SALES_DB);
    Map<String, Object> dbAttributes = db.getAttributes();
    assertEquals(QuickStartV2.SALES_DB, dbAttributes.get("name"));
    assertEquals("sales database", dbAttributes.get("description"));
}
 

private void verifyColumnsAreAddedToTable(AtlasEntity table) throws JSONException {
    Map<String, Object> tableAttributes = table.getAttributes();
    List<Map> columns = (List<Map>) tableAttributes.get("columns");
    assertEquals(4, columns.size());

    for (Map colMap : columns) {
        String colGuid = (String) colMap.get("guid");
        assertNotNull(UUID.fromString(colGuid));
    }
}
 

@Test
public void testViewIsAdded() throws AtlasServiceException, JSONException {
    Map<String, String> attributes = new HashMap<>();
    attributes.put(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, QuickStartV2.PRODUCT_DIM_VIEW);
    AtlasEntity view = atlasClientV2.getEntityByAttribute(QuickStartV2.VIEW_TYPE, attributes).getEntity();
    Map<String, Object> viewAttributes = view.getAttributes();
    assertEquals(QuickStartV2.PRODUCT_DIM_VIEW, viewAttributes.get(AtlasClient.NAME));

    String productDimId = getTable(QuickStartV2.PRODUCT_DIM_TABLE).getGuid();
    List inputTables = (List) viewAttributes.get("inputTables");
    Map inputTablesMap = (Map) inputTables.get(0);
    assertEquals(productDimId, inputTablesMap.get("guid"));
}
 

@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 void verifySimpleTableAttributes(AtlasEntity table) {
    Map<String, Object> tableAttributes = table.getAttributes();

    assertEquals(SALES_FACT_TABLE, tableAttributes.get("name"));
    assertEquals("sales fact table", tableAttributes.get("description"));
}
 

private String getAuditEventDetail(AtlasEntity entity, EntityAuditActionV2 action) {
    Map<String, Object> prunedAttributes = pruneEntityAttributesForAudit(entity);

    String auditPrefix  = getV2AuditPrefix(action);
    String auditString  = auditPrefix + AtlasType.toJson(entity);
    byte[] auditBytes   = auditString.getBytes(StandardCharsets.UTF_8);
    long   auditSize    = auditBytes != null ? auditBytes.length : 0;
    long   auditMaxSize = auditRepository.repositoryMaxSize();

    if (auditMaxSize >= 0 && auditSize > auditMaxSize) { // don't store attributes in audit
        LOG.warn("audit record too long: entityType={}, guid={}, size={}; maxSize={}. entity attribute values not stored in audit",
                entity.getTypeName(), entity.getGuid(), auditSize, auditMaxSize);

        Map<String, Object> attrValues    = entity.getAttributes();
        Map<String, Object> relAttrValues = entity.getRelationshipAttributes();

        entity.setAttributes(null);
        entity.setRelationshipAttributes(null);

        auditString = auditPrefix + AtlasType.toJson(entity);
        auditBytes  = auditString.getBytes(StandardCharsets.UTF_8); // recheck auditString size
        auditSize   = auditBytes != null ? auditBytes.length : 0;

        if (auditMaxSize >= 0 && auditSize > auditMaxSize) { // don't store classifications and meanings as well
            LOG.warn("audit record still too long: entityType={}, guid={}, size={}; maxSize={}. audit will have only summary details",
                    entity.getTypeName(), entity.getGuid(), auditSize, auditMaxSize);

            AtlasEntity shallowEntity = new AtlasEntity();

            shallowEntity.setGuid(entity.getGuid());
            shallowEntity.setTypeName(entity.getTypeName());
            shallowEntity.setCreateTime(entity.getCreateTime());
            shallowEntity.setUpdateTime(entity.getUpdateTime());
            shallowEntity.setCreatedBy(entity.getCreatedBy());
            shallowEntity.setUpdatedBy(entity.getUpdatedBy());
            shallowEntity.setStatus(entity.getStatus());
            shallowEntity.setVersion(entity.getVersion());

            auditString = auditPrefix + AtlasType.toJson(shallowEntity);
        }

        entity.setAttributes(attrValues);
        entity.setRelationshipAttributes(relAttrValues);
    }

    restoreEntityAttributes(entity, prunedAttributes);

    return auditString;
}
 

private void verifyDBIsLinkedToTable(AtlasEntity table) throws AtlasServiceException, JSONException {
    AtlasEntity db = getDB(QuickStartV2.SALES_DB);
    Map<String, Object> tableAttributes = table.getAttributes();
    Map dbFromTable = (Map) tableAttributes.get("db");
    assertEquals(db.getGuid(), dbFromTable.get("guid"));
}
 

private void verifySimpleTableAttributes(AtlasEntity table) throws JSONException {
    Map<String, Object> tableAttributes = table.getAttributes();
    assertEquals(QuickStartV2.SALES_FACT_TABLE, tableAttributes.get("name"));
    assertEquals("sales fact table", tableAttributes.get("description"));
}
 

/**
 * Fetch existing NiFiFlow entity from Atlas.
 * @param rootProcessGroupId The id of a NiFi flow root process group.
 * @param namespace The namespace of a flow.
 * @return A NiFiFlow instance filled with retrieved data from Atlas. Status objects are left blank, e.g. ProcessorStatus.
 * @throws AtlasServiceException Thrown if requesting to Atlas API failed, including when the flow is not found.
 */
public NiFiFlow fetchNiFiFlow(String rootProcessGroupId, String namespace) throws AtlasServiceException {

    final String qualifiedName = AtlasUtils.toQualifiedName(namespace, rootProcessGroupId);
    final AtlasObjectId flowId = new AtlasObjectId(TYPE_NIFI_FLOW, ATTR_QUALIFIED_NAME, qualifiedName);
    final AtlasEntity.AtlasEntityWithExtInfo nifiFlowExt = searchEntityDef(flowId);

    if (nifiFlowExt == null || nifiFlowExt.getEntity() == null) {
        return null;
    }

    final AtlasEntity nifiFlowEntity = nifiFlowExt.getEntity();
    final Map<String, AtlasEntity> nifiFlowReferredEntities = nifiFlowExt.getReferredEntities();
    final Map<String, Object> attributes = nifiFlowEntity.getAttributes();
    final NiFiFlow nifiFlow = new NiFiFlow(rootProcessGroupId);
    nifiFlow.setExEntity(nifiFlowEntity);
    nifiFlow.setFlowName(toStr(attributes.get(ATTR_NAME)));
    nifiFlow.setNamespace(namespace);
    nifiFlow.setUrl(toStr(attributes.get(ATTR_URL)));
    nifiFlow.setDescription(toStr(attributes.get(ATTR_DESCRIPTION)));

    nifiFlow.getQueues().putAll(fetchFlowComponents(TYPE_NIFI_QUEUE, nifiFlowReferredEntities));
    nifiFlow.getRootInputPortEntities().putAll(fetchFlowComponents(TYPE_NIFI_INPUT_PORT, nifiFlowReferredEntities));
    nifiFlow.getRootOutputPortEntities().putAll(fetchFlowComponents(TYPE_NIFI_OUTPUT_PORT, nifiFlowReferredEntities));

    final Map<String, NiFiFlowPath> flowPaths = nifiFlow.getFlowPaths();
    final Map<AtlasObjectId, AtlasEntity> flowPathEntities = fetchFlowComponents(TYPE_NIFI_FLOW_PATH, nifiFlowReferredEntities);

    for (AtlasEntity flowPathEntity : flowPathEntities.values()) {
        final String pathQualifiedName = toStr(flowPathEntity.getAttribute(ATTR_QUALIFIED_NAME));
        final NiFiFlowPath flowPath = new NiFiFlowPath(getComponentIdFromQualifiedName(pathQualifiedName));
        if (flowPathEntity.hasAttribute(ATTR_URL)) {
            final Matcher urlMatcher = FLOW_PATH_URL_PATTERN.matcher(toStr(flowPathEntity.getAttribute(ATTR_URL)));
            if (urlMatcher.matches()) {
                flowPath.setGroupId(urlMatcher.group(1));
            }
        }
        flowPath.setExEntity(flowPathEntity);
        flowPath.setName(toStr(flowPathEntity.getAttribute(ATTR_NAME)));
        flowPath.getInputs().addAll(toQualifiedNameIds(toAtlasObjectIds(flowPathEntity.getAttribute(ATTR_INPUTS))).keySet());
        flowPath.getOutputs().addAll(toQualifiedNameIds(toAtlasObjectIds(flowPathEntity.getAttribute(ATTR_OUTPUTS))).keySet());
        flowPath.startTrackingChanges(nifiFlow);

        flowPaths.put(flowPath.getId(), flowPath);
    }

    nifiFlow.startTrackingChanges();
    return nifiFlow;
}