Java Code Examples for org.apache.atlas.type.AtlasType#getTypeCategory()

private void mapAttribute(AtlasAttribute attribute, Object attrValue, AtlasVertex vertex, EntityOperation op, EntityMutationContext context) throws AtlasBaseException {
    boolean isDeletedEntity = context.isDeletedEntity(vertex);
    AtlasType         attrType     = attribute.getAttributeType();
    if (attrValue == null) {
        AtlasAttributeDef attributeDef = attribute.getAttributeDef();

        if (attrType.getTypeCategory() == TypeCategory.PRIMITIVE) {
            if (attributeDef.getDefaultValue() != null) {
                attrValue = attrType.createDefaultValue(attributeDef.getDefaultValue());
            } else {
                if (attribute.getAttributeDef().getIsOptional()) {
                    attrValue = attrType.createOptionalDefaultValue();
                } else {
                    attrValue = attrType.createDefaultValue();
                }
            }
        }
    }

    if (attrType.getTypeCategory() == TypeCategory.PRIMITIVE || attrType.getTypeCategory() == TypeCategory.ENUM) {
        mapPrimitiveValue(vertex, attribute, attrValue, isDeletedEntity);
    } else {
        AttributeMutationContext ctx = new AttributeMutationContext(op, vertex, attribute, attrValue);
        mapToVertexByTypeCategory(ctx, context);
    }
}
 

private void mapAttribute(AtlasAttribute attribute, Object attrValue, AtlasVertex vertex, EntityOperation op, EntityMutationContext context) throws AtlasBaseException {
    if (attrValue == null) {
        AtlasAttributeDef attributeDef = attribute.getAttributeDef();
        AtlasType attrType = attribute.getAttributeType();
        if (attrType.getTypeCategory() == TypeCategory.PRIMITIVE) {
            if (attributeDef.getDefaultValue() != null) {
                attrValue = attrType.createDefaultValue(attributeDef.getDefaultValue());
            } else {
                if (attribute.getAttributeDef().getIsOptional()) {
                    attrValue = attrType.createOptionalDefaultValue();
                } else {
                    attrValue = attrType.createDefaultValue();
                }
            }
        }
    }

    AttributeMutationContext ctx = new AttributeMutationContext(op, vertex, attribute, attrValue);

    mapToVertexByTypeCategory(ctx, context);
}
 

private static Predicate getTypePredicate(final String type) {
    return new Predicate() {
        @Override
        public boolean evaluate(Object o) {
            if (o instanceof AtlasType) {
                AtlasType atlasType = (AtlasType) o;

                switch (type.toUpperCase()) {
                    case "CLASS":
                    case "ENTITY":
                        return atlasType.getTypeCategory() == TypeCategory.ENTITY;
                    case "TRAIT":
                    case "CLASSIFICATION":
                        return atlasType.getTypeCategory() == TypeCategory.CLASSIFICATION;
                    case "STRUCT":
                        return atlasType.getTypeCategory() == TypeCategory.STRUCT;
                    case "ENUM":
                        return atlasType.getTypeCategory() == TypeCategory.ENUM;
                    case "RELATIONSHIP":
                        return atlasType.getTypeCategory() == TypeCategory.RELATIONSHIP;
                    case "BUSINESS_METADATA":
                        return atlasType.getTypeCategory() == TypeCategory.BUSINESS_METADATA;
                    default:
                        // This shouldn't have happened
                        return false;
                }
            }
            return false;
        }
    };
}
 

@Override
public AtlasEntityDef updateByGuid(String guid, AtlasEntityDef entityDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasEntityDefStoreV1.updateByGuid({})", guid);
    }

    AtlasEntityDef existingDef = typeRegistry.getEntityDefByGuid(guid);

    AtlasAuthorizationUtils.verifyAccess(new AtlasTypeAccessRequest(AtlasPrivilege.TYPE_UPDATE, existingDef), "update entity-def ", (existingDef != null ? existingDef.getName() : guid));

    validateType(entityDef);

    AtlasType type = typeRegistry.getTypeByGuid(guid);

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.ENTITY) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, entityDef.getName(), TypeCategory.CLASS.name());
    }

    AtlasVertex vertex = typeDefStore.findTypeVertexByGuidAndCategory(guid, TypeCategory.CLASS);

    if (vertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_GUID_NOT_FOUND, guid);
    }

    updateVertexPreUpdate(entityDef, (AtlasEntityType)type, vertex);
    updateVertexAddReferences(entityDef, vertex);

    AtlasEntityDef ret = toEntityDef(vertex);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasEntityDefStoreV1.updateByGuid({}): {}", guid, ret);
    }

    return ret;
}
 

@Override
public AtlasRelationshipDef updateByName(String name, AtlasRelationshipDef relationshipDef)
        throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasRelationshipDefStoreV1.updateByName({}, {})", name, relationshipDef);
    }

    validateType(relationshipDef);

    AtlasType type = typeRegistry.getType(relationshipDef.getName());

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.RELATIONSHIP) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, relationshipDef.getName(), TypeCategory.RELATIONSHIP.name());
    }

    AtlasVertex vertex = typeDefStore.findTypeVertexByNameAndCategory(name, TypeCategory.RELATIONSHIP);

    if (vertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_NOT_FOUND, name);
    }

    preUpdateCheck(relationshipDef, (AtlasRelationshipType) type, vertex);

    AtlasRelationshipDef ret = toRelationshipDef(vertex);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasRelationshipDefStoreV1.updateByName({}, {}): {}", name, relationshipDef, ret);
    }

    return ret;
}
 

@Override
public AtlasEntityDef updateByName(String name, AtlasEntityDef entityDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasEntityDefStoreV1.updateByName({}, {})", name, entityDef);
    }

    validateType(entityDef);

    AtlasType type = typeRegistry.getType(entityDef.getName());

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.ENTITY) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, entityDef.getName(), TypeCategory.CLASS.name());
    }

    AtlasVertex vertex = typeDefStore.findTypeVertexByNameAndCategory(name, TypeCategory.CLASS);

    if (vertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_NOT_FOUND, name);
    }

    updateVertexPreUpdate(entityDef, (AtlasEntityType)type, vertex);
    updateVertexAddReferences(entityDef, vertex);

    AtlasEntityDef ret = toEntityDef(vertex);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasEntityDefStoreV1.updateByName({}, {}): {}", name, entityDef, ret);
    }

    return ret;
}
 

@Override
public AtlasStructDef updateByName(String name, AtlasStructDef structDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasStructDefStoreV1.updateByName({}, {})", name, structDef);
    }

    AtlasStructDef existingDef = typeRegistry.getStructDefByName(name);

    AtlasAuthorizationUtils.verifyAccess(new AtlasTypeAccessRequest(AtlasPrivilege.TYPE_UPDATE, existingDef), "update struct-def ", name);

    validateType(structDef);

    AtlasType type = typeRegistry.getType(structDef.getName());

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.STRUCT) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, structDef.getName(), TypeCategory.STRUCT.name());
    }

    AtlasVertex vertex = typeDefStore.findTypeVertexByNameAndCategory(name, TypeCategory.STRUCT);

    if (vertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_NOT_FOUND, name);
    }

    AtlasStructDefStoreV2.updateVertexPreUpdate(structDef, (AtlasStructType)type, vertex, typeDefStore);
    AtlasStructDefStoreV2.updateVertexAddReferences(structDef, vertex, typeDefStore);

    AtlasStructDef ret = toStructDef(vertex);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasStructDefStoreV1.updateByName({}, {}): {}", name, structDef, ret);
    }

    return ret;
}
 

@Override
public AtlasVertex preCreate(AtlasStructDef structDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasStructDefStoreV1.preCreate({})", structDef);
    }

    validateType(structDef);

    AtlasType type = typeRegistry.getType(structDef.getName());

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.STRUCT) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, structDef.getName(), TypeCategory.STRUCT.name());
    }

    AtlasVertex ret = typeDefStore.findTypeVertexByName(structDef.getName());

    if (ret != null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_ALREADY_EXISTS, structDef.getName());
    }

    ret = typeDefStore.createTypeVertex(structDef);

    AtlasStructDefStoreV2.updateVertexPreCreate(structDef, (AtlasStructType)type, ret, typeDefStore);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasStructDefStoreV1.preCreate({}): {}", structDef, ret);
    }

    return ret;
}
 

@Override
public AtlasVertex preCreate(AtlasClassificationDef classificationDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasClassificationDefStoreV1.preCreate({})", classificationDef);
    }

    validateType(classificationDef);

    AtlasType type = typeRegistry.getType(classificationDef.getName());

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.CLASSIFICATION) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, classificationDef.getName(), TypeCategory.TRAIT.name());
    }

    AtlasVertex ret = typeDefStore.findTypeVertexByName(classificationDef.getName());

    if (ret != null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_ALREADY_EXISTS, classificationDef.getName());
    }

    ret = typeDefStore.createTypeVertex(classificationDef);

    updateVertexPreCreate(classificationDef, (AtlasClassificationType)type, ret);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasClassificationDefStoreV1.preCreate({}): {}", classificationDef, ret);
    }

    return ret;
}
 

@Override
public AtlasClassificationDef updateByName(String name, AtlasClassificationDef classificationDef)
    throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasClassificationDefStoreV1.updateByName({}, {})", name, classificationDef);
    }

    AtlasClassificationDef existingDef   = typeRegistry.getClassificationDefByName(name);

    AtlasAuthorizationUtils.verifyAccess(new AtlasTypeAccessRequest(AtlasPrivilege.TYPE_UPDATE, existingDef), "update classification-def ", name);

    validateType(classificationDef);

    AtlasType type = typeRegistry.getType(classificationDef.getName());

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.CLASSIFICATION) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, classificationDef.getName(), TypeCategory.TRAIT.name());
    }

    AtlasVertex vertex = typeDefStore.findTypeVertexByNameAndCategory(name, TypeCategory.TRAIT);

    if (vertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_NOT_FOUND, name);
    }

    updateVertexPreUpdate(classificationDef, (AtlasClassificationType)type, vertex);
    updateVertexAddReferences(classificationDef, vertex);

    AtlasClassificationDef ret = toClassificationDef(vertex);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasClassificationDefStoreV1.updateByName({}, {}): {}", name, classificationDef, ret);
    }

    return ret;
}
 

@Override
public AtlasStructDef updateByGuid(String guid, AtlasStructDef structDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasStructDefStoreV1.updateByGuid({})", guid);
    }

    validateType(structDef);

    AtlasType type = typeRegistry.getTypeByGuid(guid);

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.STRUCT) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, structDef.getName(), TypeCategory.STRUCT.name());
    }

    AtlasVertex vertex = typeDefStore.findTypeVertexByGuidAndCategory(guid, TypeCategory.STRUCT);

    if (vertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_GUID_NOT_FOUND, guid);
    }

    AtlasStructDefStoreV1.updateVertexPreUpdate(structDef, (AtlasStructType)type, vertex, typeDefStore);
    AtlasStructDefStoreV1.updateVertexAddReferences(structDef, vertex, typeDefStore);

    AtlasStructDef ret = toStructDef(vertex);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasStructDefStoreV1.updateByGuid({}): {}", guid, ret);
    }

    return ret;
}
 

@Override
public AtlasVertex preCreate(AtlasBusinessMetadataDef businessMetadataDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasBusinessMetadataDefStoreV2.preCreate({})", businessMetadataDef);
    }

    validateType(businessMetadataDef);

    AtlasType type = typeRegistry.getType(businessMetadataDef.getName());

    if (type.getTypeCategory() != TypeCategory.BUSINESS_METADATA) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, businessMetadataDef.getName(),
                DataTypes.TypeCategory.BUSINESS_METADATA.name());
    }

    AtlasVertex ret = typeDefStore.findTypeVertexByName(businessMetadataDef.getName());

    if (ret != null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_ALREADY_EXISTS, businessMetadataDef.getName());
    }

    ret = typeDefStore.createTypeVertex(businessMetadataDef);

    updateVertexPreCreate(businessMetadataDef, (AtlasBusinessMetadataType) type, ret);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasBusinessMetadataDefStoreV2.preCreate({}): {}", businessMetadataDef, ret);
    }

    return ret;
}
 

@Override
public AtlasClassificationDef updateByGuid(String guid, AtlasClassificationDef classificationDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasClassificationDefStoreV1.updateByGuid({})", guid);
    }

    validateType(classificationDef);

    AtlasType type = typeRegistry.getTypeByGuid(guid);

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.CLASSIFICATION) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, classificationDef.getName(), TypeCategory.TRAIT.name());
    }

    AtlasVertex vertex = typeDefStore.findTypeVertexByGuidAndCategory(guid, TypeCategory.TRAIT);

    if (vertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_GUID_NOT_FOUND, guid);
    }

    updateVertexPreUpdate(classificationDef, (AtlasClassificationType)type, vertex);
    updateVertexAddReferences(classificationDef, vertex);

    AtlasClassificationDef ret = toClassificationDef(vertex);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasClassificationDefStoreV1.updateByGuid({}): {}", guid, ret);
    }

    return ret;
}
 

@Override
public AtlasRelationshipDef updateByGuid(String guid, AtlasRelationshipDef relationshipDef)
        throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasRelationshipDefStoreV1.updateByGuid({})", guid);
    }

    AtlasRelationshipDef existingDef = typeRegistry.getRelationshipDefByGuid(guid);

    AtlasAuthorizationUtils.verifyAccess(new AtlasTypeAccessRequest(AtlasPrivilege.TYPE_UPDATE, existingDef), "update relationship-Def ", (existingDef != null ? existingDef.getName() : guid));

    validateType(relationshipDef);

    AtlasType type = typeRegistry.getTypeByGuid(guid);

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.RELATIONSHIP) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, relationshipDef.getName(), TypeCategory.RELATIONSHIP.name());
    }

    AtlasVertex vertex = typeDefStore.findTypeVertexByGuidAndCategory(guid, TypeCategory.RELATIONSHIP);

    if (vertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_GUID_NOT_FOUND, guid);
    }

    preUpdateCheck(relationshipDef, (AtlasRelationshipType) type, vertex);
    // updates should not effect the edges between the types as we do not allow updates that change the endpoints.

    AtlasRelationshipDef ret = toRelationshipDef(vertex);

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasRelationshipDefStoreV1.updateByGuid({}): {}", guid, ret);
    }

    return ret;
}
 

private void addType(AtlasType type, ExportService.ExportContext context) {
    if (type.getTypeCategory() == TypeCategory.PRIMITIVE) {
        return;
    }

    if (type instanceof AtlasArrayType) {
        AtlasArrayType arrayType = (AtlasArrayType)type;

        addType(arrayType.getElementType(), context);
    } else if (type instanceof AtlasMapType) {
        AtlasMapType mapType = (AtlasMapType)type;

        addType(mapType.getKeyType(), context);
        addType(mapType.getValueType(), context);
    } else if (type instanceof AtlasEntityType) {
        addEntityType((AtlasEntityType)type, context);
    } else if (type instanceof AtlasClassificationType) {
        addClassificationType((AtlasClassificationType)type, context);
    } else if (type instanceof AtlasRelationshipType) {
        addRelationshipType(type.getTypeName(), context);
    } else if (type instanceof AtlasBusinessMetadataType) {
        addBusinessMetadataType((AtlasBusinessMetadataType) type, context);
    } else if (type instanceof AtlasStructType) {
        addStructType((AtlasStructType)type, context);
    } else if (type instanceof AtlasEnumType) {
        addEnumType((AtlasEnumType)type, context);
    }
}
 

@Override
public AtlasVertex preCreate(AtlasRelationshipDef relationshipDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasRelationshipDefStoreV1.preCreate({})", relationshipDef);
    }

    validateType(relationshipDef);

    AtlasType type = typeRegistry.getType(relationshipDef.getName());

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.RELATIONSHIP) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, relationshipDef.getName(), TypeCategory.RELATIONSHIP.name());
    }

    AtlasVertex relationshipDefVertex = typeDefStore.findTypeVertexByName(relationshipDef.getName());

    if (relationshipDefVertex != null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_ALREADY_EXISTS, relationshipDef.getName());
    }

    relationshipDefVertex = typeDefStore.createTypeVertex(relationshipDef);

    updateVertexPreCreate(relationshipDef, (AtlasRelationshipType) type, relationshipDefVertex);

    final AtlasRelationshipEndDef endDef1        = relationshipDef.getEndDef1();
    final AtlasRelationshipEndDef endDef2        = relationshipDef.getEndDef2();
    final String                  type1          = endDef1.getType();
    final String                  type2          = endDef2.getType();
    final String                  name1          = endDef1.getName();
    final String                  name2          = endDef2.getName();
    final AtlasVertex             end1TypeVertex = typeDefStore.findTypeVertexByName(type1);
    final AtlasVertex             end2TypeVertex = typeDefStore.findTypeVertexByName(type2);

    if (end1TypeVertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_END_TYPE_NAME_NOT_FOUND, relationshipDef.getName(), type1);
    }

    if (end2TypeVertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_END_TYPE_NAME_NOT_FOUND, relationshipDef.getName(), type2);
    }

    // create an edge between the relationshipDef and each of the entityDef vertices.
    AtlasEdge edge1 = typeDefStore.getOrCreateEdge(relationshipDefVertex, end1TypeVertex, AtlasGraphUtilsV1.RELATIONSHIPTYPE_EDGE_LABEL);

    /*
    Where edge1 and edge2 have the same names and types we do not need a second edge.
    We are not invoking the equals method on the AtlasRelationshipedDef, as we only want 1 edge even if propagateTags or other properties are different.
    */

    if (type1.equals(type2) && name1.equals(name2)) {
        if (LOG.isDebugEnabled()) {
            LOG.debug("AtlasRelationshipDefStoreV1.preCreate({}): created relationshipDef vertex {}," +
                    " and one edge as {}, because end1 and end2 have the same type and name", relationshipDef, relationshipDefVertex, edge1);
        }

    } else {
        AtlasEdge edge2 = typeDefStore.getOrCreateEdge(relationshipDefVertex, end2TypeVertex, AtlasGraphUtilsV1.RELATIONSHIPTYPE_EDGE_LABEL);
        if (LOG.isDebugEnabled()) {
            LOG.debug("AtlasRelationshipDefStoreV1.preCreate({}): created relationshipDef vertex {}," +
                    " edge1 as {}, edge2 as {} ", relationshipDef, relationshipDefVertex, edge1, edge2);
        }

    }
    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasRelationshipDefStoreV1.preCreate({}): {}", relationshipDef, relationshipDefVertex);
    }
    return relationshipDefVertex;
}
 

/**
 * Deleting any type vertex. Goes over the complex attributes and removes the references
 * @param instanceVertex
 * @throws AtlasException
 */
protected void deleteTypeVertex(AtlasVertex instanceVertex, boolean force) throws AtlasBaseException {
    LOG.debug("Deleting {}", string(instanceVertex));
    String typeName = GraphHelper.getTypeName(instanceVertex);


    AtlasType parentType = typeRegistry.getType(typeName);

    if (parentType instanceof AtlasStructType) {
        AtlasStructType structType   = (AtlasStructType) parentType;
        boolean         isEntityType = (parentType instanceof AtlasEntityType);

        for (AtlasStructType.AtlasAttribute attributeInfo : structType.getAllAttributes().values()) {
            LOG.debug("Deleting attribute {} for {}", attributeInfo.getName(), string(instanceVertex));
            boolean isOwned = isEntityType && attributeInfo.isOwnedRef();

            AtlasType attrType = attributeInfo.getAttributeType();

            String edgeLabel = AtlasGraphUtilsV1.getAttributeEdgeLabel(structType, attributeInfo.getName());

            switch (attrType.getTypeCategory()) {
            case OBJECT_ID_TYPE:
                //If its class attribute, delete the reference
                deleteEdgeReference(instanceVertex, edgeLabel, attrType.getTypeCategory(), isOwned);
                break;

            case STRUCT:
                //If its struct attribute, delete the reference
                deleteEdgeReference(instanceVertex, edgeLabel, attrType.getTypeCategory(), false);
                break;

            case ARRAY:
                //For array attribute, if the element is struct/class, delete all the references
                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();
                            deleteEdgeReference(edge, elemType.getTypeCategory(), isOwned, false);
                        }
                    }
                }
                break;

            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);
                            deleteEdgeReference(instanceVertex, mapEdgeLabel, valueTypeCategory, isOwned);
                        }
                    }
                }
            }
        }
    }

    deleteVertex(instanceVertex, force);
}
 

private void validateAttribute(AtlasEntityExtInfo entityExtInfo, Object actual, Object expected, AtlasType attributeType, String attrName) throws AtlasBaseException, AtlasException {
    switch(attributeType.getTypeCategory()) {
        case OBJECT_ID_TYPE:
            Assert.assertTrue(actual instanceof AtlasObjectId);
            String guid = ((AtlasObjectId) actual).getGuid();
            Assert.assertTrue(AtlasTypeUtil.isAssignedGuid(guid), "expected assigned guid. found " + guid);
            break;

        case PRIMITIVE:
        case ENUM:
            Assert.assertEquals(actual, expected);
            break;

        case MAP:
            AtlasMapType mapType     = (AtlasMapType) attributeType;
            AtlasType    valueType   = mapType.getValueType();
            Map          actualMap   = (Map) actual;
            Map          expectedMap = (Map) expected;

            if (MapUtils.isNotEmpty(expectedMap)) {
                Assert.assertTrue(MapUtils.isNotEmpty(actualMap));

                // deleted entries are included in the attribute; hence use >=
                Assert.assertTrue(actualMap.size() >= expectedMap.size());

                for (Object key : expectedMap.keySet()) {
                    validateAttribute(entityExtInfo, actualMap.get(key), expectedMap.get(key), valueType, attrName);
                }
            }
            break;

        case ARRAY:
            AtlasArrayType arrType      = (AtlasArrayType) attributeType;
            AtlasType      elemType     = arrType.getElementType();
            List           actualList   = (List) actual;
            List           expectedList = (List) expected;

            if (CollectionUtils.isNotEmpty(expectedList)) {
                Assert.assertTrue(CollectionUtils.isNotEmpty(actualList));

                //actual list could have deleted entities. Hence size may not match.
                Assert.assertTrue(actualList.size() >= expectedList.size());

                for (int i = 0; i < expectedList.size(); i++) {
                    validateAttribute(entityExtInfo, actualList.get(i), expectedList.get(i), elemType, attrName);
                }
            }
            break;
        case STRUCT:
            AtlasStruct expectedStruct = (AtlasStruct) expected;
            AtlasStruct actualStruct   = (AtlasStruct) actual;

            validateEntity(entityExtInfo, actualStruct, expectedStruct);
            break;
        default:
            Assert.fail("Unknown type category");
    }
}
 

@Override
public AtlasVertex preCreate(AtlasRelationshipDef relationshipDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasRelationshipDefStoreV1.preCreate({})", relationshipDef);
    }

    validateType(relationshipDef);

    AtlasType type = typeRegistry.getType(relationshipDef.getName());

    if (type.getTypeCategory() != org.apache.atlas.model.TypeCategory.RELATIONSHIP) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_MATCH_FAILED, relationshipDef.getName(), TypeCategory.RELATIONSHIP.name());
    }

    AtlasVertex relationshipDefVertex = typeDefStore.findTypeVertexByName(relationshipDef.getName());

    if (relationshipDefVertex != null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_ALREADY_EXISTS, relationshipDef.getName());
    }

    relationshipDefVertex = typeDefStore.createTypeVertex(relationshipDef);

    updateVertexPreCreate(relationshipDef, (AtlasRelationshipType) type, relationshipDefVertex);

    final AtlasRelationshipEndDef endDef1        = relationshipDef.getEndDef1();
    final AtlasRelationshipEndDef endDef2        = relationshipDef.getEndDef2();
    final String                  type1          = endDef1.getType();
    final String                  type2          = endDef2.getType();
    final String                  name1          = endDef1.getName();
    final String                  name2          = endDef2.getName();
    final AtlasVertex             end1TypeVertex = typeDefStore.findTypeVertexByName(type1);
    final AtlasVertex             end2TypeVertex = typeDefStore.findTypeVertexByName(type2);

    if (end1TypeVertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_END_TYPE_NAME_NOT_FOUND, relationshipDef.getName(), type1);
    }

    if (end2TypeVertex == null) {
        throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIPDEF_END_TYPE_NAME_NOT_FOUND, relationshipDef.getName(), type2);
    }

    // create an edge between the relationshipDef and each of the entityDef vertices.
    AtlasEdge edge1 = typeDefStore.getOrCreateEdge(relationshipDefVertex, end1TypeVertex, AtlasGraphUtilsV2.RELATIONSHIPTYPE_EDGE_LABEL);

    /*
    Where edge1 and edge2 have the same names and types we do not need a second edge.
    We are not invoking the equals method on the AtlasRelationshipedDef, as we only want 1 edge even if propagateTags or other properties are different.
    */

    if (type1.equals(type2) && name1.equals(name2)) {
        if (LOG.isDebugEnabled()) {
            LOG.debug("AtlasRelationshipDefStoreV1.preCreate({}): created relationshipDef vertex {}," +
                    " and one edge as {}, because end1 and end2 have the same type and name", relationshipDef, relationshipDefVertex, edge1);
        }

    } else {
        AtlasEdge edge2 = typeDefStore.getOrCreateEdge(relationshipDefVertex, end2TypeVertex, AtlasGraphUtilsV2.RELATIONSHIPTYPE_EDGE_LABEL);
        if (LOG.isDebugEnabled()) {
            LOG.debug("AtlasRelationshipDefStoreV1.preCreate({}): created relationshipDef vertex {}," +
                    " edge1 as {}, edge2 as {} ", relationshipDef, relationshipDefVertex, edge1, edge2);
        }

    }
    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasRelationshipDefStoreV1.preCreate({}): {}", relationshipDef, relationshipDefVertex);
    }
    return relationshipDefVertex;
}
 

protected void validateAttribute(AtlasEntityExtInfo entityExtInfo, Object actual, Object expected, AtlasType attributeType, String attrName) throws AtlasBaseException, AtlasException {
    switch(attributeType.getTypeCategory()) {
        case OBJECT_ID_TYPE:
            Assert.assertTrue(actual instanceof AtlasObjectId);
            String guid = ((AtlasObjectId) actual).getGuid();
            Assert.assertTrue(AtlasTypeUtil.isAssignedGuid(guid), "expected assigned guid. found " + guid);
            break;

        case PRIMITIVE:
        case ENUM:
            Assert.assertEquals(actual, expected);
            break;

        case MAP:
            AtlasMapType mapType     = (AtlasMapType) attributeType;
            AtlasType    valueType   = mapType.getValueType();
            Map          actualMap   = (Map) actual;
            Map          expectedMap = (Map) expected;

            if (MapUtils.isNotEmpty(expectedMap)) {
                Assert.assertTrue(MapUtils.isNotEmpty(actualMap));

                // deleted entries are included in the attribute; hence use >=
                Assert.assertTrue(actualMap.size() >= expectedMap.size());

                for (Object key : expectedMap.keySet()) {
                    validateAttribute(entityExtInfo, actualMap.get(key), expectedMap.get(key), valueType, attrName);
                }
            }
            break;

        case ARRAY:
            AtlasArrayType arrType      = (AtlasArrayType) attributeType;
            AtlasType      elemType     = arrType.getElementType();
            List           actualList   = (List) actual;
            List           expectedList = (List) expected;

            if (CollectionUtils.isNotEmpty(expectedList)) {
                Assert.assertTrue(CollectionUtils.isNotEmpty(actualList));

                //actual list could have deleted entities. Hence size may not match.
                Assert.assertTrue(actualList.size() >= expectedList.size());

                for (int i = 0; i < expectedList.size(); i++) {
                    validateAttribute(entityExtInfo, actualList.get(i), expectedList.get(i), elemType, attrName);
                }
            }
            break;
        case STRUCT:
            AtlasStruct expectedStruct = (AtlasStruct) expected;
            AtlasStruct actualStruct   = (AtlasStruct) actual;

            validateEntity(entityExtInfo, actualStruct, expectedStruct);
            break;
        default:
            Assert.fail("Unknown type category");
    }
}