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/graph/GraphBackedSearchIndexer.java

@@ -76,6 +76,7 @@ import static org.apache.atlas.repository.graphdb.AtlasCardinality.LIST;
 import static org.apache.atlas.repository.graphdb.AtlasCardinality.SET;
 import static org.apache.atlas.repository.graphdb.AtlasCardinality.SINGLE;
 import static org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.isReference;
+import static org.apache.atlas.type.AtlasStructType.UNIQUE_ATTRIBUTE_SHADE_PROPERTY_PREFIX;
 import static org.apache.atlas.type.AtlasTypeUtil.isArrayType;
 import static org.apache.atlas.type.AtlasTypeUtil.isMapType;
 
@@ -107,6 +108,8 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
     private boolean     recomputeIndexedKeys = true;
     private Set<String> vertexIndexKeys      = new HashSet<>();
 
+    private enum UniqueKind { NONE, GLOBAL_UNIQUE, PER_TYPE_UNIQUE }
+
     @Inject
     public GraphBackedSearchIndexer(AtlasTypeRegistry typeRegistry) throws AtlasException {
         this(new AtlasGraphProvider(), ApplicationProperties.get(), typeRegistry);
@@ -261,21 +264,21 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
             }
 
             // create vertex indexes
-            createVertexIndex(management, GUID_PROPERTY_KEY, String.class, true, SINGLE, true, false);
-            createVertexIndex(management, ENTITY_TYPE_PROPERTY_KEY, String.class, false, SINGLE, true, false);
-            createVertexIndex(management, SUPER_TYPES_PROPERTY_KEY, String.class, false, SET, true, false);
-            createVertexIndex(management, TIMESTAMP_PROPERTY_KEY, Long.class, false, SINGLE, false, false);
-            createVertexIndex(management, MODIFICATION_TIMESTAMP_PROPERTY_KEY, Long.class, false, SINGLE, false, false);
-            createVertexIndex(management, STATE_PROPERTY_KEY, String.class, false, SINGLE, false, false);
-            createVertexIndex(management, CREATED_BY_KEY, String.class, false, SINGLE, true, true);
-            createVertexIndex(management, MODIFIED_BY_KEY, String.class, false, SINGLE, true, true);
-            createVertexIndex(management, TRAIT_NAMES_PROPERTY_KEY, String.class, false, SET, true, true);
-            createVertexIndex(management, PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, String.class, false, LIST, true, true);
-            createVertexIndex(management, TYPENAME_PROPERTY_KEY, String.class, true, SINGLE, true, false);
-            createVertexIndex(management, TYPESERVICETYPE_PROPERTY_KEY, String.class, false, SINGLE, true, false);
-            createVertexIndex(management, VERTEX_TYPE_PROPERTY_KEY, String.class, false, SINGLE, true, true);
-            createVertexIndex(management, CLASSIFICATION_ENTITY_GUID, String.class, false, SINGLE, true, false);
-            createVertexIndex(management, VERTEX_ID_IN_IMPORT_KEY, Long.class, false, SINGLE, true, false);
+            createVertexIndex(management, GUID_PROPERTY_KEY, UniqueKind.GLOBAL_UNIQUE, String.class, SINGLE, true, false);
+            createVertexIndex(management, TYPENAME_PROPERTY_KEY, UniqueKind.GLOBAL_UNIQUE, String.class, SINGLE, true, false);
+            createVertexIndex(management, TYPESERVICETYPE_PROPERTY_KEY, UniqueKind.NONE, String.class, SINGLE, true, false);
+            createVertexIndex(management, VERTEX_TYPE_PROPERTY_KEY, UniqueKind.NONE, String.class, SINGLE, true, false);
+            createVertexIndex(management, VERTEX_ID_IN_IMPORT_KEY, UniqueKind.NONE, Long.class, SINGLE, true, false);
+
+            createVertexIndex(management, ENTITY_TYPE_PROPERTY_KEY, UniqueKind.NONE, String.class, SINGLE, true, false);
+            createVertexIndex(management, SUPER_TYPES_PROPERTY_KEY, UniqueKind.NONE, String.class, SET, true, false);
+            createVertexIndex(management, TIMESTAMP_PROPERTY_KEY, UniqueKind.NONE, Long.class, SINGLE, false, false);
+            createVertexIndex(management, MODIFICATION_TIMESTAMP_PROPERTY_KEY, UniqueKind.NONE, Long.class, SINGLE, false, false);
+            createVertexIndex(management, STATE_PROPERTY_KEY, UniqueKind.NONE, String.class, SINGLE, false, false);
+            createVertexIndex(management, CREATED_BY_KEY, UniqueKind.NONE, String.class, SINGLE, false, false);
+            createVertexIndex(management, MODIFIED_BY_KEY, UniqueKind.NONE, String.class, SINGLE, false, false);
+            createVertexIndex(management, TRAIT_NAMES_PROPERTY_KEY, UniqueKind.NONE, String.class, SET, true, true);
+            createVertexIndex(management, PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, UniqueKind.NONE, String.class, LIST, true, true);
 
             // create vertex-centric index
             createVertexCentricIndex(management, CLASSIFICATION_LABEL, AtlasEdgeDirection.BOTH, CLASSIFICATION_EDGE_NAME_PROPERTY_KEY, String.class, SINGLE);
@@ -325,6 +328,8 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
         boolean          isBuiltInType  = AtlasTypeUtil.isBuiltInType(attribTypeName);
         boolean          isArrayType    = isArrayType(attribTypeName);
         boolean          isMapType      = isMapType(attribTypeName);
+        final String     uniqPropName   = isUnique ? AtlasGraphUtilsV2.encodePropertyKey(typeName + "." + UNIQUE_ATTRIBUTE_SHADE_PROPERTY_PREFIX + attributeDef.getName()) : null;
+
 
         try {
             AtlasType atlasType     = typeRegistry.getType(typeName);
@@ -363,13 +368,21 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
                 if (isRelationshipType(atlasType)) {
                     createEdgeIndex(management, propertyName, getPrimitiveClass(attribTypeName), cardinality, false);
                 } else {
-                    createVertexIndex(management, propertyName, getPrimitiveClass(attribTypeName), isUnique, cardinality, false, isIndexable);
+                    createVertexIndex(management, propertyName, UniqueKind.NONE, getPrimitiveClass(attribTypeName), cardinality, isIndexable, false);
+
+                    if (uniqPropName != null) {
+                        createVertexIndex(management, uniqPropName, UniqueKind.PER_TYPE_UNIQUE, getPrimitiveClass(attribTypeName), cardinality, isIndexable, true);
+                    }
                 }
             } else if (isEnumType(attributeType)) {
                 if (isRelationshipType(atlasType)) {
                     createEdgeIndex(management, propertyName, String.class, cardinality, false);
                 } else {
-                    createVertexIndex(management, propertyName, String.class, isUnique, cardinality, false, isIndexable);
+                    createVertexIndex(management, propertyName, UniqueKind.NONE, String.class, cardinality, isIndexable, false);
+
+                    if (uniqPropName != null) {
+                        createVertexIndex(management, uniqPropName, UniqueKind.PER_TYPE_UNIQUE, String.class, cardinality, isIndexable, true);
+                    }
                 }
             } else if (isStructType(attributeType)) {
                 AtlasStructDef structDef = typeRegistry.getStructDefByName(attribTypeName);
@@ -479,9 +492,9 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
         return propertyKey;
     }
 
-    private AtlasPropertyKey createVertexIndex(AtlasGraphManagement management, String propertyName, Class propertyClass,
-                                               boolean isUnique, AtlasCardinality cardinality, boolean createCompositeIndex,
-                                               boolean createCompositeIndexWithTypeAndSuperTypes) {
+    private void createVertexIndex(AtlasGraphManagement management, String propertyName, UniqueKind uniqueKind, Class propertyClass,
+                                   AtlasCardinality cardinality, boolean createCompositeIndex, boolean createCompositeIndexWithTypeAndSuperTypes) {
+        if (propertyName != null) {
             AtlasPropertyKey propertyKey = management.getPropertyKey(propertyName);
 
             if (propertyKey == null) {
@@ -499,19 +512,18 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
             }
 
             if (propertyKey != null) {
-            if (createCompositeIndex) {
-                createVertexCompositeIndex(management, propertyClass, propertyKey, isUnique);
+                if (createCompositeIndex || uniqueKind == UniqueKind.GLOBAL_UNIQUE || uniqueKind == UniqueKind.PER_TYPE_UNIQUE) {
+                    createVertexCompositeIndex(management, propertyClass, propertyKey, uniqueKind == UniqueKind.GLOBAL_UNIQUE);
                 }
 
                 if (createCompositeIndexWithTypeAndSuperTypes) {
-                createVertexCompositeIndexWithTypeName(management, propertyClass, propertyKey);
+                    createVertexCompositeIndexWithTypeName(management, propertyClass, propertyKey, uniqueKind == UniqueKind.PER_TYPE_UNIQUE);
                     createVertexCompositeIndexWithSuperTypeName(management, propertyClass, propertyKey);
                 }
             } else {
                 LOG.warn("Index not created for {}: propertyKey is null", propertyName);
             }
-
-        return propertyKey;
+        }
     }
 
     private void createVertexCentricIndex(AtlasGraphManagement management, String edgeLabel, AtlasEdgeDirection edgeDirection,
@@ -563,8 +575,9 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
     }
 
 
-    private AtlasPropertyKey createEdgeIndex(AtlasGraphManagement management, String propertyName, Class propertyClass,
+    private void createEdgeIndex(AtlasGraphManagement management, String propertyName, Class propertyClass,
                                  AtlasCardinality cardinality, boolean createCompositeIndex) {
+        if (propertyName != null) {
             AtlasPropertyKey propertyKey = management.getPropertyKey(propertyName);
 
             if (propertyKey == null) {
@@ -588,8 +601,7 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
             } else {
                 LOG.warn("Index not created for {}: propertyKey is null", propertyName);
             }
-
-        return propertyKey;
+        }
     }
 
     private AtlasPropertyKey createFullTextIndex(AtlasGraphManagement management, String propertyName, Class propertyClass,
@@ -648,16 +660,16 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
         }
     }
 
-    private void createVertexCompositeIndexWithTypeName(AtlasGraphManagement management, Class propertyClass, AtlasPropertyKey propertyKey) {
-        createVertexCompositeIndexWithSystemProperty(management, propertyClass, propertyKey, ENTITY_TYPE_PROPERTY_KEY, SINGLE);
+    private void createVertexCompositeIndexWithTypeName(AtlasGraphManagement management, Class propertyClass, AtlasPropertyKey propertyKey, boolean isUnique) {
+        createVertexCompositeIndexWithSystemProperty(management, propertyClass, propertyKey, ENTITY_TYPE_PROPERTY_KEY, SINGLE, isUnique);
     }
 
     private void createVertexCompositeIndexWithSuperTypeName(AtlasGraphManagement management, Class propertyClass, AtlasPropertyKey propertyKey) {
-        createVertexCompositeIndexWithSystemProperty(management, propertyClass, propertyKey, SUPER_TYPES_PROPERTY_KEY, SET);
+        createVertexCompositeIndexWithSystemProperty(management, propertyClass, propertyKey, SUPER_TYPES_PROPERTY_KEY, SET, false);
     }
 
     private void createVertexCompositeIndexWithSystemProperty(AtlasGraphManagement management, Class propertyClass, AtlasPropertyKey propertyKey,
-                                                              final String systemPropertyKey, AtlasCardinality cardinality) {
+                                                              final String systemPropertyKey, AtlasCardinality cardinality, boolean isUnique) {
         if (LOG.isDebugEnabled()) {
             LOG.debug("Creating composite index for property {} of type {} and {}", propertyKey.getName(), propertyClass.getName(),  systemPropertyKey);
         }
@@ -672,9 +684,9 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
 
         if (existingIndex == null) {
             List<AtlasPropertyKey> keys = new ArrayList<>(2);
-            keys.add(propertyKey);
             keys.add(typePropertyKey);
-            management.createVertexCompositeIndex(indexName, false, keys);
+            keys.add(propertyKey);
+            management.createVertexCompositeIndex(indexName, isUnique, keys);
 
             LOG.info("Created composite index for property {} of type {} and {}", propertyKey.getName(), propertyClass.getName(), systemPropertyKey);
         }

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;
     }