Examples with EntityPartialUpdateRequest org.apache.atlas.notification.hook.HookNotification.EntityPartialUpdateRequest used on opensource projects

Example 1 with EntityPartialUpdateRequest

use of org.apache.atlas.notification.hook.HookNotification.EntityPartialUpdateRequest in project nifi by apache.

the class NiFiAtlasHook method commitMessages.

public void commitMessages() {
    final Map<Boolean, List<HookNotificationMessage>> partialNiFiFlowPathUpdateAndOthers = messages.stream().collect(Collectors.groupingBy(msg -> ENTITY_PARTIAL_UPDATE.equals(msg.getType()) && TYPE_NIFI_FLOW_PATH.equals(((EntityPartialUpdateRequest) msg).getTypeName()) && ATTR_QUALIFIED_NAME.equals(((EntityPartialUpdateRequest) msg).getAttribute())));
    final List<HookNotificationMessage> otherMessages = partialNiFiFlowPathUpdateAndOthers.computeIfAbsent(false, k -> Collections.emptyList());
    final List<HookNotificationMessage> partialNiFiFlowPathUpdates = partialNiFiFlowPathUpdateAndOthers.computeIfAbsent(true, k -> Collections.emptyList());
    logger.info("Commit messages: {} partialNiFiFlowPathUpdate and {} other messages.", partialNiFiFlowPathUpdates.size(), otherMessages.size());
    final Metrics metrics = new Metrics();
    metrics.totalMessages = messages.size();
    metrics.partialNiFiFlowPathUpdates = partialNiFiFlowPathUpdates.size();
    metrics.otherMessages = otherMessages.size();
    try {
        // Notify other messages first.
        notifyEntities(otherMessages);
        // De-duplicate messages.
        final List<HookNotificationMessage> deduplicatedMessages = partialNiFiFlowPathUpdates.stream().map(msg -> (EntityPartialUpdateRequest) msg).collect(Collectors.groupingBy(EntityPartialUpdateRequest::getAttributeValue)).entrySet().stream().map(entry -> {
            final String flowPathQualifiedName = entry.getKey();
            final Map<String, Referenceable> distinctInputs;
            final Map<String, Referenceable> distinctOutputs;
            final String flowPathGuid;
            try {
                // Fetch existing nifi_flow_path and its inputs/ouputs.
                metrics.flowPathSearched++;
                final AtlasEntity.AtlasEntityWithExtInfo flowPathExt = atlasClient.searchEntityDef(new AtlasObjectId(TYPE_NIFI_FLOW_PATH, ATTR_QUALIFIED_NAME, flowPathQualifiedName));
                final AtlasEntity flowPathEntity = flowPathExt.getEntity();
                flowPathGuid = flowPathEntity.getGuid();
                distinctInputs = toReferenceables(flowPathEntity.getAttribute(ATTR_INPUTS), metrics);
                distinctOutputs = toReferenceables(flowPathEntity.getAttribute(ATTR_OUTPUTS), metrics);
            } catch (AtlasServiceException e) {
                if (ClientResponse.Status.NOT_FOUND.equals(e.getStatus())) {
                    logger.debug("nifi_flow_path was not found for qualifiedName {}", flowPathQualifiedName);
                } else {
                    logger.warn("Failed to retrieve nifi_flow_path with qualifiedName {} due to {}", flowPathQualifiedName, e, e);
                }
                return null;
            }
            // Merge all inputs and outputs for this nifi_flow_path.
            for (EntityPartialUpdateRequest msg : entry.getValue()) {
                fromReferenceable(msg.getEntity().get(ATTR_INPUTS), metrics).entrySet().stream().filter(ref -> !distinctInputs.containsKey(ref.getKey())).forEach(ref -> distinctInputs.put(ref.getKey(), ref.getValue()));
                fromReferenceable(msg.getEntity().get(ATTR_OUTPUTS), metrics).entrySet().stream().filter(ref -> !distinctOutputs.containsKey(ref.getKey())).forEach(ref -> distinctOutputs.put(ref.getKey(), ref.getValue()));
            }
            // Consolidate messages into one.
            final Referenceable flowPathRef = new Referenceable(flowPathGuid, TYPE_NIFI_FLOW_PATH, null);
            // NOTE: distinctInputs.values() returns HashMap$Values, which causes following error. To avoid that, wrap with ArrayList:
            // org.json4s.package$MappingException: Can't find ScalaSig for class org.apache.atlas.typesystem.Referenceable
            flowPathRef.set(ATTR_INPUTS, new ArrayList<>(distinctInputs.values()));
            flowPathRef.set(ATTR_OUTPUTS, new ArrayList<>(distinctOutputs.values()));
            return new EntityPartialUpdateRequest(NIFI_USER, TYPE_NIFI_FLOW_PATH, ATTR_QUALIFIED_NAME, flowPathQualifiedName, flowPathRef);
        }).filter(Objects::nonNull).collect(Collectors.toList());
        metrics.dedupedPartialNiFiFlowPathUpdates = deduplicatedMessages.size();
        notifyEntities(deduplicatedMessages);
    } finally {
        metrics.log("Committed");
        messages.clear();
    }
}