ATLAS-3043: added unique-attributes constraints at store-level

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

@@ -43,6 +43,8 @@ import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasConstraintDef.C
 public class AtlasStructType extends AtlasType {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasStructType.class);
 
+    public static final String UNIQUE_ATTRIBUTE_SHADE_PROPERTY_PREFIX = "__u_";
+
     private final AtlasStructDef structDef;
 
     protected Map<String, AtlasAttribute> allAttributes  = Collections.emptyMap();
@@ -697,6 +699,7 @@ public class AtlasStructType extends AtlasType {
         private final AtlasAttributeDef        attributeDef;
         private final String                   qualifiedName;
         private final String                   vertexPropertyName;
+        private final String                   vertexUniquePropertyName;
         private final boolean                  isOwnedRef;
         private final String                   inverseRefAttributeName;
         private AtlasAttribute                 inverseRefAttribute;
@@ -709,6 +712,7 @@ public class AtlasStructType extends AtlasType {
             this.attributeType            = attributeType.getTypeForAttribute();
             this.qualifiedName            = getQualifiedAttributeName(definedInType.getStructDef(), attributeDef.getName());
             this.vertexPropertyName       = encodePropertyKey(this.qualifiedName);
+            this.vertexUniquePropertyName = attrDef.getIsUnique() ? encodePropertyKey(getQualifiedAttributeName(definedInType.getStructDef(), UNIQUE_ATTRIBUTE_SHADE_PROPERTY_PREFIX + attributeDef.getName())) : null;
             this.relationshipEdgeLabel    = getRelationshipEdgeLabel(relationshipLabel);
             boolean isOwnedRef            = false;
             String  inverseRefAttribute   = null;
@@ -758,6 +762,8 @@ public class AtlasStructType extends AtlasType {
 
         public String getVertexPropertyName() { return vertexPropertyName; }
 
+        public String getVertexUniquePropertyName() { return vertexUniquePropertyName; }
+
         public boolean isOwnedRef() { return isOwnedRef; }
 
         public String getInverseRefAttributeName() { return inverseRefAttributeName; }

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

@@ -794,6 +794,12 @@ public abstract class DeleteHandlerV1 {
                             }
                         }
                      break;
+
+                    case PRIMITIVE:
+                        if (attributeInfo.getVertexUniquePropertyName() != null) {
+                            instanceVertex.removeProperty(attributeInfo.getVertexUniquePropertyName());
+                        }
+                    break;
                 }
             }
         }

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasGraphUtilsV2.java

@@ -72,6 +72,7 @@ public class AtlasGraphUtilsV2 {
     public static final String VERTEX_TYPE                 = "typeSystem";
 
     private static boolean USE_INDEX_QUERY_TO_FIND_ENTITY_BY_UNIQUE_ATTRIBUTES = false;
+    private static boolean USE_UNIQUE_INDEX_PROPERTY_TO_FIND_ENTITY            = false;
     private static String  INDEX_SEARCH_PREFIX;
 
     static {
@@ -79,6 +80,7 @@ public class AtlasGraphUtilsV2 {
             Configuration conf = ApplicationProperties.get();
 
             USE_INDEX_QUERY_TO_FIND_ENTITY_BY_UNIQUE_ATTRIBUTES = conf.getBoolean("atlas.use.index.query.to.find.entity.by.unique.attributes", USE_INDEX_QUERY_TO_FIND_ENTITY_BY_UNIQUE_ATTRIBUTES);
+            USE_UNIQUE_INDEX_PROPERTY_TO_FIND_ENTITY            = conf.getBoolean("atlas.unique.index.property.to.find.entity", USE_UNIQUE_INDEX_PROPERTY_TO_FIND_ENTITY);
             INDEX_SEARCH_PREFIX                                 = conf.getString(INDEX_SEARCH_VERTEX_PREFIX_PROPERTY, INDEX_SEARCH_VERTEX_PREFIX_DEFAULT);
         } catch (Exception excp) {
             LOG.error("Error reading configuration", excp);
@@ -287,6 +289,14 @@ public class AtlasGraphUtilsV2 {
                 if (canUseIndexQuery(entityType, attribute.getName())) {
                     vertex = AtlasGraphUtilsV2.getAtlasVertexFromIndexQuery(entityType, attribute, attrValue);
                 } else {
+                    if (USE_UNIQUE_INDEX_PROPERTY_TO_FIND_ENTITY && attribute.getVertexUniquePropertyName() != null) {
+                        vertex = AtlasGraphUtilsV2.findByTypeAndUniquePropertyName(entityType.getTypeName(), attribute.getVertexUniquePropertyName(), attrValue);
+
+                        // if no instance of given typeName is found, try to find an instance of type's sub-type
+                        if (vertex == null && !entityType.getAllSubTypes().isEmpty()) {
+                            vertex = AtlasGraphUtilsV2.findBySuperTypeAndUniquePropertyName(entityType.getTypeName(), attribute.getVertexUniquePropertyName(), attrValue);
+                        }
+                    } else {
                         vertex = AtlasGraphUtilsV2.findByTypeAndPropertyName(entityType.getTypeName(), attribute.getVertexPropertyName(), attrValue);
 
                         // if no instance of given typeName is found, try to find an instance of type's sub-type
@@ -295,6 +305,8 @@ public class AtlasGraphUtilsV2 {
                         }
                     }
 
+                }
+
                 if (vertex != null) {
                     if (LOG.isDebugEnabled()) {
                         LOG.debug("findByUniqueAttributes(type={}, attrName={}, attrValue={}: found vertex {}",
@@ -350,6 +362,38 @@ public class AtlasGraphUtilsV2 {
         return hasInstanceVertex;
     }
 
+    public static AtlasVertex findByTypeAndUniquePropertyName(String typeName, String propertyName, Object attrVal) {
+        MetricRecorder metric = RequestContext.get().startMetricRecord("findByTypeAndUniquePropertyName");
+
+        AtlasGraphQuery query = AtlasGraphProvider.getGraphInstance().query()
+                                                    .has(Constants.ENTITY_TYPE_PROPERTY_KEY, typeName)
+                                                    .has(propertyName, attrVal);
+
+        Iterator<AtlasVertex> results = query.vertices().iterator();
+
+        AtlasVertex vertex = results.hasNext() ? results.next() : null;
+
+        RequestContext.get().endMetricRecord(metric);
+
+        return vertex;
+    }
+
+    public static AtlasVertex findBySuperTypeAndUniquePropertyName(String typeName, String propertyName, Object attrVal) {
+        MetricRecorder metric = RequestContext.get().startMetricRecord("findBySuperTypeAndUniquePropertyName");
+
+        AtlasGraphQuery query = AtlasGraphProvider.getGraphInstance().query()
+                                                    .has(Constants.SUPER_TYPES_PROPERTY_KEY, typeName)
+                                                    .has(propertyName, attrVal);
+
+        Iterator<AtlasVertex> results = query.vertices().iterator();
+
+        AtlasVertex vertex = results.hasNext() ? results.next() : null;
+
+        RequestContext.get().endMetricRecord(metric);
+
+        return vertex;
+    }
+
     public static AtlasVertex findByTypeAndPropertyName(String typeName, String propertyName, Object attrVal) {
         MetricRecorder metric = RequestContext.get().startMetricRecord("findByTypeAndPropertyName");
 

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java

@@ -705,6 +705,16 @@ public class EntityGraphMapper {
 
         AtlasGraphUtilsV2.setEncodedProperty(ctx.getReferringVertex(), ctx.getVertexProperty(), ret);
 
+        String uniqPropName = ctx.getAttribute() != null ? ctx.getAttribute().getVertexUniquePropertyName() : null;
+
+        if (uniqPropName != null) {
+            if (AtlasGraphUtilsV2.getState(ctx.getReferringVertex()) == DELETED) {
+                ctx.getReferringVertex().removeProperty(uniqPropName);
+            } else {
+                AtlasGraphUtilsV2.setEncodedProperty(ctx.getReferringVertex(), uniqPropName, ret);
+            }
+        }
+
         return ret;
     }