use of org.apache.nifi.atlas.NiFiFlow.EntityChangeType.CREATED in project nifi by apache.
the class NiFiAtlasClient method registerDataSetEntities.
/**
* Register DataSet within specified NiFiFlow.
* @return Set of registered Atlas type names and its remaining entities without deleted ones.
*/
private Map<String, List<AtlasEntity>> registerDataSetEntities(final NiFiFlow nifiFlow) throws AtlasServiceException {
final Map<NiFiFlow.EntityChangeType, List<AtlasEntity>> changedEntities = nifiFlow.getChangedDataSetEntities();
if (changedEntities.containsKey(CREATED)) {
final List<AtlasEntity> createdEntities = changedEntities.get(CREATED);
final AtlasEntity.AtlasEntitiesWithExtInfo atlasEntities = new AtlasEntity.AtlasEntitiesWithExtInfo(createdEntities);
final EntityMutationResponse mutationResponse = atlasClient.createEntities(atlasEntities);
logger.debug("Created DataSet entities mutation response={}", mutationResponse);
final Map<String, String> guidAssignments = mutationResponse.getGuidAssignments();
for (AtlasEntity entity : createdEntities) {
final String guid = guidAssignments.get(entity.getGuid());
final String qualifiedName = toStr(entity.getAttribute(ATTR_QUALIFIED_NAME));
if (StringUtils.isEmpty(guid)) {
logger.warn("GUID was not assigned for {}::{} for some reason.", entity.getTypeName(), qualifiedName);
continue;
}
final Map<AtlasObjectId, AtlasEntity> entityMap;
switch(entity.getTypeName()) {
case TYPE_NIFI_INPUT_PORT:
entityMap = nifiFlow.getRootInputPortEntities();
break;
case TYPE_NIFI_OUTPUT_PORT:
entityMap = nifiFlow.getRootOutputPortEntities();
break;
case TYPE_NIFI_QUEUE:
entityMap = nifiFlow.getQueues();
break;
default:
throw new RuntimeException(entity.getTypeName() + " is not expected.");
}
// In order to replace the id, remove current id which does not have GUID.
findIdByQualifiedName(entityMap.keySet(), qualifiedName).ifPresent(entityMap::remove);
entity.setGuid(guid);
final AtlasObjectId idWithGuid = new AtlasObjectId(guid, entity.getTypeName(), Collections.singletonMap(ATTR_QUALIFIED_NAME, qualifiedName));
entityMap.put(idWithGuid, entity);
}
}
if (changedEntities.containsKey(UPDATED)) {
final List<AtlasEntity> updatedEntities = changedEntities.get(UPDATED);
final AtlasEntity.AtlasEntitiesWithExtInfo atlasEntities = new AtlasEntity.AtlasEntitiesWithExtInfo(updatedEntities);
final EntityMutationResponse mutationResponse = atlasClient.updateEntities(atlasEntities);
logger.debug("Updated DataSet entities mutation response={}", mutationResponse);
}
final Set<String> changedTypeNames = changedEntities.entrySet().stream().filter(entry -> !AS_IS.equals(entry.getKey())).flatMap(entry -> entry.getValue().stream()).map(AtlasEntity::getTypeName).collect(Collectors.toSet());
// NOTE: Cascading DELETE will be performed when parent NiFiFlow is updated without removed DataSet entities.
final Map<String, List<AtlasEntity>> remainingEntitiesByType = changedEntities.entrySet().stream().filter(entry -> !DELETED.equals(entry.getKey())).flatMap(entry -> entry.getValue().stream()).filter(entity -> changedTypeNames.contains(entity.getTypeName())).collect(Collectors.groupingBy(AtlasEntity::getTypeName));
// If all entities are deleted for a type (e.g. nifi_intput_port), then remainingEntitiesByType will not contain such key.
// If the returning map does not contain anything for a type, then the corresponding attribute will not be updated.
// To empty an attribute when all of its elements are deleted, add empty list for a type.
changedTypeNames.forEach(changedTypeName -> remainingEntitiesByType.computeIfAbsent(changedTypeName, k -> Collections.emptyList()));
return remainingEntitiesByType;
}
use of org.apache.nifi.atlas.NiFiFlow.EntityChangeType.CREATED in project nifi by apache.
the class NiFiAtlasClient method registerFlowPathEntities.
private Set<AtlasObjectId> registerFlowPathEntities(final NiFiFlow nifiFlow) throws AtlasServiceException {
final Map<NiFiFlow.EntityChangeType, List<AtlasEntity>> changedEntities = nifiFlow.getChangedFlowPathEntities();
if (changedEntities.containsKey(CREATED)) {
final List<AtlasEntity> createdEntities = changedEntities.get(CREATED);
final AtlasEntity.AtlasEntitiesWithExtInfo atlasEntities = new AtlasEntity.AtlasEntitiesWithExtInfo(createdEntities);
final EntityMutationResponse mutationResponse = atlasClient.createEntities(atlasEntities);
logger.debug("Created FlowPath entities mutation response={}", mutationResponse);
final Map<String, String> guidAssignments = mutationResponse.getGuidAssignments();
createdEntities.forEach(entity -> {
final String guid = entity.getGuid();
entity.setGuid(guidAssignments.get(guid));
final String pathId = getComponentIdFromQualifiedName(toStr(entity.getAttribute(ATTR_QUALIFIED_NAME)));
final NiFiFlowPath path = nifiFlow.getFlowPaths().get(pathId);
path.setExEntity(entity);
});
}
if (changedEntities.containsKey(UPDATED)) {
final List<AtlasEntity> updatedEntities = changedEntities.get(UPDATED);
final AtlasEntity.AtlasEntitiesWithExtInfo atlasEntities = new AtlasEntity.AtlasEntitiesWithExtInfo(updatedEntities);
final EntityMutationResponse mutationResponse = atlasClient.updateEntities(atlasEntities);
logger.debug("Updated FlowPath entities mutation response={}", mutationResponse);
updatedEntities.forEach(entity -> {
final String pathId = getComponentIdFromQualifiedName(toStr(entity.getAttribute(ATTR_QUALIFIED_NAME)));
final NiFiFlowPath path = nifiFlow.getFlowPaths().get(pathId);
path.setExEntity(entity);
});
}
if (NiFiFlow.EntityChangeType.containsChange(changedEntities.keySet())) {
return changedEntities.entrySet().stream().filter(entry -> !DELETED.equals(entry.getKey())).flatMap(entry -> entry.getValue().stream()).map(path -> new AtlasObjectId(path.getGuid(), TYPE_NIFI_FLOW_PATH, Collections.singletonMap(ATTR_QUALIFIED_NAME, path.getAttribute(ATTR_QUALIFIED_NAME)))).collect(Collectors.toSet());
}
return null;
}
Aggregations