Examples with TYPE_NIFI_INPUT_PORT org.apache.nifi.atlas.NiFiTypes.TYPE_NIFI_INPUT_PORT used on opensource projects

Example 1 with TYPE_NIFI_INPUT_PORT

use of org.apache.nifi.atlas.NiFiTypes.TYPE_NIFI_INPUT_PORT 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;
}