Search in sources :

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();
    }
}
Also used : ClientResponse(com.sun.jersey.api.client.ClientResponse) LoggerFactory(org.slf4j.LoggerFactory) HashMap(java.util.HashMap) LineageContext(org.apache.nifi.atlas.provenance.lineage.LineageContext) ATTR_QUALIFIED_NAME(org.apache.nifi.atlas.NiFiTypes.ATTR_QUALIFIED_NAME) ArrayList(java.util.ArrayList) LinkedHashMap(java.util.LinkedHashMap) ATTR_OUTPUTS(org.apache.nifi.atlas.NiFiTypes.ATTR_OUTPUTS) EntityPartialUpdateRequest(org.apache.atlas.notification.hook.HookNotification.EntityPartialUpdateRequest) TYPE_NIFI_FLOW_PATH(org.apache.nifi.atlas.NiFiTypes.TYPE_NIFI_FLOW_PATH) HookNotificationMessage(org.apache.atlas.notification.hook.HookNotification.HookNotificationMessage) AtlasObjectId(org.apache.atlas.model.instance.AtlasObjectId) Map(java.util.Map) Id(org.apache.atlas.typesystem.persistence.Id) AtlasEntity(org.apache.atlas.model.instance.AtlasEntity) ATTR_INPUTS(org.apache.nifi.atlas.NiFiTypes.ATTR_INPUTS) Logger(org.slf4j.Logger) AtlasUtils.toTypedQualifiedName(org.apache.nifi.atlas.AtlasUtils.toTypedQualifiedName) AtlasServiceException(org.apache.atlas.AtlasServiceException) Collectors(java.util.stream.Collectors) ATTR_GUID(org.apache.nifi.atlas.NiFiTypes.ATTR_GUID) Objects(java.util.Objects) List(java.util.List) Tuple(org.apache.nifi.util.Tuple) ENTITY_PARTIAL_UPDATE(org.apache.atlas.notification.hook.HookNotification.HookNotificationType.ENTITY_PARTIAL_UPDATE) Referenceable(org.apache.atlas.typesystem.Referenceable) ATTR_TYPENAME(org.apache.nifi.atlas.NiFiTypes.ATTR_TYPENAME) Collections(java.util.Collections) AtlasHook(org.apache.atlas.hook.AtlasHook) AtlasObjectId(org.apache.atlas.model.instance.AtlasObjectId) EntityPartialUpdateRequest(org.apache.atlas.notification.hook.HookNotification.EntityPartialUpdateRequest) Referenceable(org.apache.atlas.typesystem.Referenceable) AtlasServiceException(org.apache.atlas.AtlasServiceException) AtlasEntity(org.apache.atlas.model.instance.AtlasEntity) HookNotificationMessage(org.apache.atlas.notification.hook.HookNotification.HookNotificationMessage) ArrayList(java.util.ArrayList) List(java.util.List)

Aggregations

ClientResponse (com.sun.jersey.api.client.ClientResponse)1 ArrayList (java.util.ArrayList)1 Collections (java.util.Collections)1 HashMap (java.util.HashMap)1 LinkedHashMap (java.util.LinkedHashMap)1 List (java.util.List)1 Map (java.util.Map)1 Objects (java.util.Objects)1 Collectors (java.util.stream.Collectors)1 AtlasServiceException (org.apache.atlas.AtlasServiceException)1 AtlasHook (org.apache.atlas.hook.AtlasHook)1 AtlasEntity (org.apache.atlas.model.instance.AtlasEntity)1 AtlasObjectId (org.apache.atlas.model.instance.AtlasObjectId)1 EntityPartialUpdateRequest (org.apache.atlas.notification.hook.HookNotification.EntityPartialUpdateRequest)1 HookNotificationMessage (org.apache.atlas.notification.hook.HookNotification.HookNotificationMessage)1 ENTITY_PARTIAL_UPDATE (org.apache.atlas.notification.hook.HookNotification.HookNotificationType.ENTITY_PARTIAL_UPDATE)1 Referenceable (org.apache.atlas.typesystem.Referenceable)1 Id (org.apache.atlas.typesystem.persistence.Id)1 AtlasUtils.toTypedQualifiedName (org.apache.nifi.atlas.AtlasUtils.toTypedQualifiedName)1 ATTR_GUID (org.apache.nifi.atlas.NiFiTypes.ATTR_GUID)1