org.apache.atlas.type.AtlasClassificationType Java Examples

public static TypesDef toTypesDef(AtlasType type, AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
    final TypesDef ret;

    if (type instanceof AtlasEnumType) {
        ret = TypeConverterUtil.enumToTypesDef((AtlasEnumType) type);
    } else if (type instanceof AtlasEntityType) {
        ret = TypeConverterUtil.entityToTypesDef((AtlasEntityType) type, typeRegistry);
    } else if (type instanceof AtlasClassificationType) {
        ret = TypeConverterUtil.classificationToTypesDef((AtlasClassificationType) type, typeRegistry);
    } else if (type instanceof AtlasStructType) {
        ret = TypeConverterUtil.structToTypesDef((AtlasStructType) type, typeRegistry);
    } else {
        ret = new TypesDef();
    }

    return ret;
}
 

@Test(dataProvider = "hdfs_path1", expectedExceptions = AtlasBaseException.class)
public void importHdfs_path1(InputStream inputStream) throws IOException, AtlasBaseException {
    loadBaseModel();
    loadFsModel();
    loadModelFromResourcesJson("tag1.json", typeDefStore, typeRegistry);

    try {
        runImportWithNoParameters(importService, inputStream);
    } catch (AtlasBaseException e) {
        assertEquals(e.getAtlasErrorCode(), AtlasErrorCode.INVALID_IMPORT_ATTRIBUTE_TYPE_CHANGED);
        AtlasClassificationType tag1 = typeRegistry.getClassificationTypeByName("tag1");
        assertNotNull(tag1);
        assertEquals(tag1.getAllAttributes().size(), 2);
        throw e;
    }
}
 

/**
 * Validate if classification is not already associated with the entities
 *
 * @param guid            unique entity id
 * @param classifications list of classifications to be associated
 */
private void validateEntityAssociations(String guid, List<AtlasClassification> classifications) throws AtlasBaseException {
    List<String>    entityClassifications = getClassificationNames(guid);
    String          entityTypeName        = AtlasGraphUtilsV2.getTypeNameFromGuid(graph, guid);
    AtlasEntityType entityType            = typeRegistry.getEntityTypeByName(entityTypeName);

    for (AtlasClassification classification : classifications) {
        String newClassification = classification.getTypeName();

        if (CollectionUtils.isNotEmpty(entityClassifications) && entityClassifications.contains(newClassification)) {
            throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "entity: " + guid +
                    ", already associated with classification: " + newClassification);
        }

        // for each classification, check whether there are entities it should be restricted to
        AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(newClassification);

        if (!classificationType.canApplyToEntityType(entityType)) {
            throw new AtlasBaseException(AtlasErrorCode.INVALID_ENTITY_FOR_CLASSIFICATION, guid, entityTypeName, newClassification);
        }
    }
}
 

private void validateAndNormalize(AtlasClassification classification) throws AtlasBaseException {
    AtlasClassificationType type = typeRegistry.getClassificationTypeByName(classification.getTypeName());

    if (type == null) {
        throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATION_NOT_FOUND, classification.getTypeName());
    }

    List<String> messages = new ArrayList<>();

    type.validateValue(classification, classification.getTypeName(), messages);

    if (!messages.isEmpty()) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, messages);
    }

    type.getNormalizedValue(classification);
}
 

public void validateAndNormalizeForUpdate(AtlasClassification classification) throws AtlasBaseException {
    AtlasClassificationType type = typeRegistry.getClassificationTypeByName(classification.getTypeName());

    if (type == null) {
        throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATION_NOT_FOUND, classification.getTypeName());
    }

    List<String> messages = new ArrayList<>();

    type.validateValueForUpdate(classification, classification.getTypeName(), messages);

    if (!messages.isEmpty()) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, messages);
    }

    type.getNormalizedValueForUpdate(classification);
}
 

private void addType(AtlasType type, 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 AtlasStructType) {
        addStructType((AtlasStructType)type, context);
    } else if (type instanceof AtlasEnumType) {
        addEnumType((AtlasEnumType)type, context);
    }
}
 

public static TypesDef toTypesDef(AtlasType type, AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
    final TypesDef ret;

    if (type instanceof AtlasEnumType) {
        ret = TypeConverterUtil.enumToTypesDef((AtlasEnumType) type);
    } else if (type instanceof AtlasEntityType) {
        ret = TypeConverterUtil.entityToTypesDef((AtlasEntityType) type, typeRegistry);
    } else if (type instanceof AtlasClassificationType) {
        ret = TypeConverterUtil.classificationToTypesDef((AtlasClassificationType) type, typeRegistry);
    } else if (type instanceof AtlasStructType) {
        ret = TypeConverterUtil.structToTypesDef((AtlasStructType) type, typeRegistry);
    } else {
        ret = new TypesDef();
    }

    return ret;
}
 

private static TypesDef classificationToTypesDef(AtlasClassificationType classificationType,
                                                 AtlasTypeRegistry registry) throws AtlasBaseException {
    String                typeName    = classificationType.getClassificationDef().getName();
    String                typeDesc    = classificationType.getClassificationDef().getDescription();
    String                typeVersion = classificationType.getClassificationDef().getTypeVersion();
    ImmutableSet          superTypes  = ImmutableSet.copyOf(classificationType.getClassificationDef().getSuperTypes());
    AttributeDefinition[] attributes  = getAttributes(classificationType, registry);

    HierarchicalTypeDefinition traitType = TypesUtil.createTraitTypeDef(typeName, typeDesc, typeVersion, superTypes,
                                                                         attributes);
    TypesDef ret = TypesUtil.getTypesDef(ImmutableList.<EnumTypeDefinition>of(),
                                         ImmutableList.<StructTypeDefinition>of(),
                                         ImmutableList.<HierarchicalTypeDefinition<TraitType>>of(traitType),
                                         ImmutableList.<HierarchicalTypeDefinition<ClassType>>of());
    return ret;
}
 

private void validateAndNormalize(AtlasClassification classification) throws AtlasBaseException {
    AtlasClassificationType type = typeRegistry.getClassificationTypeByName(classification.getTypeName());

    if (type == null) {
        throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATION_NOT_FOUND, classification.getTypeName());
    }

    List<String> messages = new ArrayList<>();

    type.validateValue(classification, classification.getTypeName(), messages);

    if (!messages.isEmpty()) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, messages);
    }

    type.getNormalizedValue(classification);
}
 

public String getTagAttributeType(String tagTypeName, String tagAttributeName) {
    String ret = DEFAULT_TAG_ATTRIBUTE_TYPE;

    if (typeRegistry != null) {
        AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(tagTypeName);
        if (classificationType != null) {
            AtlasStructType.AtlasAttribute atlasAttribute = classificationType.getAttribute(tagAttributeName);

            if (atlasAttribute != null) {
                ret = atlasAttribute.getTypeName();
            }
        }
    }

    return ret;
}
 

private Set<String> getClassificationNames(String classification) throws AtlasBaseException {
    Set<String> classificationNames = new HashSet<>();

    if (StringUtils.isNotEmpty(classification)) {
        String[] types                  = classification.split(TYPENAME_DELIMITER);
        Set<String> names = new HashSet<>(Arrays.asList(types));

        names.forEach(name -> {
            AtlasClassificationType type = getClassificationType(name);
            if (type != null || name.contains(WILDCARD_CLASSIFICATIONS)) {
                classificationNames.add(name);
            }
        });

        // Validate if the classification exists
        if (CollectionUtils.isEmpty(classificationNames) || classificationNames.size() != names.size()) {
            if (CollectionUtils.isNotEmpty(classificationNames)) {
                names.removeAll(classificationNames);
            }
            throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_CLASSIFICATION, String.join(TYPENAME_DELIMITER, names));
        }
    }

    return classificationNames;
}
 

private AtlasClassificationType getClassificationType(String classificationName) {
    AtlasClassificationType ret;

    if (StringUtils.equals(classificationName, MATCH_ALL_WILDCARD_CLASSIFICATION.getTypeName())) {
        ret = MATCH_ALL_WILDCARD_CLASSIFICATION;
    } else if (StringUtils.equals(classificationName, MATCH_ALL_CLASSIFIED.getTypeName())) {
        ret = MATCH_ALL_CLASSIFIED;
    } else if (StringUtils.equals(classificationName, MATCH_ALL_NOT_CLASSIFIED.getTypeName())) {
        ret = MATCH_ALL_NOT_CLASSIFIED;
    } else if (StringUtils.equals(classificationName, ALL_CLASSIFICATION_TYPES)){
        ret = MATCH_ALL_CLASSIFICATION_TYPES;
    } else {
        ret = typeRegistry.getClassificationTypeByName(classificationName);
    }

    return ret;
}
 

private void validateAndNormalizeForUpdate(AtlasClassification classification) throws AtlasBaseException {
    AtlasClassificationType type = typeRegistry.getClassificationTypeByName(classification.getTypeName());

    if (type == null) {
        throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATION_NOT_FOUND, classification.getTypeName());
    }

    List<String> messages = new ArrayList<>();

    type.validateValueForUpdate(classification, classification.getTypeName(), messages);

    if (!messages.isEmpty()) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, messages);
    }

    type.getNormalizedValueForUpdate(classification);
}
 

@Override
public AtlasClassificationDef updateByName(String name, AtlasClassificationDef classificationDef)
    throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasClassificationDefStoreV1.updateByName({}, {})", name, 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 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;
}
 

public static AtlasClassification newClassification(AtlasClassificationDef classificationDef,
                                                    AtlasTypeRegistry typesRegistry) {
    AtlasClassification ret = null;

    AtlasClassificationType classificationType = typesRegistry.getClassificationTypeByName(classificationDef.getName());

    if (classificationType != null) {
        ret = classificationType.createDefaultValue();
    } else {
        LOG.error("failed to get classification-type {}", classificationDef.getName());
    }

    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;
}
 

@Test
public void testGetOnAllClassificationTypes() throws AtlasBaseException {
    AtlasClassificationDef classificationTypeDef = typeDefStore.getClassificationDefByName("_ALL_CLASSIFICATION_TYPES");

    assertNotNull(classificationTypeDef);
    assertEquals(classificationTypeDef, AtlasClassificationType.getClassificationRoot().getClassificationDef());
}
 

private AtlasClassificationType ensureClassificationType(String typeName) throws AtlasBaseException {
    AtlasClassificationType ret = typeRegistry.getClassificationTypeByName(typeName);

    if (ret == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.CLASSIFICATION.name(), typeName);
    }

    return ret;
}
 

public static AtlasClassification newClassification(AtlasClassificationDef classificationDef,
                                                    AtlasTypeRegistry typesRegistry) {
    AtlasClassification ret = null;

    AtlasClassificationType classificationType = typesRegistry.getClassificationTypeByName(classificationDef.getName());

    if (classificationType != null) {
        ret = classificationType.createDefaultValue();
    } else {
        LOG.error("failed to get classification-type {}", classificationDef.getName());
    }

    return ret;
}
 

public Set<String> getClassificationTypeAndAllSuperTypes(String classificationName, AtlasTypeRegistry typeRegistry) {
    final Set<String> ret;

    if (classificationName == null) {
        ret = Collections.emptySet();
    } else if (typeRegistry == null) {
        ret = Collections.singleton(classificationName);
    } else {
        AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(classificationName);

        return classificationType != null ? classificationType.getTypeAndAllSuperTypes() : Collections.singleton(classificationName);
    }

    return ret;
}
 

@Test
public void testGetAllTraitsSuperTraits() throws Exception {
    AtlasTypeRegistry       typeRegistry        = mock(AtlasTypeRegistry.class);
    String                  traitName           = "MyTrait";
    Struct                  myTrait             = new Struct(traitName);
    String                  superTraitName      = "MySuperTrait";
    AtlasClassificationType traitType           = mock(AtlasClassificationType.class);
    Set<String>             superTypeNames      = Collections.singleton(superTraitName);
    AtlasClassificationType superTraitType      = mock(AtlasClassificationType.class);
    Set<String>             superSuperTypeNames = Collections.emptySet();
    Referenceable           entity              = getEntity("id", myTrait);

    when(typeRegistry.getClassificationTypeByName(traitName)).thenReturn(traitType);
    when(typeRegistry.getClassificationTypeByName(superTraitName)).thenReturn(superTraitType);

    when(traitType.getAllSuperTypes()).thenReturn(superTypeNames);
    when(superTraitType.getAllSuperTypes()).thenReturn(superSuperTypeNames);

    EntityNotificationV1 entityNotification = new EntityNotificationV1(entity, OperationType.TRAIT_ADD, typeRegistry);

    List<Struct> allTraits = entityNotification.getAllTraits();

    assertEquals(2, allTraits.size());

    for (Struct trait : allTraits) {
        String typeName = trait.getTypeName();

        assertTrue(typeName.equals(traitName) || typeName.equals(superTraitName));
    }
}
 

private AtlasClassificationType ensureClassificationType(String typeName) throws AtlasBaseException {
    AtlasClassificationType ret = typeRegistry.getClassificationTypeByName(typeName);

    if (ret == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.CLASSIFICATION.name(), typeName);
    }

    return ret;
}
 

private void addClassificationType(String typeName, ExportContext context) {
    if (!context.classificationTypes.contains(typeName)) {
        AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(typeName);

        addClassificationType(classificationType, context);
    }
}
 

private void addClassificationType(AtlasClassificationType classificationType, ExportContext context) {
    if (!context.classificationTypes.contains(classificationType.getTypeName())) {
        context.classificationTypes.add(classificationType.getTypeName());

        addAttributeTypes(classificationType, context);

        if (CollectionUtils.isNotEmpty(classificationType.getAllSuperTypes())) {
            for (String superType : classificationType.getAllSuperTypes()) {
                addClassificationType(superType, context);
            }
        }
    }
}
 

private static List<Struct> getAllTraits(Referenceable entityDefinition, AtlasTypeRegistry typeRegistry) {
    List<Struct> ret = new LinkedList<>();

    for (String traitName : entityDefinition.getTraitNames()) {
        Struct                  trait          = entityDefinition.getTrait(traitName);
        AtlasClassificationType traitType      = typeRegistry.getClassificationTypeByName(traitName);
        Set<String>             superTypeNames = traitType != null ? traitType.getAllSuperTypes() : null;

        ret.add(trait);

        if (CollectionUtils.isNotEmpty(superTypeNames)) {
            for (String superTypeName : superTypeNames) {
                Struct superTypeTrait = new Struct(superTypeName);

                if (MapUtils.isNotEmpty(trait.getValues())) {
                    AtlasClassificationType superType = typeRegistry.getClassificationTypeByName(superTypeName);

                    if (superType != null && MapUtils.isNotEmpty(superType.getAllAttributes())) {
                        Map<String, Object> superTypeTraitAttributes = new HashMap<>();

                        for (Map.Entry<String, Object> attrEntry : trait.getValues().entrySet()) {
                            String attrName = attrEntry.getKey();

                            if (superType.getAllAttributes().containsKey(attrName)) {
                                superTypeTraitAttributes.put(attrName, attrEntry.getValue());
                            }
                        }

                        superTypeTrait.setValues(superTypeTraitAttributes);
                    }
                }

                ret.add(superTypeTrait);
            }
        }
    }

    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;
}
 

public AtlasClassification getClassification(IStruct classification) throws AtlasBaseException {
    AtlasFormatConverter converter          = instanceFormatters.getConverter(TypeCategory.CLASSIFICATION);
    AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(classification.getTypeName());
    if (classificationType == null) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.CLASSIFICATION.name(), classification.getTypeName());
    }
    AtlasClassification  ret                = (AtlasClassification)converter.fromV1ToV2(classification, classificationType, new AtlasFormatConverter.ConverterContext());

    return ret;
}
 

private static TypesDef classificationToTypesDef(AtlasClassificationType classificationType, AtlasTypeRegistry registry) {
    String                    typeName    = classificationType.getClassificationDef().getName();
    String                    typeDesc    = classificationType.getClassificationDef().getDescription();
    String                    typeVersion = classificationType.getClassificationDef().getTypeVersion();
    Set<String>               superTypes  = new HashSet<>(classificationType.getClassificationDef().getSuperTypes());
    List<AttributeDefinition> attributes  = getAttributes(classificationType, registry);

    TraitTypeDefinition traitTypeDef = new TraitTypeDefinition(typeName, typeDesc, typeVersion, attributes, superTypes);

    TypesDef ret = new TypesDef(null, null, Arrays.asList(traitTypeDef), null);

    return ret;
}
 

@VisibleForTesting
public static List<Struct> getAllTraits(Referenceable entityDefinition, AtlasTypeRegistry typeRegistry) throws AtlasException {
    List<Struct> ret = new ArrayList<>();

    for (String traitName : entityDefinition.getTraitNames()) {
        Struct                  trait          = entityDefinition.getTrait(traitName);
        AtlasClassificationType traitType      = typeRegistry.getClassificationTypeByName(traitName);
        Set<String>             superTypeNames = traitType != null ? traitType.getAllSuperTypes() : null;

        ret.add(trait);

        if (CollectionUtils.isNotEmpty(superTypeNames)) {
            for (String superTypeName : superTypeNames) {
                Struct superTypeTrait = new Struct(superTypeName);

                if (MapUtils.isNotEmpty(trait.getValues())) {
                    AtlasClassificationType superType = typeRegistry.getClassificationTypeByName(superTypeName);

                    if (superType != null && MapUtils.isNotEmpty(superType.getAllAttributes())) {
                        Map<String, Object> superTypeTraitAttributes = new HashMap<>();

                        for (Map.Entry<String, Object> attrEntry : trait.getValues().entrySet()) {
                            String attrName = attrEntry.getKey();

                            if (superType.getAllAttributes().containsKey(attrName)) {
                                superTypeTraitAttributes.put(attrName, attrEntry.getValue());
                            }
                        }

                        superTypeTrait.setValues(superTypeTraitAttributes);
                    }
                }

                ret.add(superTypeTrait);
            }
        }
    }

    return ret;
}
 

private Set<AtlasClassificationType> getClassificationTypes(Set<String> classificationNames) {
    if (CollectionUtils.isNotEmpty(classificationNames)) {
        return classificationNames.stream().map(n ->
                getClassificationType(n)).filter(Objects::nonNull).collect(Collectors.toSet());
    }

    return null;
}