ATLAS-1490: added methods to get sub-types of entity and classification types

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

@@ -24,6 +24,7 @@ import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.typedef.AtlasClassificationDef;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -44,8 +45,9 @@ public class AtlasClassificationType extends AtlasStructType {
 
     private final AtlasClassificationDef classificationDef;
 
-    private List<AtlasClassificationType>  superTypes        = Collections.emptyList();
-    private Set<String>                    allSuperTypes     = Collections.emptySet();
+    private List<AtlasClassificationType>  superTypes    = Collections.emptyList();
+    private Set<String>                    allSuperTypes = Collections.emptySet();
+    private Set<String>                    allSubTypes   = Collections.emptySet();
 
     public AtlasClassificationType(AtlasClassificationDef classificationDef) {
         super(classificationDef);
@@ -68,9 +70,9 @@ public class AtlasClassificationType extends AtlasStructType {
     public void resolveReferences(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
         super.resolveReferences(typeRegistry);
 
-        List<AtlasClassificationType>  s    = new ArrayList<>();
-        Set<String>                    allS = new HashSet<>();
-        Map<String, AtlasAttribute> allA    = new HashMap<>();
+        List<AtlasClassificationType> s    = new ArrayList<>();
+        Set<String>                   allS = new HashSet<>();
+        Map<String, AtlasAttribute>   allA = new HashMap<>();
 
         getTypeHierarchyInfo(typeRegistry, allS, allA);
 
@@ -85,9 +87,24 @@ public class AtlasClassificationType extends AtlasStructType {
             }
         }
 
-        this.superTypes        = Collections.unmodifiableList(s);
-        this.allSuperTypes     = Collections.unmodifiableSet(allS);
-        this.allAttributes     = Collections.unmodifiableMap(allA);
+        this.superTypes    = Collections.unmodifiableList(s);
+        this.allSuperTypes = Collections.unmodifiableSet(allS);
+        this.allAttributes = Collections.unmodifiableMap(allA);
+        this.allSubTypes   = new HashSet<>(); // this will be populated in resolveReferencesPhase2()
+    }
+
+    @Override
+    public void resolveReferencesPhase2(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
+        super.resolveReferencesPhase2(typeRegistry);
+
+        for (String superTypeName : allSuperTypes) {
+            AtlasClassificationType superType = typeRegistry.getClassificationTypeByName(superTypeName);
+            superType.addSubType(this);
+        }
+    }
+
+    private void addSubType(AtlasClassificationType subType) {
+        allSubTypes.add(subType.getTypeName());
     }
 
     public Set<String> getSuperTypes() {
@@ -96,14 +113,26 @@ public class AtlasClassificationType extends AtlasStructType {
 
     public Set<String> getAllSuperTypes() { return allSuperTypes; }
 
+    public Set<String> getAllSubTypes() {
+        return allSubTypes;
+    }
+
     public boolean isSuperTypeOf(AtlasClassificationType classificationType) {
-        return classificationType != null && classificationType.getAllSuperTypes().contains(this.getTypeName());
+        return classificationType != null && allSubTypes.contains(classificationType.getTypeName());
+    }
+
+    public boolean isSuperTypeOf(String classificationName) {
+        return StringUtils.isNotEmpty(classificationName) && allSubTypes.contains(classificationName);
     }
 
     public boolean isSubTypeOf(AtlasClassificationType classificationType) {
         return classificationType != null && allSuperTypes.contains(classificationType.getTypeName());
     }
 
+    public boolean isSubTypeOf(String classificationName) {
+        return StringUtils.isNotEmpty(classificationName) && allSuperTypes.contains(classificationName);
+    }
+
     @Override
     public AtlasClassification createDefaultValue() {
         AtlasClassification ret = new AtlasClassification(classificationDef.getName());
@@ -217,11 +246,9 @@ public class AtlasClassificationType extends AtlasStructType {
         if (CollectionUtils.isNotEmpty(classificationDef.getSuperTypes())) {
             visitedTypes.add(classificationDef.getName());
             for (String superTypeName : classificationDef.getSuperTypes()) {
-                AtlasType type = typeRegistry.getType(superTypeName);
-
-                if (type instanceof AtlasClassificationType) {
-                    AtlasClassificationType superType = (AtlasClassificationType) type;
+                AtlasClassificationType superType = typeRegistry.getClassificationTypeByName(superTypeName);
 
+                if (superType != null) {
                     superType.collectTypeHierarchyInfo(typeRegistry, allSuperTypeNames, allAttributes, visitedTypes);
                 }
             }

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

@@ -45,8 +45,9 @@ public class AtlasEntityType extends AtlasStructType {
 
     private final AtlasEntityDef entityDef;
 
-    private List<AtlasEntityType>          superTypes        = Collections.emptyList();
-    private Set<String>                    allSuperTypes     = Collections.emptySet();
+    private List<AtlasEntityType> superTypes    = Collections.emptyList();
+    private Set<String>           allSuperTypes = Collections.emptySet();
+    private Set<String>           allSubTypes   = Collections.emptySet();
 
     public AtlasEntityType(AtlasEntityDef entityDef) {
         super(entityDef);
@@ -68,8 +69,8 @@ public class AtlasEntityType extends AtlasStructType {
     public void resolveReferences(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
         super.resolveReferences(typeRegistry);
 
-        List<AtlasEntityType>          s    = new ArrayList<>();
-        Set<String>                    allS = new HashSet<>();
+        List<AtlasEntityType>       s    = new ArrayList<>();
+        Set<String>                 allS = new HashSet<>();
         Map<String, AtlasAttribute> allA    = new HashMap<>();
 
         getTypeHierarchyInfo(typeRegistry, allS, allA);
@@ -84,9 +85,24 @@ public class AtlasEntityType extends AtlasStructType {
             }
         }
 
-        this.superTypes        = Collections.unmodifiableList(s);
-        this.allSuperTypes     = Collections.unmodifiableSet(allS);
-        this.allAttributes     = Collections.unmodifiableMap(allA);
+        this.superTypes    = Collections.unmodifiableList(s);
+        this.allSuperTypes = Collections.unmodifiableSet(allS);
+        this.allAttributes = Collections.unmodifiableMap(allA);
+        this.allSubTypes   = new HashSet<>(); // this will be populated in resolveReferencesPhase2()
+    }
+
+    @Override
+    public void resolveReferencesPhase2(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
+        super.resolveReferencesPhase2(typeRegistry);
+
+        for (String superTypeName : allSuperTypes) {
+            AtlasEntityType superType = typeRegistry.getEntityTypeByName(superTypeName);
+            superType.addSubType(this);
+        }
+    }
+
+    private void addSubType(AtlasEntityType subType) {
+        allSubTypes.add(subType.getTypeName());
     }
 
     public Set<String> getSuperTypes() {
@@ -97,8 +113,16 @@ public class AtlasEntityType extends AtlasStructType {
         return allSuperTypes;
     }
 
+    public Set<String> getAllSubTypes() {
+        return allSubTypes;
+    }
+
     public boolean isSuperTypeOf(AtlasEntityType entityType) {
-        return entityType != null && entityType.getAllSuperTypes().contains(this.getTypeName());
+        return entityType != null && allSubTypes.contains(entityType.getTypeName());
+    }
+
+    public boolean isSuperTypeOf(String entityTypeName) {
+        return StringUtils.isNotEmpty(entityTypeName) && allSubTypes.contains(entityTypeName);
     }
 
     public boolean isSubTypeOf(AtlasEntityType entityType) {
@@ -238,10 +262,9 @@ public class AtlasEntityType extends AtlasStructType {
         if (CollectionUtils.isNotEmpty(entityDef.getSuperTypes())) {
             visitedTypes.add(entityDef.getName());
             for (String superTypeName : entityDef.getSuperTypes()) {
-                AtlasType type = typeRegistry.getType(superTypeName);
+                AtlasEntityType superType = typeRegistry.getEntityTypeByName(superTypeName);
 
-                if (type instanceof AtlasEntityType) {
-                    AtlasEntityType superType = (AtlasEntityType) type;
+                if (superType != null) {
                     superType.collectTypeHierarchyInfo(typeRegistry, allSuperTypeNames, allAttributes, visitedTypes);
                 }
             }
@@ -263,11 +286,7 @@ public class AtlasEntityType extends AtlasStructType {
             return false;
         } else {
             String typeName = objId.getTypeName();
-            if (!typeName.equals(getTypeName())) {
-                //TODO - Enable below after enabling subType check
-//                        if ( !isSuperTypeOf(typeName)) {
-//                            return false;
-//                        }
+            if (!typeName.equals(getTypeName()) && !isSuperTypeOf(typeName)) {
                 return false;
             }
         }

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

@@ -50,6 +50,9 @@ public abstract class AtlasType {
     public void resolveReferences(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
     }
 
+    public void resolveReferencesPhase2(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
+    }
+
     public String getTypeName() { return typeName; }
 
     public TypeCategory getTypeCategory() { return typeCategory; }

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

@@ -64,6 +64,8 @@ public class AtlasTypeRegistry {
 
     public Collection<String> getAllTypeNames() { return registryData.allTypes.getAllTypeNames(); }
 
+    public Collection<AtlasType> getAllTypes() { return registryData.allTypes.getAllTypes(); }
+
     public boolean isRegisteredType(String typeName) {
         return registryData.allTypes.isKnownType(typeName);
     }
@@ -120,6 +122,7 @@ public class AtlasTypeRegistry {
 
     public AtlasBaseTypeDef getTypeDefByGuid(String guid) { return registryData.getTypeDefByGuid(guid); }
 
+
     public Collection<AtlasEnumDef> getAllEnumDefs() { return registryData.enumDefs.getAll(); }
 
     public AtlasEnumDef getEnumDefByGuid(String guid) {
@@ -130,6 +133,12 @@ public class AtlasTypeRegistry {
         return registryData.enumDefs.getTypeDefByName(name);
     }
 
+    public Collection<String> getAllEnumDefNames() { return registryData.enumDefs.getAllNames(); }
+
+    public Collection<AtlasEnumType> getAllEnumTypes() { return registryData.enumDefs.getAllTypes(); }
+
+    public AtlasEnumType getEnumTypeByName(String name) { return registryData.enumDefs.getTypeByName(name); }
+
 
     public Collection<AtlasStructDef> getAllStructDefs() { return registryData.structDefs.getAll(); }
 
@@ -139,6 +148,12 @@ public class AtlasTypeRegistry {
 
     public AtlasStructDef getStructDefByName(String name) { return registryData.structDefs.getTypeDefByName(name); }
 
+    public Collection<String> getAllStructDefNames() { return registryData.structDefs.getAllNames(); }
+
+    public Collection<AtlasStructType> getAllStructTypes() { return registryData.structDefs.getAllTypes(); }
+
+    public AtlasStructType getStructTypeByName(String name) { return registryData.structDefs.getTypeByName(name); }
+
 
     public Collection<AtlasClassificationDef> getAllClassificationDefs() {
         return registryData.classificationDefs.getAll();
@@ -152,6 +167,16 @@ public class AtlasTypeRegistry {
         return registryData.classificationDefs.getTypeDefByName(name);
     }
 
+    public Collection<String> getAllClassificationDefNames() { return registryData.classificationDefs.getAllNames(); }
+
+    public Collection<AtlasClassificationType> getAllClassificationTypes() {
+        return registryData.classificationDefs.getAllTypes();
+    }
+
+    public AtlasClassificationType getClassificationTypeByName(String name) {
+        return registryData.classificationDefs.getTypeByName(name);
+    }
+
 
     public Collection<AtlasEntityDef> getAllEntityDefs() { return registryData.entityDefs.getAll(); }
 
@@ -163,6 +188,13 @@ public class AtlasTypeRegistry {
         return registryData.entityDefs.getTypeDefByName(name);
     }
 
+    public Collection<String> getAllEntityDefNames() { return registryData.entityDefs.getAllNames(); }
+
+    public Collection<AtlasEntityType> getAllEntityTypes() { return registryData.entityDefs.getAllTypes(); }
+
+    public AtlasEntityType getEntityTypeByName(String name) { return registryData.entityDefs.getTypeByName(name); }
+
+
     public AtlasTransientTypeRegistry createTransientTypeRegistry() {
         return new AtlasTransientTypeRegistry(this);
     }
@@ -172,12 +204,12 @@ public class AtlasTypeRegistry {
     }
 
     static class RegistryData {
-        final TypeCache                            allTypes;
-        final TypeDefCache<AtlasEnumDef>           enumDefs;
-        final TypeDefCache<AtlasStructDef>         structDefs;
-        final TypeDefCache<AtlasClassificationDef> classificationDefs;
-        final TypeDefCache<AtlasEntityDef>         entityDefs;
-        final TypeDefCache<? extends AtlasBaseTypeDef>[] allDefCaches;
+        final TypeCache                                                       allTypes;
+        final TypeDefCache<AtlasEnumDef, AtlasEnumType>                       enumDefs;
+        final TypeDefCache<AtlasStructDef, AtlasStructType>                   structDefs;
+        final TypeDefCache<AtlasClassificationDef, AtlasClassificationType>   classificationDefs;
+        final TypeDefCache<AtlasEntityDef, AtlasEntityType>                   entityDefs;
+        final TypeDefCache<? extends AtlasBaseTypeDef, ? extends AtlasType>[] allDefCaches;
 
         RegistryData() {
             allTypes           = new TypeCache();
@@ -284,6 +316,10 @@ public class AtlasTypeRegistry {
             for (AtlasType type : registryData.allTypes.getAllTypes()) {
                 type.resolveReferences(this);
             }
+
+            for (AtlasType type : registryData.allTypes.getAllTypes()) {
+                type.resolveReferencesPhase2(this);
+            }
         }
 
         public void addType(AtlasBaseTypeDef typeDef) throws AtlasBaseException {
@@ -730,25 +766,28 @@ class TypeCache {
     }
 }
 
-class TypeDefCache<T extends AtlasBaseTypeDef> {
+class TypeDefCache<T1 extends AtlasBaseTypeDef, T2 extends AtlasType> {
     private static final Logger LOG = LoggerFactory.getLogger(TypeDefCache.class);
-    private final TypeCache      typeCache;
-    private final Map<String, T> typeDefGuidMap;
-    private final Map<String, T> typeDefNameMap;
+    private final TypeCache       typeCache;
+    private final Map<String, T1> typeDefGuidMap;
+    private final Map<String, T1> typeDefNameMap;
+    private final Map<String, T2> typeNameMap;
 
     public TypeDefCache(TypeCache typeCache) {
         this.typeCache      = typeCache;
         this.typeDefGuidMap = new ConcurrentHashMap<>();
         this.typeDefNameMap = new ConcurrentHashMap<>();
+        this.typeNameMap    = new ConcurrentHashMap<>();
     }
 
     public TypeDefCache(TypeDefCache other, TypeCache typeCache) {
         this.typeCache      = typeCache;
         this.typeDefGuidMap = new ConcurrentHashMap<>(other.typeDefGuidMap);
         this.typeDefNameMap = new ConcurrentHashMap<>(other.typeDefNameMap);
+        this.typeNameMap    = new ConcurrentHashMap<>(other.typeNameMap);
     }
 
-    public void addType(T typeDef, AtlasType type) {
+    public void addType(T1 typeDef, T2 type) {
         if (typeDef != null && type != null) {
             if (StringUtils.isNotEmpty(typeDef.getGuid())) {
                 typeDefGuidMap.put(typeDef.getGuid(), typeDef);
@@ -756,28 +795,38 @@ class TypeDefCache<T extends AtlasBaseTypeDef> {
 
             if (StringUtils.isNotEmpty(typeDef.getName())) {
                 typeDefNameMap.put(typeDef.getName(), typeDef);
+                typeNameMap.put(typeDef.getName(), type);
             }
 
             typeCache.addType(typeDef, type);
         }
     }
 
-    public Collection<T> getAll() {
+    public Collection<T1> getAll() {
         return Collections.unmodifiableCollection(typeDefNameMap.values());
     }
 
-    public T getTypeDefByGuid(String guid) {
+    public Collection<String> getAllNames() { return Collections.unmodifiableCollection(typeDefNameMap.keySet()); }
+
+    public T1 getTypeDefByGuid(String guid) {
         return guid != null ? typeDefGuidMap.get(guid) : null;
     }
 
-    public T getTypeDefByName(String name) {
-
+    public T1 getTypeDefByName(String name) {
         return name != null ? typeDefNameMap.get(name) : null;
     }
 
+    public Collection<T2> getAllTypes() {
+        return Collections.unmodifiableCollection(typeNameMap.values());
+    }
+
+    public T2 getTypeByName(String name) {
+        return name != null ? typeNameMap.get(name) : null;
+    }
+
     public void updateGuid(String typeName, String newGuid) {
         if (typeName != null) {
-            T typeDef = typeDefNameMap.get(typeName);
+            T1 typeDef = typeDefNameMap.get(typeName);
 
             if (typeDef != null) {
                 String currGuid = typeDef.getGuid();
@@ -807,7 +856,7 @@ class TypeDefCache<T extends AtlasBaseTypeDef> {
 
     public void removeTypeDefByGuid(String guid) {
         if (guid != null) {
-            T typeDef = typeDefGuidMap.remove(guid);
+            T1 typeDef = typeDefGuidMap.remove(guid);
 
             typeCache.removeTypeByGuid(guid);
 
@@ -815,6 +864,7 @@ class TypeDefCache<T extends AtlasBaseTypeDef> {
 
             if (name != null) {
                 typeDefNameMap.remove(name);
+                typeNameMap.remove(name);
                 typeCache.removeTypeByName(name);
             }
 
@@ -823,8 +873,9 @@ class TypeDefCache<T extends AtlasBaseTypeDef> {
 
     public void removeTypeDefByName(String name) {
         if (name != null) {
-            T typeDef = typeDefNameMap.remove(name);
+            T1 typeDef = typeDefNameMap.remove(name);
 
+            typeNameMap.remove(name);
             typeCache.removeTypeByName(name);
 
             String guid = typeDef != null ? typeDef.getGuid() : null;

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

@@ -302,14 +302,12 @@ public final class  ModelTestUtil {
     public static AtlasEntity newEntity(AtlasEntityDef entityDef, AtlasTypeRegistry typesRegistry) {
         AtlasEntity ret = null;
 
-        try {
-            AtlasType dataType = typesRegistry.getType(entityDef.getName());
+        AtlasEntityType entityType = typesRegistry.getEntityTypeByName(entityDef.getName());
 
-            if (dataType instanceof AtlasEntityType) {
-                ret = ((AtlasEntityType) dataType).createDefaultValue();
-            }
-        } catch (AtlasBaseException excp) {
-            LOG.error("failed to get entity-type {}", entityDef.getName(), excp);
+        if (entityType != null) {
+            ret = entityType.createDefaultValue();
+        } else {
+            LOG.error("failed to get entity-type {}", entityDef.getName());
         }
 
         return ret;
@@ -322,14 +320,12 @@ public final class  ModelTestUtil {
     public static AtlasStruct newStruct(AtlasStructDef structDef, AtlasTypeRegistry typesRegistry) {
         AtlasStruct ret = null;
 
-        try {
-            AtlasType dataType = typesRegistry.getType(structDef.getName());
+        AtlasStructType structType = typesRegistry.getStructTypeByName(structDef.getName());
 
-            if (dataType instanceof AtlasStructType) {
-                ret = ((AtlasStructType)dataType).createDefaultValue();
-            }
-        } catch (AtlasBaseException excp) {
-            LOG.error("failed to get struct-type {}", structDef.getName(), excp);
+        if (structType != null) {
+            ret = structType.createDefaultValue();
+        } else {
+            LOG.error("failed to get struct-type {}", structDef.getName());
         }
 
         return ret;
@@ -343,14 +339,12 @@ public final class  ModelTestUtil {
                                                         AtlasTypeRegistry typesRegistry) {
         AtlasClassification ret = null;
 
-        try {
-            AtlasType dataType = typesRegistry.getType(classificationDef.getName());
+        AtlasClassificationType classificationType = typesRegistry.getClassificationTypeByName(classificationDef.getName());
 
-            if (dataType instanceof AtlasClassificationType) {
-                ret = ((AtlasClassificationType)dataType).createDefaultValue();
-            }
-        } catch (AtlasBaseException excp) {
-            LOG.error("failed to get classification-type {}", classificationDef.getName(), excp);
+        if (classificationType != null) {
+            ret = classificationType.createDefaultValue();
+        } else {
+            LOG.error("failed to get classification-type {}", classificationDef.getName());
         }
 
         return ret;

intg/src/test/java/org/apache/atlas/model/instance/TestAtlasClassification.java

@@ -26,6 +26,7 @@ import org.apache.atlas.type.AtlasTypeRegistry;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.assertTrue;
 
 
@@ -33,11 +34,11 @@ public class TestAtlasClassification {
 
     @Test
     public void testClassificationSerDe() throws AtlasBaseException {
-        AtlasClassificationDef classificationDef = ModelTestUtil.getClassificationDef();
-        AtlasTypeRegistry      typeRegistry      = ModelTestUtil.getTypesRegistry();
-        AtlasType              dataType          = typeRegistry.getType(classificationDef.getName());
+        AtlasClassificationDef  classificationDef  = ModelTestUtil.getClassificationDef();
+        AtlasTypeRegistry       typeRegistry       = ModelTestUtil.getTypesRegistry();
+        AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(classificationDef.getName());
 
-        assertTrue(dataType instanceof  AtlasClassificationType);
+        assertNotNull(classificationType);
 
         AtlasClassification ent1 = ModelTestUtil.newClassification(classificationDef, typeRegistry);
 
@@ -45,45 +46,45 @@ public class TestAtlasClassification {
 
         AtlasClassification ent2 = AtlasType.fromJson(jsonString, AtlasClassification.class);
 
-        ((AtlasClassificationType)dataType).normalizeAttributeValues(ent2);
+        classificationType.normalizeAttributeValues(ent2);
 
         assertEquals(ent2, ent1, "Incorrect serialization/deserialization of AtlasClassification");
     }
 
     @Test
     public void testClassificationSerDeWithSuperType() throws AtlasBaseException {
-        AtlasClassificationDef classificationDef = ModelTestUtil.getClassificationDefWithSuperType();
-        AtlasTypeRegistry      typeRegistry      = ModelTestUtil.getTypesRegistry();
-        AtlasType              dataType          = typeRegistry.getType(classificationDef.getName());
+        AtlasClassificationDef  classificationDef  = ModelTestUtil.getClassificationDefWithSuperType();
+        AtlasTypeRegistry       typeRegistry       = ModelTestUtil.getTypesRegistry();
+        AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(classificationDef.getName());
 
-        assertTrue(dataType instanceof AtlasClassificationType);
+        assertNotNull(classificationType);
 
-        AtlasClassification ent1 =  ((AtlasClassificationType)dataType).createDefaultValue();
+        AtlasClassification ent1 =  classificationType.createDefaultValue();
 
         String jsonString = AtlasType.toJson(ent1);
 
         AtlasClassification ent2 = AtlasType.fromJson(jsonString, AtlasClassification.class);
 
-        ((AtlasClassificationType)dataType).normalizeAttributeValues(ent2);
+        classificationType.normalizeAttributeValues(ent2);
 
         assertEquals(ent2, ent1, "Incorrect serialization/deserialization of AtlasClassification with superType");
     }
 
     @Test
     public void testClassificationSerDeWithSuperTypes() throws AtlasBaseException {
-        AtlasClassificationDef classificationDef = ModelTestUtil.getClassificationDefWithSuperTypes();
-        AtlasTypeRegistry      typeRegistry      = ModelTestUtil.getTypesRegistry();
-        AtlasType              dataType          = typeRegistry.getType(classificationDef.getName());
+        AtlasClassificationDef  classificationDef  = ModelTestUtil.getClassificationDefWithSuperTypes();
+        AtlasTypeRegistry       typeRegistry       = ModelTestUtil.getTypesRegistry();
+        AtlasClassificationType classificationType = typeRegistry.getClassificationTypeByName(classificationDef.getName());
 
-        assertTrue(dataType instanceof  AtlasClassificationType);
+        assertNotNull(classificationType);
 
-        AtlasClassification ent1 =  ((AtlasClassificationType)dataType).createDefaultValue();
+        AtlasClassification ent1 =  classificationType.createDefaultValue();
 
         String jsonString = AtlasType.toJson(ent1);
 
         AtlasClassification ent2 = AtlasType.fromJson(jsonString, AtlasClassification.class);
 
-        ((AtlasClassificationType)dataType).normalizeAttributeValues(ent2);
+        classificationType.normalizeAttributeValues(ent2);
 
         assertEquals(ent2, ent1, "Incorrect serialization/deserialization of AtlasClassification with superTypes");
     }

intg/src/test/java/org/apache/atlas/model/instance/TestAtlasEntity.java

@@ -26,6 +26,7 @@ import org.apache.atlas.type.AtlasTypeRegistry;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.assertTrue;
 
 
@@ -35,17 +36,17 @@ public class TestAtlasEntity {
     public void testEntitySerDe() throws AtlasBaseException {
         AtlasEntityDef    entityDef    = ModelTestUtil.getEntityDef();
         AtlasTypeRegistry typeRegistry = ModelTestUtil.getTypesRegistry();
-        AtlasType         dataType     = typeRegistry.getType(entityDef.getName());
+        AtlasEntityType   entityType   = typeRegistry.getEntityTypeByName(entityDef.getName());
 
-        assertTrue(dataType instanceof  AtlasEntityType);
+        assertNotNull(entityType);
 
-        AtlasEntity ent1 =  ((AtlasEntityType)dataType).createDefaultValue();
+        AtlasEntity ent1 =  entityType.createDefaultValue();
 
         String jsonString = AtlasType.toJson(ent1);
 
         AtlasEntity ent2 = AtlasType.fromJson(jsonString, AtlasEntity.class);
 
-        ((AtlasEntityType)dataType).normalizeAttributeValues(ent2);
+        entityType.normalizeAttributeValues(ent2);
 
         assertEquals(ent2, ent1, "Incorrect serialization/deserialization of AtlasEntity");
     }
@@ -54,17 +55,17 @@ public class TestAtlasEntity {
     public void testEntitySerDeWithSuperType() throws AtlasBaseException {
         AtlasEntityDef    entityDef    = ModelTestUtil.getEntityDefWithSuperType();
         AtlasTypeRegistry typeRegistry = ModelTestUtil.getTypesRegistry();
-        AtlasType         dataType     = typeRegistry.getType(entityDef.getName());
+        AtlasEntityType   entityType   = typeRegistry.getEntityTypeByName(entityDef.getName());
 
-        assertTrue(dataType instanceof  AtlasEntityType);
+        assertNotNull(entityType);
 
-        AtlasEntity ent1 =  ((AtlasEntityType)dataType).createDefaultValue();
+        AtlasEntity ent1 =  entityType.createDefaultValue();
 
         String jsonString = AtlasType.toJson(ent1);
 
         AtlasEntity ent2 = AtlasType.fromJson(jsonString, AtlasEntity.class);
 
-        ((AtlasEntityType)dataType).normalizeAttributeValues(ent2);
+        entityType.normalizeAttributeValues(ent2);
 
         assertEquals(ent2, ent1, "Incorrect serialization/deserialization of AtlasEntity with superType");
     }
@@ -73,17 +74,17 @@ public class TestAtlasEntity {
     public void testEntitySerDeWithSuperTypes() throws AtlasBaseException {
         AtlasEntityDef    entityDef    = ModelTestUtil.getEntityDefWithSuperTypes();
         AtlasTypeRegistry typeRegistry = ModelTestUtil.getTypesRegistry();
-        AtlasType         dataType     = typeRegistry.getType(entityDef.getName());
+        AtlasEntityType   entityType   = typeRegistry.getEntityTypeByName(entityDef.getName());
 
-        assertTrue(dataType instanceof  AtlasEntityType);
+        assertNotNull(entityType);
 
-        AtlasEntity ent1 =  ((AtlasEntityType)dataType).createDefaultValue();
+        AtlasEntity ent1 =  entityType.createDefaultValue();
 
         String jsonString = AtlasType.toJson(ent1);
 
         AtlasEntity ent2 = AtlasType.fromJson(jsonString, AtlasEntity.class);
 
-        ((AtlasEntityType)dataType).normalizeAttributeValues(ent2);
+        entityType.normalizeAttributeValues(ent2);
 
         assertEquals(ent2, ent1, "Incorrect serialization/deserialization of AtlasEntity with superTypes");
     }

intg/src/test/java/org/apache/atlas/type/TestAtlasTypeRegistry.java

@@ -46,7 +46,7 @@ public class TestAtlasTypeRegistry {
      *   L2_1  L2_2   L2_3   L2_4
      */
     @Test
-    public void testClassificationDefValidSuperTypes() {
+    public void testClassificationDefValidHierarchy() {
         AtlasClassificationDef classifiL0   = new AtlasClassificationDef("L0");
         AtlasClassificationDef classifiL1_1 = new AtlasClassificationDef("L1-1");
         AtlasClassificationDef classifiL1_2 = new AtlasClassificationDef("L1-2");
@@ -102,6 +102,14 @@ public class TestAtlasTypeRegistry {
         validateSuperTypes(typeRegistry, "L2-3", new HashSet<>(Arrays.asList("L1-1", "L0", "L1-2")));
         validateSuperTypes(typeRegistry, "L2-4", new HashSet<>(Arrays.asList("L1-2", "L0")));
 
+        validateSubTypes(typeRegistry, "L0", new HashSet<>(Arrays.asList("L1-1", "L1-2", "L2-1", "L2-2", "L2-3", "L2-4")));
+        validateSubTypes(typeRegistry, "L1-1", new HashSet<>(Arrays.asList("L2-1", "L2-2", "L2-3")));
+        validateSubTypes(typeRegistry, "L1-2", new HashSet<>(Arrays.asList("L2-3", "L2-4")));
+        validateSubTypes(typeRegistry, "L2-1", new HashSet<String>());
+        validateSubTypes(typeRegistry, "L2-2", new HashSet<String>());
+        validateSubTypes(typeRegistry, "L2-3", new HashSet<String>());
+        validateSubTypes(typeRegistry, "L2-4", new HashSet<String>());
+
         validateAttributeNames(typeRegistry, "L0", new HashSet<>(Arrays.asList("L0_a1")));
         validateAttributeNames(typeRegistry, "L1-1", new HashSet<>(Arrays.asList("L0_a1", "L1-1_a1")));
         validateAttributeNames(typeRegistry, "L1-2", new HashSet<>(Arrays.asList("L0_a1", "L1-2_a1")));
@@ -112,7 +120,7 @@ public class TestAtlasTypeRegistry {
     }
 
     @Test
-    public void testClassificationDefInvalidSuperTypes_Self() {
+    public void testClassificationDefInvalidHierarchy_Self() {
         AtlasClassificationDef classifiDef1 = new AtlasClassificationDef("classifiDef-1");
 
         classifiDef1.addSuperType(classifiDef1.getName());
@@ -141,7 +149,7 @@ public class TestAtlasTypeRegistry {
      *   L2_1  L2_2   L2_3   L2_4
      */
     @Test
-    public void testClassificationDefInvalidSuperTypes_CircularRef() {
+    public void testClassificationDefInvalidHierarchy_CircularRef() {
         AtlasClassificationDef classifiL0   = new AtlasClassificationDef("L0");
         AtlasClassificationDef classifiL1_1 = new AtlasClassificationDef("L1-1");
         AtlasClassificationDef classifiL1_2 = new AtlasClassificationDef("L1-2");
@@ -191,7 +199,7 @@ public class TestAtlasTypeRegistry {
      *   L2_1  L2_2   L2_3   L2_4
      */
     @Test
-    public void testEntityDefValidSuperTypes() {
+    public void testEntityDefValidHierarchy() {
         AtlasEntityDef entL0   = new AtlasEntityDef("L0");
         AtlasEntityDef entL1_1 = new AtlasEntityDef("L1-1");
         AtlasEntityDef entL1_2 = new AtlasEntityDef("L1-2");
@@ -247,6 +255,14 @@ public class TestAtlasTypeRegistry {
         validateSuperTypes(typeRegistry, "L2-3", new HashSet<>(Arrays.asList("L1-1", "L0", "L1-2")));
         validateSuperTypes(typeRegistry, "L2-4", new HashSet<>(Arrays.asList("L1-2", "L0")));
 
+        validateSubTypes(typeRegistry, "L0", new HashSet<>(Arrays.asList("L1-1", "L1-2", "L2-1", "L2-2", "L2-3", "L2-4")));
+        validateSubTypes(typeRegistry, "L1-1", new HashSet<>(Arrays.asList("L2-1", "L2-2", "L2-3")));
+        validateSubTypes(typeRegistry, "L1-2", new HashSet<>(Arrays.asList("L2-3", "L2-4")));
+        validateSubTypes(typeRegistry, "L2-1", new HashSet<String>());
+        validateSubTypes(typeRegistry, "L2-2", new HashSet<String>());
+        validateSubTypes(typeRegistry, "L2-3", new HashSet<String>());
+        validateSubTypes(typeRegistry, "L2-4", new HashSet<String>());
+
         validateAttributeNames(typeRegistry, "L0", new HashSet<>(Arrays.asList("L0_a1")));
         validateAttributeNames(typeRegistry, "L1-1", new HashSet<>(Arrays.asList("L0_a1", "L1-1_a1")));
         validateAttributeNames(typeRegistry, "L1-2", new HashSet<>(Arrays.asList("L0_a1", "L1-2_a1")));
@@ -257,7 +273,7 @@ public class TestAtlasTypeRegistry {
     }
 
     @Test
-    public void testEntityDefInvalidSuperTypes_Self() {
+    public void testEntityDefInvalidHierarchy_Self() {
         AtlasEntityDef entDef1 = new AtlasEntityDef("entDef-1");
 
         entDef1.addSuperType(entDef1.getName());
@@ -286,7 +302,7 @@ public class TestAtlasTypeRegistry {
      *   L2_1  L2_2   L2_3   L2_4
      */
     @Test
-    public void testEntityDefInvalidSuperTypes_CircularRef() {
+    public void testEntityDefInvalidHierarchy_CircularRef() {
         AtlasEntityDef entL0   = new AtlasEntityDef("L0");
         AtlasEntityDef entL1_1 = new AtlasEntityDef("L1-1");
         AtlasEntityDef entL1_2 = new AtlasEntityDef("L1-2");
@@ -347,6 +363,27 @@ public class TestAtlasTypeRegistry {
         assertEquals(superTypes, expectedSuperTypes);
     }
 
+    private void validateSubTypes(AtlasTypeRegistry typeRegistry, String typeName, Set<String> expectedSubTypes) {
+        AtlasType type = null;
+
+        try {
+            type = typeRegistry.getType(typeName);
+        } catch (AtlasBaseException excp) {
+        }
+
+        Set<String> subTypes = null;
+
+        if (type != null) {
+            if (type instanceof AtlasEntityType) {
+                subTypes = ((AtlasEntityType) type).getAllSubTypes();
+            } else if (type instanceof AtlasClassificationType) {
+                subTypes = ((AtlasClassificationType) type).getAllSubTypes();
+            }
+        }
+
+        assertEquals(subTypes, expectedSubTypes);
+    }
+
     private void validateAttributeNames(AtlasTypeRegistry typeRegistry, String typeName, Set<String> attributeNames) {
         AtlasType type = null;
 

release-log.txt

@@ -9,6 +9,7 @@ ATLAS-1060 Add composite indexes for exact match performance improvements for al
 ATLAS-1127 Modify creation and modification timestamps to Date instead of Long(sumasai)
 
 ALL CHANGES:
+ATLAS-1490 added methods to get sub-types of entity and classification types (mneethiraj)
 ATLAS-1437 UI update to disallow tag association changes to deleted entities (Kalyanikashikar via mneethiraj)
 ATLAS-1352 fix for error in redirecting to Knox gateway URL (nixonrodrigues via mneethiraj)
 ATLAS-1467 instance create/full-Update implementation (sumasai via mneethiraj)