Repository: incubator-atlas Updated Branches: refs/heads/master f74e43c2b -> 02e4e86b5
http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/02e4e86b/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 157f8cd..1282be5 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 @@ -45,6 +45,7 @@ import org.apache.atlas.type.AtlasEntityType; import org.apache.atlas.type.AtlasMapType; import org.apache.atlas.type.AtlasStructType; import org.apache.atlas.type.AtlasStructType.AtlasAttribute; +import org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection; import org.apache.atlas.type.AtlasType; import org.apache.atlas.type.AtlasTypeRegistry; import org.apache.atlas.type.AtlasTypeUtil; @@ -63,7 +64,11 @@ import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.DE 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.repository.Constants.STATE_PROPERTY_KEY; +import static org.apache.atlas.repository.graph.GraphHelper.getTypeName; +import static org.apache.atlas.repository.graph.GraphHelper.isRelationshipEdge; import static org.apache.atlas.repository.graph.GraphHelper.string; +import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.getIdFromVertex; +import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection.IN; @Component public class EntityGraphMapper { @@ -144,6 +149,8 @@ public class EntityGraphMapper { mapAttributes(createdEntity, vertex, CREATE, context); + mapRelationshipAttributes(createdEntity, vertex, CREATE, context); + resp.addEntity(CREATE, constructHeader(createdEntity, entityType, vertex)); addClassifications(context, guid, createdEntity.getClassifications()); } @@ -157,6 +164,8 @@ public class EntityGraphMapper { mapAttributes(updatedEntity, vertex, UPDATE, context); + mapRelationshipAttributes(updatedEntity, vertex, UPDATE, context); + if (isPartialUpdate) { resp.addEntity(PARTIAL_UPDATE, constructHeader(updatedEntity, entityType, vertex)); } else { @@ -238,6 +247,7 @@ public class EntityGraphMapper { mapAttribute(attribute, attrValue, vertex, op, context); } + } else if (op.equals(UPDATE)) { for (String attrName : struct.getAttributes().keySet()) { AtlasAttribute attribute = structType.getAttribute(attrName); @@ -260,6 +270,41 @@ public class EntityGraphMapper { } } + private void mapRelationshipAttributes(AtlasEntity entity, AtlasVertex vertex, EntityOperation op, + EntityMutationContext context) throws AtlasBaseException { + if (LOG.isDebugEnabled()) { + LOG.debug("==> mapRelationshipAttributes({}, {})", op, entity.getTypeName()); + } + + if (MapUtils.isNotEmpty(entity.getRelationshipAttributes())) { + AtlasEntityType entityType = getEntityType(entity.getTypeName()); + + if (op.equals(CREATE)) { + for (AtlasAttribute attribute : entityType.getRelationshipAttributes().values()) { + Object attrValue = entity.getRelationshipAttribute(attribute.getName()); + + mapAttribute(attribute, attrValue, vertex, op, context); + } + + } else if (op.equals(UPDATE)) { + // relationship attributes mapping + for (AtlasAttribute attribute : entityType.getRelationshipAttributes().values()) { + if (attribute != null && entity.hasRelationshipAttribute(attribute.getName())) { + Object attrValue = entity.getRelationshipAttribute(attribute.getName()); + + mapAttribute(attribute, attrValue, vertex, op, context); + } + } + } + + updateModificationMetadata(vertex); + } + + if (LOG.isDebugEnabled()) { + LOG.debug("<== mapRelationshipAttributes({}, {})", op, entity.getTypeName()); + } + } + private void mapAttribute(AtlasAttribute attribute, Object attrValue, AtlasVertex vertex, EntityOperation op, EntityMutationContext context) throws AtlasBaseException { if (attrValue == null) { AtlasAttributeDef attributeDef = attribute.getAttributeDef(); @@ -309,14 +354,15 @@ public class EntityGraphMapper { } case OBJECT_ID_TYPE: { - String edgeLabel = ctx.getAttribute().getRelationshipEdgeLabel(); + String edgeLabel = ctx.getAttribute().getRelationshipEdgeLabel(); + AtlasRelationshipEdgeDirection edgeDirection = ctx.getAttribute().getRelationshipEdgeDirection(); // legacy case - if relationship attribute doesn't exist, use legacy edge label. if (StringUtils.isEmpty(edgeLabel)) { edgeLabel = AtlasGraphUtilsV1.getEdgeLabel(ctx.getVertexProperty()); } - AtlasEdge currentEdge = graphHelper.getEdgeForLabel(ctx.getReferringVertex(), edgeLabel); + AtlasEdge currentEdge = graphHelper.getEdgeForLabel(ctx.getReferringVertex(), edgeLabel, edgeDirection); AtlasEdge newEdge = null; if (ctx.getValue() != null) { @@ -328,14 +374,41 @@ public class EntityGraphMapper { newEdge = mapObjectIdValueUsingRelationship(ctx, context); - if (newEdge != null && ctx.getAttribute().getInverseRefAttribute() != null) { + // legacy case update inverse attribute + if (ctx.getAttribute().getInverseRefAttribute() != null) { // Update the inverse reference using relationship on the target entity addInverseReference(ctx.getAttribute().getInverseRefAttribute(), newEdge); } } + // created new relationship, + // record entity update on both vertices of the new relationship + if (currentEdge == null && newEdge != null) { + + // based on relationship edge direction record update only on attribute vertex + if (edgeDirection == IN) { + recordEntityUpdate(newEdge.getOutVertex()); + + } else { + recordEntityUpdate(newEdge.getInVertex()); + } + } + + // update references, if current and new edge don't match + // record entity update on new reference and delete(edge) old reference. if (currentEdge != null && !currentEdge.equals(newEdge)) { - deleteHandler.deleteEdgeReference(currentEdge, ctx.getAttrType().getTypeCategory(), ctx.getAttribute().isOwnedRef(), true); + + //record entity update on new edge + if (isRelationshipEdge(newEdge)) { + AtlasVertex attrVertex = context.getDiscoveryContext().getResolvedEntityVertex(getGuid(ctx.getValue())); + + recordEntityUpdate(attrVertex); + updateModificationMetadata(attrVertex); + } + + //delete old reference + deleteHandler.deleteEdgeReference(currentEdge, ctx.getAttrType().getTypeCategory(), ctx.getAttribute().isOwnedRef(), + true, ctx.getAttribute().getRelationshipEdgeDirection()); } return newEdge; @@ -402,7 +475,7 @@ public class EntityGraphMapper { if (inverseUpdated) { updateModificationMetadata(inverseVertex); - AtlasObjectId inverseEntityId = new AtlasObjectId(AtlasGraphUtilsV1.getIdFromVertex(inverseVertex), inverseType.getTypeName()); + AtlasObjectId inverseEntityId = new AtlasObjectId(getIdFromVertex(inverseVertex), inverseType.getTypeName()); RequestContextV1.get().recordEntityUpdate(inverseEntityId); } } @@ -424,7 +497,7 @@ public class EntityGraphMapper { if (entityType.hasRelationshipAttribute(inverseAttributeName)) { String relationshipName = graphHelper.getRelationshipDefName(inverseVertex, entityType, inverseAttributeName); - ret = getOrCreateRelationship(inverseVertex, vertex, relationshipName); + ret = getOrCreateRelationship(inverseVertex, vertex, relationshipName, inverseAttribute); } else { if (LOG.isDebugEnabled()) { @@ -576,24 +649,47 @@ public class EntityGraphMapper { String attributeName = ctx.getAttribute().getName(); AtlasType type = typeRegistry.getType(AtlasGraphUtilsV1.getTypeName(entityVertex)); + AtlasRelationshipEdgeDirection edgeDirection = ctx.getAttribute().getRelationshipEdgeDirection(); + String edgeLabel = ctx.getAttribute().getRelationshipEdgeLabel(); + if (type instanceof AtlasEntityType) { AtlasEntityType entityType = (AtlasEntityType) type; // use relationship to create/update edges if (entityType.hasRelationshipAttribute(attributeName)) { if (ctx.getCurrentEdge() != null) { - ret = updateRelationship(ctx.getCurrentEdge(), attributeVertex); + ret = updateRelationship(ctx.getCurrentEdge(), attributeVertex, edgeDirection, ctx.getAttribute()); + + recordEntityUpdate(attributeVertex); } else { - String relationshipName = graphHelper.getRelationshipDefName(entityVertex, entityType, attributeName); - ret = getOrCreateRelationship(entityVertex, attributeVertex, relationshipName); - } + String relationshipName = graphHelper.getRelationshipDefName(entityVertex, entityType, attributeName); + AtlasVertex fromVertex; + AtlasVertex toVertex; + + if (edgeDirection == IN) { + fromVertex = attributeVertex; + toVertex = entityVertex; + + } else { + fromVertex = entityVertex; + toVertex = attributeVertex; + } + boolean relationshipExists = isRelationshipExists(fromVertex, toVertex, edgeLabel); + + ret = getOrCreateRelationship(fromVertex, toVertex, relationshipName, ctx.getAttribute()); + // if relationship did not exist before and new relationship was created + // record entity update on both relationship vertices + if (!relationshipExists) { + recordEntityUpdate(attributeVertex); + } + } } else { // use legacy way to create/update edges if (LOG.isDebugEnabled()) { - LOG.debug("No RelationshipDef defined between {} and {} on attribute: {}", AtlasGraphUtilsV1.getTypeName(entityVertex), - AtlasGraphUtilsV1.getTypeName(attributeVertex), attributeName); + LOG.debug("No RelationshipDef defined between {} and {} on attribute: {}", getTypeName(entityVertex), + getTypeName(attributeVertex), attributeName); } ret = mapObjectIdValue(ctx, context); @@ -728,6 +824,21 @@ 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); @@ -767,6 +878,16 @@ public class EntityGraphMapper { return (AtlasStructType)objType; } + private AtlasEntityType getEntityType(String typeName) throws AtlasBaseException { + AtlasType objType = typeRegistry.getType(typeName); + + if (!(objType instanceof AtlasEntityType)) { + throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, typeName); + } + + return (AtlasEntityType)objType; + } + private Object mapCollectionElementsToVertex(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException { switch(ctx.getAttrType().getTypeCategory()) { case PRIMITIVE: @@ -918,8 +1039,8 @@ public class EntityGraphMapper { // Update edge if it exists AtlasVertex currentVertex = currentEdge.getInVertex(); - String currentEntityId = AtlasGraphUtilsV1.getIdFromVertex(currentVertex); - String newEntityId = AtlasGraphUtilsV1.getIdFromVertex(entityVertex); + String currentEntityId = getIdFromVertex(currentVertex); + String newEntityId = getIdFromVertex(entityVertex); AtlasEdge newEdge = currentEdge; if (!currentEntityId.equals(newEntityId)) { @@ -936,16 +1057,25 @@ public class EntityGraphMapper { return newEdge; } - private AtlasEdge updateRelationship(AtlasEdge currentEdge, final AtlasVertex entityVertex) throws AtlasBaseException { + private AtlasEdge updateRelationship(AtlasEdge currentEdge, final AtlasVertex newEntityVertex, + AtlasRelationshipEdgeDirection edgeDirection, AtlasAttribute attribute) + throws AtlasBaseException { if (LOG.isDebugEnabled()) { - LOG.debug("Updating entity reference using relationship {} for reference attribute {}", AtlasGraphUtilsV1.getTypeName(entityVertex)); + LOG.debug("Updating entity reference using relationship {} for reference attribute {}", getTypeName(newEntityVertex)); } - String currentEntityId = AtlasGraphUtilsV1.getIdFromVertex(currentEdge.getInVertex()); - String newEntityId = AtlasGraphUtilsV1.getIdFromVertex(entityVertex); - AtlasEdge ret = currentEdge; + // Max's manager updated from Jane to Julius (Max.manager --> Jane.subordinates) + // manager attribute (OUT direction), current manager vertex (Jane) (IN vertex) - if (!currentEntityId.equals(newEntityId)) { + // Max's mentor updated from John to Jane (John.mentee --> Max.mentor) + // mentor attribute (IN direction), current mentee vertex (John) (OUT vertex) + String currentEntityId = (edgeDirection == IN) ? getIdFromVertex(currentEdge.getOutVertex()) : + getIdFromVertex(currentEdge.getInVertex()); + + String newEntityId = getIdFromVertex(newEntityVertex); + AtlasEdge ret = currentEdge; + + if (!currentEntityId.equals(newEntityId) && newEntityVertex != null) { // create a new relationship edge to the new attribute vertex from the instance String relationshipName = AtlasGraphUtilsV1.getTypeName(currentEdge); @@ -953,7 +1083,8 @@ public class EntityGraphMapper { relationshipName = currentEdge.getLabel(); } - ret = getOrCreateRelationship(currentEdge.getOutVertex(), entityVertex, relationshipName); + ret = (edgeDirection == IN) ? getOrCreateRelationship(newEntityVertex, currentEdge.getInVertex(), relationshipName, attribute) : + getOrCreateRelationship(currentEdge.getOutVertex(), newEntityVertex, relationshipName, attribute); } return ret; @@ -983,8 +1114,7 @@ public class EntityGraphMapper { //Removes unused edges from the old collection, compared to the new collection private List<AtlasEdge> removeUnusedArrayEntries(AtlasAttribute attribute, List<AtlasEdge> currentEntries, List<AtlasEdge> newEntries) throws AtlasBaseException { if (CollectionUtils.isNotEmpty(currentEntries)) { - AtlasStructType entityType = attribute.getDefinedInType(); - AtlasType entryType = ((AtlasArrayType)attribute.getAttributeType()).getElementType(); + AtlasType entryType = ((AtlasArrayType) attribute.getAttributeType()).getElementType(); if (AtlasGraphUtilsV1.isReference(entryType)) { Collection<AtlasEdge> edgesToRemove = CollectionUtils.subtract(currentEntries, newEntries); @@ -993,7 +1123,8 @@ public class EntityGraphMapper { List<AtlasEdge> additionalElements = new ArrayList<>(); for (AtlasEdge edge : edgesToRemove) { - boolean deleted = deleteHandler.deleteEdgeReference(edge, entryType.getTypeCategory(), attribute.isOwnedRef(), true); + boolean deleted = deleteHandler.deleteEdgeReference(edge, entryType.getTypeCategory(), attribute.isOwnedRef(), + true, attribute.getRelationshipEdgeDirection()); if (!deleted) { additionalElements.add(edge); @@ -1021,7 +1152,7 @@ public class EntityGraphMapper { private AtlasEntityHeader constructHeader(AtlasEntity entity, final AtlasEntityType type, AtlasVertex vertex) { AtlasEntityHeader header = new AtlasEntityHeader(entity.getTypeName()); - header.setGuid(AtlasGraphUtilsV1.getIdFromVertex(vertex)); + header.setGuid(getIdFromVertex(vertex)); for (AtlasAttribute attribute : type.getUniqAttributes().values()) { header.setAttribute(attribute.getName(), entity.getAttribute(attribute.getName())); @@ -1148,7 +1279,7 @@ public class EntityGraphMapper { for (String classificationName : classificationNames) { try { - final String entityTypeName = GraphHelper.getTypeName(instanceVertex); + final String entityTypeName = getTypeName(instanceVertex); String relationshipLabel = GraphHelper.getTraitLabel(entityTypeName, classificationName); AtlasEdge edge = graphHelper.getEdgeForLabel(instanceVertex, relationshipLabel); if (edge != null) { @@ -1182,20 +1313,20 @@ public class EntityGraphMapper { } } - private AtlasEdge getOrCreateRelationship(AtlasVertex end1Vertex, AtlasVertex end2Vertex, String relationshipName) throws AtlasBaseException { - AtlasEdge ret = null; - AtlasObjectId end1 = new AtlasObjectId(AtlasGraphUtilsV1.getIdFromVertex(end1Vertex), AtlasGraphUtilsV1.getTypeName(end1Vertex)); - AtlasObjectId end2 = new AtlasObjectId(AtlasGraphUtilsV1.getIdFromVertex(end2Vertex), AtlasGraphUtilsV1.getTypeName(end2Vertex)); - AtlasRelationship relationship = relationshipStore.getOrCreate(new AtlasRelationship(relationshipName, end1, end2)); + private AtlasEdge getOrCreateRelationship(AtlasVertex end1Vertex, AtlasVertex end2Vertex, String relationshipName, AtlasAttribute attribute) throws AtlasBaseException { + AtlasEdge ret = null; + AtlasObjectId end1 = new AtlasObjectId(getIdFromVertex(end1Vertex), AtlasGraphUtilsV1.getTypeName(end1Vertex)); + AtlasObjectId end2 = new AtlasObjectId(getIdFromVertex(end2Vertex), AtlasGraphUtilsV1.getTypeName(end2Vertex)); + AtlasRelationship relationship = relationshipStore.getOrCreate(new AtlasRelationship(relationshipName, end1, end2)); // return newly created AtlasEdge - // if multiple edges are returned, compare using id to pick the right one + // if multiple edges are returned, compare using guid to pick the right one Iterator<AtlasEdge> outEdges = graphHelper.getOutGoingEdgesByLabel(end1Vertex, relationship.getLabel()); while (outEdges.hasNext()) { AtlasEdge edge = outEdges.next(); - if (AtlasGraphUtilsV1.getIdFromVertex(end2Vertex).equals(AtlasGraphUtilsV1.getIdFromVertex(edge.getInVertex()))) { + if (getIdFromVertex(end2Vertex).equals(getIdFromVertex(edge.getInVertex()))) { ret = edge; break; } @@ -1203,4 +1334,47 @@ public class EntityGraphMapper { return ret; } + + private boolean isRelationshipExists(AtlasVertex fromVertex, AtlasVertex toVertex, String edgeLabel) { + boolean ret = false; + Iterator<AtlasEdge> edges = graphHelper.getOutGoingEdgesByLabel(fromVertex, edgeLabel); + + while (edges != null && edges.hasNext()) { + AtlasEdge edge = edges.next(); + AtlasVertex inVertex = edge.getInVertex(); + + if (inVertex != null && StringUtils.equals(getIdFromVertex(inVertex), getIdFromVertex(toVertex))) { + ret = true; + } + } + + return ret; + } + + private void recordEntityUpdate(AtlasVertex vertex) { + AtlasObjectId objectId = new AtlasObjectId(GraphHelper.getGuid(vertex), GraphHelper.getTypeName(vertex)); + RequestContextV1 req = RequestContextV1.get(); + + if (!objectIdsContain(req.getUpdatedEntityIds(), objectId)) { + req.recordEntityUpdate(objectId); + } + } + + private boolean objectIdsContain(Collection<AtlasObjectId> objectIds, AtlasObjectId objectId) { + boolean ret = false; + + if (objectIds != null && objectIds.isEmpty()) { + ret = false; + + } else { + for (AtlasObjectId id : objectIds) { + if (StringUtils.equals(id.getGuid(), objectId.getGuid())) { + ret = true; + break; + } + } + } + + return ret; + } } http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/02e4e86b/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 f4257be..31fc837 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 @@ -69,6 +69,7 @@ import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_SHORT; import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_STRING; import static org.apache.atlas.repository.graph.GraphHelper.EDGE_LABEL_PREFIX; import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.getIdFromVertex; +import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection; public final class EntityGraphRetriever { @@ -136,6 +137,10 @@ public final class EntityGraphRetriever { return toAtlasEntityHeader(entityVertex, Collections.<String>emptySet()); } + public AtlasEntityHeader toAtlasEntityHeader(AtlasVertex atlasVertex, Set<String> attributes) throws AtlasBaseException { + return atlasVertex != null ? mapVertexToAtlasEntityHeader(atlasVertex, attributes) : null; + } + private AtlasVertex getEntityVertex(String guid) throws AtlasBaseException { AtlasVertex ret = AtlasGraphUtilsV1.findByGuid(guid); @@ -188,7 +193,7 @@ public final class EntityGraphRetriever { mapAttributes(entityVertex, entity, entityExtInfo); - mapRelationshipAttributes(entityVertex, entity, entityExtInfo); + mapRelationshipAttributes(entityVertex, entity); mapClassifications(entityVertex, entity, entityExtInfo); } @@ -300,23 +305,6 @@ public final class EntityGraphRetriever { } } - private void mapRelationshipAttributes(AtlasVertex entityVertex, AtlasEntity entity, AtlasEntityExtInfo entityExtInfo) throws AtlasBaseException { - AtlasType objType = typeRegistry.getType(entity.getTypeName()); - - if (!(objType instanceof AtlasEntityType)) { - throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, entity.getTypeName()); - } - - AtlasEntityType entityType = (AtlasEntityType) objType; - - for (AtlasAttribute attribute : entityType.getRelationshipAttributes().values()) { - - Object attrValue = mapVertexToRelationshipAttribute(entityVertex, entityType, attribute, entityExtInfo); - - entity.addRelationshipAttribute(attribute.getName(), attrValue); - } - } - public List<AtlasClassification> getClassifications(String guid) throws AtlasBaseException { AtlasVertex instanceVertex = AtlasGraphUtilsV1.findByGuid(guid); @@ -401,6 +389,7 @@ public final class EntityGraphRetriever { String vertexPropertyName = attribute.getQualifiedName(); String edgeLabel = EDGE_LABEL_PREFIX + vertexPropertyName; boolean isOwnedAttribute = attribute.isOwnedRef(); + AtlasRelationshipEdgeDirection edgeDirection = attribute.getRelationshipEdgeDirection(); if (LOG.isDebugEnabled()) { LOG.debug("Mapping vertex {} to atlas entity {}.{}", entityVertex, attribute.getDefinedInDef().getName(), attribute.getName()); @@ -417,13 +406,13 @@ public final class EntityGraphRetriever { ret = mapVertexToStruct(entityVertex, edgeLabel, null, entityExtInfo); break; case OBJECT_ID_TYPE: - ret = mapVertexToObjectId(entityVertex, edgeLabel, null, entityExtInfo, isOwnedAttribute); + ret = mapVertexToObjectId(entityVertex, edgeLabel, null, entityExtInfo, isOwnedAttribute, edgeDirection); break; case ARRAY: - ret = mapVertexToArray(entityVertex, (AtlasArrayType) attrType, vertexPropertyName, entityExtInfo, isOwnedAttribute); + ret = mapVertexToArray(entityVertex, (AtlasArrayType) attrType, vertexPropertyName, entityExtInfo, isOwnedAttribute, edgeDirection); break; case MAP: - ret = mapVertexToMap(entityVertex, (AtlasMapType) attrType, vertexPropertyName, entityExtInfo, isOwnedAttribute); + ret = mapVertexToMap(entityVertex, (AtlasMapType) attrType, vertexPropertyName, entityExtInfo, isOwnedAttribute, edgeDirection); break; case CLASSIFICATION: // do nothing @@ -433,51 +422,10 @@ public final class EntityGraphRetriever { return ret; } - private Object mapVertexToRelationshipAttribute(AtlasVertex entityVertex, AtlasEntityType entityType, AtlasAttribute attribute, - AtlasEntityExtInfo entityExtInfo) throws AtlasBaseException { - Object ret = null; - AtlasRelationshipDef relationshipDef = graphHelper.getRelationshipDef(entityVertex, entityType, attribute.getName()); - - if (relationshipDef == null) { - throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_INVALID, "relationshipDef is null"); - } - - AtlasRelationshipEndDef endDef1 = relationshipDef.getEndDef1(); - AtlasRelationshipEndDef endDef2 = relationshipDef.getEndDef2(); - AtlasEntityType endDef1Type = typeRegistry.getEntityTypeByName(endDef1.getType()); - AtlasEntityType endDef2Type = typeRegistry.getEntityTypeByName(endDef2.getType()); - AtlasRelationshipEndDef attributeEndDef = null; - - if (endDef1Type.isTypeOrSuperTypeOf(entityType.getTypeName()) && StringUtils.equals(endDef1.getName(), attribute.getName())) { - attributeEndDef = endDef1; - - } else if (endDef2Type.isTypeOrSuperTypeOf(entityType.getTypeName()) && StringUtils.equals(endDef2.getName(), attribute.getName())) { - attributeEndDef = endDef2; - } - - if (attributeEndDef == null) { - throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_INVALID, relationshipDef.toString()); - } - - String relationshipLabel = attribute.getRelationshipEdgeLabel(); - - switch (attributeEndDef.getCardinality()) { - case SINGLE: - ret = mapVertexToObjectId(entityVertex, relationshipLabel, null, entityExtInfo, attributeEndDef.getIsContainer()); - break; - - case LIST: - case SET: - ret = mapVertexToRelationshipArrayAttribute(entityVertex, (AtlasArrayType) attribute.getAttributeType(), relationshipLabel, - entityExtInfo, attributeEndDef.getIsContainer()); - break; - } - - return ret; - } - private Map<String, Object> mapVertexToMap(AtlasVertex entityVertex, AtlasMapType atlasMapType, final String propertyName, - AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute) throws AtlasBaseException { + AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute, + AtlasRelationshipEdgeDirection edgeDirection) throws AtlasBaseException { + List<String> mapKeys = GraphHelper.getListProperty(entityVertex, propertyName); if (CollectionUtils.isEmpty(mapKeys)) { @@ -496,7 +444,9 @@ public final class EntityGraphRetriever { final String edgeLabel = EDGE_LABEL_PREFIX + keyPropertyName; final Object keyValue = GraphHelper.getMapValueProperty(mapValueType, entityVertex, keyPropertyName); - Object mapValue = mapVertexToCollectionEntry(entityVertex, mapValueType, keyValue, edgeLabel, entityExtInfo, isOwnedAttribute); + Object mapValue = mapVertexToCollectionEntry(entityVertex, mapValueType, keyValue, edgeLabel, + entityExtInfo, isOwnedAttribute, edgeDirection); + if (mapValue != null) { ret.put(mapKey, mapValue); } @@ -506,7 +456,9 @@ public final class EntityGraphRetriever { } private List<Object> mapVertexToArray(AtlasVertex entityVertex, AtlasArrayType arrayType, String propertyName, - AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute) throws AtlasBaseException { + AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute, + AtlasRelationshipEdgeDirection edgeDirection) throws AtlasBaseException { + AtlasType arrayElementType = arrayType.getElementType(); List<Object> arrayElements = GraphHelper.getArrayElementsProperty(arrayElementType, entityVertex, propertyName); @@ -522,8 +474,8 @@ public final class EntityGraphRetriever { String edgeLabel = EDGE_LABEL_PREFIX + propertyName; for (Object element : arrayElements) { - Object arrValue = mapVertexToCollectionEntry(entityVertex, arrayElementType, element, - edgeLabel, entityExtInfo, isOwnedAttribute); + Object arrValue = mapVertexToCollectionEntry(entityVertex, arrayElementType, element, edgeLabel, + entityExtInfo, isOwnedAttribute, edgeDirection); if (arrValue != null) { arrValues.add(arrValue); @@ -533,42 +485,9 @@ public final class EntityGraphRetriever { return arrValues; } - private List<Object> mapVertexToRelationshipArrayAttribute(AtlasVertex entityVertex, AtlasArrayType arrayType, - String relationshipName, AtlasEntityExtInfo entityExtInfo, - boolean isContainer) throws AtlasBaseException { - - Iterator<AtlasEdge> relationshipEdges = graphHelper.getBothEdgesByLabel(entityVertex, relationshipName); - AtlasType arrayElementType = arrayType.getElementType(); - List<AtlasEdge> arrayElements = new ArrayList<>(); - - if (LOG.isDebugEnabled()) { - LOG.debug("Mapping array attribute {} for vertex {}", arrayElementType.getTypeName(), entityVertex); - } - - while (relationshipEdges.hasNext()) { - arrayElements.add(relationshipEdges.next()); - } - - if (CollectionUtils.isEmpty(arrayElements)) { - return null; - } - - List arrValues = new ArrayList(arrayElements.size()); - - for (Object element : arrayElements) { - Object arrValue = mapVertexToCollectionEntry(entityVertex, arrayElementType, element, relationshipName, - entityExtInfo, isContainer); - - if (arrValue != null) { - arrValues.add(arrValue); - } - } - - return arrValues; - } - - private Object mapVertexToCollectionEntry(AtlasVertex entityVertex, AtlasType arrayElement, Object value, String edgeLabel, - AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute) throws AtlasBaseException { + private Object mapVertexToCollectionEntry(AtlasVertex entityVertex, AtlasType arrayElement, Object value, + String edgeLabel, AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute, + AtlasRelationshipEdgeDirection edgeDirection) throws AtlasBaseException { Object ret = null; switch (arrayElement.getTypeCategory()) { @@ -587,7 +506,7 @@ public final class EntityGraphRetriever { break; case OBJECT_ID_TYPE: - ret = mapVertexToObjectId(entityVertex, edgeLabel, (AtlasEdge) value, entityExtInfo, isOwnedAttribute); + ret = mapVertexToObjectId(entityVertex, edgeLabel, (AtlasEdge) value, entityExtInfo, isOwnedAttribute, edgeDirection); break; default: @@ -646,11 +565,12 @@ public final class EntityGraphRetriever { } private AtlasObjectId mapVertexToObjectId(AtlasVertex entityVertex, String edgeLabel, AtlasEdge edge, - AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute) throws AtlasBaseException { + AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute, + AtlasRelationshipEdgeDirection edgeDirection) throws AtlasBaseException { AtlasObjectId ret = null; if (edge == null) { - edge = graphHelper.getEdgeForLabel(entityVertex, edgeLabel); + edge = graphHelper.getEdgeForLabel(entityVertex, edgeLabel, edgeDirection); } if (GraphHelper.elementExists(edge)) { @@ -697,7 +617,102 @@ public final class EntityGraphRetriever { return vertex != null && attribute != null ? mapVertexToAttribute(vertex, attribute, null) : null; } - public AtlasEntityHeader toAtlasEntityHeader(AtlasVertex atlasVertex, Set<String> attributes) throws AtlasBaseException { - return atlasVertex != null ? mapVertexToAtlasEntityHeader(atlasVertex, attributes) : null; + private void mapRelationshipAttributes(AtlasVertex entityVertex, AtlasEntity entity) throws AtlasBaseException { + AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName()); + + if (entityType == null) { + throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, entity.getTypeName()); + } + + for (AtlasAttribute attribute : entityType.getRelationshipAttributes().values()) { + Object attrValue = mapVertexToRelationshipAttribute(entityVertex, entityType, attribute); + + entity.setRelationshipAttribute(attribute.getName(), attrValue); + } + } + + private Object mapVertexToRelationshipAttribute(AtlasVertex entityVertex, AtlasEntityType entityType, AtlasAttribute attribute) throws AtlasBaseException { + Object ret = null; + AtlasRelationshipDef relationshipDef = graphHelper.getRelationshipDef(entityVertex, entityType, attribute.getName()); + + if (relationshipDef == null) { + throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_INVALID, "relationshipDef is null"); + } + + AtlasRelationshipEndDef endDef1 = relationshipDef.getEndDef1(); + AtlasRelationshipEndDef endDef2 = relationshipDef.getEndDef2(); + AtlasEntityType endDef1Type = typeRegistry.getEntityTypeByName(endDef1.getType()); + AtlasEntityType endDef2Type = typeRegistry.getEntityTypeByName(endDef2.getType()); + AtlasRelationshipEndDef attributeEndDef = null; + + if (endDef1Type.isTypeOrSuperTypeOf(entityType.getTypeName()) && StringUtils.equals(endDef1.getName(), attribute.getName())) { + attributeEndDef = endDef1; + } else if (endDef2Type.isTypeOrSuperTypeOf(entityType.getTypeName()) && StringUtils.equals(endDef2.getName(), attribute.getName())) { + attributeEndDef = endDef2; + } + + if (attributeEndDef == null) { + throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_INVALID, relationshipDef.toString()); + } + + switch (attributeEndDef.getCardinality()) { + case SINGLE: + ret = mapRelatedVertexToObjectId(entityVertex, attribute); + break; + + case LIST: + case SET: + ret = mapRelationshipArrayAttribute(entityVertex, attribute); + break; + } + + return ret; + } + + private AtlasObjectId mapRelatedVertexToObjectId(AtlasVertex entityVertex, AtlasAttribute attribute) throws AtlasBaseException { + AtlasEdge edge = graphHelper.getEdgeForLabel(entityVertex, attribute.getRelationshipEdgeLabel(), attribute.getRelationshipEdgeDirection()); + + return mapRelatedVertexToObjectId(entityVertex, edge); + } + + private List<AtlasObjectId> mapRelationshipArrayAttribute(AtlasVertex entityVertex, AtlasAttribute attribute) throws AtlasBaseException { + List<AtlasObjectId> ret = new ArrayList<>(); + Iterator<AtlasEdge> edges = null; + + if (attribute.getRelationshipEdgeDirection() == AtlasRelationshipEdgeDirection.IN) { + edges = graphHelper.getIncomingEdgesByLabel(entityVertex, attribute.getRelationshipEdgeLabel()); + } else if (attribute.getRelationshipEdgeDirection() == AtlasRelationshipEdgeDirection.OUT) { + edges = graphHelper.getOutGoingEdgesByLabel(entityVertex, attribute.getRelationshipEdgeLabel()); + } + + if (edges != null) { + while (edges.hasNext()) { + AtlasEdge relationshipEdge = edges.next(); + + AtlasObjectId objectId = mapRelatedVertexToObjectId(entityVertex, relationshipEdge); + + ret.add(objectId); + } + } + + return ret; + } + + private AtlasObjectId mapRelatedVertexToObjectId(AtlasVertex entityVertex, AtlasEdge edge) throws AtlasBaseException { + AtlasObjectId ret = null; + + if (GraphHelper.elementExists(edge)) { + AtlasVertex referenceVertex = edge.getInVertex(); + + if (StringUtils.equals(getIdFromVertex(referenceVertex), getIdFromVertex(entityVertex))) { + referenceVertex = edge.getOutVertex(); + } + + if (referenceVertex != null) { + ret = new AtlasObjectId(GraphHelper.getGuid(referenceVertex), GraphHelper.getTypeName(referenceVertex)); + } + } + + return ret; } } http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/02e4e86b/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java ---------------------------------------------------------------------- diff --git a/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java b/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java index 404225c..7210799 100644 --- a/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java +++ b/repository/src/test/java/org/apache/atlas/repository/impexp/ImportServiceTest.java @@ -18,6 +18,7 @@ package org.apache.atlas.repository.impexp; import com.google.inject.Inject; +import org.apache.atlas.RequestContextV1; import org.apache.atlas.TestModules; import org.apache.atlas.exception.AtlasBaseException; import org.apache.atlas.model.impexp.AtlasImportRequest; http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/02e4e86b/repository/src/test/java/org/apache/atlas/repository/impexp/ZipFileResourceTestUtils.java ---------------------------------------------------------------------- diff --git a/repository/src/test/java/org/apache/atlas/repository/impexp/ZipFileResourceTestUtils.java b/repository/src/test/java/org/apache/atlas/repository/impexp/ZipFileResourceTestUtils.java index d901731..82692cf 100644 --- a/repository/src/test/java/org/apache/atlas/repository/impexp/ZipFileResourceTestUtils.java +++ b/repository/src/test/java/org/apache/atlas/repository/impexp/ZipFileResourceTestUtils.java @@ -18,6 +18,7 @@ package org.apache.atlas.repository.impexp; import com.google.common.collect.Sets; +import org.apache.atlas.RequestContextV1; import org.apache.atlas.exception.AtlasBaseException; import org.apache.atlas.model.impexp.AtlasExportResult; import org.apache.atlas.model.impexp.AtlasImportRequest; @@ -149,6 +150,8 @@ public class ZipFileResourceTestUtils { AtlasExportResult exportResult = zipSource.getExportResult(); List<String> creationOrder = zipSource.getCreationOrder(); + RequestContextV1.clear(); + AtlasImportRequest request = getDefaultImportRequest(); AtlasImportResult result = runImportWithParameters(importService, request, zipSource); http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/02e4e86b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreHardDeleteV1Test.java ---------------------------------------------------------------------- diff --git a/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreHardDeleteV1Test.java b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreHardDeleteV1Test.java new file mode 100644 index 0000000..2c31140 --- /dev/null +++ b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreHardDeleteV1Test.java @@ -0,0 +1,54 @@ +/** + * 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 com.google.common.collect.ImmutableList; +import org.apache.atlas.TestModules; +import org.apache.atlas.model.instance.AtlasEntity; +import org.testng.annotations.Guice; + +import static org.apache.atlas.type.AtlasTypeUtil.getAtlasObjectId; + +/** + * Inverse reference update test with {@link HardDeleteHandlerV1} + */ +@Guice(modules = TestModules.HardDeleteModule.class) +public class AtlasRelationshipStoreHardDeleteV1Test extends AtlasRelationshipStoreV1Test { + + @Override + protected void verifyRelationshipAttributeUpdate_NonComposite_OneToMany(AtlasEntity jane) throws Exception { + // Max should have been removed from the subordinates list, leaving only John. + verifyRelationshipAttributeList(jane, "subordinates", ImmutableList.of(employeeNameIdMap.get("John"))); + } + + @Override + protected void verifyRelationshipAttributeUpdate_NonComposite_ManyToOne(AtlasEntity a1, AtlasEntity a2, + AtlasEntity a3, AtlasEntity b) { + + verifyRelationshipAttributeValue(a1, "oneB", null); + + verifyRelationshipAttributeValue(a2, "oneB", null); + + verifyRelationshipAttributeList(b, "manyA", ImmutableList.of(getAtlasObjectId(a3))); + } + + @Override + protected void verifyRelationshipAttributeUpdate_NonComposite_OneToOne(AtlasEntity a1, AtlasEntity b) { + verifyRelationshipAttributeValue(a1, "b", null); + } +} http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/02e4e86b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreSoftDeleteV1Test.java ---------------------------------------------------------------------- diff --git a/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreSoftDeleteV1Test.java b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreSoftDeleteV1Test.java new file mode 100644 index 0000000..33ef8c0 --- /dev/null +++ b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreSoftDeleteV1Test.java @@ -0,0 +1,55 @@ +/** + * 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 com.google.common.collect.ImmutableList; +import org.apache.atlas.TestModules; +import org.apache.atlas.model.instance.AtlasEntity; +import org.testng.annotations.Guice; + +import static org.apache.atlas.type.AtlasTypeUtil.getAtlasObjectId; + + +/** + * Inverse reference update test with {@link SoftDeleteHandlerV1} + */ +@Guice(modules = TestModules.SoftDeleteModule.class) +public class AtlasRelationshipStoreSoftDeleteV1Test extends AtlasRelationshipStoreV1Test { + + @Override + protected void verifyRelationshipAttributeUpdate_NonComposite_OneToMany(AtlasEntity jane) throws Exception { + // Max is still in the subordinates list, as the edge still exists with state DELETED + verifyRelationshipAttributeList(jane, "subordinates", ImmutableList.of(employeeNameIdMap.get("John"), employeeNameIdMap.get("Max"))); + } + + @Override + protected void verifyRelationshipAttributeUpdate_NonComposite_ManyToOne(AtlasEntity a1, AtlasEntity a2, + AtlasEntity a3, AtlasEntity b) { + + verifyRelationshipAttributeValue(a1, "oneB", b.getGuid()); + + verifyRelationshipAttributeValue(a2, "oneB", b.getGuid()); + + verifyRelationshipAttributeList(b, "manyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2), getAtlasObjectId(a3))); + } + + @Override + protected void verifyRelationshipAttributeUpdate_NonComposite_OneToOne(AtlasEntity a1, AtlasEntity b) { + verifyRelationshipAttributeValue(a1, "b", b.getGuid()); + } +} http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/02e4e86b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreV1Test.java ---------------------------------------------------------------------- diff --git a/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreV1Test.java b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreV1Test.java index 6770223..31efe86 100644 --- a/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreV1Test.java +++ b/repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasRelationshipStoreV1Test.java @@ -17,18 +17,25 @@ */ package org.apache.atlas.repository.store.graph.v1; +import com.google.common.collect.ImmutableList; import org.apache.atlas.RequestContextV1; import org.apache.atlas.TestModules; -import org.apache.atlas.TestUtilsV2; +import org.apache.atlas.exception.AtlasBaseException; +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.EntityMutationResponse; import org.apache.atlas.model.typedef.AtlasTypesDef; import org.apache.atlas.repository.graph.AtlasGraphProvider; import org.apache.atlas.repository.graph.GraphBackedSearchIndexer; -import org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer; import org.apache.atlas.repository.store.graph.AtlasEntityStore; import org.apache.atlas.repository.store.graph.AtlasRelationshipStore; import org.apache.atlas.store.AtlasTypeDefStore; +import org.apache.atlas.type.AtlasEntityType; import org.apache.atlas.type.AtlasTypeRegistry; +import org.apache.commons.collections.CollectionUtils; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.BeforeTest; @@ -36,11 +43,25 @@ import org.testng.annotations.Guice; import org.testng.annotations.Test; import javax.inject.Inject; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import static org.apache.atlas.TestRelationshipUtilsV2.EMPLOYEE_TYPE; +import static org.apache.atlas.TestRelationshipUtilsV2.getDepartmentEmployeeInstances; +import static org.apache.atlas.TestRelationshipUtilsV2.getDepartmentEmployeeTypes; +import static org.apache.atlas.TestRelationshipUtilsV2.getInverseReferenceTestTypes; +import static org.apache.atlas.TestUtils.NAME; +import static org.apache.atlas.type.AtlasTypeUtil.getAtlasObjectId; import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; @Guice(modules = TestModules.TestOnlyModule.class) -public class AtlasRelationshipStoreV1Test { +public abstract class AtlasRelationshipStoreV1Test { @Inject AtlasTypeRegistry typeRegistry; @@ -56,32 +77,28 @@ public class AtlasRelationshipStoreV1Test { AtlasEntityStore entityStore; AtlasRelationshipStore relationshipStore; - AtlasEntityWithExtInfo dbEntity; - AtlasEntityWithExtInfo tblEntity; AtlasEntityChangeNotifier mockChangeNotifier = mock(AtlasEntityChangeNotifier.class); + protected Map<String, AtlasObjectId> employeeNameIdMap = new HashMap<>(); + @BeforeClass public void setUp() throws Exception { new GraphBackedSearchIndexer(typeRegistry); - AtlasTypesDef[] testTypesDefs = new AtlasTypesDef[] { TestUtilsV2.defineDeptEmployeeTypes(), - TestUtilsV2.defineHiveTypes() }; + // create employee relationship types + AtlasTypesDef employeeTypes = getDepartmentEmployeeTypes(); + typeDefStore.createTypesDef(employeeTypes); - for (AtlasTypesDef typesDef : testTypesDefs) { - AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(typesDef, typeRegistry); + AtlasEntitiesWithExtInfo employeeInstances = getDepartmentEmployeeInstances(); + EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(employeeInstances), false); - if (!typesToCreate.isEmpty()) { - typeDefStore.createTypesDef(typesToCreate); - } + for (AtlasEntityHeader entityHeader : response.getCreatedEntities()) { + employeeNameIdMap.put((String) entityHeader.getAttribute(NAME), getAtlasObjectId(entityHeader)); } - dbEntity = TestUtilsV2.createDBEntityV2(); - tblEntity = TestUtilsV2.createTableEntityV2(dbEntity.getEntity()); - } - - @AfterClass - public void clear() { - AtlasGraphProvider.cleanup(); + init(); + AtlasTypesDef testTypes = getInverseReferenceTestTypes(); + typeDefStore.createTypesDef(testTypes); } @BeforeTest @@ -92,8 +109,384 @@ public class AtlasRelationshipStoreV1Test { RequestContextV1.clear(); } + @AfterClass + public void clear() { + AtlasGraphProvider.cleanup(); + } + + @Test + public void testDepartmentEmployeeEntitiesUsingRelationship() throws Exception { + AtlasObjectId hrId = employeeNameIdMap.get("hr"); + AtlasObjectId maxId = employeeNameIdMap.get("Max"); + AtlasObjectId johnId = employeeNameIdMap.get("John"); + AtlasObjectId juliusId = employeeNameIdMap.get("Julius"); + AtlasObjectId janeId = employeeNameIdMap.get("Jane"); + + AtlasEntity hrDept = getEntityFromStore(hrId.getGuid()); + AtlasEntity max = getEntityFromStore(maxId.getGuid()); + AtlasEntity john = getEntityFromStore(johnId.getGuid()); + AtlasEntity julius = getEntityFromStore(juliusId.getGuid()); + AtlasEntity jane = getEntityFromStore(janeId.getGuid()); + + // Department relationship attributes + List<AtlasObjectId> deptEmployees = toAtlasObjectIds(hrDept.getRelationshipAttribute("employees")); + assertNotNull(deptEmployees); + assertEquals(deptEmployees.size(), 4); + assertObjectIdsContains(deptEmployees, maxId); + assertObjectIdsContains(deptEmployees, johnId); + assertObjectIdsContains(deptEmployees, juliusId); + assertObjectIdsContains(deptEmployees, janeId); + + // Max employee validation + AtlasObjectId maxDepartmentId = toAtlasObjectId(max.getRelationshipAttribute("department")); + assertNotNull(maxDepartmentId); + assertObjectIdEquals(maxDepartmentId, hrId); + + AtlasObjectId maxManagerId = toAtlasObjectId(max.getRelationshipAttribute("manager")); + assertNotNull(maxManagerId); + assertObjectIdEquals(maxManagerId, janeId); + + AtlasObjectId maxMentorId = toAtlasObjectId(max.getRelationshipAttribute("mentor")); + assertNotNull(maxMentorId); + assertObjectIdEquals(maxMentorId, juliusId); + + List<AtlasObjectId> maxMenteesId = toAtlasObjectIds(max.getRelationshipAttribute("mentees")); + assertNotNull(maxMenteesId); + assertEquals(maxMenteesId.size(), 1); + assertObjectIdEquals(maxMenteesId.get(0), johnId); + + // John Employee validation + AtlasObjectId johnDepartmentId = toAtlasObjectId(john.getRelationshipAttribute("department")); + assertNotNull(johnDepartmentId); + assertObjectIdEquals(johnDepartmentId, hrId); + + AtlasObjectId johnManagerId = toAtlasObjectId(john.getRelationshipAttribute("manager")); + assertNotNull(johnManagerId); + assertObjectIdEquals(johnManagerId, janeId); + + AtlasObjectId johnMentorId = toAtlasObjectId(john.getRelationshipAttribute("mentor")); + assertNotNull(johnMentorId); + assertObjectIdEquals(johnMentorId, maxId); + + List<AtlasObjectId> johnMenteesId = toAtlasObjectIds(john.getRelationshipAttribute("mentees")); + assertNull(johnMenteesId); + + // Jane Manager validation + AtlasObjectId janeDepartmentId = toAtlasObjectId(jane.getRelationshipAttribute("department")); + assertNotNull(janeDepartmentId); + assertObjectIdEquals(janeDepartmentId, hrId); + + AtlasObjectId janeManagerId = toAtlasObjectId(jane.getRelationshipAttribute("manager")); + assertNull(janeManagerId); + + AtlasObjectId janeMentorId = toAtlasObjectId(jane.getRelationshipAttribute("mentor")); + assertNull(janeMentorId); + + List<AtlasObjectId> janeMenteesId = toAtlasObjectIds(jane.getRelationshipAttribute("mentees")); + assertNull(janeMenteesId); + + List<AtlasObjectId> janeSubordinateIds = toAtlasObjectIds(jane.getRelationshipAttribute("subordinates")); + assertNotNull(janeSubordinateIds); + assertEquals(janeSubordinateIds.size(), 2); + assertObjectIdsContains(janeSubordinateIds, maxId); + assertObjectIdsContains(janeSubordinateIds, johnId); + + // Julius Manager validation + AtlasObjectId juliusDepartmentId = toAtlasObjectId(julius.getRelationshipAttribute("department")); + assertNotNull(juliusDepartmentId); + assertObjectIdEquals(juliusDepartmentId, hrId); + + AtlasObjectId juliusManagerId = toAtlasObjectId(julius.getRelationshipAttribute("manager")); + assertNull(juliusManagerId); + + AtlasObjectId juliusMentorId = toAtlasObjectId(julius.getRelationshipAttribute("mentor")); + assertNull(juliusMentorId); + + List<AtlasObjectId> juliusMenteesId = toAtlasObjectIds(julius.getRelationshipAttribute("mentees")); + assertNotNull(juliusMenteesId); + assertEquals(juliusMenteesId.size(), 1); + assertObjectIdsContains(juliusMenteesId, maxId); + + List<AtlasObjectId> juliusSubordinateIds = toAtlasObjectIds(julius.getRelationshipAttribute("subordinates")); + assertNull(juliusSubordinateIds); + } + + @Test + public void testRelationshipAttributeUpdate_NonComposite_OneToMany() throws Exception { + AtlasObjectId maxId = employeeNameIdMap.get("Max"); + AtlasObjectId juliusId = employeeNameIdMap.get("Julius"); + AtlasObjectId janeId = employeeNameIdMap.get("Jane"); + + // Change Max's Employee.manager reference to Julius and apply the change as a partial update. + // This should also update Julius to add Max to the inverse Manager.subordinates reference. + AtlasEntity maxEntityForUpdate = new AtlasEntity(EMPLOYEE_TYPE); + maxEntityForUpdate.setRelationshipAttribute("manager", juliusId); + + AtlasEntityType employeeType = typeRegistry.getEntityTypeByName(EMPLOYEE_TYPE); + Map<String, Object> uniqAttributes = Collections.<String, Object>singletonMap("name", "Max"); + EntityMutationResponse updateResponse = entityStore.updateByUniqueAttributes(employeeType, uniqAttributes , new AtlasEntityWithExtInfo(maxEntityForUpdate)); + + List<AtlasEntityHeader> partialUpdatedEntities = updateResponse.getPartialUpdatedEntities(); + assertEquals(partialUpdatedEntities.size(), 3); + // 3 entities should have been updated: + // * Max to change the Employee.manager reference + // * Julius to add Max to Manager.subordinates + // * Jane to remove Max from Manager.subordinates + + AtlasEntitiesWithExtInfo updatedEntities = entityStore.getByIds(ImmutableList.of(maxId.getGuid(), juliusId.getGuid(), janeId.getGuid())); + + // Max's manager updated as Julius + AtlasEntity maxEntity = updatedEntities.getEntity(maxId.getGuid()); + verifyRelationshipAttributeValue(maxEntity, "manager", juliusId.getGuid()); + + // Max added to the subordinate list of Julius + AtlasEntity juliusEntity = updatedEntities.getEntity(juliusId.getGuid()); + verifyRelationshipAttributeList(juliusEntity, "subordinates", ImmutableList.of(maxId)); + + // Max removed from the subordinate list of Julius + AtlasEntity janeEntity = updatedEntities.getEntity(janeId.getGuid()); + + // Jane's subordinates list includes John and Max for soft delete + // Jane's subordinates list includes only John for hard delete + verifyRelationshipAttributeUpdate_NonComposite_OneToMany(janeEntity); + } + + @Test + public void testRelationshipAttributeUpdate_NonComposite_ManyToOne() throws Exception { + AtlasEntity a1 = new AtlasEntity("A"); + a1.setAttribute(NAME, "a1_name"); + + AtlasEntity a2 = new AtlasEntity("A"); + a2.setAttribute(NAME, "a2_name"); + + AtlasEntity a3 = new AtlasEntity("A"); + a3.setAttribute(NAME, "a3_name"); + + AtlasEntity b = new AtlasEntity("B"); + b.setAttribute(NAME, "b_name"); + + AtlasEntitiesWithExtInfo entitiesWithExtInfo = new AtlasEntitiesWithExtInfo(); + entitiesWithExtInfo.addEntity(a1); + entitiesWithExtInfo.addEntity(a2); + entitiesWithExtInfo.addEntity(a3); + entitiesWithExtInfo.addEntity(b); + entityStore.createOrUpdate(new AtlasEntityStream(entitiesWithExtInfo) , false); + + AtlasEntity bPartialUpdate = new AtlasEntity("B"); + bPartialUpdate.setRelationshipAttribute("manyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2))); + + init(); + EntityMutationResponse response = entityStore.updateByUniqueAttributes(typeRegistry.getEntityTypeByName("B"), + Collections.singletonMap(NAME, b.getAttribute(NAME)), + new AtlasEntityWithExtInfo(bPartialUpdate)); + // Verify 3 entities were updated: + // * set b.manyA reference to a1 and a2 + // * set inverse a1.oneB reference to b + // * set inverse a2.oneB reference to b + assertEquals(response.getPartialUpdatedEntities().size(), 3); + AtlasEntitiesWithExtInfo updatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), b.getGuid())); + + AtlasEntity a1Entity = updatedEntities.getEntity(a1.getGuid()); + verifyRelationshipAttributeValue(a1Entity, "oneB", b.getGuid()); + + AtlasEntity a2Entity = updatedEntities.getEntity(a2.getGuid()); + verifyRelationshipAttributeValue(a2Entity, "oneB", b.getGuid()); + + AtlasEntity bEntity = updatedEntities.getEntity(b.getGuid()); + verifyRelationshipAttributeList(bEntity, "manyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2))); + + + bPartialUpdate.setRelationshipAttribute("manyA", ImmutableList.of(getAtlasObjectId(a3))); + init(); + response = entityStore.updateByUniqueAttributes(typeRegistry.getEntityTypeByName("B"), + Collections.singletonMap(NAME, b.getAttribute(NAME)), + new AtlasEntityWithExtInfo(bPartialUpdate)); + // Verify 4 entities were updated: + // * set b.manyA reference to a3 + // * set inverse a3.oneB reference to b + // * disconnect inverse a1.oneB reference to b + // * disconnect inverse a2.oneB reference to b + assertEquals(response.getPartialUpdatedEntities().size(), 4); + init(); + + updatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), a3.getGuid(), b.getGuid())); + a1Entity = updatedEntities.getEntity(a1.getGuid()); + a2Entity = updatedEntities.getEntity(a2.getGuid()); + bEntity = updatedEntities.getEntity(b.getGuid()); + + AtlasEntity a3Entity = updatedEntities.getEntity(a3.getGuid()); + verifyRelationshipAttributeValue(a3Entity, "oneB", b.getGuid()); + + verifyRelationshipAttributeUpdate_NonComposite_ManyToOne(a1Entity, a2Entity, a3Entity, bEntity); + } + + @Test + public void testRelationshipAttributeUpdate_NonComposite_OneToOne() throws Exception { + AtlasEntity a1 = new AtlasEntity("A"); + a1.setAttribute(NAME, "a1_name"); + + AtlasEntity a2 = new AtlasEntity("A"); + a2.setAttribute(NAME, "a2_name"); + + AtlasEntity b = new AtlasEntity("B"); + b.setAttribute(NAME, "b_name"); + + AtlasEntitiesWithExtInfo entitiesWithExtInfo = new AtlasEntitiesWithExtInfo(); + entitiesWithExtInfo.addEntity(a1); + entitiesWithExtInfo.addEntity(a2); + entitiesWithExtInfo.addEntity(b); + + EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesWithExtInfo) , false); + + AtlasEntity partialUpdateB = new AtlasEntity("B"); + partialUpdateB.setRelationshipAttribute("a", getAtlasObjectId(a1)); + + init(); + AtlasEntityType bType = typeRegistry.getEntityTypeByName("B"); + + response = entityStore.updateByUniqueAttributes(bType, Collections.singletonMap(NAME, b.getAttribute(NAME)), new AtlasEntityWithExtInfo(partialUpdateB)); + List<AtlasEntityHeader> partialUpdatedEntitiesHeader = response.getPartialUpdatedEntities(); + // Verify 2 entities were updated: + // * set b.a reference to a1 + // * set inverse a1.b reference to b + assertEquals(partialUpdatedEntitiesHeader.size(), 2); + AtlasEntitiesWithExtInfo partialUpdatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), b.getGuid())); + + AtlasEntity a1Entity = partialUpdatedEntities.getEntity(a1.getGuid()); + verifyRelationshipAttributeValue(a1Entity, "b", b.getGuid()); + + AtlasEntity bEntity = partialUpdatedEntities.getEntity(b.getGuid()); + verifyRelationshipAttributeValue(bEntity, "a", a1.getGuid()); + + init(); + + // Update b.a to reference a2. + partialUpdateB.setRelationshipAttribute("a", getAtlasObjectId(a2)); + response = entityStore.updateByUniqueAttributes(bType, Collections.<String, Object>singletonMap(NAME, b.getAttribute(NAME)), new AtlasEntityWithExtInfo(partialUpdateB)); + partialUpdatedEntitiesHeader = response.getPartialUpdatedEntities(); + // Verify 3 entities were updated: + // * set b.a reference to a2 + // * set a2.b reference to b + // * disconnect a1.b reference + assertEquals(partialUpdatedEntitiesHeader.size(), 3); + partialUpdatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), b.getGuid())); + + bEntity = partialUpdatedEntities.getEntity(b.getGuid()); + verifyRelationshipAttributeValue(bEntity, "a", a2.getGuid()); + + AtlasEntity a2Entity = partialUpdatedEntities.getEntity(a2.getGuid()); + verifyRelationshipAttributeValue(a2Entity, "b", b.getGuid()); + + a1Entity = partialUpdatedEntities.getEntity(a1.getGuid()); + verifyRelationshipAttributeUpdate_NonComposite_OneToOne(a1Entity, bEntity); + } + @Test - public void testDbTableRelationship() throws Exception { - // Add tests - in progress + public void testRelationshipAttributeUpdate_NonComposite_ManyToMany() throws Exception { + AtlasEntity a1 = new AtlasEntity("A"); + a1.setAttribute(NAME, "a1_name"); + + AtlasEntity a2 = new AtlasEntity("A"); + a2.setAttribute(NAME, "a2_name"); + + AtlasEntity a3 = new AtlasEntity("A"); + a3.setAttribute(NAME, "a3_name"); + + AtlasEntity b1 = new AtlasEntity("B"); + b1.setAttribute(NAME, "b1_name"); + + AtlasEntity b2 = new AtlasEntity("B"); + b2.setAttribute(NAME, "b2_name"); + + AtlasEntitiesWithExtInfo entitiesWithExtInfo = new AtlasEntitiesWithExtInfo(); + entitiesWithExtInfo.addEntity(a1); + entitiesWithExtInfo.addEntity(a2); + entitiesWithExtInfo.addEntity(a3); + entitiesWithExtInfo.addEntity(b1); + entitiesWithExtInfo.addEntity(b2); + entityStore.createOrUpdate(new AtlasEntityStream(entitiesWithExtInfo) , false); + + AtlasEntity b1PartialUpdate = new AtlasEntity("B"); + b1PartialUpdate.setRelationshipAttribute("manyToManyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2))); + + init(); + EntityMutationResponse response = entityStore.updateByUniqueAttributes(typeRegistry.getEntityTypeByName("B"), + Collections.singletonMap(NAME, b1.getAttribute(NAME)), + new AtlasEntityWithExtInfo(b1PartialUpdate)); + + List<AtlasEntityHeader> updatedEntityHeaders = response.getPartialUpdatedEntities(); + assertEquals(updatedEntityHeaders.size(), 3); + + AtlasEntitiesWithExtInfo updatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), b1.getGuid())); + + AtlasEntity b1Entity = updatedEntities.getEntity(b1.getGuid()); + verifyRelationshipAttributeList(b1Entity, "manyToManyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2))); + + AtlasEntity a1Entity = updatedEntities.getEntity(a1.getGuid()); + verifyRelationshipAttributeList(a1Entity, "manyB", ImmutableList.of(getAtlasObjectId(b1))); + + AtlasEntity a2Entity = updatedEntities.getEntity(a2.getGuid()); + verifyRelationshipAttributeList(a2Entity, "manyB", ImmutableList.of(getAtlasObjectId(b1))); + } + + protected abstract void verifyRelationshipAttributeUpdate_NonComposite_OneToOne(AtlasEntity a1, AtlasEntity b); + + protected abstract void verifyRelationshipAttributeUpdate_NonComposite_OneToMany(AtlasEntity entity) throws Exception; + + protected abstract void verifyRelationshipAttributeUpdate_NonComposite_ManyToOne(AtlasEntity a1, AtlasEntity a2, AtlasEntity a3, AtlasEntity b); + + private static void assertObjectIdsContains(List<AtlasObjectId> objectIds, AtlasObjectId objectId) { + assertTrue(CollectionUtils.isNotEmpty(objectIds)); + assertTrue(objectIds.contains(objectId)); + } + + private static void assertObjectIdEquals(AtlasObjectId objId1, AtlasObjectId objId2) { + assertTrue(objId1.equals(objId2)); + } + + private static List<AtlasObjectId> toAtlasObjectIds(Object objectIds) { + if (objectIds instanceof List) { + return (List<AtlasObjectId>) objectIds; + } + + return null; + } + + private static AtlasObjectId toAtlasObjectId(Object objectId) { + if (objectId instanceof AtlasObjectId) { + return (AtlasObjectId) objectId; + } + + return null; + } + + private AtlasEntity getEntityFromStore(String guid) throws AtlasBaseException { + AtlasEntityWithExtInfo entity = guid != null ? entityStore.getById(guid) : null; + + return entity != null ? entity.getEntity() : null; + } + + protected static void verifyRelationshipAttributeList(AtlasEntity entity, String relationshipAttrName, List<AtlasObjectId> expectedValues) { + Object refValue = entity.getRelationshipAttribute(relationshipAttrName); + assertTrue(refValue instanceof List); + + List<AtlasObjectId> refList = (List<AtlasObjectId>) refValue; + assertEquals(refList.size(), expectedValues.size()); + + if (expectedValues.size() > 0) { + assertTrue(refList.containsAll(expectedValues)); + } + } + + protected static void verifyRelationshipAttributeValue(AtlasEntity entity, String relationshipAttrName, String expectedGuid) { + Object refValue = entity.getRelationshipAttribute(relationshipAttrName); + if (expectedGuid == null) { + assertNull(refValue); + } + else { + assertTrue(refValue instanceof AtlasObjectId); + AtlasObjectId referencedObjectId = (AtlasObjectId) refValue; + assertEquals(referencedObjectId.getGuid(), expectedGuid); + } } -} +} \ No newline at end of file
