ATLAS-1227: Added support for attribute constraints in the API

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

@@ -20,9 +20,11 @@ package org.apache.atlas.model.typedef;
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.ListIterator;
+import java.util.Map;
 import java.util.Set;
 
 import javax.xml.bind.annotation.XmlAccessType;
@@ -264,15 +266,17 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
         private int                      valuesMaxCount;
         private boolean                  isUnique;
         private boolean                  isIndexable;
+        private List<AtlasConstraintDef> constraintDefs;
 
         public AtlasAttributeDef() { this(null, null); }
 
         public AtlasAttributeDef(String name, String typeName) {
-            this(name, typeName, false, Cardinality.SINGLE, 1, 1, false, false);
+            this(name, typeName, false, Cardinality.SINGLE, 1, 1, false, false, null);
         }
 
         public AtlasAttributeDef(String name, String typeName, boolean isOptional, Cardinality cardinality,
-                                 int valuesMinCount, int valuesMaxCount, boolean isUnique, boolean isIndexable) {
+                                 int valuesMinCount, int valuesMaxCount, boolean isUnique, boolean isIndexable,
+                                 List<AtlasConstraintDef> constraintDefs) {
             setName(name);
             setTypeName(typeName);
             setOptional(isOptional);
@@ -281,6 +285,7 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
             setValuesMaxCount(valuesMaxCount);
             setUnique(isUnique);
             setIndexable(isIndexable);
+            setConstraintDefs(constraintDefs);
         }
 
         public AtlasAttributeDef(AtlasAttributeDef other) {
@@ -293,6 +298,7 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
                 setValuesMaxCount(other.getValuesMaxCount());
                 setUnique(other.isUnique());
                 setIndexable(other.isIndexable());
+                setConstraintDefs(other.getConstraintDefs());
             }
         }
 
@@ -358,6 +364,33 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
             isIndexable = idexable;
         }
 
+        public List<AtlasConstraintDef> getConstraintDefs() { return constraintDefs; }
+
+        public void setConstraintDefs(List<AtlasConstraintDef> constraintDefs) {
+            if (this.constraintDefs != null && this.constraintDefs == constraintDefs) {
+                return;
+            }
+
+            if (CollectionUtils.isEmpty(constraintDefs)) {
+                this.constraintDefs = null;
+            } else {
+                this.constraintDefs = new ArrayList<AtlasConstraintDef>(constraintDefs);
+            }
+        }
+
+        public void addConstraint(AtlasConstraintDef constraintDef) {
+            List<AtlasConstraintDef> cDefs = constraintDefs;
+
+            if (cDefs == null) {
+                cDefs = new ArrayList<>();
+            } else {
+                cDefs = new ArrayList<>(cDefs);
+            }
+
+            cDefs.add(constraintDef);
+
+            this.constraintDefs = cDefs;
+        }
         public StringBuilder toString(StringBuilder sb) {
             if (sb == null) {
                 sb = new StringBuilder();
@@ -372,6 +405,18 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
             sb.append(", valuesMaxCount=").append(valuesMaxCount);
             sb.append(", isUnique=").append(isUnique);
             sb.append(", isIndexable=").append(isIndexable);
+            sb.append(", constraintDefs=[");
+            if (CollectionUtils.isNotEmpty(constraintDefs)) {
+                int i = 0;
+                for (AtlasConstraintDef constraintDef : constraintDefs) {
+                    constraintDef.toString(sb);
+                    if (i > 0) {
+                        sb.append(", ");
+                    }
+                    i++;
+                }
+            }
+            sb.append("]");
             sb.append('}');
 
             return sb;
@@ -394,6 +439,9 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
             if (valuesMaxCount != that.valuesMaxCount) { return false; }
             if (isUnique != that.isUnique) { return false; }
             if (isIndexable != that.isIndexable) { return false; }
+            if (constraintDefs != null ? !constraintDefs.equals(that.constraintDefs) : that.constraintDefs != null) {
+                return false;
+            }
 
             return true;
         }
@@ -408,6 +456,7 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
             result = 31 * result + valuesMaxCount;
             result = 31 * result + (isUnique ? 1 : 0);
             result = 31 * result + (isIndexable ? 1 : 0);
+            result = 31 * result + (constraintDefs != null ? constraintDefs.hashCode() : 0);
             return result;
         }
 
@@ -418,6 +467,104 @@ public class AtlasStructDef extends AtlasBaseTypeDef implements Serializable {
     }
 
     /**
+     * class that captures details of a constraint.
+     */
+    @JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
+    @JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
+    @JsonIgnoreProperties(ignoreUnknown=true)
+    @XmlRootElement
+    @XmlAccessorType(XmlAccessType.PROPERTY)
+    public static class AtlasConstraintDef implements Serializable {
+        private static final long serialVersionUID = 1L;
+
+        public static final String CONSTRAINT_TYPE_FOREIGN_KEY     = "foreignKey";
+        public static final String CONSTRAINT_TYPE_MAPPED_FROM_REF = "mappedFromRef";
+        public static final String CONSTRAINT_PARAM_REF_ATTRIBUTE  = "refAttribute";
+        public static final String CONSTRAINT_PARAM_ON_DELETE      = "onDelete";
+        public static final String CONSTRAINT_PARAM_VAL_CASCADE    = "cascade";
+
+        private String              type;   // foreignKey/mappedFromRef/valueInRange
+        private Map<String, Object> params; // onDelete=cascade/refAttribute=attr2/min=0,max=23
+
+
+        public AtlasConstraintDef() { }
+
+        public AtlasConstraintDef(String type) {
+            this(type, null);
+        }
+
+        public AtlasConstraintDef(String type, Map<String, Object> params) {
+            this.type = type;
+
+            if (params != null) {
+                this.params = new HashMap<String, Object>(params);
+            }
+        }
+
+        public AtlasConstraintDef(AtlasConstraintDef that) {
+            if (that != null) {
+                this.type = that.type;
+
+                if (that.params != null) {
+                    this.params = new HashMap<String, Object>(that.params);
+                }
+            }
+        }
+
+        public String getType() {
+            return type;
+        }
+
+        public void setType(String type) {
+            this.type = type;
+        }
+
+        public Map<String, Object> getParams() {
+            return params;
+        }
+
+        public void setParams(Map<String, Object> params) {
+            this.params = params;
+        }
+
+        public StringBuilder toString(StringBuilder sb) {
+            if (sb == null) {
+                sb = new StringBuilder();
+            }
+
+            sb.append("AtlasConstraintDef{");
+            sb.append("type='").append(type).append('\'');
+            sb.append(", params='").append(params).append('\'');
+            sb.append('}');
+
+            return sb;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+
+            AtlasConstraintDef that = (AtlasConstraintDef) o;
+
+            if (type != null ? !type.equals(that.type) : that.type != null) return false;
+            if (params != null ? !params.equals(that.params) : that.params != null) return false;
+
+            return true;
+        }
+
+        @Override
+        public int hashCode() {
+            int result = type != null ? type.hashCode() : 0;
+            result = 31 * result + (params != null ? params.hashCode() : 0);
+            return result;
+        }
+
+        @Override
+        public String toString() { return toString(new StringBuilder()).toString(); }
+    }
+
+    /**
      * REST serialization friendly list.
      */
     @JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)

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

@@ -22,8 +22,14 @@ import java.util.*;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasStruct;
 import org.apache.atlas.model.typedef.AtlasStructDef;
+import org.apache.atlas.model.typedef.AtlasStructDef.AtlasConstraintDef;
+import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasConstraintDef.CONSTRAINT_PARAM_REF_ATTRIBUTE;
+import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasConstraintDef.CONSTRAINT_TYPE_FOREIGN_KEY;
+import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasConstraintDef.CONSTRAINT_TYPE_MAPPED_FROM_REF;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef.Cardinality;
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -37,6 +43,8 @@ public class AtlasStructType extends AtlasType {
     private final AtlasStructDef structDef;
 
     private Map<String, AtlasType>         attrTypes               = Collections.emptyMap();
+    private Set<String>                    foreignKeyAttributes    = new HashSet<>();
+    private Map<String, TypeAttributePair> mappedFromRefAttributes = new HashMap<String, TypeAttributePair>();
 
 
     public AtlasStructType(AtlasStructDef structDef) {
@@ -57,6 +65,31 @@ public class AtlasStructType extends AtlasType {
 
     public AtlasAttributeDef getAttributeDef(String attributeName) { return structDef.getAttribute(attributeName); }
 
+    public boolean isForeignKeyAttribute(String attributeName) {
+        return foreignKeyAttributes.contains(attributeName);
+    }
+
+    public boolean isMappedFromRefAttribute(String attributeName) {
+        return mappedFromRefAttributes.containsKey(attributeName);
+    }
+
+    public String getMappedFromRefAttribute(String typeName, String attribName) {
+        String ret = null;
+
+        for (Map.Entry<String, TypeAttributePair> e : mappedFromRefAttributes.entrySet()) {
+            String refTypeName   = e.getValue().typeName;
+            String refAttribName = e.getValue().attributeName;
+
+            if(StringUtils.equals(refTypeName, typeName) && StringUtils.equals(refAttribName, attribName)) {
+                ret = e.getKey();
+
+                break;
+            }
+        }
+
+        return ret;
+    }
+
     @Override
     public void resolveReferences(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
         Map<String, AtlasType> a = new HashMap<String, AtlasType>();
@@ -69,6 +102,8 @@ public class AtlasStructType extends AtlasType {
                                              + structDef.getName() + "." + attributeDef.getName());
             }
 
+            resolveConstraints(attributeDef, attrType);
+
             Cardinality cardinality = attributeDef.getCardinality();
 
             if (cardinality == Cardinality.LIST || cardinality == Cardinality.SET) {
@@ -290,4 +325,111 @@ public class AtlasStructType extends AtlasType {
 
         return null;
     }
+
+    private void resolveConstraints(AtlasAttributeDef attribDef, AtlasType attribType) throws AtlasBaseException {
+        if (attribDef == null || CollectionUtils.isEmpty(attribDef.getConstraintDefs()) || attribType == null) {
+            return;
+        }
+
+        for (AtlasStructDef.AtlasConstraintDef constraintDef : attribDef.getConstraintDefs()) {
+            String constraintType = constraintDef != null ? constraintDef.getType() : null;
+
+            if (StringUtils.isBlank(constraintType)) {
+                continue;
+            }
+
+            if (constraintType.equals(CONSTRAINT_TYPE_FOREIGN_KEY)) {
+                resolveForeignKeyConstraint(attribDef, constraintDef, attribType);
+            } else if (constraintType.equals(CONSTRAINT_TYPE_MAPPED_FROM_REF)) {
+                resolveMappedFromRefConstraint(attribDef, constraintDef, attribType);
+            } else {
+                throw new AtlasBaseException(getTypeName() + "." + attribDef.getName()
+                                             + ": unknown constraint " + constraintType);
+            }
+        }
+    }
+
+    /*
+     * valid conditions for foreign-key constraint:
+     *  - supported only in entity-type
+     *  - attribute should be an entity-type or an array of entity-type
+     */
+    private void resolveForeignKeyConstraint(AtlasAttributeDef attribDef, AtlasConstraintDef constraintDef,
+                                             AtlasType attribType) throws AtlasBaseException {
+        if (!(this instanceof AtlasEntityType)) {
+            throw new AtlasBaseException(getTypeName() + "." + attribDef.getName() + ": "
+                         + AtlasStructDef.AtlasConstraintDef.CONSTRAINT_TYPE_FOREIGN_KEY + " constraint not supported");
+        }
+
+        if (attribType instanceof AtlasArrayType) {
+            attribType = ((AtlasArrayType)attribType).getElementType();
+        }
+
+        if (!(attribType instanceof AtlasEntityType)) {
+            throw new AtlasBaseException(getTypeName() + "." + attribDef.getName() + ": "
+                                      + AtlasConstraintDef.CONSTRAINT_TYPE_FOREIGN_KEY + " incompatible attribute type "
+                                      + attribType.getTypeName());
+        }
+
+        foreignKeyAttributes.add(attribDef.getName());
+    }
+
+    /*
+     * valid conditions for mapped-from-ref constraint:
+     *  - supported only in entity-type
+     *  - attribute should be an entity-type or an array of entity-type
+     *  - attribute's entity-type should have a foreign-key constraint to this type
+     */
+    private void resolveMappedFromRefConstraint(AtlasAttributeDef attribDef, AtlasConstraintDef constraintDef,
+                                                AtlasType attribType) throws AtlasBaseException {
+        if (!(this instanceof AtlasEntityType)) {
+            throw new AtlasBaseException(getTypeName() + "." + attribDef.getName() + ": "
+                                    + CONSTRAINT_TYPE_MAPPED_FROM_REF + " constraint not supported");
+        }
+
+        if (attribType instanceof AtlasArrayType) {
+            attribType = ((AtlasArrayType)attribType).getElementType();
+        }
+
+        if (!(attribType instanceof AtlasEntityType)) {
+            throw new AtlasBaseException(getTypeName() + "." + attribDef.getName() + ": "
+                                  + CONSTRAINT_TYPE_MAPPED_FROM_REF + " incompatible attribute type "
+                                  + attribType.getTypeName());
+        }
+
+        String refAttribName = AtlasTypeUtil.getStringValue(constraintDef.getParams(), CONSTRAINT_PARAM_REF_ATTRIBUTE);
+
+        if (StringUtils.isBlank(refAttribName)) {
+            throw new AtlasBaseException(getTypeName() + "." + attribDef.getName() + ": "
+                        + CONSTRAINT_TYPE_MAPPED_FROM_REF + " invalid constraint. missing parameter "
+                        + CONSTRAINT_PARAM_REF_ATTRIBUTE);
+        }
+
+        AtlasStructType   structType = (AtlasStructType)attribType;
+        AtlasAttributeDef refAttrib  = structType.getAttributeDef(refAttribName);
+
+        if (refAttrib == null) {
+            throw new AtlasBaseException(getTypeName() + "." + attribDef.getName() + ": invalid constraint. "
+                              + CONSTRAINT_PARAM_REF_ATTRIBUTE + " " + structType.getTypeName() + "." + refAttribName
+                              + " does not exist");
+        }
+
+        if (!StringUtils.equals(getTypeName(), refAttrib.getTypeName())) {
+            throw new AtlasBaseException(getTypeName() + "." + attribDef.getName() + ": invalid constraint. Datatype of"
+                    + CONSTRAINT_PARAM_REF_ATTRIBUTE + " " + structType.getTypeName() + "." + refAttribName
+                    + " should be " + getTypeName() + ", but found " + refAttrib.getTypeName());
+        }
+
+        mappedFromRefAttributes.put(attribDef.getName(), new TypeAttributePair(refAttrib.getTypeName(), refAttribName));
+    }
+
+    private class TypeAttributePair {
+        public final String typeName;
+        public final String attributeName;
+
+        public TypeAttributePair(String typeName, String attributeName) {
+            this.typeName      = typeName;
+            this.attributeName = attributeName;
+        }
+    }
 }

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

@@ -26,6 +26,7 @@ import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_MAP_KEY
 import org.apache.commons.lang.StringUtils;
 
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -52,6 +53,22 @@ public class AtlasTypeUtil {
         return ATLAS_BUILTIN_TYPENAMES.contains(typeName);
     }
 
+    public static boolean isArrayType(String typeName) {
+        return StringUtils.startsWith(typeName, ATLAS_TYPE_ARRAY_PREFIX)
+            && StringUtils.endsWith(typeName, ATLAS_TYPE_ARRAY_SUFFIX);
+    }
+
+    public static boolean isMapType(String typeName) {
+        return StringUtils.startsWith(typeName, ATLAS_TYPE_MAP_PREFIX)
+                && StringUtils.endsWith(typeName, ATLAS_TYPE_MAP_SUFFIX);
+    }
+
+    public static String getStringValue(Map map, Object key) {
+        Object ret = map != null ? map.get(key) : null;
+
+        return ret != null ? ret.toString() : null;
+    }
+
     private static void getReferencedTypeNames(String typeName, Set<String> referencedTypeNames) {
         if (StringUtils.isNotBlank(typeName) && !referencedTypeNames.contains(typeName)) {
             if (typeName.startsWith(ATLAS_TYPE_ARRAY_PREFIX) && typeName.endsWith(ATLAS_TYPE_ARRAY_SUFFIX)) {

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

@@ -24,6 +24,8 @@ import org.apache.atlas.model.ModelTestUtil;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.model.typedef.AtlasEntityDef;
+import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
+import org.apache.atlas.model.typedef.AtlasStructDef.AtlasConstraintDef;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.*;
@@ -115,6 +117,74 @@ public class TestAtlasEntityType {
         }
     }
 
+    @Test
+    public void testForeignKeyConstraintValid() {
+        AtlasTypeRegistry    typeRegistry = new AtlasTypeRegistry();
+        List<AtlasEntityDef> entityDefs   = new ArrayList<>();
+        String               failureMsg   = null;
+
+        entityDefs.add(createTableEntityDef());
+        entityDefs.add(createColumnEntityDef());
+
+        try {
+            typeRegistry.addTypesWithNoRefResolve(entityDefs);
+            typeRegistry.resolveReferences();
+        } catch (AtlasBaseException excp) {
+            failureMsg = excp.getMessage();
+        }
+        assertNull(failureMsg, "failed to create types my_table and my_column");
+    }
+
+    @Test
+    public void testForeignKeyConstraintInValidMappedFromRef() {
+        AtlasTypeRegistry    typeRegistry = new AtlasTypeRegistry();
+        List<AtlasEntityDef> entityDefs   = new ArrayList<>();
+        String               failureMsg   = null;
+
+        entityDefs.add(createTableEntityDef());
+
+        try {
+            typeRegistry.addTypes(entityDefs);
+        } catch (AtlasBaseException excp) {
+            failureMsg = excp.getMessage();
+        }
+        assertNotNull(failureMsg, "expected invalid constraint failure - unknown attribute in mappedFromRef");
+    }
+
+    @Test
+    public void testForeignKeyConstraintInValidMappedFromRef2() {
+        AtlasTypeRegistry    typeRegistry = new AtlasTypeRegistry();
+        List<AtlasEntityDef> entityDefs   = new ArrayList<>();
+        String               failureMsg   = null;
+
+        entityDefs.add(createTableEntityDefWithMissingRefAttribute());
+        entityDefs.add(createColumnEntityDef());
+
+        try {
+            typeRegistry.addTypesWithNoRefResolve(entityDefs);
+            typeRegistry.resolveReferences();
+        } catch (AtlasBaseException excp) {
+            failureMsg = excp.getMessage();
+        }
+        assertNotNull(failureMsg, "expected invalid constraint failure - missing refAttribute in mappedFromRef");
+    }
+
+    @Test
+    public void testForeignKeyConstraintInValidForeignKey() {
+        AtlasTypeRegistry    typeRegistry = new AtlasTypeRegistry();
+        List<AtlasEntityDef> entityDefs   = new ArrayList<>();
+        String               failureMsg   = null;
+
+        entityDefs.add(createColumnEntityDef());
+
+        try {
+            typeRegistry.addTypes(entityDefs);
+        } catch (AtlasBaseException excp) {
+            failureMsg = excp.getMessage();
+        }
+        assertNotNull(failureMsg, "expected invalid constraint failure - unknown attribute in foreignKey");
+    }
+
     private static AtlasEntityType getEntityType(AtlasEntityDef entityDef) {
         try {
             return new AtlasEntityType(entityDef, ModelTestUtil.getTypesRegistry());
@@ -122,4 +192,39 @@ public class TestAtlasEntityType {
             return null;
         }
     }
+
+    private AtlasEntityDef createTableEntityDef() {
+        AtlasEntityDef    table       = new AtlasEntityDef("my_table");
+        AtlasAttributeDef attrColumns = new AtlasAttributeDef("columns",
+                                                              AtlasBaseTypeDef.getArrayTypeName("my_column"));
+
+        Map<String, Object> params = new HashMap<>();
+        params.put(AtlasConstraintDef.CONSTRAINT_PARAM_REF_ATTRIBUTE, "table");
+
+        attrColumns.addConstraint(new AtlasConstraintDef(AtlasConstraintDef.CONSTRAINT_TYPE_MAPPED_FROM_REF, params));
+        table.addAttribute(attrColumns);
+
+        return table;
+    }
+
+    private AtlasEntityDef createTableEntityDefWithMissingRefAttribute() {
+        AtlasEntityDef    table       = new AtlasEntityDef("my_table");
+        AtlasAttributeDef attrColumns = new AtlasAttributeDef("columns",
+                AtlasBaseTypeDef.getArrayTypeName("my_column"));
+
+        attrColumns.addConstraint(new AtlasConstraintDef(AtlasConstraintDef.CONSTRAINT_TYPE_MAPPED_FROM_REF, null));
+        table.addAttribute(attrColumns);
+
+        return table;
+    }
+
+    private AtlasEntityDef createColumnEntityDef() {
+        AtlasEntityDef    column    = new AtlasEntityDef("my_column");
+        AtlasAttributeDef attrTable = new AtlasAttributeDef("table", "my_table");
+
+        attrTable.addConstraint(new AtlasConstraintDef(AtlasConstraintDef.CONSTRAINT_TYPE_FOREIGN_KEY));
+        column.addAttribute(attrTable);
+
+        return column;
+    }
 }

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-1227 Added support for attribute constraints in the API
 ATLAS-1225 Optimize AtlasTypeDefGraphStore to use typeRegistry (mneethiraj via sumasai)
 ATLAS-1221 build failure in windows SyntaxError: invalid syntax (zhangqiang2 via shwethags)
 ATLAS-1226 Servlet init-params in web.xml are unused (mneethiraj via shwethags)

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

@@ -345,7 +345,7 @@ public class AtlasClassificationDefStoreV1 implements AtlasClassificationDefStor
         typeDefStore.createSuperTypeEdges(vertex, classificationDef.getSuperTypes());
     }
 
-    private AtlasClassificationDef toClassificationDef(AtlasVertex vertex) {
+    private AtlasClassificationDef toClassificationDef(AtlasVertex vertex) throws AtlasBaseException {
         AtlasClassificationDef ret = null;
 
         if (vertex != null && typeDefStore.isTypeVertex(vertex, TypeCategory.TRAIT)) {

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

@@ -344,7 +344,7 @@ public class AtlasEntityDefStoreV1 implements AtlasEntityDefStore {
         typeDefStore.createSuperTypeEdges(vertex, entityDef.getSuperTypes());
     }
 
-    private AtlasEntityDef toEntityDef(AtlasVertex vertex) {
+    private AtlasEntityDef toEntityDef(AtlasVertex vertex) throws AtlasBaseException {
         AtlasEntityDef ret = null;
 
         if (vertex != null && typeDefStore.isTypeVertex(vertex, TypeCategory.CLASS)) {