| /* |
| * Licensed to the Apache Software Foundation (ASF) under one or more |
| * contributor license agreements. See the NOTICE file distributed with |
| * this work for additional information regarding copyright ownership. |
| * The ASF licenses this file to You under the Apache License, Version 2.0 |
| * (the "License"); you may not use this file except in compliance with |
| * the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.nifi.atlas; |
| |
| import com.sun.jersey.api.client.UniformInterfaceException; |
| import com.sun.jersey.core.util.MultivaluedMapImpl; |
| import org.apache.atlas.AtlasClientV2; |
| import org.apache.atlas.AtlasErrorCode; |
| import org.apache.atlas.AtlasServiceException; |
| import org.apache.atlas.model.SearchFilter; |
| import org.apache.atlas.model.instance.AtlasEntity; |
| import org.apache.atlas.model.instance.AtlasObjectId; |
| import org.apache.atlas.model.instance.EntityMutationResponse; |
| import org.apache.atlas.model.typedef.AtlasEntityDef; |
| import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef; |
| import org.apache.atlas.model.typedef.AtlasTypesDef; |
| import org.apache.nifi.util.StringUtils; |
| import org.apache.nifi.util.Tuple; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| import javax.ws.rs.core.MultivaluedMap; |
| import java.util.ArrayList; |
| import java.util.Collections; |
| import java.util.HashMap; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Objects; |
| import java.util.Set; |
| import java.util.concurrent.atomic.AtomicBoolean; |
| import java.util.function.Function; |
| import java.util.regex.Matcher; |
| import java.util.regex.Pattern; |
| import java.util.stream.Collectors; |
| |
| import static org.apache.nifi.atlas.AtlasUtils.findIdByQualifiedName; |
| import static org.apache.nifi.atlas.AtlasUtils.getComponentIdFromQualifiedName; |
| import static org.apache.nifi.atlas.AtlasUtils.toStr; |
| import static org.apache.nifi.atlas.NiFiFlow.EntityChangeType.AS_IS; |
| import static org.apache.nifi.atlas.NiFiFlow.EntityChangeType.CREATED; |
| import static org.apache.nifi.atlas.NiFiFlow.EntityChangeType.DELETED; |
| import static org.apache.nifi.atlas.NiFiFlow.EntityChangeType.UPDATED; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_DESCRIPTION; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_FLOW_PATHS; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_GUID; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_INPUTS; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_INPUT_PORTS; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_NAME; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_OUTPUTS; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_OUTPUT_PORTS; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_QUALIFIED_NAME; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_QUEUES; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_TYPENAME; |
| import static org.apache.nifi.atlas.NiFiTypes.ATTR_URL; |
| import static org.apache.nifi.atlas.NiFiTypes.ENTITIES; |
| import static org.apache.nifi.atlas.NiFiTypes.TYPE_NIFI_FLOW; |
| import static org.apache.nifi.atlas.NiFiTypes.TYPE_NIFI_FLOW_PATH; |
| import static org.apache.nifi.atlas.NiFiTypes.TYPE_NIFI_INPUT_PORT; |
| import static org.apache.nifi.atlas.NiFiTypes.TYPE_NIFI_OUTPUT_PORT; |
| import static org.apache.nifi.atlas.NiFiTypes.TYPE_NIFI_QUEUE; |
| |
| public class NiFiAtlasClient { |
| |
| private static final Logger logger = LoggerFactory.getLogger(NiFiAtlasClient.class); |
| |
| private final AtlasClientV2 atlasClient; |
| |
| public NiFiAtlasClient(AtlasClientV2 atlasClient) { |
| this.atlasClient = atlasClient; |
| } |
| |
| /** |
| * This is an utility method to delete unused types. |
| * Should be used during development or testing only. |
| * @param typeNames to delete |
| */ |
| void deleteTypeDefs(String ... typeNames) throws AtlasServiceException { |
| final AtlasTypesDef existingTypeDef = getTypeDefs(typeNames); |
| try { |
| atlasClient.deleteAtlasTypeDefs(existingTypeDef); |
| } catch (UniformInterfaceException e) { |
| if (e.getResponse().getStatus() == 204) { |
| // 204 is a successful response. |
| // NOTE: However after executing this, Atlas should be restarted to work properly. |
| logger.info("Deleted type defs: {}", existingTypeDef); |
| } else { |
| throw e; |
| } |
| } |
| } |
| |
| /** |
| * @return True when required NiFi types are already created. |
| */ |
| public boolean isNiFiTypeDefsRegistered() throws AtlasServiceException { |
| final Set<String> typeNames = ENTITIES.keySet(); |
| final Map<String, AtlasEntityDef> existingDefs = getTypeDefs(typeNames.toArray(new String[typeNames.size()])).getEntityDefs().stream() |
| .collect(Collectors.toMap(AtlasEntityDef::getName, Function.identity())); |
| return typeNames.stream().allMatch(existingDefs::containsKey); |
| } |
| |
| /** |
| * Create or update NiFi types in Atlas type system. |
| * @param update If false, doesn't perform anything if there is existing type def for the name. |
| */ |
| public void registerNiFiTypeDefs(boolean update) throws AtlasServiceException { |
| final Set<String> typeNames = ENTITIES.keySet(); |
| final Map<String, AtlasEntityDef> existingDefs = getTypeDefs(typeNames.toArray(new String[typeNames.size()])).getEntityDefs().stream() |
| .collect(Collectors.toMap(AtlasEntityDef::getName, Function.identity())); |
| |
| |
| final AtomicBoolean shouldUpdate = new AtomicBoolean(false); |
| |
| final AtlasTypesDef type = new AtlasTypesDef(); |
| |
| typeNames.stream().filter(typeName -> { |
| final AtlasEntityDef existingDef = existingDefs.get(typeName); |
| if (existingDef != null) { |
| // type is already defined. |
| if (!update) { |
| return false; |
| } |
| shouldUpdate.set(true); |
| } |
| return true; |
| }).forEach(typeName -> { |
| final NiFiTypes.EntityDefinition def = ENTITIES.get(typeName); |
| |
| final AtlasEntityDef entity = new AtlasEntityDef(); |
| type.getEntityDefs().add(entity); |
| |
| entity.setName(typeName); |
| |
| Set<String> superTypes = new HashSet<>(); |
| List<AtlasAttributeDef> attributes = new ArrayList<>(); |
| |
| def.define(entity, superTypes, attributes); |
| |
| entity.setSuperTypes(superTypes); |
| entity.setAttributeDefs(attributes); |
| }); |
| |
| // Create or Update. |
| final AtlasTypesDef atlasTypeDefsResult = shouldUpdate.get() |
| ? atlasClient.updateAtlasTypeDefs(type) |
| : atlasClient.createAtlasTypeDefs(type); |
| logger.debug("Result={}", atlasTypeDefsResult); |
| } |
| |
| private AtlasTypesDef getTypeDefs(String ... typeNames) throws AtlasServiceException { |
| final AtlasTypesDef typeDefs = new AtlasTypesDef(); |
| for (int i = 0; i < typeNames.length; i++) { |
| final MultivaluedMap<String, String> searchParams = new MultivaluedMapImpl(); |
| searchParams.add(SearchFilter.PARAM_NAME, typeNames[i]); |
| final AtlasTypesDef typeDef = atlasClient.getAllTypeDefs(new SearchFilter(searchParams)); |
| typeDefs.getEntityDefs().addAll(typeDef.getEntityDefs()); |
| } |
| logger.debug("typeDefs={}", typeDefs); |
| return typeDefs; |
| } |
| |
| private Pattern FLOW_PATH_URL_PATTERN = Pattern.compile("^http.+processGroupId=([0-9a-z\\-]+).*$"); |
| /** |
| * Fetch existing NiFiFlow entity from Atlas. |
| * @param rootProcessGroupId The id of a NiFi flow root process group. |
| * @param clusterName The cluster name 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 clusterName) throws AtlasServiceException { |
| |
| final String qualifiedName = AtlasUtils.toQualifiedName(clusterName, 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.setClusterName(clusterName); |
| 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; |
| } |
| |
| /** |
| * Retrieves the flow components of type {@code componentType} from Atlas server. |
| * Deleted components will be filtered out before calling Atlas. |
| * Atlas object ids will be initialized with all the attributes (guid, type, unique attributes) in order to be able |
| * to match ids retrieved from Atlas (having guid) and ids created by the reporting task (not having guid yet). |
| * |
| * @param componentType Atlas type of the flow component (nifi_flow_path, nifi_queue, nifi_input_port, nifi_output_port) |
| * @param referredEntities referred entities of the flow entity (returned when the flow fetched) containing the basic data (id, status) of the flow components |
| * @return flow component entities mapped to their object ids |
| */ |
| private Map<AtlasObjectId, AtlasEntity> fetchFlowComponents(String componentType, Map<String, AtlasEntity> referredEntities) { |
| return referredEntities.values().stream() |
| .filter(referredEntity -> referredEntity.getTypeName().equals(componentType)) |
| .filter(referredEntity -> referredEntity.getStatus() == AtlasEntity.Status.ACTIVE) |
| .map(referredEntity -> { |
| final Map<String, Object> uniqueAttributes = Collections.singletonMap(ATTR_QUALIFIED_NAME, referredEntity.getAttribute(ATTR_QUALIFIED_NAME)); |
| final AtlasObjectId id = new AtlasObjectId(referredEntity.getGuid(), componentType, uniqueAttributes); |
| try { |
| final AtlasEntity.AtlasEntityWithExtInfo fetchedEntityExt = searchEntityDef(id); |
| return new Tuple<>(id, fetchedEntityExt.getEntity()); |
| } catch (AtlasServiceException e) { |
| logger.warn("Failed to search entity by id {}, due to {}", id, e); |
| return null; |
| } |
| }) |
| .filter(Objects::nonNull) |
| .collect(Collectors.toMap(Tuple::getKey, Tuple::getValue)); |
| } |
| |
| @SuppressWarnings("unchecked") |
| private List<AtlasObjectId> toAtlasObjectIds(Object _references) { |
| if (_references == null) { |
| return Collections.emptyList(); |
| } |
| List<Map<String, Object>> references = (List<Map<String, Object>>) _references; |
| return references.stream() |
| .map(ref -> new AtlasObjectId(toStr(ref.get(ATTR_GUID)), toStr(ref.get(ATTR_TYPENAME)), ref)) |
| .collect(Collectors.toList()); |
| } |
| |
| /** |
| * <p>AtlasObjectIds returned from Atlas have GUID, but do not have qualifiedName, while ones created by the reporting task |
| * do not have GUID, but qualifiedName. AtlasObjectId.equals returns false for this combination. |
| * In order to match ids correctly, this method converts fetches actual entities from ids to get qualifiedName attribute.</p> |
| * |
| * <p>Also, AtlasObjectIds returned from Atlas does not have entity state. |
| * If Atlas is configured to use soft-delete (default), deleted ids are still returned. |
| * Fetched entities are used to determine whether an AtlasObjectId is still active or deleted. |
| * Deleted entities will not be included in the result of this method. |
| * </p> |
| * @param ids to convert |
| * @return AtlasObjectIds with qualifiedName |
| */ |
| private Map<AtlasObjectId, AtlasEntity> toQualifiedNameIds(List<AtlasObjectId> ids) { |
| if (ids == null) { |
| return Collections.emptyMap(); |
| } |
| |
| return ids.stream().distinct().map(id -> { |
| try { |
| final AtlasEntity.AtlasEntityWithExtInfo entityExt = searchEntityDef(id); |
| final AtlasEntity entity = entityExt.getEntity(); |
| if (AtlasEntity.Status.DELETED.equals(entity.getStatus())) { |
| return null; |
| } |
| final Map<String, Object> uniqueAttrs = Collections.singletonMap(ATTR_QUALIFIED_NAME, entity.getAttribute(ATTR_QUALIFIED_NAME)); |
| return new Tuple<>(new AtlasObjectId(id.getGuid(), id.getTypeName(), uniqueAttrs), entity); |
| } catch (AtlasServiceException e) { |
| logger.warn("Failed to search entity by id {}, due to {}", id, e); |
| return null; |
| } |
| }).filter(Objects::nonNull).collect(Collectors.toMap(Tuple::getKey, Tuple::getValue)); |
| } |
| |
| public void registerNiFiFlow(NiFiFlow nifiFlow) throws AtlasServiceException { |
| |
| // Create parent flow entity, so that common properties are taken over. |
| final AtlasEntity flowEntity = registerNiFiFlowEntity(nifiFlow); |
| |
| // Create DataSet entities those are created by this NiFi flow. |
| final Map<String, List<AtlasEntity>> updatedDataSetEntities = registerDataSetEntities(nifiFlow); |
| |
| // Create path entities. |
| final Set<AtlasObjectId> remainingPathIds = registerFlowPathEntities(nifiFlow); |
| |
| // Update these attributes only if anything is created, updated or removed. |
| boolean shouldUpdateNiFiFlow = nifiFlow.isMetadataUpdated(); |
| if (remainingPathIds != null) { |
| flowEntity.setAttribute(ATTR_FLOW_PATHS, remainingPathIds); |
| shouldUpdateNiFiFlow = true; |
| } |
| if (updatedDataSetEntities.containsKey(TYPE_NIFI_QUEUE)) { |
| flowEntity.setAttribute(ATTR_QUEUES, updatedDataSetEntities.get(TYPE_NIFI_QUEUE)); |
| shouldUpdateNiFiFlow = true; |
| } |
| if (updatedDataSetEntities.containsKey(TYPE_NIFI_INPUT_PORT)) { |
| flowEntity.setAttribute(ATTR_INPUT_PORTS, updatedDataSetEntities.get(TYPE_NIFI_INPUT_PORT)); |
| shouldUpdateNiFiFlow = true; |
| } |
| if (updatedDataSetEntities.containsKey(TYPE_NIFI_OUTPUT_PORT)) { |
| flowEntity.setAttribute(ATTR_OUTPUT_PORTS, updatedDataSetEntities.get(TYPE_NIFI_OUTPUT_PORT)); |
| shouldUpdateNiFiFlow = true; |
| } |
| |
| if (logger.isDebugEnabled()) { |
| logger.debug("### NiFi Flow Audit Logs START"); |
| nifiFlow.getUpdateAudit().forEach(logger::debug); |
| nifiFlow.getFlowPaths().forEach((k, v) -> { |
| logger.debug("--- NiFiFlowPath Audit Logs: {}", k); |
| v.getUpdateAudit().forEach(logger::debug); |
| }); |
| logger.debug("### NiFi Flow Audit Logs END"); |
| } |
| |
| if (shouldUpdateNiFiFlow) { |
| // Send updated entities. |
| final List<AtlasEntity> entities = new ArrayList<>(); |
| final AtlasEntity.AtlasEntitiesWithExtInfo atlasEntities = new AtlasEntity.AtlasEntitiesWithExtInfo(entities); |
| entities.add(flowEntity); |
| try { |
| final EntityMutationResponse mutationResponse = atlasClient.createEntities(atlasEntities); |
| logger.debug("mutation response={}", mutationResponse); |
| } catch (AtlasServiceException e) { |
| if (e.getStatus().getStatusCode() == AtlasErrorCode.INSTANCE_NOT_FOUND.getHttpCode().getStatusCode() |
| && e.getMessage().contains(AtlasErrorCode.INSTANCE_NOT_FOUND.getErrorCode())) { |
| // NOTE: If previously existed nifi_flow_path entity is removed because the path is removed from NiFi, |
| // then Atlas respond with 404 even though the entity is successfully updated. |
| // Following exception is thrown in this case. Just log it. |
| // org.apache.atlas.AtlasServiceException: |
| // Metadata service API org.apache.atlas.AtlasBaseClient$APIInfo@45a37759 |
| // failed with status 404 (Not Found) Response Body |
| // ({"errorCode":"ATLAS-404-00-00B","errorMessage":"Given instance is invalid/not found: |
| // Could not find entities in the repository with guids: [96d24487-cd66-4795-b552-f00b426fed26]"}) |
| logger.debug("Received error response from Atlas but it should be stored." + e); |
| } else { |
| throw e; |
| } |
| } |
| } |
| } |
| |
| private AtlasEntity registerNiFiFlowEntity(final NiFiFlow nifiFlow) throws AtlasServiceException { |
| final List<AtlasEntity> entities = new ArrayList<>(); |
| final AtlasEntity.AtlasEntitiesWithExtInfo atlasEntities = new AtlasEntity.AtlasEntitiesWithExtInfo(entities); |
| |
| if (!nifiFlow.isMetadataUpdated()) { |
| // Nothing has been changed, return existing entity. |
| return nifiFlow.getExEntity(); |
| } |
| |
| // Create parent flow entity using existing NiFiFlow entity if available, so that common properties are taken over. |
| final AtlasEntity flowEntity = nifiFlow.getExEntity() != null ? new AtlasEntity(nifiFlow.getExEntity()) : new AtlasEntity(); |
| flowEntity.setTypeName(TYPE_NIFI_FLOW); |
| flowEntity.setVersion(1L); |
| flowEntity.setAttribute(ATTR_NAME, nifiFlow.getFlowName()); |
| flowEntity.setAttribute(ATTR_QUALIFIED_NAME, nifiFlow.toQualifiedName(nifiFlow.getRootProcessGroupId())); |
| flowEntity.setAttribute(ATTR_URL, nifiFlow.getUrl()); |
| flowEntity.setAttribute(ATTR_DESCRIPTION, nifiFlow.getDescription()); |
| |
| // If flowEntity is not persisted yet, then store nifi_flow entity to make nifiFlowId available for other entities. |
| if (flowEntity.getGuid().startsWith("-")) { |
| entities.add(flowEntity); |
| final EntityMutationResponse mutationResponse = atlasClient.createEntities(atlasEntities); |
| logger.debug("Registered a new nifi_flow entity, mutation response={}", mutationResponse); |
| final String assignedNiFiFlowGuid = mutationResponse.getGuidAssignments().get(flowEntity.getGuid()); |
| flowEntity.setGuid(assignedNiFiFlowGuid); |
| nifiFlow.setAtlasGuid(assignedNiFiFlowGuid); |
| } |
| |
| return flowEntity; |
| } |
| |
| /** |
| * 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; |
| } |
| |
| 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; |
| } |
| |
| public AtlasEntity.AtlasEntityWithExtInfo searchEntityDef(AtlasObjectId id) throws AtlasServiceException { |
| final String guid = id.getGuid(); |
| if (!StringUtils.isEmpty(guid)) { |
| return atlasClient.getEntityByGuid(guid, true, false); |
| } |
| final Map<String, String> attributes = new HashMap<>(); |
| id.getUniqueAttributes().entrySet().stream().filter(entry -> entry.getValue() != null) |
| .forEach(entry -> attributes.put(entry.getKey(), entry.getValue().toString())); |
| return atlasClient.getEntityByAttribute(id.getTypeName(), attributes, true, false); |
| } |
| |
| } |
