ATLAS-2029: Restrict entities, classifications can be applied to

intg/src/main/java/org/apache/atlas/AtlasErrorCode.java

@@ -95,6 +95,11 @@ public enum AtlasErrorCode {
     INVALID_RELATIONSHIP_ATTRIBUTE(400, "ATLAS-400-00-048", "Expected attribute {0} to be a relationship but found type {1}"),
     INVALID_RELATIONSHIP_TYPE(400, "ATLAS-400-00-049", "Invalid entity type '{0}', guid '{1}' in relationship search"),
     INVALID_IMPORT_ATTRIBUTE_TYPE_CHANGED(400, "ATLAS-400-00-050", "Attribute {0}.{1} is of type {2}. Import has this attribute type as {3}"),
+    ENTITYTYPE_REMOVAL_NOT_SUPPORTED(400, "ATLAS-400-00-051", "EntityTypes cannot be removed from ClassificationDef ‘{0}‘"),
+    CLASSIFICATIONDEF_INVALID_ENTITYTYPES(400, "ATLAS-400-00-052", "ClassificationDef ‘{0}‘ has invalid ‘{1}‘ in entityTypes"),
+    CLASSIFICATIONDEF_PARENTS_ENTITYTYPES_DISJOINT(400, "ATLAS-400-00-053", "ClassificationDef ‘{0}‘ has supertypes whose entityTypes are disjoint; e.g. 2 supertypes that are not related by inheritance specify different non empty entityType lists. This means the child cannot honour the restrictions specified in both parents."),
+    CLASSIFICATIONDEF_ENTITYTYPES_NOT_PARENTS_SUBSET(400, "ATLAS-400-00-054", "ClassificationDef ‘{0}‘ has entityTypes ‘{1}‘ which are not subsets of it's supertypes entityTypes"),
+    INVALID_ENTITY_FOR_CLASSIFICATION (400, "ATLAS-400-00-055", "Entity (guid=‘{0}‘,typename=‘{1}‘) cannot be classified by Classification ‘{2}‘, because ‘{1}‘ is not in the ClassificationDef's restrictions."),
 
     // All Not found enums go here
     TYPE_NAME_NOT_FOUND(404, "ATLAS-404-00-001", "Given typename {0} was invalid"),

intg/src/main/java/org/apache/atlas/model/typedef/AtlasClassificationDef.java

@@ -51,6 +51,7 @@ public class AtlasClassificationDef extends AtlasStructDef implements java.io.Se
     private static final long serialVersionUID = 1L;
 
     private Set<String> superTypes;
+    private Set<String> entityTypes;
 
 
     public AtlasClassificationDef() {
@@ -82,9 +83,16 @@ public class AtlasClassificationDef extends AtlasStructDef implements java.io.Se
     public AtlasClassificationDef(String name, String description, String typeVersion,
                                   List<AtlasAttributeDef> attributeDefs, Set<String> superTypes,
                                   Map<String, String> options) {
+        this(name, description, typeVersion, attributeDefs, superTypes, null, options);
+    }
+
+    public AtlasClassificationDef(String name, String description, String typeVersion,
+                                  List<AtlasAttributeDef> attributeDefs, Set<String> superTypes,
+                                  Set<String> entityTypes, Map<String, String> options) {
         super(TypeCategory.CLASSIFICATION, name, description, typeVersion, attributeDefs, options);
 
         setSuperTypes(superTypes);
+        setEntityTypes(entityTypes);
     }
 
     public AtlasClassificationDef(AtlasClassificationDef other) {
@@ -141,6 +149,66 @@ public class AtlasClassificationDef extends AtlasStructDef implements java.io.Se
         return superTypes != null && typeName != null && superTypes.contains(typeName);
     }
 
+    /**
+     * Specifying a list of entityType names in the classificationDef, ensures that classifications can
+     * only be applied to those entityTypes.
+     * <ul>
+     * <li>Any subtypes of the entity types inherit the restriction</li>
+     * <li>Any classificationDef subtypes inherit the parents entityTypes restrictions</li>
+     * <li>Any classificationDef subtypes can further restrict the parents entityTypes restrictions by specifying a subset of the entityTypes</li>
+     * <li>An empty entityTypes list when there are no parent restrictions means there are no restrictions</li>
+     * <li>An empty entityTypes list when there are parent restrictions means that the subtype picks up the parents restrictions</li>
+     * <li>If a list of entityTypes are supplied, where one inherits from another, this will be rejected. This should encourage cleaner classificationsDefs</li>
+     * </ul>
+     */
+    public Set<String> getEntityTypes() {
+        return entityTypes;
+    }
+
+    public void setEntityTypes(Set<String> entityTypes) {
+        if (entityTypes != null && this.entityTypes == entityTypes) {
+            return;
+        }
+
+        if (CollectionUtils.isEmpty(entityTypes)) {
+            this.entityTypes = new HashSet<>();
+        } else {
+            this.entityTypes = new HashSet<>(entityTypes);
+        }
+    }
+
+    public boolean hasEntityType(String typeName) {
+        return hasEntityType(entityTypes, typeName);
+    }
+
+    public void addEntityType(String typeName) {
+        Set<String> s = this.entityTypes;
+
+        if (!hasEntityType(s, typeName)) {
+            s = new HashSet<>(s);
+
+            s.add(typeName);
+
+            this.entityTypes = s;
+        }
+    }
+
+    public void removeEntityType(String typeName) {
+        Set<String> s = this.entityTypes;
+
+        if (hasEntityType(s, typeName)) {
+            s = new HashSet<>(s);
+
+            s.remove(typeName);
+
+            this.entityTypes = s;
+        }
+    }
+
+    private static boolean hasEntityType(Set<String> entityTypes, String typeName) {
+        return entityTypes != null && typeName != null && entityTypes.contains(typeName);
+    }
+
     @Override
     public StringBuilder toString(StringBuilder sb) {
         if (sb == null) {
@@ -151,6 +219,8 @@ public class AtlasClassificationDef extends AtlasStructDef implements java.io.Se
         super.toString(sb);
         sb.append(", superTypes=[");
         dumpObjects(superTypes, sb);
+        sb.append("], entityTypes=[");
+        dumpObjects(entityTypes, sb);
         sb.append("]");
         sb.append('}');
 
@@ -164,7 +234,8 @@ public class AtlasClassificationDef extends AtlasStructDef implements java.io.Se
         if (!super.equals(o)) { return false; }
 
         AtlasClassificationDef that = (AtlasClassificationDef) o;
-        return Objects.equals(superTypes, that.superTypes);
+
+        return Objects.equals(superTypes, that.superTypes) && Objects.equals(entityTypes,that.entityTypes);
     }
 
     @Override

intg/src/main/java/org/apache/atlas/type/AtlasClassificationType.java

@@ -28,13 +28,7 @@ import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 
 /**
@@ -51,12 +45,28 @@ public class AtlasClassificationType extends AtlasStructType {
     private Set<String>                   typeAndAllSubTypes       = Collections.emptySet();
     private String                        typeAndAllSubTypesQryStr = "";
 
+    // we need to store the entityTypes specified in our supertypes. i.e. our parent classificationDefs may specify more entityTypes
+    // that we also need to allow
+    private Set<String> entityTypes = Collections.emptySet();
+
+    /**
+     * Note this constructor does NOT run resolveReferences - so some fields are not setup.
+     * @param classificationDef
+     */
     public AtlasClassificationType(AtlasClassificationDef classificationDef) {
         super(classificationDef);
 
         this.classificationDef = classificationDef;
     }
 
+    /**
+     * ClassificationType needs to be constructed with a type registry so that is can resolve references
+     * at constructor time. This is only used by junits.
+     *
+     * @param classificationDef
+     * @param typeRegistry
+     * @throws AtlasBaseException
+     */
     public AtlasClassificationType(AtlasClassificationDef classificationDef, AtlasTypeRegistry typeRegistry)
         throws AtlasBaseException {
         super(classificationDef);
@@ -95,6 +105,7 @@ public class AtlasClassificationType extends AtlasStructType {
         this.uniqAttributes     = getUniqueAttributes(this.allAttributes);
         this.allSubTypes        = new HashSet<>(); // this will be populated in resolveReferencesPhase2()
         this.typeAndAllSubTypes = new HashSet<>(); // this will be populated in resolveReferencesPhase2()
+        this.entityTypes        = new HashSet<>(); // this will be populated in resolveReferencesPhase3()
 
         this.typeAndAllSubTypes.add(this.getTypeName());
     }
@@ -109,11 +120,92 @@ public class AtlasClassificationType extends AtlasStructType {
         }
     }
 
+    /**
+     * This method processes the entityTypes to ensure they are valid, using the following principles:
+     * - entityTypes are supplied on the classificationDef to restrict the types of entities that this classification can be applied to
+     * - Any subtypes of the specified entity type can also have this classification applied
+     * - Any subtypes of the classificationDef inherit the parents entityTypes restrictions
+     * - Any subtypes of the classificationDef can further restrict the parents entityTypes restrictions
+     * - An empty entityTypes list when there are no parent restrictions means there are no restrictions
+     * - An empty entityTypes list when there are parent restrictions means that the subtype picks up the parents restrictions
+     *
+     * This method validates that these priniciples are adhered to.
+     *
+     * Note that if duplicate Strings in the entityTypes are specified on an add / update, the duplicates are ignored - as Java Sets cannot have duplicates.
+     * Note if an entityType is supplied in the list that is a subtype of one of the other supplied entityTypes, we are not policing this case as invalid.
+     *
+     * @param typeRegistry
+     * @throws AtlasBaseException
+     */
     @Override
     void resolveReferencesPhase3(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
         allSubTypes              = Collections.unmodifiableSet(allSubTypes);
         typeAndAllSubTypes       = Collections.unmodifiableSet(typeAndAllSubTypes);
         typeAndAllSubTypesQryStr = ""; // will be computed on next access
+
+        /*
+          Add any entityTypes defined in our parents as restrictions.
+         */
+        Set<String> superTypeEntityTypes = null;
+
+        final Set<String> classificationDefEntityTypes = classificationDef.getEntityTypes();
+
+        // Here we find the intersection of the entityTypes specified in all our supertypes; in this way we will honour our parents restrictions.
+        // This following logic assumes typeAndAllSubTypes is populated so needs to be run after resolveReferencesPhase2().
+
+        for (String superType : this.allSuperTypes) {
+            AtlasClassificationDef superTypeDef    = typeRegistry.getClassificationDefByName(superType);
+            Set<String>            entityTypeNames = superTypeDef.getEntityTypes();
+
+            if (CollectionUtils.isEmpty(entityTypeNames)) { // no restrictions specified
+                continue;
+            }
+
+            // classification is applicable for specified entityTypes and their sub-entityTypes
+            Set<String> typesAndSubEntityTypes = AtlasEntityType.getEntityTypesAndAllSubTypes(entityTypeNames, typeRegistry);
+
+            if (superTypeEntityTypes == null) {
+                superTypeEntityTypes = new HashSet<>(typesAndSubEntityTypes);
+            } else {
+                // retain only the intersections.
+                superTypeEntityTypes.retainAll(typesAndSubEntityTypes);
+            }
+            if (superTypeEntityTypes.isEmpty()) {
+                // if we have no intersections then we are disjoint - so no need to check other supertypes
+                break;
+            }
+        }
+
+        if (superTypeEntityTypes == null) {  // no supertype restrictions; use current classification restrictions
+            this.entityTypes = AtlasEntityType.getEntityTypesAndAllSubTypes(classificationDefEntityTypes, typeRegistry);
+        } else {                             // restrictions are specified in super-types
+            if (CollectionUtils.isEmpty(superTypeEntityTypes)) {
+                /*
+                 Restrictions in superTypes are disjoint! This means that the child cannot exist as it cannot be a restriction of it's parents.
+
+                 For example:
+                  parent1 specifies entityTypes ["EntityA"]
+                  parent2 specifies entityTypes ["EntityB"]
+
+                  In order to be a valid child of Parent1 the child could only be applied to EntityAs.
+                  In order to be a valid child of Parent2 the child could only be applied to EntityBs.
+
+                  Reject the creation of the classificationDef - as it would compromise Atlas's integrity.
+                 */
+                throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATIONDEF_PARENTS_ENTITYTYPES_DISJOINT, this.classificationDef.getName());
+            }
+
+            if (CollectionUtils.isEmpty(classificationDefEntityTypes)) { // no restriction specified; use the restrictions from super-types
+                this.entityTypes = superTypeEntityTypes;
+            } else {
+                this.entityTypes = AtlasEntityType.getEntityTypesAndAllSubTypes(classificationDefEntityTypes,typeRegistry);
+                // Compatible parents and entityTypes, now check whether the specified entityTypes are the same as the effective entityTypes due to our parents or a subset.
+                // Only allowed to restrict our parents.
+                if (!superTypeEntityTypes.containsAll(this.entityTypes)) {
+                    throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATIONDEF_ENTITYTYPES_NOT_PARENTS_SUBSET, classificationDef.getName(), classificationDefEntityTypes.toString());
+                }
+            }
+        }
     }
 
     private void addSubType(AtlasClassificationType subType) {
@@ -155,6 +247,16 @@ public class AtlasClassificationType extends AtlasStructType {
         return StringUtils.isNotEmpty(classificationName) && allSuperTypes.contains(classificationName);
     }
 
+    /**
+     * List of all the entity type names that are valid for this classification type.
+     *
+     * An empty list means there are no restrictions on which entities can be classified by these classifications.
+     * @return
+     */
+    public Set<String> getEntityTypes() {
+        return entityTypes;
+    }
+
     @Override
     public AtlasClassification createDefaultValue() {
         AtlasClassification ret = new AtlasClassification(classificationDef.getName());
@@ -313,6 +415,27 @@ public class AtlasClassificationType extends AtlasStructType {
         }
     }
 
+    /**
+     * Check whether the supplied entityType can be applied to this classification.
+     *
+     * We can apply this classification to the supplied entityType if
+     * - we have no restrictions (entityTypes empty including null)
+     * or
+     * - the entityType is in our list of restricted entityTypes (which includes our parent classification restrictions)
+     *
+     * @param entityType
+     * @return whether can apply
+     */
+    /**
+     * Check whether the supplied entityType can be applied to this classification.
+     *
+     * @param entityType
+     * @return whether can apply
+     */
+    public boolean canApplyToEntityType(AtlasEntityType entityType) {
+        return CollectionUtils.isEmpty(this.entityTypes) || this.entityTypes.contains(entityType.getTypeName());
+    }
+
     private void getTypeHierarchyInfo(AtlasTypeRegistry              typeRegistry,
                                       Set<String>                    allSuperTypeNames,
                                       Map<String, AtlasAttribute>    allAttributes) throws AtlasBaseException {

intg/src/main/java/org/apache/atlas/type/AtlasEntityType.java

@@ -522,6 +522,28 @@ public class AtlasEntityType extends AtlasStructType {
         return true;
     }
 
+    /**
+     * Takes a set of entityType names and a registry and returns a set of the entitytype names and the names of all their subTypes.
+     *
+     * @param entityTypes
+     * @param typeRegistry
+     * @return set of strings of the types and their subtypes.
+     */
+    static public Set<String> getEntityTypesAndAllSubTypes(Set<String> entityTypes, AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
+        Set<String> ret = new HashSet<>();
+
+        for (String typeName : entityTypes) {
+            AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
+            if (entityType == null) {
+                throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_NOT_FOUND, typeName);
+            }
+
+            ret.addAll(entityType.getTypeAndAllSubTypes());
+        }
+
+        return ret;
+    }
+
     private boolean isAssignableValue(Object value, AtlasAttributeDef attributeDef) {
         boolean ret = true;
 

intg/src/main/java/org/apache/atlas/type/AtlasTypeUtil.java

@@ -220,6 +220,10 @@ public class AtlasTypeUtil {
         return new AtlasClassificationDef(name, description, version, Arrays.asList(attrDefs), superTypes);
     }
 
+    public static AtlasClassificationDef createAtlasClassificationDef(String name, String description, String version, ImmutableSet<String> superTypes, ImmutableSet<String> entityTypes, AtlasAttributeDef... attrDefs) {
+        return new AtlasClassificationDef(name, description, version, Arrays.asList(attrDefs), superTypes, entityTypes, null);
+    }
+
     public static AtlasStructDef createStructTypeDef(String name, AtlasAttributeDef... attrDefs) {
         return createStructTypeDef(name, null, attrDefs);
     }

intg/src/test/java/org/apache/atlas/TestUtilsV2.java

@@ -1004,6 +1004,29 @@ public final class TestUtilsV2 {
         return ret;
     }
 
+    public static List<AtlasClassificationDef> getClassificationWithName(String name) {
+        AtlasClassificationDef classificationTypeDef =
+                AtlasTypeUtil.createTraitTypeDef(name, "s_description", ImmutableSet.<String>of(),
+                        AtlasTypeUtil.createRequiredAttrDef("level", "int"));
+
+
+        List<AtlasClassificationDef> ret = Arrays.asList(classificationTypeDef);
+
+        populateSystemAttributes(ret);
+
+        return ret;
+    }
+
+    public static AtlasClassificationDef getSingleClassificationWithName(String name) {
+        AtlasClassificationDef classificaitonTypeDef =
+                AtlasTypeUtil.createTraitTypeDef(name, "s_description", ImmutableSet.<String>of(),
+                        AtlasTypeUtil.createRequiredAttrDef("level", "int"));
+
+        populateSystemAttributes(classificaitonTypeDef);
+
+        return classificaitonTypeDef;
+    }
+
     public static List<AtlasClassificationDef> getClassificationWithValidAttribute(){
         return getClassificationWithValidSuperType();
     }
@@ -1021,6 +1044,28 @@ public final class TestUtilsV2 {
         return ret;
     }
 
+    public static List<AtlasEntityDef> getEntityWithName(String name) {
+        AtlasEntityDef developerTypeDef = AtlasTypeUtil.createClassTypeDef(name, "Developer_description", ImmutableSet.<String>of(),
+                new AtlasAttributeDef("language", String.format("array<%s>", "string"), false, AtlasAttributeDef.Cardinality.SET,
+                        1, 10, false, false,
+                        Collections.<AtlasConstraintDef>emptyList()));
+
+        List<AtlasEntityDef> ret = Arrays.asList(developerTypeDef);
+
+        populateSystemAttributes(ret);
+
+        return ret;
+    }
+
+    public static AtlasEntityDef getSingleEntityWithName(String name) {
+        AtlasEntityDef developerTypeDef = AtlasTypeUtil.createClassTypeDef(name, "Developer_description", ImmutableSet.<String>of(),
+                new AtlasAttributeDef("language", String.format("array<%s>", "string"), false, AtlasAttributeDef.Cardinality.SET,
+                        1, 10, false, false,
+                        Collections.<AtlasConstraintDef>emptyList()));
+
+        return developerTypeDef;
+    }
+
     public static List<AtlasEntityDef> getEntityWithValidAttribute() {
         List<AtlasEntityDef> entityDefs = getEntityWithValidSuperType();
         entityDefs.get(1).getSuperTypes().clear();

intg/src/test/java/org/apache/atlas/model/ModelTestUtil.java

@@ -216,7 +216,7 @@ public final class  ModelTestUtil {
     }
 
     public static AtlasEntityDef newEntityDef(AtlasTypeRegistry typesRegistry) {
-        return newEntityDef(getTypesRegistry(), null);
+        return newEntityDef(typesRegistry, null);
     }
 
     public static AtlasEntityDef newEntityDef(AtlasTypeRegistry typesRegistry, AtlasEntityDef[] superTypes) {

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasClassificationDefStoreV1.java

@@ -316,6 +316,8 @@ class AtlasClassificationDefStoreV1 extends AtlasAbstractDefStoreV1<AtlasClassif
         AtlasStructDefStoreV1.updateVertexAddReferences(classificationDef, vertex, typeDefStore);
 
         typeDefStore.createSuperTypeEdges(vertex, classificationDef.getSuperTypes(), TypeCategory.TRAIT);
+        // create edges from this vertex to entity Type vertices with the supplied entity type names
+        typeDefStore.createEntityTypeEdges(vertex, classificationDef.getEntityTypes());
     }
 
     private AtlasClassificationDef toClassificationDef(AtlasVertex vertex) throws AtlasBaseException {
@@ -327,6 +329,7 @@ class AtlasClassificationDefStoreV1 extends AtlasAbstractDefStoreV1<AtlasClassif
             AtlasStructDefStoreV1.toStructDef(vertex, ret, typeDefStore);
 
             ret.setSuperTypes(typeDefStore.getSuperTypeNames(vertex));
+            ret.setEntityTypes(typeDefStore.getEntityTypeNames(vertex));
         }
 
         return ret;

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java

@@ -729,6 +729,8 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
      */
     private void validateEntityAssociations(String guid, List<AtlasClassification> classifications) throws AtlasBaseException {
         List<String>    entityClassifications = getClassificationNames(guid);
+        String          entityTypeName        = AtlasGraphUtilsV1.getTypeNameFromGuid(guid);
+        AtlasEntityType entityType            = typeRegistry.getEntityTypeByName(entityTypeName);
 
         for (AtlasClassification classification : classifications) {
             String newClassification = classification.getTypeName();
@@ -737,6 +739,13 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
                 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);
+            }
         }
     }
 

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasGraphUtilsV1.java

@@ -57,6 +57,7 @@ public class AtlasGraphUtilsV1 {
 
     public static final String PROPERTY_PREFIX      = Constants.INTERNAL_PROPERTY_KEY_PREFIX + "type.";
     public static final String SUPERTYPE_EDGE_LABEL = PROPERTY_PREFIX + ".supertype";
+    public static final String ENTITYTYPE_EDGE_LABEL = PROPERTY_PREFIX + ".entitytype";
     public static final String VERTEX_TYPE          = "typeSystem";
     public static final String RELATIONSHIPTYPE_EDGE_LABEL = PROPERTY_PREFIX + ".relationshipType";
 

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasTypeDefGraphStoreV1.java

@@ -380,9 +380,43 @@ public class AtlasTypeDefGraphStoreV1 extends AtlasTypeDefGraphStore {
         }
     }
 
+    public void createEntityTypeEdges(AtlasVertex classificationVertex, Set<String> entityTypes) throws AtlasBaseException {
+        Set<String> currentEntityTypes     = getEntityTypeNames(classificationVertex);
+        String      classificationTypeName = classificationVertex.getProperty(Constants.TYPENAME_PROPERTY_KEY, String.class);
+
+        if (CollectionUtils.isNotEmpty(entityTypes)) {
+            if (!entityTypes.containsAll(currentEntityTypes)) {
+                throw new AtlasBaseException(AtlasErrorCode.ENTITYTYPE_REMOVAL_NOT_SUPPORTED, classificationTypeName);
+            }
+
+            for (String entityType : entityTypes) {
+                AtlasVertex entityTypeVertex = findTypeVertexByNameAndCategory(entityType, TypeCategory.CLASS);
+                if (entityTypeVertex == null) {
+                    throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATIONDEF_INVALID_ENTITYTYPES, classificationTypeName,entityType);
+
+                }
+                getOrCreateEdge(classificationVertex, entityTypeVertex, AtlasGraphUtilsV1.ENTITYTYPE_EDGE_LABEL);
+            }
+        }
+    }
+
     Set<String>  getSuperTypeNames(AtlasVertex vertex) {
+        return getTypeNamesFromEdges(vertex,AtlasGraphUtilsV1.SUPERTYPE_EDGE_LABEL);
+    }
+
+    Set<String>  getEntityTypeNames(AtlasVertex vertex) {
+        return getTypeNamesFromEdges(vertex,AtlasGraphUtilsV1.ENTITYTYPE_EDGE_LABEL);
+    }
+
+    /**
+     * Get the typename properties from the edges, that are associated with the vertex and have the supplied edge label.
+     * @param vertex
+     * @param edgeLabel
+     * @return set of type names
+     */
+    private Set<String> getTypeNamesFromEdges(AtlasVertex vertex,String edgeLabel) {
         Set<String>    ret   = new HashSet<>();
-        Iterable<AtlasEdge> edges = vertex.getEdges(AtlasEdgeDirection.OUT, AtlasGraphUtilsV1.SUPERTYPE_EDGE_LABEL);
+        Iterable<AtlasEdge> edges = vertex.getEdges(AtlasEdgeDirection.OUT, edgeLabel);
 
         for (AtlasEdge edge : edges) {
             ret.add(edge.getInVertex().getProperty(Constants.TYPENAME_PROPERTY_KEY, String.class));

repository/src/test/java/org/apache/atlas/repository/store/graph/AtlasTypeDefGraphStoreTest.java

@@ -18,17 +18,13 @@
 package org.apache.atlas.repository.store.graph;
 
 import com.google.inject.Inject;
+import org.apache.atlas.RequestContextV1;
 import org.apache.atlas.TestModules;
 import org.apache.atlas.TestUtilsV2;
-import org.apache.atlas.RequestContextV1;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.SearchFilter;
-import org.apache.atlas.model.typedef.AtlasClassificationDef;
-import org.apache.atlas.model.typedef.AtlasEntityDef;
-import org.apache.atlas.model.typedef.AtlasEnumDef;
-import org.apache.atlas.model.typedef.AtlasStructDef;
+import org.apache.atlas.model.typedef.*;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
-import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.store.AtlasTypeDefStore;
 import org.apache.atlas.type.AtlasType;
 import org.slf4j.Logger;
@@ -38,11 +34,12 @@ import org.testng.annotations.BeforeTest;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
-
-import java.util.Arrays;
 import java.util.Collections;
-import java.util.Date;
 import java.util.List;
+import java.util.Set;
+import java.util.HashSet;
+import java.util.Arrays;
+import java.util.Date;
 
 import static org.testng.Assert.*;
 
@@ -389,6 +386,146 @@ public class AtlasTypeDefGraphStoreTest {
     }
 
     @Test(dependsOnMethods = "testGet")
+    public void testCreateClassificationDefWithValidEntityType(){
+        final String entityTypeName ="testCreateClassificationDefWithValidEntityTypeEntity1";
+        final String classificationTypeName ="testCreateClassificationDefWithValidEntityTypeClassification1";
+
+        List<AtlasEntityDef> entityDefs = TestUtilsV2.getEntityWithName(entityTypeName);
+
+        // Test Classification with entitytype
+        List<AtlasClassificationDef> classificationDefs = TestUtilsV2.getClassificationWithName(classificationTypeName);
+
+        Set<String> entityTypeNames =  new HashSet<String>();
+        entityTypeNames.add(entityTypeName);
+        classificationDefs.get(0).setEntityTypes(entityTypeNames);
+        AtlasTypesDef toCreate = new AtlasTypesDef(Collections.<AtlasEnumDef>emptyList(),
+                Collections.<AtlasStructDef>emptyList(),
+                classificationDefs,
+                entityDefs);
+        try {
+            AtlasTypesDef created = typeDefStore.createTypesDef(toCreate);
+            assertEquals(created.getClassificationDefs(), toCreate.getClassificationDefs(),
+                    "Classification creation with valid entitytype should've succeeded");
+        } catch (AtlasBaseException e) {
+            fail("Classification creation with valid entitytype should've succeeded. Failed with " + e.getMessage());
+        }
+    }
+
+    @Test(dependsOnMethods = "testGet")
+    public void testCreateWithInvalidEntityType(){
+        final String classificationTypeName ="testCreateClassificationDefWithInvalidEntityTypeClassification1";
+        // Test Classification with entitytype
+        List<AtlasClassificationDef> classificationDefs = TestUtilsV2.getClassificationWithName(classificationTypeName);
+
+        Set<String> entityTypeNames =  new HashSet<String>();
+        entityTypeNames.add("cccc");
+        classificationDefs.get(0).setEntityTypes(entityTypeNames);
+        AtlasTypesDef toCreate = new AtlasTypesDef(Collections.<AtlasEnumDef>emptyList(),
+                Collections.<AtlasStructDef>emptyList(),
+                classificationDefs,
+                Collections.<AtlasEntityDef>emptyList());
+        try {
+            AtlasTypesDef created = typeDefStore.createTypesDef(toCreate);
+            fail("Classification creation with invalid entitytype should've failed");
+        } catch (AtlasBaseException e) {
+
+        }
+    }
+
+    /**
+     * test that specifying an entitytype in a child classificationDef when then parent has unrestricted entityTypes fails.
+     */
+    @Test(dependsOnMethods = "testGet")
+    public void testCreateWithInvalidEntityType2(){
+        final String classificationTypeName1 ="testCreateClassificationDefWithInvalidEntityType2Classification1";
+        final String classificationTypeName2 ="testCreateClassificationDefWithInvalidEntityType2Classification2";
+        final String entityTypeName1 ="testCreateClassificationDefWithInvalidEntityType2Entity1";
+
+        // Test Classification with entitytype
+        AtlasClassificationDef classificationDef1 = TestUtilsV2.getSingleClassificationWithName(classificationTypeName1);
+        AtlasClassificationDef classificationDef2 = TestUtilsV2.getSingleClassificationWithName(classificationTypeName2);
+        List<AtlasEntityDef> entityDefs = TestUtilsV2.getEntityWithName(entityTypeName1);
+
+
+        Set<String> entityTypeNames =  new HashSet<String>();
+        entityTypeNames.add(entityTypeName1);
+
+        Set<String> superTypes =  new HashSet<String>();
+        superTypes.add(classificationTypeName1);
+
+        classificationDef2.setSuperTypes(superTypes);
+        classificationDef1.setEntityTypes(entityTypeNames);
+
+        TestUtilsV2.populateSystemAttributes(classificationDef1);
+        TestUtilsV2.populateSystemAttributes(classificationDef2);
+
+        List<AtlasClassificationDef>  classificationDefs = Arrays.asList(classificationDef1,classificationDef2);
+
+        AtlasTypesDef toCreate = new AtlasTypesDef(Collections.<AtlasEnumDef>emptyList(),
+                Collections.<AtlasStructDef>emptyList(),
+                classificationDefs,
+                Collections.<AtlasEntityDef>emptyList());
+        try {
+            AtlasTypesDef created = typeDefStore.createTypesDef(toCreate);
+            fail("Classification creation with invalid entitytype should've failed");
+        } catch (AtlasBaseException e) {
+
+        }
+    }
+
+    /**
+     * test that specifying an entitytype in a child classificationDef which is not in the parent fails
+     */
+    @Test(dependsOnMethods = "testGet")
+    public void testCreateWithInvalidEntityType3(){
+        final String classificationTypeName1 ="testCreateClassificationDefWithInvalidEntityType3Classification1";
+        final String classificationTypeName2 ="testCreateClassificationDefWithInvalidEntityType3Classification2";
+        final String entityTypeName1 ="testCreateClassificationDefWithInvalidEntityType3Entity1";
+        final String entityTypeName2 ="testCreateClassificationDefWithInvalidEntityType3Entity2";
+
+
+        // Test Classification with entitytype
+        AtlasClassificationDef classificationDef1 = TestUtilsV2.getSingleClassificationWithName(classificationTypeName1);
+        AtlasClassificationDef classificationDef2 = TestUtilsV2.getSingleClassificationWithName(classificationTypeName2);
+        AtlasEntityDef entityDef1 = TestUtilsV2.getSingleEntityWithName(entityTypeName1);
+        AtlasEntityDef entityDef2 = TestUtilsV2.getSingleEntityWithName(entityTypeName2);
+
+        Set<String> entityTypeNames1 =  new HashSet<String>();
+        entityTypeNames1.add(entityTypeName1);
+
+        Set<String> entityTypeNames2 =  new HashSet<String>();
+        entityTypeNames2.add(entityTypeName2);
+
+        Set<String> superTypes =  new HashSet<String>();
+        superTypes.add(classificationTypeName1);
+
+        classificationDef1.setEntityTypes(entityTypeNames1);
+
+
+        classificationDef2.setSuperTypes(superTypes);
+        classificationDef2.setEntityTypes(entityTypeNames2);
+
+        TestUtilsV2.populateSystemAttributes(classificationDef1);
+        TestUtilsV2.populateSystemAttributes(classificationDef2);
+        TestUtilsV2.populateSystemAttributes(entityDef1);
+        TestUtilsV2.populateSystemAttributes(entityDef2);
+
+        List<AtlasClassificationDef>  classificationDefs = Arrays.asList(classificationDef1,classificationDef2);
+        List<AtlasEntityDef>  entityDefs = Arrays.asList(entityDef1,entityDef2);
+
+        AtlasTypesDef toCreate = new AtlasTypesDef(Collections.<AtlasEnumDef>emptyList(),
+                Collections.<AtlasStructDef>emptyList(),
+                classificationDefs,
+                entityDefs);
+        try {
+            AtlasTypesDef created = typeDefStore.createTypesDef(toCreate);
+            fail("Classification creation with invalid entitytype should've failed");
+        } catch (AtlasBaseException e) {
+
+        }
+    }
+
+    @Test(dependsOnMethods = "testGet")
     public void testSearchFunctionality() {
         SearchFilter searchFilter = new SearchFilter();
         searchFilter.setParam(SearchFilter.PARAM_SUPERTYPE, "Person");

repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1Test.java

@@ -941,7 +941,7 @@ public class AtlasEntityStoreV1Test {
             String guid = createdEntity.get(0).getGuid();
             entityStore.addClassification(Arrays.asList(guid), new AtlasClassification(aTag.getName(), "testAttribute", "test-string"));
         } catch (AtlasBaseException e) {
-            fail("DB entity creation should've succeeded");
+            fail("DB entity creation should've succeeded, e.getMessage() => " + e.getMessage());
         }
 
     }