http://git-wip-us.apache.org/repos/asf/atlas/blob/6e7aa6ed/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java index f28d771..d285e69 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java @@ -51,6 +51,11 @@ import org.slf4j.LoggerFactory; import java.util.*; +import static org.apache.atlas.model.TypeCategory.ARRAY; +import static org.apache.atlas.model.TypeCategory.CLASSIFICATION; +import static org.apache.atlas.model.TypeCategory.MAP; +import static org.apache.atlas.model.TypeCategory.OBJECT_ID_TYPE; +import static org.apache.atlas.model.TypeCategory.STRUCT; import static org.apache.atlas.model.instance.AtlasEntity.Status.DELETED; import static org.apache.atlas.model.typedef.AtlasRelationshipDef.PropagateTags.ONE_TO_TWO; import static org.apache.atlas.repository.Constants.CLASSIFICATION_EDGE_NAME_PROPERTY_KEY; @@ -58,7 +63,6 @@ import static org.apache.atlas.repository.Constants.CLASSIFICATION_LABEL; import static org.apache.atlas.repository.Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY; import static org.apache.atlas.repository.Constants.RELATIONSHIP_GUID_PROPERTY_KEY; import static org.apache.atlas.repository.Constants.TRAIT_NAMES_PROPERTY_KEY; -import static org.apache.atlas.repository.graph.GraphHelper.EDGE_LABEL_PREFIX; import static org.apache.atlas.repository.graph.GraphHelper.addToPropagatedTraitNames; import static org.apache.atlas.repository.graph.GraphHelper.getAllClassificationEdges; import static org.apache.atlas.repository.graph.GraphHelper.getAssociatedEntityVertex; @@ -66,7 +70,9 @@ import static org.apache.atlas.repository.graph.GraphHelper.getClassificationEdg import static org.apache.atlas.repository.graph.GraphHelper.getClassificationEntityGuid; import static org.apache.atlas.repository.graph.GraphHelper.getClassificationName; import static org.apache.atlas.repository.graph.GraphHelper.getClassificationVertices; +import static org.apache.atlas.repository.graph.GraphHelper.getCollectionElementsUsingRelationship; import static org.apache.atlas.repository.graph.GraphHelper.getGuid; +import static org.apache.atlas.repository.graph.GraphHelper.getMapValuesUsingRelationship; import static org.apache.atlas.repository.graph.GraphHelper.getPropagatedClassificationEdge; import static org.apache.atlas.repository.graph.GraphHelper.getPropagatedEdges; import static org.apache.atlas.repository.graph.GraphHelper.getPropagationEnabledClassificationVertices; @@ -77,7 +83,9 @@ import static org.apache.atlas.repository.graph.GraphHelper.isPropagatedClassifi import static org.apache.atlas.repository.graph.GraphHelper.string; import static org.apache.atlas.repository.graph.GraphHelper.updateModificationMetadata; import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.getIdFromEdge; +import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.getQualifiedAttributePropertyKey; import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.getState; +import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.isReference; public abstract class DeleteHandlerV1 { public static final Logger LOG = LoggerFactory.getLogger(DeleteHandlerV1.class); @@ -198,63 +206,43 @@ public abstract class DeleteHandlerV1 { continue; } - String edgeLabel = AtlasGraphUtilsV1.getAttributeEdgeLabel(entityType, attributeInfo.getName()); - AtlasType attrType = attributeInfo.getAttributeType(); + String edgeLabel = attributeInfo.getRelationshipEdgeLabel(); + AtlasType attrType = attributeInfo.getAttributeType(); + TypeCategory typeCategory = attrType.getTypeCategory(); - switch (attrType.getTypeCategory()) { - case OBJECT_ID_TYPE: { - AtlasEdge edge = graphHelper.getEdgeForLabel(vertex, edgeLabel); + if (typeCategory == OBJECT_ID_TYPE) { + AtlasEdge edge = graphHelper.getEdgeForLabel(vertex, edgeLabel); - if (edge != null && getState(edge) == AtlasEntity.Status.ACTIVE) { - vertices.push(edge.getInVertex()); - } + if (edge == null || getState(edge) == DELETED) { + continue; } - break; - - case ARRAY: { - AtlasArrayType arrType = (AtlasArrayType) attrType; - - if (arrType.getElementType().getTypeCategory() != TypeCategory.OBJECT_ID_TYPE) { - continue; - } - Iterator<AtlasEdge> edges = graphHelper.getOutGoingEdgesByLabel(vertex, edgeLabel); + vertices.push(edge.getInVertex()); - if (edges != null) { - while (edges.hasNext()) { - AtlasEdge edge = edges.next(); + } else if (typeCategory == ARRAY || typeCategory == MAP) { + TypeCategory elementType = null; - if (edge != null && getState(edge) == AtlasEntity.Status.ACTIVE) { - vertices.push(edge.getInVertex()); - } - } - } + if (typeCategory == ARRAY) { + elementType = ((AtlasArrayType) attrType).getElementType().getTypeCategory(); + } else if (typeCategory == MAP) { + elementType = ((AtlasMapType) attrType).getValueType().getTypeCategory(); } - break; - - case MAP: { - AtlasMapType mapType = (AtlasMapType) attrType; - TypeCategory valueTypeCategory = mapType.getValueType().getTypeCategory(); - - if (valueTypeCategory != TypeCategory.OBJECT_ID_TYPE) { - continue; - } - String propertyName = AtlasGraphUtilsV1.getQualifiedAttributePropertyKey(entityType, attributeInfo.getName()); - List<String> keys = vertex.getProperty(propertyName, List.class); + if (elementType != OBJECT_ID_TYPE) { + continue; + } - if (keys != null) { - for (String key : keys) { - String mapEdgeLabel = GraphHelper.getQualifiedNameForMapKey(edgeLabel, key); - AtlasEdge edge = graphHelper.getEdgeForLabel(vertex, mapEdgeLabel); + List<AtlasEdge> edges = getCollectionElementsUsingRelationship(vertex, attributeInfo); - if (edge != null && getState(edge) == AtlasEntity.Status.ACTIVE) { - vertices.push(edge.getInVertex()); - } + if (CollectionUtils.isNotEmpty(edges)) { + for (AtlasEdge edge : edges) { + if (edge == null || getState(edge) == DELETED) { + continue; } + + vertices.push(edge.getInVertex()); } } - break; } } } @@ -284,14 +272,14 @@ public abstract class DeleteHandlerV1 { } boolean isInternalType = isInternalType(entityVertex); - boolean forceDelete = (typeCategory == TypeCategory.STRUCT || typeCategory == TypeCategory.CLASSIFICATION) + boolean forceDelete = (typeCategory == STRUCT || typeCategory == CLASSIFICATION) && (forceDeleteStructTrait || isInternalType); if (LOG.isDebugEnabled()) { LOG.debug("isInternal = {}, forceDelete = {}", isInternalType, forceDelete); } - if (typeCategory == TypeCategory.STRUCT || typeCategory == TypeCategory.CLASSIFICATION || (typeCategory == TypeCategory.OBJECT_ID_TYPE && isOwned)) { + if (typeCategory == STRUCT || typeCategory == CLASSIFICATION || (typeCategory == OBJECT_ID_TYPE && isOwned)) { if (LOG.isDebugEnabled()) { LOG.debug("Processing delete for typeCategory={}, isOwned={}", typeCategory, isOwned); } @@ -681,7 +669,7 @@ public abstract class DeleteHandlerV1 { boolean isOwned = isEntityType && attributeInfo.isOwnedRef(); AtlasType attrType = attributeInfo.getAttributeType(); - String edgeLabel = AtlasGraphUtilsV1.getAttributeEdgeLabel(structType, attributeInfo.getName()); + String edgeLabel = attributeInfo.getRelationshipEdgeLabel(); switch (attrType.getTypeCategory()) { case OBJECT_ID_TYPE: @@ -699,13 +687,11 @@ public abstract class DeleteHandlerV1 { AtlasArrayType arrType = (AtlasArrayType) attrType; AtlasType elemType = arrType.getElementType(); - if (AtlasGraphUtilsV1.isReference(elemType.getTypeCategory())) { - Iterator<AtlasEdge> edges = graphHelper.getOutGoingEdgesByLabel(instanceVertex, edgeLabel); - - if (edges != null) { - while (edges.hasNext()) { - AtlasEdge edge = edges.next(); + if (isReference(elemType.getTypeCategory())) { + List<AtlasEdge> edges = getCollectionElementsUsingRelationship(instanceVertex, attributeInfo); + if (CollectionUtils.isNotEmpty(edges)) { + for (AtlasEdge edge : edges) { deleteEdgeReference(edge, elemType.getTypeCategory(), isOwned, false, instanceVertex); } } @@ -715,19 +701,13 @@ public abstract class DeleteHandlerV1 { case MAP: //For map attribute, if the value type is struct/class, delete all the references AtlasMapType mapType = (AtlasMapType) attrType; - AtlasType keyType = mapType.getKeyType(); TypeCategory valueTypeCategory = mapType.getValueType().getTypeCategory(); - String propertyName = AtlasGraphUtilsV1.getQualifiedAttributePropertyKey(structType, attributeInfo.getName()); - - if (AtlasGraphUtilsV1.isReference(valueTypeCategory)) { - List<Object> keys = EntityGraphMapper.getArrayElementsProperty(keyType, instanceVertex, propertyName); - if (keys != null) { - for (Object key : keys) { - String mapEdgeLabel = GraphHelper.getQualifiedNameForMapKey(edgeLabel, (String) key); + if (isReference(valueTypeCategory)) { + List<AtlasEdge> edges = getMapValuesUsingRelationship(instanceVertex, attributeInfo); - deleteEdgeReference(instanceVertex, mapEdgeLabel, valueTypeCategory, isOwned); - } + for (AtlasEdge edge : edges) { + deleteEdgeReference(edge, valueTypeCategory, isOwned, false, instanceVertex); } } break; @@ -772,8 +752,8 @@ public abstract class DeleteHandlerV1 { } AtlasStructType parentType = (AtlasStructType) typeRegistry.getType(typeName); - String propertyName = AtlasGraphUtilsV1.getQualifiedAttributePropertyKey(parentType, attribute.getName()); - String edgeLabel = EDGE_LABEL_PREFIX + propertyName; + String propertyName = getQualifiedAttributePropertyKey(parentType, attribute.getName()); + String edgeLabel = attribute.getRelationshipEdgeLabel(); AtlasEdge edge = null; AtlasAttributeDef attrDef = attribute.getAttributeDef(); AtlasType attrType = attribute.getAttributeType(); @@ -796,14 +776,12 @@ public abstract class DeleteHandlerV1 { case ARRAY: { //If its array attribute, find the right edge between the two vertices and update array property - List<String> elements = GraphHelper.getListProperty(outVertex, propertyName); - - if (elements != null) { - elements = new ArrayList<>(elements); + List<AtlasEdge> elementEdges = getCollectionElementsUsingRelationship(outVertex, attribute); - for (String elementEdgeId : elements) { - AtlasEdge elementEdge = graphHelper.getEdgeByEdgeId(outVertex, edgeLabel, elementEdgeId); + if (elementEdges != null) { + elementEdges = new ArrayList<>(elementEdges); + for (AtlasEdge elementEdge : elementEdges) { if (elementEdge == null) { continue; } @@ -814,26 +792,10 @@ public abstract class DeleteHandlerV1 { edge = elementEdge; //TODO element.size includes deleted items as well. should exclude - if (!attrDef.getIsOptional() && elements.size() <= attrDef.getValuesMinCount()) { + if (!attrDef.getIsOptional() && elementEdges.size() <= attrDef.getValuesMinCount()) { // Deleting this edge would violate the attribute's lower bound. throw new AtlasBaseException("Cannot remove array element from required attribute " + propertyName + " on " + GraphHelper.getVertexDetails(outVertex) + " " + GraphHelper.getEdgeDetails(elementEdge)); } - - if (shouldUpdateInverseReferences) { - //if composite attribute, remove the reference as well. else, just remove the edge - //for example, when table is deleted, process still references the table - //but when column is deleted, table will not reference the deleted column - if (LOG.isDebugEnabled()) { - LOG.debug("Removing edge {} from the array attribute {}", string(elementEdge), attribute.getName()); - } - - // Remove all occurrences of the edge ID from the list. - // This prevents dangling edge IDs (i.e. edge IDs for deleted edges) - // from the remaining in the list if there are duplicates. - elements.removeAll(Collections.singletonList(elementEdge.getId().toString())); - GraphHelper.setProperty(outVertex, propertyName, elements); - break; - } } } } @@ -842,37 +804,22 @@ public abstract class DeleteHandlerV1 { case MAP: { //If its map attribute, find the right edge between two vertices and update map property - List<String> keys = GraphHelper.getListProperty(outVertex, propertyName); + List<AtlasEdge> mapEdges = getMapValuesUsingRelationship(outVertex, attribute); - if (keys != null) { - keys = new ArrayList<>(keys); - - for (String key : keys) { - String keyPropertyName = GraphHelper.getQualifiedNameForMapKey(propertyName, key); - String mapEdgeId = GraphHelper.getSingleValuedProperty(outVertex, keyPropertyName, String.class); - AtlasEdge mapEdge = graphHelper.getEdgeByEdgeId(outVertex, keyPropertyName, mapEdgeId); + if (mapEdges != null) { + mapEdges = new ArrayList<>(mapEdges); + for (AtlasEdge mapEdge : mapEdges) { if (mapEdge != null) { AtlasVertex mapVertex = mapEdge.getInVertex(); if (mapVertex.getId().toString().equals(inVertex.getId().toString())) { //TODO keys.size includes deleted items as well. should exclude - if (attrDef.getIsOptional() || keys.size() > attrDef.getValuesMinCount()) { + if (attrDef.getIsOptional() || mapEdges.size() > attrDef.getValuesMinCount()) { edge = mapEdge; } else { // Deleting this entry would violate the attribute's lower bound. - throw new AtlasBaseException("Cannot remove map entry " + keyPropertyName + " from required attribute " + propertyName + " on " + GraphHelper.getVertexDetails(outVertex) + " " + GraphHelper.getEdgeDetails(mapEdge)); - } - - if (shouldUpdateInverseReferences) { - //remove this key - if (LOG.isDebugEnabled()) { - LOG.debug("Removing edge {}, key {} from the map attribute {}", string(mapEdge), key, attribute.getName()); - } - - keys.remove(key); - GraphHelper.setProperty(outVertex, propertyName, keys); - GraphHelper.setProperty(outVertex, keyPropertyName, null); + throw new AtlasBaseException("Cannot remove map entry " + propertyName + " from required attribute " + propertyName + " on " + GraphHelper.getVertexDetails(outVertex) + " " + GraphHelper.getEdgeDetails(mapEdge)); } break; } @@ -968,7 +915,7 @@ public abstract class DeleteHandlerV1 { removeTagPropagation(classificationVertex); - deleteEdgeReference(edge, TypeCategory.CLASSIFICATION, false, false, instanceVertex); + deleteEdgeReference(edge, CLASSIFICATION, false, false, instanceVertex); } }
http://git-wip-us.apache.org/repos/asf/atlas/blob/6e7aa6ed/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java index 1728549..cd00639 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java @@ -56,7 +56,6 @@ import org.apache.atlas.type.AtlasTypeRegistry; import org.apache.atlas.type.AtlasTypeUtil; import org.apache.atlas.utils.AtlasJson; import org.apache.commons.collections.CollectionUtils; -import org.apache.commons.collections.IteratorUtils; import org.apache.commons.collections.MapUtils; import org.apache.commons.lang3.StringUtils; import org.slf4j.Logger; @@ -68,26 +67,36 @@ import java.util.*; import java.util.stream.Collectors; import static org.apache.atlas.model.TypeCategory.CLASSIFICATION; +import static org.apache.atlas.model.instance.AtlasEntity.Status.DELETED; import static org.apache.atlas.model.instance.AtlasRelatedObjectId.KEY_RELATIONSHIP_ATTRIBUTES; import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.CREATE; import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.DELETE; import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.PARTIAL_UPDATE; import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.UPDATE; import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef.Cardinality.SET; +import static org.apache.atlas.repository.Constants.ATTRIBUTE_KEY_PROPERTY_KEY; import static org.apache.atlas.repository.Constants.CLASSIFICATION_LABEL; import static org.apache.atlas.repository.Constants.STATE_PROPERTY_KEY; import static org.apache.atlas.repository.Constants.TRAIT_NAMES_PROPERTY_KEY; +import static org.apache.atlas.repository.Constants.ATTRIBUTE_INDEX_PROPERTY_KEY; +import static org.apache.atlas.repository.graph.GraphHelper.getCollectionElementsUsingRelationship; import static org.apache.atlas.repository.graph.GraphHelper.getClassificationEdge; import static org.apache.atlas.repository.graph.GraphHelper.getClassificationVertex; +import static org.apache.atlas.repository.graph.GraphHelper.getListProperty; +import static org.apache.atlas.repository.graph.GraphHelper.getMapElementsProperty; +import static org.apache.atlas.repository.graph.GraphHelper.getQualifiedNameForMapKey; +import static org.apache.atlas.repository.graph.GraphHelper.getStatus; import static org.apache.atlas.repository.graph.GraphHelper.getTraitLabel; import static org.apache.atlas.repository.graph.GraphHelper.getTraitNames; import static org.apache.atlas.repository.graph.GraphHelper.getTypeName; import static org.apache.atlas.repository.graph.GraphHelper.getTypeNames; import static org.apache.atlas.repository.graph.GraphHelper.isPropagationEnabled; import static org.apache.atlas.repository.graph.GraphHelper.isRelationshipEdge; +import static org.apache.atlas.repository.graph.GraphHelper.setListProperty; import static org.apache.atlas.repository.graph.GraphHelper.string; import static org.apache.atlas.repository.graph.GraphHelper.updateModificationMetadata; import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.getIdFromVertex; +import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.isReference; import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection.IN; import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection.OUT; @@ -794,64 +803,59 @@ public class EntityGraphMapper { LOG.debug("==> mapMapValue({})", ctx); } - @SuppressWarnings("unchecked") - Map<Object, Object> newVal = (Map<Object, Object>) ctx.getValue(); - Map<String, Object> newMap = new HashMap<>(); - AtlasMapType mapType = (AtlasMapType) ctx.getAttrType(); + Map<Object, Object> newVal = (Map<Object, Object>) ctx.getValue(); + Map<String, Object> newMap = new HashMap<>(); + AtlasMapType mapType = (AtlasMapType) ctx.getAttrType(); + AtlasAttribute attribute = ctx.getAttribute(); + Map<String, Object> currentMap = getMapElementsProperty(mapType, ctx.getReferringVertex(), ctx.getVertexProperty(), attribute); + boolean isReference = isReference(mapType.getValueType()); - try { - AtlasAttribute attribute = ctx.getAttribute(); - List<String> currentKeys = GraphHelper.getListProperty(ctx.getReferringVertex(), ctx.getVertexProperty()); - Map<String, Object> currentMap = new HashMap<>(); - - if (CollectionUtils.isNotEmpty(currentKeys)) { - for (String key : currentKeys) { - String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(ctx.getVertexProperty(), GraphHelper.encodePropertyKey(key)); - Object propertyValueForKey = getMapValueProperty(mapType.getValueType(), ctx.getReferringVertex(), propertyNameForKey); + if (MapUtils.isNotEmpty(newVal)) { + AtlasAttribute inverseRefAttribute = attribute.getInverseRefAttribute(); - currentMap.put(key, propertyValueForKey); - } - } + for (Map.Entry<Object, Object> entry : newVal.entrySet()) { + String key = entry.getKey().toString(); + String propertyName = (isReference) ? ctx.getVertexProperty() : getQualifiedNameForMapKey(ctx.getVertexProperty(), GraphHelper.encodePropertyKey(key)); + AtlasEdge existingEdge = getEdgeIfExists(mapType, currentMap, key); - if (MapUtils.isNotEmpty(newVal)) { - boolean isReference = AtlasGraphUtilsV1.isReference(mapType.getValueType()); - AtlasAttribute inverseRefAttribute = attribute.getInverseRefAttribute(); - for (Map.Entry<Object, Object> entry : newVal.entrySet()) { - String key = entry.getKey().toString(); - String propertyName = GraphHelper.getQualifiedNameForMapKey(ctx.getVertexProperty(), GraphHelper.encodePropertyKey(key)); - AtlasEdge existingEdge = getEdgeIfExists(mapType, currentMap, key); + AttributeMutationContext mapCtx = new AttributeMutationContext(ctx.getOp(), ctx.getReferringVertex(), attribute, entry.getValue(), + propertyName, mapType.getValueType(), existingEdge); - AttributeMutationContext mapCtx = new AttributeMutationContext(ctx.getOp(), ctx.getReferringVertex(), attribute, entry.getValue(), propertyName, mapType.getValueType(), existingEdge); + //Add/Update/Remove property value + Object newEntry = mapCollectionElementsToVertex(mapCtx, context); - //Add/Update/Remove property value - Object newEntry = mapCollectionElementsToVertex(mapCtx, context); - setMapValueProperty(mapType.getValueType(), ctx.getReferringVertex(), propertyName, newEntry); + if (isReference) { + AtlasEdge edge = (AtlasEdge) newEntry; + edge.setProperty(ATTRIBUTE_KEY_PROPERTY_KEY, key); + } else { + ctx.getReferringVertex().setProperty(propertyName, newEntry); + } - newMap.put(key, newEntry); + newMap.put(key, newEntry); - // If value type indicates this attribute is a reference, and the attribute has an inverse reference attribute, - // update the inverse reference value. - if (isReference && newEntry instanceof AtlasEdge && inverseRefAttribute != null) { - AtlasEdge newEdge = (AtlasEdge) newEntry; + // If value type indicates this attribute is a reference, and the attribute has an inverse reference attribute, + // update the inverse reference value. + if (isReference && newEntry instanceof AtlasEdge && inverseRefAttribute != null) { + AtlasEdge newEdge = (AtlasEdge) newEntry; - addInverseReference(context, inverseRefAttribute, newEdge, getRelationshipAttributes(ctx.getValue())); - } + addInverseReference(context, inverseRefAttribute, newEdge, getRelationshipAttributes(ctx.getValue())); } } + } - Map<String, Object> finalMap = removeUnusedMapEntries(attribute, ctx.getReferringVertex(), ctx.getVertexProperty(), currentMap, newMap); + Map<String, Object> finalMap = removeUnusedMapEntries(attribute, ctx.getReferringVertex(), ctx.getVertexProperty(), currentMap, newMap); - for (Object newEntry : newMap.values()) { - updateInConsistentOwnedMapVertices(ctx, mapType, newEntry); - } + for (Object newEntry : newMap.values()) { + updateInConsistentOwnedMapVertices(ctx, mapType, newEntry); + } + // for dereference on way out for primitive map types + if (!isReference) { Set<String> newKeys = new LinkedHashSet<>(newMap.keySet()); + newKeys.addAll(finalMap.keySet()); - // for dereference on way out - GraphHelper.setListProperty(ctx.getReferringVertex(), ctx.getVertexProperty(), new ArrayList<>(newKeys)); - } catch (AtlasException e) { - throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, e); + setListProperty(ctx.getReferringVertex(), ctx.getVertexProperty(), new ArrayList<>(newKeys)); } if (LOG.isDebugEnabled()) { @@ -870,16 +874,16 @@ public class EntityGraphMapper { List newElements = (List) ctx.getValue(); AtlasArrayType arrType = (AtlasArrayType) attribute.getAttributeType(); AtlasType elementType = arrType.getElementType(); - boolean isReference = AtlasGraphUtilsV1.isReference(elementType); + boolean isReference = isReference(elementType); AtlasAttribute inverseRefAttribute = attribute.getInverseRefAttribute(); Cardinality cardinality = attribute.getAttributeDef().getCardinality(); List<Object> newElementsCreated = new ArrayList<>(); List<Object> currentElements; - if (isRelationshipAttribute(attribute)) { - currentElements = getArrayElementsUsingRelationship(ctx.getReferringVertex(), attribute, elementType); + if (isReference) { + currentElements = (List) getCollectionElementsUsingRelationship(ctx.getReferringVertex(), attribute); } else { - currentElements = getArrayElementsProperty(elementType, ctx.getReferringVertex(), ctx.getVertexProperty()); + currentElements = (List) getArrayElementsProperty(elementType, ctx.getReferringVertex(), ctx.getVertexProperty()); } if (CollectionUtils.isNotEmpty(newElements)) { @@ -912,6 +916,15 @@ public class EntityGraphMapper { newElementsCreated.addAll(additionalEdges); } + // add index to attributes of array type + for (int index = 0; index < newElementsCreated.size(); index++) { + Object element = newElementsCreated.get(index); + + if (element instanceof AtlasEdge) { + AtlasGraphUtilsV1.setProperty((AtlasEdge) element, ATTRIBUTE_INDEX_PROPERTY_KEY, index); + } + } + // for dereference on way out setArrayElementsProperty(elementType, ctx.getReferringVertex(), ctx.getVertexProperty(), newElementsCreated); @@ -922,22 +935,6 @@ public class EntityGraphMapper { return newElementsCreated; } - private boolean isRelationshipAttribute(AtlasAttribute attribute) { - boolean ret = false; - - if (attribute != null) { - AtlasStructType structType = attribute.getDefinedInType(); - String attributeName = attribute.getName(); - - if (structType instanceof AtlasEntityType) { - ret = ((AtlasEntityType) structType).hasRelationshipAttribute(attributeName); - } - } - - return ret; - } - - private AtlasEdge createVertex(AtlasStruct struct, AtlasVertex referringVertex, String edgeLabel, EntityMutationContext context) throws AtlasBaseException { AtlasVertex vertex = createStructVertex(struct); @@ -1089,7 +1086,7 @@ public class EntityGraphMapper { public static Object getMapValueProperty(AtlasType elementType, AtlasVertex vertex, String vertexPropertyName) { - if (AtlasGraphUtilsV1.isReference(elementType)) { + if (isReference(elementType)) { return vertex.getProperty(vertexPropertyName, AtlasEdge.class); } else if (elementType instanceof AtlasArrayType) { return vertex.getProperty(vertexPropertyName, List.class); @@ -1102,7 +1099,7 @@ public class EntityGraphMapper { } private static void setMapValueProperty(AtlasType elementType, AtlasVertex vertex, String vertexPropertyName, Object value) { - if (AtlasGraphUtilsV1.isReference(elementType)) { + if (isReference(elementType)) { vertex.setPropertyFromElementId(vertexPropertyName, (AtlasEdge)value); } else { @@ -1113,16 +1110,18 @@ public class EntityGraphMapper { //Remove unused entries from map private Map<String, Object> removeUnusedMapEntries(AtlasAttribute attribute, AtlasVertex vertex, String propertyName, Map<String, Object> currentMap, Map<String, Object> newMap) - throws AtlasException, AtlasBaseException { + throws AtlasBaseException { + AtlasMapType mapType = (AtlasMapType) attribute.getAttributeType(); + boolean isReference = isReference(mapType.getValueType()); Map<String, Object> additionalMap = new HashMap<>(); for (String currentKey : currentMap.keySet()) { boolean shouldDeleteKey = !newMap.containsKey(currentKey); - if (AtlasGraphUtilsV1.isReference(mapType.getValueType())) { + if (isReference) { //Delete the edge reference if its not part of new edges created/updated - AtlasEdge currentEdge = (AtlasEdge)currentMap.get(currentKey); + AtlasEdge currentEdge = (AtlasEdge) currentMap.get(currentKey); if (!newMap.values().contains(currentEdge)) { boolean deleted = deleteHandler.deleteEdgeReference(currentEdge, mapType.getValueType().getTypeCategory(), attribute.isOwnedRef(), true, vertex); @@ -1134,8 +1133,9 @@ public class EntityGraphMapper { } } - if (shouldDeleteKey) { - String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(propertyName, GraphHelper.encodePropertyKey(currentKey)); + if (!isReference && shouldDeleteKey) { + String propertyNameForKey = getQualifiedNameForMapKey(propertyName, GraphHelper.encodePropertyKey(currentKey)); + GraphHelper.setProperty(vertex, propertyNameForKey, null); } } @@ -1146,7 +1146,7 @@ public class EntityGraphMapper { private static AtlasEdge getEdgeIfExists(AtlasMapType mapType, Map<String, Object> currentMap, String keyStr) { AtlasEdge ret = null; - if (AtlasGraphUtilsV1.isReference(mapType.getValueType())) { + if (isReference(mapType.getValueType())) { Object val = currentMap.get(keyStr); if (val != null) { @@ -1232,7 +1232,7 @@ public class EntityGraphMapper { } public static List<Object> getArrayElementsProperty(AtlasType elementType, AtlasVertex vertex, String vertexPropertyName) { - if (AtlasGraphUtilsV1.isReference(elementType)) { + if (isReference(elementType)) { return (List)vertex.getListProperty(vertexPropertyName, AtlasEdge.class); } else { @@ -1240,25 +1240,10 @@ public class EntityGraphMapper { } } - public static List<Object> getArrayElementsUsingRelationship(AtlasVertex vertex, AtlasAttribute attribute, AtlasType elementType) { - List<Object> ret = null; - - if (AtlasGraphUtilsV1.isReference(elementType)) { - - AtlasRelationshipEdgeDirection edgeDirection = attribute.getRelationshipEdgeDirection(); - String edgeLabel = attribute.getRelationshipEdgeLabel(); - - Iterator<AtlasEdge> edgesForLabel = GraphHelper.getEdgesForLabel(vertex, edgeLabel, edgeDirection); - - ret = IteratorUtils.toList(edgesForLabel); - } - return ret; - } - private AtlasEdge getEdgeAt(List<Object> currentElements, int index, AtlasType elemType) { AtlasEdge ret = null; - if (AtlasGraphUtilsV1.isReference(elemType)) { + if (isReference(elemType)) { if (currentElements != null && index < currentElements.size()) { ret = (AtlasEdge) currentElements.get(index); } @@ -1273,7 +1258,7 @@ public class EntityGraphMapper { if (CollectionUtils.isNotEmpty(currentEntries)) { AtlasType entryType = ((AtlasArrayType) attribute.getAttributeType()).getElementType(); - if (AtlasGraphUtilsV1.isReference(entryType)) { + if (isReference(entryType)) { Collection<AtlasEdge> edgesToRemove = CollectionUtils.subtract(currentEntries, newEntries); if (CollectionUtils.isNotEmpty(edgesToRemove)) { @@ -1296,15 +1281,11 @@ public class EntityGraphMapper { return Collections.emptyList(); } private void setArrayElementsProperty(AtlasType elementType, AtlasVertex vertex, String vertexPropertyName, List<Object> values) { - if (AtlasGraphUtilsV1.isReference(elementType)) { - GraphHelper.setListPropertyFromElementIds(vertex, vertexPropertyName, (List) values); - } - else { + if (!isReference(elementType)) { GraphHelper.setProperty(vertex, vertexPropertyName, values); } } - private AtlasEntityHeader constructHeader(AtlasEntity entity, final AtlasEntityType type, AtlasVertex vertex) { AtlasEntityHeader header = new AtlasEntityHeader(entity.getTypeName()); @@ -1324,9 +1305,9 @@ public class EntityGraphMapper { private void updateInConsistentOwnedMapVertices(AttributeMutationContext ctx, AtlasMapType mapType, Object val) { if (mapType.getValueType().getTypeCategory() == TypeCategory.OBJECT_ID_TYPE) { AtlasEdge edge = (AtlasEdge) val; - if (ctx.getAttribute().isOwnedRef() && - GraphHelper.getStatus(edge) == AtlasEntity.Status.DELETED && - GraphHelper.getStatus(edge.getInVertex()) == AtlasEntity.Status.DELETED) { + + if (ctx.getAttribute().isOwnedRef() && getStatus(edge) == DELETED && getStatus(edge.getInVertex()) == DELETED) { + //Resurrect the vertex and edge to ACTIVE state GraphHelper.setProperty(edge, STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name()); GraphHelper.setProperty(edge.getInVertex(), STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name()); @@ -1813,6 +1794,7 @@ public class EntityGraphMapper { currentEntityId = getIdFromOutVertex(currentEdge); } + return currentEntityId; } http://git-wip-us.apache.org/repos/asf/atlas/blob/6e7aa6ed/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java index 0d65471..183a2f6 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java @@ -54,6 +54,7 @@ import org.apache.atlas.type.AtlasTypeRegistry; import org.apache.atlas.type.AtlasTypeUtil; import org.apache.atlas.utils.AtlasJson; import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections.MapUtils; import org.apache.commons.lang3.StringUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -97,15 +98,19 @@ import static org.apache.atlas.repository.graph.GraphHelper.getAdjacentEdgesByLa import static org.apache.atlas.repository.graph.GraphHelper.getAllClassificationEdges; import static org.apache.atlas.repository.graph.GraphHelper.getAllTraitNames; import static org.apache.atlas.repository.graph.GraphHelper.getBlockedClassificationIds; +import static org.apache.atlas.repository.graph.GraphHelper.getArrayElementsProperty; import static org.apache.atlas.repository.graph.GraphHelper.getClassificationVertices; import static org.apache.atlas.repository.graph.GraphHelper.getGuid; import static org.apache.atlas.repository.graph.GraphHelper.getIncomingEdgesByLabel; +import static org.apache.atlas.repository.graph.GraphHelper.getPrimitiveMap; +import static org.apache.atlas.repository.graph.GraphHelper.getReferenceMap; import static org.apache.atlas.repository.graph.GraphHelper.getOutGoingEdgesByLabel; import static org.apache.atlas.repository.graph.GraphHelper.getPropagateTags; import static org.apache.atlas.repository.graph.GraphHelper.getRelationshipGuid; import static org.apache.atlas.repository.graph.GraphHelper.getTypeName; import static org.apache.atlas.repository.graph.GraphHelper.isPropagationEnabled; import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.getIdFromVertex; +import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.isReference; import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection; import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection.BOTH; import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection.IN; @@ -614,10 +619,10 @@ public final class EntityGraphRetriever { ret = mapVertexToObjectId(entityVertex, edgeLabel, null, entityExtInfo, isOwnedAttribute, edgeDirection); break; case ARRAY: - ret = mapVertexToArray(entityVertex, (AtlasArrayType) attrType, vertexPropertyName, entityExtInfo, isOwnedAttribute, edgeDirection); + ret = mapVertexToArray(entityVertex, vertexPropertyName, entityExtInfo, isOwnedAttribute, attribute); break; case MAP: - ret = mapVertexToMap(entityVertex, (AtlasMapType) attrType, vertexPropertyName, entityExtInfo, isOwnedAttribute, edgeDirection); + ret = mapVertexToMap(entityVertex, vertexPropertyName, entityExtInfo, isOwnedAttribute, attribute); break; case CLASSIFICATION: // do nothing @@ -627,45 +632,50 @@ public final class EntityGraphRetriever { return ret; } - private Map<String, Object> mapVertexToMap(AtlasVertex entityVertex, AtlasMapType atlasMapType, final String propertyName, - AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute, - AtlasRelationshipEdgeDirection edgeDirection) throws AtlasBaseException { + private Map<String, Object> mapVertexToMap(AtlasVertex entityVertex, final String propertyName, AtlasEntityExtInfo entityExtInfo, + boolean isOwnedAttribute, AtlasAttribute attribute) throws AtlasBaseException { - List<String> mapKeys = GraphHelper.getListProperty(entityVertex, propertyName); - - if (CollectionUtils.isEmpty(mapKeys)) { - return null; - } + Map<String, Object> ret = null; + AtlasMapType mapType = (AtlasMapType) attribute.getAttributeType(); + AtlasType mapValueType = mapType.getValueType(); if (LOG.isDebugEnabled()) { - LOG.debug("Mapping map attribute {} for vertex {}", atlasMapType.getTypeName(), entityVertex); + LOG.debug("Mapping map attribute {} for vertex {}", mapType.getTypeName(), entityVertex); } - Map<String, Object> ret = new HashMap<>(mapKeys.size()); - AtlasType mapValueType = atlasMapType.getValueType(); + if (isReference(mapValueType)) { + Map<String, Object> currentMap = getReferenceMap(entityVertex, attribute); - for (String mapKey : mapKeys) { - final String keyPropertyName = propertyName + "." + mapKey; - final String edgeLabel = EDGE_LABEL_PREFIX + keyPropertyName; - final Object keyValue = GraphHelper.getMapValueProperty(mapValueType, entityVertex, keyPropertyName); + if (MapUtils.isNotEmpty(currentMap)) { + ret = new HashMap<>(); - Object mapValue = mapVertexToCollectionEntry(entityVertex, mapValueType, keyValue, edgeLabel, - entityExtInfo, isOwnedAttribute, edgeDirection); - - if (mapValue != null) { - ret.put(mapKey, mapValue); + for (Map.Entry<String, Object> entry : currentMap.entrySet()) { + String mapKey = entry.getKey(); + Object keyValue = entry.getValue(); + Object mapValue = mapVertexToCollectionEntry(entityVertex, mapValueType, keyValue, attribute.getRelationshipEdgeLabel(), + entityExtInfo, isOwnedAttribute, attribute.getRelationshipEdgeDirection()); + if (mapValue != null) { + ret.put(mapKey, mapValue); + } + } } + } else { + ret = getPrimitiveMap(entityVertex, propertyName, mapValueType); + } + + if (MapUtils.isEmpty(ret)) { + ret = null; } return ret; } - private List<Object> mapVertexToArray(AtlasVertex entityVertex, AtlasArrayType arrayType, String propertyName, - AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute, - AtlasRelationshipEdgeDirection edgeDirection) throws AtlasBaseException { + private List<Object> mapVertexToArray(AtlasVertex entityVertex, String propertyName, AtlasEntityExtInfo entityExtInfo, + boolean isOwnedAttribute, AtlasAttribute attribute) throws AtlasBaseException { - AtlasType arrayElementType = arrayType.getElementType(); - List<Object> arrayElements = GraphHelper.getArrayElementsProperty(arrayElementType, entityVertex, propertyName); + AtlasArrayType arrayType = (AtlasArrayType) attribute.getAttributeType(); + AtlasType arrayElementType = arrayType.getElementType(); + List<Object> arrayElements = getArrayElementsProperty(arrayElementType, entityVertex, propertyName, attribute); if (CollectionUtils.isEmpty(arrayElements)) { return null; @@ -675,8 +685,9 @@ public final class EntityGraphRetriever { LOG.debug("Mapping array attribute {} for vertex {}", arrayElementType.getTypeName(), entityVertex); } - List arrValues = new ArrayList(arrayElements.size()); - String edgeLabel = EDGE_LABEL_PREFIX + propertyName; + List arrValues = new ArrayList(arrayElements.size()); + String edgeLabel = attribute.getRelationshipEdgeLabel(); + AtlasRelationshipEdgeDirection edgeDirection = attribute.getRelationshipEdgeDirection(); for (Object element : arrayElements) { // When internal types are deleted, sometimes the collection type attribute will contain a null value http://git-wip-us.apache.org/repos/asf/atlas/blob/6e7aa6ed/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasComplexAttributesTest.java ---------------------------------------------------------------------- diff --git a/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasComplexAttributesTest.java b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasComplexAttributesTest.java new file mode 100644 index 0000000..e494728 --- /dev/null +++ b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasComplexAttributesTest.java @@ -0,0 +1,396 @@ +/** + * 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.atlas.repository.store.graph.v1; + +import org.apache.atlas.TestModules; +import org.apache.atlas.TestUtilsV2; +import org.apache.atlas.model.instance.AtlasEntity; +import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo; +import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo; +import org.apache.atlas.model.instance.AtlasEntityHeader; +import org.apache.atlas.model.instance.AtlasObjectId; +import org.apache.atlas.model.instance.AtlasStruct; +import org.apache.atlas.model.instance.EntityMutationResponse; +import org.apache.atlas.model.typedef.AtlasTypesDef; +import org.apache.atlas.repository.graphdb.AtlasEdge; +import org.apache.atlas.repository.graphdb.AtlasEdgeDirection; +import org.apache.atlas.repository.graphdb.AtlasVertex; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Guice; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; + +import static org.apache.atlas.TestUtilsV2.ENTITY_TYPE; +import static org.apache.atlas.TestUtilsV2.ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR; +import static org.apache.atlas.TestUtilsV2.ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR_DELETE; +import static org.apache.atlas.TestUtilsV2.NAME; +import static org.apache.atlas.repository.graph.GraphHelper.getStatus; +import static org.apache.atlas.type.AtlasTypeUtil.getAtlasObjectId; +import static org.testng.Assert.assertNull; +import static org.testng.AssertJUnit.assertEquals; + +@Guice(modules = TestModules.TestOnlyModule.class) +public class AtlasComplexAttributesTest extends AtlasEntityTestBase { + private AtlasEntityWithExtInfo complexCollectionAttrEntity; + private AtlasEntityWithExtInfo complexCollectionAttrEntityForDelete; + + @BeforeClass + public void setUp() throws Exception { + super.setUp(); + + // create typeDefs + AtlasTypesDef[] testTypesDefs = new AtlasTypesDef[] { TestUtilsV2.defineTypeWithComplexCollectionAttributes() }; + createTypesDef(testTypesDefs); + + // create entity + complexCollectionAttrEntity = TestUtilsV2.createComplexCollectionAttrEntity(); + complexCollectionAttrEntityForDelete = TestUtilsV2.createComplexCollectionAttrEntity(); + } + + @Test + public void testCreateComplexAttributeEntity() throws Exception { + init(); + + EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(complexCollectionAttrEntity), false); + AtlasEntityHeader entityCreated = response.getFirstCreatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + + validateEntity(complexCollectionAttrEntity, getEntityFromStore(entityCreated)); + } + + @Test(dependsOnMethods = "testCreateComplexAttributeEntity") + public void testStructArray() throws Exception { + init(); + AtlasEntity complexEntity = getEntityFromStore(complexCollectionAttrEntity.getEntity().getGuid()); + AtlasEntitiesWithExtInfo complexEntitiesInfo = new AtlasEntitiesWithExtInfo(complexEntity); + + // Modify array of structs + List<AtlasStruct> structList = new ArrayList<>(Arrays.asList(new AtlasStruct("struct_type", "name", "structArray00"), + new AtlasStruct("struct_type", "name", "structArray11"), + new AtlasStruct("struct_type", "name", "structArray22"))); + complexEntity.setAttribute("listOfStructs", structList); + + EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + AtlasEntityHeader updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // add a new element to array of struct + init(); + structList.add(new AtlasStruct("struct_type", "name", "structArray33")); + complexEntity.setAttribute("listOfStructs", structList); + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // remove one of the struct values - structArray00 + init(); + structList.remove(0); + complexEntity.setAttribute("listOfStructs", structList); + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // Update struct value within array of struct + init(); + structList.get(0).setAttribute(NAME, "structArray11-edit"); + complexEntity.setAttribute("listOfStructs", structList); + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // add a repeated element to array of struct + init(); + structList.add(new AtlasStruct("struct_type", "name", "structArray33")); + complexEntity.setAttribute("listOfStructs", structList); + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // Remove all elements. Should set array attribute to null + init(); + structList.clear(); + complexEntity.setAttribute("listOfStructs", structList); + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + } + + @Test(dependsOnMethods = "testStructArray") + public void testEntityArray() throws Exception { + init(); + AtlasEntity complexEntity = getEntityFromStore(complexCollectionAttrEntity.getEntity().getGuid()); + AtlasEntitiesWithExtInfo complexEntitiesInfo = new AtlasEntitiesWithExtInfo(complexEntity); + + // Modify array of entities + AtlasEntity e0Array = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityArray00"); put("isReplicated", true); }}); + AtlasEntity e1Array = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityArray11"); put("isReplicated", false); }}); + AtlasEntity e2Array = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityArray22"); put("isReplicated", true); }}); + + List<AtlasObjectId> entityList = new ArrayList<>(Arrays.asList(getAtlasObjectId(e0Array), getAtlasObjectId(e1Array), getAtlasObjectId(e2Array))); + + complexEntity.setAttribute("listOfEntities", entityList); + complexEntitiesInfo.addReferredEntity(e0Array); + complexEntitiesInfo.addReferredEntity(e1Array); + complexEntitiesInfo.addReferredEntity(e2Array); + + init(); + EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + AtlasEntityHeader updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // add a new element to array of entities + init(); + AtlasEntity e3Array = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityArray33"); put("isReplicated", true); }}); + entityList.add(getAtlasObjectId(e3Array)); + complexEntity.setAttribute("listOfEntities", entityList); + complexEntitiesInfo.addReferredEntity(e3Array); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // remove one of the entity values - entityArray00 + init(); + entityList.remove(0); + complexEntity.setAttribute("listOfEntities", entityList); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // Update entity value within array of struct + init(); + e1Array.setAttribute(NAME, "entityArray11-edit"); + complexEntity.setAttribute("listOfEntities", entityList); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // add a repeated element to array of struct + init(); + AtlasEntity e3Array_duplicate = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityArray33"); put("isReplicated", true); }}); + entityList.add(getAtlasObjectId(e3Array_duplicate)); + complexEntity.setAttribute("listOfEntities", entityList); + complexEntitiesInfo.addReferredEntity(e3Array_duplicate); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // Remove all elements. Should set array attribute to null + init(); + entityList.clear(); + complexEntity.setAttribute("listOfEntities", entityList); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + } + + @Test(dependsOnMethods = "testEntityArray") + public void testStructMap() throws Exception { + init(); + AtlasEntity complexEntity = getEntityFromStore(complexCollectionAttrEntity.getEntity().getGuid()); + AtlasEntitiesWithExtInfo complexEntitiesInfo = new AtlasEntitiesWithExtInfo(complexEntity); + + // Modify map of structs + HashMap<String, AtlasStruct> structMap = new HashMap<String, AtlasStruct>() {{ + put("key00", new AtlasStruct("struct_type", "name", "structMap00")); + put("key11", new AtlasStruct("struct_type", "name", "structMap11")); + put("key22", new AtlasStruct("struct_type", "name", "structMap22")); }}; + + complexEntity.setAttribute("mapOfStructs", structMap); + + EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + AtlasEntityHeader updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // add a new element to map of struct - structMap6 + init(); + structMap.put("key33", new AtlasStruct("struct_type", "name", "structMap33")); + complexEntity.setAttribute("mapOfStructs", structMap); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // remove one of the entity values - structMap + init(); + structMap.remove("key00"); + complexEntity.setAttribute("mapOfStructs", structMap); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // Update struct value within map of struct + init(); + structMap.get("key11").setAttribute("name", "structMap11-edit"); + complexEntity.setAttribute("mapOfStructs", structMap); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // add a repeated element to array of struct + init(); + structMap.put("key33", new AtlasStruct("struct_type", "name", "structMap33")); + complexEntity.setAttribute("mapOfStructs", structMap); + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + // no update since duplicate entry + assertNull(updatedComplexEntity); + + // Remove all elements. Should set array attribute to null + init(); + structMap.clear(); + complexEntity.setAttribute("mapOfStructs", structMap); + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + } + + @Test(dependsOnMethods = "testStructMap") + public void testEntityMap() throws Exception { + init(); + AtlasEntity complexEntity = getEntityFromStore(complexCollectionAttrEntity.getEntity().getGuid()); + AtlasEntitiesWithExtInfo complexEntitiesInfo = new AtlasEntitiesWithExtInfo(complexEntity); + + // Modify map of entities + AtlasEntity e0MapValue = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityMapValue00"); put("isReplicated", false); }}); + AtlasEntity e1MapValue = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityMapValue11"); put("isReplicated", true); }}); + AtlasEntity e2MapValue = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityMapValue22"); put("isReplicated", false); }}); + + HashMap<String, Object> entityMap = new HashMap<String, Object>() {{ put("key00", getAtlasObjectId(e0MapValue)); + put("key11", getAtlasObjectId(e1MapValue)); + put("key22", getAtlasObjectId(e2MapValue)); }}; + complexEntity.setAttribute("mapOfEntities", entityMap); + complexEntitiesInfo.addReferredEntity(e0MapValue); + complexEntitiesInfo.addReferredEntity(e1MapValue); + complexEntitiesInfo.addReferredEntity(e2MapValue); + + init(); + EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + AtlasEntityHeader updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // add a new element to map of entities + init(); + AtlasEntity e3MapValue = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityMapValue33"); put("isReplicated", false); }}); + entityMap.put("key33", getAtlasObjectId(e3MapValue)); + complexEntity.setAttribute("mapOfEntities", entityMap); + complexEntitiesInfo.addReferredEntity(e3MapValue); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // remove one of the entity values - [key00 : entityMapValue00] + init(); + entityMap.remove("key00"); + complexEntity.setAttribute("mapOfEntities", entityMap); + complexEntitiesInfo.addReferredEntity(e3MapValue); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // Update entity value within map of entities + init(); + + AtlasEntity e1MapValueEdit = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityMapValue11-edit"); put("isReplicated", false); }}); + entityMap.clear(); + entityMap.put("key11", getAtlasObjectId(e1MapValueEdit)); + entityMap.put("key22", getAtlasObjectId(e2MapValue)); + entityMap.put("key33", getAtlasObjectId(e3MapValue)); + complexEntity.setAttribute("mapOfEntities", entityMap); + complexEntitiesInfo.addReferredEntity(e1MapValueEdit); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // add a repeated element to map of entities + init(); + e3MapValue = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityMapValue33"); put("isReplicated", false); }}); + entityMap.put("key33", getAtlasObjectId(e3MapValue)); + complexEntity.setAttribute("mapOfEntities", entityMap); + complexEntitiesInfo.addReferredEntity(e3MapValue); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + + // Remove all elements. Should set map attribute to null + init(); + entityMap.clear(); + complexEntity.setAttribute("mapOfEntities", entityMap); + complexEntitiesInfo.addReferredEntity(e3MapValue); + + response = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false); + updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity)); + } + + @Test(dependsOnMethods = "testEntityMap") + public void testDeleteEntityRemoveReferences() throws Exception { + init(); + + complexCollectionAttrEntityForDelete.getEntity().setAttribute(NAME, ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR_DELETE); + + EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(complexCollectionAttrEntityForDelete), false); + AtlasEntityHeader entityCreated = response.getFirstCreatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + + validateEntity(complexCollectionAttrEntityForDelete, getEntityFromStore(entityCreated)); + + // delete entity and check if referenced complex attribute edges are deleted + response = entityStore.deleteById(entityCreated.getGuid()); + + AtlasEntityHeader entityDeleted = response.getFirstDeletedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR); + + AtlasEntityWithExtInfo deletedEntityWithExtInfo = entityStore.getById(entityDeleted.getGuid()); + AtlasVertex deletedEntityVertex = AtlasGraphUtilsV1.findByGuid(entityDeleted.getGuid()); + Iterator<AtlasEdge> edges = deletedEntityVertex.getEdges(AtlasEdgeDirection.OUT).iterator(); + + // validate all attribute edges are deleted + while (edges != null && edges.hasNext()) { + assertEquals(getStatus(edges.next()), AtlasEntity.Status.DELETED); + } + + AtlasEntity deletedEntity = deletedEntityWithExtInfo.getEntity(); + List<AtlasObjectId> listOfEntities = (List<AtlasObjectId>) deletedEntity.getAttribute("listOfEntities"); + Map<String, AtlasObjectId> mapOfEntities = (Map<String, AtlasObjectId>) deletedEntity.getAttribute("mapOfEntities"); + + // validate entity attributes are deleted + for (AtlasObjectId o : listOfEntities) { + AtlasEntity entity = deletedEntityWithExtInfo.getEntity(o.getGuid()); + assertEquals(entity.getStatus(), AtlasEntity.Status.DELETED); + } + + for (AtlasObjectId o : mapOfEntities.values()) { + AtlasEntity entity = deletedEntityWithExtInfo.getEntity(o.getGuid()); + assertEquals(entity.getStatus(), AtlasEntity.Status.DELETED); + } + } +} \ No newline at end of file
