ATLAS-1283: added attribute AtlasBaseTypeDef.catagory

intg/src/main/java/org/apache/atlas/model/TypeCategory.java

@@ -18,5 +18,5 @@
 package org.apache.atlas.model;
 
 public enum TypeCategory {
-    PRIMITIVE, ARRAY, MAP, ENTITY, STRUCT, CLASSIFICATION, OBJECT_ID_TYPE
+    PRIMITIVE, OBJECT_ID_TYPE, ENUM, STRUCT, CLASSIFICATION, ENTITY, ARRAY, MAP
 }

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

@@ -27,6 +27,7 @@ import javax.xml.bind.annotation.XmlAccessType;
 import javax.xml.bind.annotation.XmlAccessorType;
 import javax.xml.bind.annotation.XmlRootElement;
 
+import org.apache.atlas.model.TypeCategory;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.collections.MapUtils;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
@@ -103,6 +104,7 @@ public abstract class AtlasBaseTypeDef implements java.io.Serializable {
     public static final String     SERIALIZED_DATE_FORMAT_STR = "yyyy-MM-dd'T'HH:mm:ss.SSS'Z'";
     public static final DateFormat DATE_FORMATTER             = new SimpleDateFormat(SERIALIZED_DATE_FORMAT_STR);
 
+    private final TypeCategory category;
     private String  guid       = null;
     private String  createdBy  = null;
     private String  updatedBy  = null;
@@ -113,21 +115,11 @@ public abstract class AtlasBaseTypeDef implements java.io.Serializable {
     private String  description;
     private String  typeVersion;
 
-    public AtlasBaseTypeDef() {
-        this(null, null, null);
-    }
-
-    public AtlasBaseTypeDef(String name) {
-        this(name, null, null);
-    }
-
-    public AtlasBaseTypeDef(String name, String description) {
-        this(name, description, null);
-    }
-
-    public AtlasBaseTypeDef(String name, String description, String typeVersion) {
+    protected AtlasBaseTypeDef(TypeCategory category, String name, String description, String typeVersion) {
         super();
 
+        this.category = category;
+
         setGuid(null);
         setCreatedBy(null);
         setUpdatedBy(null);
@@ -139,8 +131,10 @@ public abstract class AtlasBaseTypeDef implements java.io.Serializable {
         setTypeVersion(typeVersion);
     }
 
-    public AtlasBaseTypeDef(AtlasBaseTypeDef other) {
+    protected AtlasBaseTypeDef(AtlasBaseTypeDef other) {
         if (other != null) {
+            this.category = other.category;
+
             setGuid(other.getGuid());
             setCreatedBy(other.getCreatedBy());
             setUpdatedBy(other.getUpdatedBy());
@@ -151,6 +145,8 @@ public abstract class AtlasBaseTypeDef implements java.io.Serializable {
             setDescription(other.getDescription());
             setTypeVersion(other.getTypeVersion());
         } else {
+            this.category = TypeCategory.PRIMITIVE;
+
             setGuid(null);
             setCreatedBy(null);
             setUpdatedBy(null);
@@ -163,6 +159,8 @@ public abstract class AtlasBaseTypeDef implements java.io.Serializable {
         }
     }
 
+    public TypeCategory getCategory() { return category; }
+
     public String getGuid() {
         return guid;
     }
@@ -242,7 +240,8 @@ public abstract class AtlasBaseTypeDef implements java.io.Serializable {
         }
 
         sb.append("AtlasBaseTypeDef{");
-        sb.append("guid='").append(guid).append('\'');
+        sb.append("category='").append(category).append('\'');
+        sb.append(", guid='").append(guid).append('\'');
         sb.append(", createdBy='").append(createdBy).append('\'');
         sb.append(", updatedBy='").append(updatedBy).append('\'');
         dumpDateField(", createTime=", createTime, sb);
@@ -263,6 +262,7 @@ public abstract class AtlasBaseTypeDef implements java.io.Serializable {
 
         AtlasBaseTypeDef that = (AtlasBaseTypeDef) o;
 
+        if (category != null ? !category.equals(that.category) : that.category != null) { return false; }
         if (guid != null ? !guid.equals(that.guid) : that.guid != null) { return false; }
         if (createdBy != null ? !createdBy.equals(that.createdBy) : that.createdBy != null) { return false; }
         if (updatedBy != null ? !updatedBy.equals(that.updatedBy) : that.updatedBy != null) { return false; }
@@ -279,7 +279,8 @@ public abstract class AtlasBaseTypeDef implements java.io.Serializable {
 
     @Override
     public int hashCode() {
-        int result = guid != null ? guid.hashCode() : 0;
+        int result = category != null ? category.hashCode() : 0;
+        result = 31 * result + (guid != null ? guid.hashCode() : 0);
         result = 31 * result + (createdBy != null ? createdBy.hashCode() : 0);
         result = 31 * result + (updatedBy != null ? updatedBy.hashCode() : 0);
         result = 31 * result + (createTime != null ? createTime.hashCode() : 0);

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

@@ -28,6 +28,7 @@ import javax.xml.bind.annotation.XmlSeeAlso;
 
 import org.apache.atlas.model.PList;
 import org.apache.atlas.model.SearchFilter.SortType;
+import org.apache.atlas.model.TypeCategory;
 import org.apache.commons.collections.CollectionUtils;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
 import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONLY;
@@ -51,9 +52,7 @@ public class AtlasClassificationDef extends AtlasStructDef implements java.io.Se
 
 
     public AtlasClassificationDef() {
-        super();
-
-        setSuperTypes(null);
+        this(null, null, null, null, null);
     }
 
     public AtlasClassificationDef(String name) {
@@ -75,7 +74,7 @@ public class AtlasClassificationDef extends AtlasStructDef implements java.io.Se
 
     public AtlasClassificationDef(String name, String description, String typeVersion,
                                   List<AtlasAttributeDef> attributeDefs, Set<String> superTypes) {
-        super(name, description, typeVersion, attributeDefs);
+        super(TypeCategory.CLASSIFICATION, name, description, typeVersion, attributeDefs);
 
         setSuperTypes(superTypes);
     }

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

@@ -28,6 +28,7 @@ import javax.xml.bind.annotation.XmlSeeAlso;
 
 import org.apache.atlas.model.PList;
 import org.apache.atlas.model.SearchFilter.SortType;
+import org.apache.atlas.model.TypeCategory;
 import org.apache.commons.collections.CollectionUtils;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
 import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONLY;
@@ -72,7 +73,7 @@ public class AtlasEntityDef extends AtlasStructDef implements java.io.Serializab
 
     public AtlasEntityDef(String name, String description, String typeVersion, List<AtlasAttributeDef> attributeDefs,
                           Set<String> superTypes) {
-        super(name, description, typeVersion, attributeDefs);
+        super(TypeCategory.ENTITY, name, description, typeVersion, attributeDefs);
 
         setSuperTypes(superTypes);
     }

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

@@ -27,6 +27,7 @@ import javax.xml.bind.annotation.XmlSeeAlso;
 
 import org.apache.atlas.model.PList;
 import org.apache.atlas.model.SearchFilter.SortType;
+import org.apache.atlas.model.TypeCategory;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.hadoop.util.StringUtils;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
@@ -51,10 +52,7 @@ public class AtlasEnumDef extends AtlasBaseTypeDef implements Serializable {
     private String                    defaultValue;
 
     public AtlasEnumDef() {
-        super();
-
-        setElementDefs(null);
-        setDefaultValue(null);
+        this(null, null, null, null, null);
     }
 
     public AtlasEnumDef(String name) {
@@ -79,7 +77,7 @@ public class AtlasEnumDef extends AtlasBaseTypeDef implements Serializable {
 
     public AtlasEnumDef(String name, String description, String typeVersion, List<AtlasEnumElementDef> elementDefs,
                         String defaultValue) {
-        super(name, description, typeVersion);
+        super(TypeCategory.ENUM, name, description, typeVersion);
 
         setElementDefs(elementDefs);
         setDefaultValue(defaultValue);

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

@@ -34,6 +34,7 @@ import javax.xml.bind.annotation.XmlSeeAlso;
 
 import org.apache.atlas.model.PList;
 import org.apache.atlas.model.SearchFilter.SortType;
+import org.apache.atlas.model.TypeCategory;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.hadoop.util.StringUtils;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
@@ -75,7 +76,11 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
     }
 
     public AtlasStructDef(String name, String description, String typeVersion, List<AtlasAttributeDef> attributeDefs) {
-        super(name, description, typeVersion);
+        this(TypeCategory.STRUCT, name, description, typeVersion, attributeDefs);
+    }
+
+    protected AtlasStructDef(TypeCategory category, String name, String description, String typeVersion, List<AtlasAttributeDef> attributeDefs) {
+        super(category, name, description, typeVersion);
 
         setAttributeDefs(attributeDefs);
     }

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

@@ -34,19 +34,25 @@ import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONL
 @JsonIgnoreProperties(ignoreUnknown=true)
 @XmlRootElement
 @XmlAccessorType(XmlAccessType.PROPERTY)
-public class AtlasTypeDefHeader {
+public class AtlasTypeDefHeader implements java.io.Serializable {
+    private static final long serialVersionUID = 1L;
+
     private String guid;
     private String name;
     private TypeCategory category;
 
+    public AtlasTypeDefHeader() {
+        this(null, null, null);
+    }
+
     public AtlasTypeDefHeader(String guid, String name, TypeCategory category) {
         this.guid = guid;
         this.name = name;
         this.category = category;
     }
 
-    public AtlasTypeDefHeader() {
-        this(null, null, null);
+    public AtlasTypeDefHeader(AtlasBaseTypeDef typeDef) {
+        this(typeDef.getGuid(), typeDef.getName(), typeDef.getCategory());
     }
 
     public AtlasTypeDefHeader(AtlasTypeDefHeader other) {

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

@@ -45,12 +45,13 @@ public class AtlasClassificationType extends AtlasStructType {
 
     private final AtlasClassificationDef classificationDef;
 
-    private List<AtlasClassificationType>  superTypes       = Collections.emptyList();
-    private Set<String>                    allSuperTypes    = Collections.emptySet();
-    private Map<String, AtlasAttributeDef> allAttributeDefs = Collections.emptyMap();
+    private List<AtlasClassificationType>  superTypes        = Collections.emptyList();
+    private Set<String>                    allSuperTypes     = Collections.emptySet();
+    private Map<String, AtlasAttributeDef> allAttributeDefs  = Collections.emptyMap();
+    private Map<String, AtlasType>         allAttributeTypes = new HashMap<>();
 
     public AtlasClassificationType(AtlasClassificationDef classificationDef) {
-        super(classificationDef, TypeCategory.CLASSIFICATION);
+        super(classificationDef);
 
         this.classificationDef = classificationDef;
     }
@@ -64,6 +65,8 @@ public class AtlasClassificationType extends AtlasStructType {
         resolveReferences(typeRegistry);
     }
 
+    public AtlasClassificationDef getClassificationDef() { return classificationDef; }
+
     @Override
     public void resolveReferences(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
         super.resolveReferences(typeRegistry);
@@ -85,9 +88,10 @@ public class AtlasClassificationType extends AtlasStructType {
             }
         }
 
-        this.superTypes       = Collections.unmodifiableList(s);
-        this.allSuperTypes    = Collections.unmodifiableSet(allS);
-        this.allAttributeDefs = Collections.unmodifiableMap(allA);
+        this.superTypes        = Collections.unmodifiableList(s);
+        this.allSuperTypes     = Collections.unmodifiableSet(allS);
+        this.allAttributeDefs  = Collections.unmodifiableMap(allA);
+        this.allAttributeTypes = new HashMap<>(); // this will be rebuilt on calls to getAttributeType()
     }
 
     public Set<String> getSuperTypes() {
@@ -98,6 +102,49 @@ public class AtlasClassificationType extends AtlasStructType {
 
     public Map<String, AtlasAttributeDef> getAllAttributeDefs() { return allAttributeDefs; }
 
+    @Override
+    public AtlasType getAttributeType(String attributeName) {
+        AtlasType ret = allAttributeTypes.get(attributeName);
+
+        if (ret == null) {
+            ret = super.getAttributeType(attributeName);
+
+            if (ret == null) {
+                for (AtlasClassificationType superType : superTypes) {
+                    ret = superType.getAttributeType(attributeName);
+
+                    if (ret != null) {
+                        break;
+                    }
+                }
+            }
+
+            if (ret != null) {
+                allAttributeTypes.put(attributeName, ret);
+            }
+        }
+
+        return ret;
+    }
+
+
+    @Override
+    public AtlasAttributeDef getAttributeDef(String attributeName) {
+        AtlasAttributeDef ret = super.getAttributeDef(attributeName);
+
+        if (ret == null) {
+            for (AtlasClassificationType superType : superTypes) {
+                ret = superType.getAttributeDef(attributeName);
+
+                if (ret != null) {
+                    break;
+                }
+            }
+        }
+
+        return ret;
+    }
+
     public boolean isSuperTypeOf(AtlasClassificationType classificationType) {
         return classificationType != null ? classificationType.getAllSuperTypes().contains(this.getTypeName()) : false;
     }

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

@@ -44,12 +44,13 @@ public class AtlasEntityType extends AtlasStructType {
 
     private final AtlasEntityDef entityDef;
 
-    private List<AtlasEntityType>          superTypes       = Collections.emptyList();
-    private Set<String>                    allSuperTypes    = Collections.emptySet();
-    private Map<String, AtlasAttributeDef> allAttributeDefs = Collections.emptyMap();
+    private List<AtlasEntityType>          superTypes        = Collections.emptyList();
+    private Set<String>                    allSuperTypes     = Collections.emptySet();
+    private Map<String, AtlasAttributeDef> allAttributeDefs  = Collections.emptyMap();
+    private Map<String, AtlasType>         allAttributeTypes = new HashMap<>();
 
     public AtlasEntityType(AtlasEntityDef entityDef) {
-        super(entityDef, TypeCategory.ENTITY);
+        super(entityDef);
 
         this.entityDef = entityDef;
     }
@@ -62,6 +63,8 @@ public class AtlasEntityType extends AtlasStructType {
         resolveReferences(typeRegistry);
     }
 
+    public AtlasEntityDef getEntityDef() { return entityDef; }
+
     @Override
     public void resolveReferences(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
         super.resolveReferences(typeRegistry);
@@ -78,14 +81,14 @@ public class AtlasEntityType extends AtlasStructType {
             if (superType instanceof AtlasEntityType) {
                 s.add((AtlasEntityType)superType);
             } else {
-                throw new AtlasBaseException(AtlasErrorCode.INCOMPATIBLE_SUPERTYPE, superTypeName,
-                        entityDef.getName());
+                throw new AtlasBaseException(AtlasErrorCode.INCOMPATIBLE_SUPERTYPE, superTypeName, entityDef.getName());
             }
         }
 
-        this.superTypes       = Collections.unmodifiableList(s);
-        this.allSuperTypes    = Collections.unmodifiableSet(allS);
-        this.allAttributeDefs = Collections.unmodifiableMap(allA);
+        this.superTypes        = Collections.unmodifiableList(s);
+        this.allSuperTypes     = Collections.unmodifiableSet(allS);
+        this.allAttributeDefs  = Collections.unmodifiableMap(allA);
+        this.allAttributeTypes = new HashMap<>(); // this will be rebuilt on calls to getAttributeType()
     }
 
     public Set<String> getSuperTypes() {
@@ -98,6 +101,48 @@ public class AtlasEntityType extends AtlasStructType {
 
     public Map<String, AtlasAttributeDef> getAllAttributeDefs() { return allAttributeDefs; }
 
+    @Override
+    public AtlasType getAttributeType(String attributeName) {
+        AtlasType ret = allAttributeTypes.get(attributeName);
+
+        if (ret == null) {
+            ret = super.getAttributeType(attributeName);
+
+            if (ret == null) {
+                for (AtlasEntityType superType : superTypes) {
+                    ret = superType.getAttributeType(attributeName);
+
+                    if (ret != null) {
+                        break;
+                    }
+                }
+            }
+
+            if (ret != null) {
+                allAttributeTypes.put(attributeName, ret);
+            }
+        }
+
+        return ret;
+    }
+
+    @Override
+    public AtlasAttributeDef getAttributeDef(String attributeName) {
+        AtlasAttributeDef ret = super.getAttributeDef(attributeName);
+
+        if (ret == null) {
+            for (AtlasEntityType superType : superTypes) {
+                ret = superType.getAttributeDef(attributeName);
+
+                if (ret != null) {
+                    break;
+                }
+            }
+        }
+
+        return ret;
+    }
+
     public boolean isSuperTypeOf(AtlasEntityType entityType) {
         return entityType != null ? entityType.getAllSuperTypes().contains(this.getTypeName()) : false;
     }

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

@@ -37,7 +37,7 @@ public class AtlasEnumType extends AtlasType {
     private final String                           defaultValue;
 
     public AtlasEnumType(AtlasEnumDef enumDef) {
-        super(enumDef.getName(), TypeCategory.PRIMITIVE);
+        super(enumDef);
 
         Map<String, AtlasEnumElementDef> e = new HashMap<String, AtlasEnumElementDef>();
 
@@ -60,6 +60,8 @@ public class AtlasEnumType extends AtlasType {
         this.defaultValue = d;
     }
 
+    public AtlasEnumDef getEnumDef() { return enumDef; }
+
     @Override
     public void resolveReferences(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
     }

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

@@ -55,25 +55,21 @@ public class AtlasStructType extends AtlasType {
 
 
     public AtlasStructType(AtlasStructDef structDef) {
-        super(structDef.getName(), TypeCategory.STRUCT);
-
-        this.structDef = structDef;
-    }
-
-    public AtlasStructType(AtlasStructDef structDef, TypeCategory category) {
-        super(structDef.getName(), category);
+        super(structDef);
 
         this.structDef = structDef;
     }
 
     public AtlasStructType(AtlasStructDef structDef, AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
-        super(structDef.getName(), TypeCategory.STRUCT);
+        super(structDef);
 
         this.structDef = structDef;
 
         this.resolveReferences(typeRegistry);
     }
 
+    public AtlasStructDef getStructDef() { return structDef; }
+
     public AtlasType getAttributeType(String attributeName) { return attrTypes.get(attributeName); }
 
     public AtlasAttributeDef getAttributeDef(String attributeName) { return structDef.getAttribute(attributeName); }
@@ -442,8 +438,4 @@ public class AtlasStructType extends AtlasType {
             this.attributeName = attributeName;
         }
     }
-
-    public AtlasStructDef getStructDefinition() {
-        return structDef;
-    }
 }

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

@@ -35,13 +35,16 @@ public abstract class AtlasType {
     private static final Gson GSON =
             new GsonBuilder().setDateFormat(AtlasBaseTypeDef.SERIALIZED_DATE_FORMAT_STR).create();
 
-    private final String typeName;
-
+    private final String       typeName;
     private final TypeCategory typeCategory;
 
-    protected AtlasType(String typeName, TypeCategory category) {
-        this.typeName = typeName;
-        this.typeCategory = category;
+    protected AtlasType(AtlasBaseTypeDef typeDef) {
+        this(typeDef.getName(), typeDef.getCategory());
+    }
+
+    protected AtlasType(String typeName, TypeCategory typeCategory) {
+        this.typeName     = typeName;
+        this.typeCategory = typeCategory;
     }
 
     public void resolveReferences(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
@@ -49,16 +52,14 @@ public abstract class AtlasType {
 
     public String getTypeName() { return typeName; }
 
+    public TypeCategory getTypeCategory() { return typeCategory; }
+
     public abstract Object createDefaultValue();
 
     public abstract boolean isValidValue(Object obj);
 
     public abstract Object getNormalizedValue(Object obj);
 
-    public TypeCategory getTypeCategory() {
-        return typeCategory;
-    }
-
     public boolean validateValue(Object obj, String objName, List<String> messages) {
         boolean ret = isValidValue(obj);
 

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

@@ -29,7 +29,6 @@ import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef.Cardinality;
 import org.apache.atlas.model.typedef.AtlasTypeDefHeader;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
-import org.apache.atlas.model.TypeCategory;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang.StringUtils;
 
@@ -202,23 +201,22 @@ public class AtlasTypeUtil {
         List<AtlasTypeDefHeader> headerList = new LinkedList<>();
         if (CollectionUtils.isNotEmpty(typesDef.getEnumDefs())) {
             for (AtlasEnumDef enumDef : typesDef.getEnumDefs()) {
-                headerList.add(new AtlasTypeDefHeader(enumDef.getGuid(), enumDef.getName(), TypeCategory.PRIMITIVE));
+                headerList.add(new AtlasTypeDefHeader(enumDef));
             }
         }
         if (CollectionUtils.isNotEmpty(typesDef.getStructDefs())) {
             for (AtlasStructDef structDef : typesDef.getStructDefs()) {
-                headerList.add(new AtlasTypeDefHeader(structDef.getGuid(), structDef.getName(), TypeCategory.STRUCT));
+                headerList.add(new AtlasTypeDefHeader(structDef));
             }
         }
         if (CollectionUtils.isNotEmpty(typesDef.getClassificationDefs())) {
             for (AtlasClassificationDef classificationDef : typesDef.getClassificationDefs()) {
-                headerList.add(new AtlasTypeDefHeader(classificationDef.getGuid(), classificationDef.getName(),
-                        TypeCategory.CLASSIFICATION));
+                headerList.add(new AtlasTypeDefHeader(classificationDef));
             }
         }
         if (CollectionUtils.isNotEmpty(typesDef.getEntityDefs())) {
             for (AtlasEntityDef entityDef : typesDef.getEntityDefs()) {
-                headerList.add(new AtlasTypeDefHeader(entityDef.getGuid(), entityDef.getName(), TypeCategory.ENTITY));
+                headerList.add(new AtlasTypeDefHeader(entityDef));
             }
         }
 

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

@@ -367,19 +367,21 @@ public class AtlasTypeDefGraphStoreV1 extends AtlasTypeDefGraphStore {
     }
 
     TypeCategory getTypeCategory(AtlasBaseTypeDef typeDef) {
-        TypeCategory ret = null;
-
-        if (typeDef instanceof AtlasEntityDef) {
-            ret = TypeCategory.CLASS;
-        } else if (typeDef instanceof AtlasClassificationDef) {
-            ret = TypeCategory.TRAIT;
-        } else if (typeDef instanceof AtlasStructDef) {
-            ret = TypeCategory.STRUCT;
-        } else if (typeDef instanceof AtlasEnumDef) {
-            ret = TypeCategory.ENUM;
+        switch (typeDef.getCategory()) {
+            case ENTITY:
+                return TypeCategory.CLASS;
+
+            case CLASSIFICATION:
+                return TypeCategory.TRAIT;
+
+            case STRUCT:
+                return TypeCategory.STRUCT;
+
+            case ENUM:
+                return TypeCategory.ENUM;
         }
 
-        return ret;
+        return null;
     }
 
     /*

repository/src/main/java/org/apache/atlas/repository/util/FilterUtil.java

@@ -18,6 +18,7 @@
 package org.apache.atlas.repository.util;
 
 import org.apache.atlas.model.SearchFilter;
+import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.model.typedef.AtlasClassificationDef;
 import org.apache.atlas.model.typedef.AtlasEntityDef;
@@ -81,17 +82,19 @@ public class FilterUtil {
             @Override
             public boolean evaluate(Object o) {
                 if (o instanceof AtlasBaseTypeDef) {
+                    AtlasBaseTypeDef typeDef = (AtlasBaseTypeDef)o;
+
                     switch (type.toUpperCase()) {
                         case "CLASS":
                         case "ENTITY":
-                            return o instanceof AtlasEntityDef;
+                            return typeDef.getCategory() == TypeCategory.ENTITY;
                         case "TRAIT":
                         case "CLASSIFICATION":
-                            return o instanceof AtlasClassificationDef;
+                            return typeDef.getCategory() == TypeCategory.CLASSIFICATION;
                         case "STRUCT":
-                            return o instanceof AtlasStructDef;
+                            return typeDef.getCategory() == TypeCategory.STRUCT;
                         case "ENUM":
-                            return o instanceof AtlasEnumDef;
+                            return typeDef.getCategory() == TypeCategory.ENUM;
                         default:
                             // This shouldn't have happened
                             return false;