ATLAS-3054: updated notification pre-process to handle updates to ownedRef attributes - #3

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

@@ -22,7 +22,6 @@ import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasObjectId;
-import org.apache.atlas.model.instance.AtlasStruct;
 import org.apache.atlas.model.typedef.AtlasEntityDef;
 import org.apache.atlas.model.typedef.AtlasEntityDef.AtlasRelationshipAttributeDef;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
@@ -544,7 +543,9 @@ public class AtlasEntityType extends AtlasStructType {
                 superType.normalizeAttributeValues(ent);
             }
 
-            normalizeValues(ent);
+            super.normalizeAttributeValues(ent);
+
+            normalizeRelationshipAttributeValues(ent, false);
         }
     }
 
@@ -555,6 +556,8 @@ public class AtlasEntityType extends AtlasStructType {
             }
 
             super.normalizeAttributeValuesForUpdate(ent);
+
+            normalizeRelationshipAttributeValues(ent, true);
         }
     }
 
@@ -565,7 +568,9 @@ public class AtlasEntityType extends AtlasStructType {
                 superType.normalizeAttributeValues(obj);
             }
 
-            normalizeValues(obj);
+            super.normalizeAttributeValues(obj);
+
+            normalizeRelationshipAttributeValues(obj, false);
         }
     }
 
@@ -576,6 +581,8 @@ public class AtlasEntityType extends AtlasStructType {
             }
 
             super.normalizeAttributeValuesForUpdate(obj);
+
+            normalizeRelationshipAttributeValues(obj, true);
         }
     }
 
@@ -743,67 +750,56 @@ public class AtlasEntityType extends AtlasStructType {
         return ret;
     }
 
-    private void normalizeRelationshipAttributeValues(AtlasStruct obj) {
-        if (obj != null && obj instanceof AtlasEntity) {
-            AtlasEntity entityObj = (AtlasEntity) obj;
-
+    private void normalizeRelationshipAttributeValues(AtlasEntity entity, boolean isUpdate) {
+        if (entity != null) {
             for (String attributeName : relationshipAttributes.keySet()) {
-                if (entityObj.hasRelationshipAttribute(attributeName)) {
-                    Object            attributeValue   = entityObj.getRelationshipAttribute(attributeName);
-                    String            relationshipType = AtlasEntityUtil.getRelationshipType(attributeValue);
-                    AtlasAttribute    attribute        = getRelationshipAttribute(attributeName, relationshipType);
-                    AtlasAttributeDef attributeDef     = attribute.getAttributeDef();
+                if (entity.hasRelationshipAttribute(attributeName)) {
+                    Object         attributeValue   = entity.getRelationshipAttribute(attributeName);
+                    String         relationshipType = AtlasEntityUtil.getRelationshipType(attributeValue);
+                    AtlasAttribute attribute        = getRelationshipAttribute(attributeName, relationshipType);
 
-                    attributeValue = getNormalizedValue(attributeValue, attributeDef);
+                    if (attribute != null) {
+                        AtlasType attrType = attribute.getAttributeType();
+
+                        if (isValidRelationshipType(attrType)) {
+                            if (isUpdate) {
+                                attributeValue = attrType.getNormalizedValueForUpdate(attributeValue);
+                            } else {
+                                attributeValue = attrType.getNormalizedValue(attributeValue);
+                            }
 
-                    entityObj.setRelationshipAttribute(attributeName, attributeValue);
+                            entity.setRelationshipAttribute(attributeName, attributeValue);
+                        }
+                    }
                 }
             }
         }
     }
 
-    public void normalizeRelationshipAttributeValues(Map<String, Object> obj) {
+    public void normalizeRelationshipAttributeValues(Map<String, Object> obj, boolean isUpdate) {
         if (obj != null) {
             for (String attributeName : relationshipAttributes.keySet()) {
                 if (obj.containsKey(attributeName)) {
-                    Object            attributeValue   = obj.get(attributeName);
-                    String            relationshipType = AtlasEntityUtil.getRelationshipType(attributeValue);
-                    AtlasAttribute    attribute        = getRelationshipAttribute(attributeName, relationshipType);
-                    AtlasAttributeDef attributeDef     = attribute.getAttributeDef();
-
-                    attributeValue = getNormalizedValue(attributeValue, attributeDef);
-
-                    obj.put(attributeName, attributeValue);
-                }
-            }
-        }
-    }
+                    Object         attributeValue   = obj.get(attributeName);
+                    String         relationshipType = AtlasEntityUtil.getRelationshipType(attributeValue);
+                    AtlasAttribute attribute        = getRelationshipAttribute(attributeName, relationshipType);
 
-    private Object getNormalizedValue(Object value, AtlasAttributeDef attributeDef) {
-        String         relationshipType = AtlasEntityUtil.getRelationshipType(value);
-        AtlasAttribute attribute        = getRelationshipAttribute(attributeDef.getName(), relationshipType);
+                    if (attribute != null) {
+                        AtlasType attrType = attribute.getAttributeType();
 
-        if (attribute != null) {
-            AtlasType attrType = attribute.getAttributeType();
+                        if (isValidRelationshipType(attrType)) {
+                            if (isUpdate) {
+                                attributeValue = attrType.getNormalizedValueForUpdate(attributeValue);
+                            } else {
+                                attributeValue = attrType.getNormalizedValue(attributeValue);
+                            }
 
-            if (isValidRelationshipType(attrType) && value != null) {
-                return attrType.getNormalizedValue(value);
+                            obj.put(attributeName, attributeValue);
+                        }
+                    }
+                }
             }
         }
-
-        return null;
-    }
-
-    private void normalizeValues(AtlasEntity ent) {
-        super.normalizeAttributeValues(ent);
-
-        normalizeRelationshipAttributeValues(ent);
-    }
-
-    private void normalizeValues(Map<String, Object> obj) {
-        super.normalizeAttributeValues(obj);
-
-        normalizeRelationshipAttributeValues(obj);
     }
 
     private boolean validateRelationshipAttributes(Object obj, String objName, List<String> messages) {

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

@@ -346,8 +346,15 @@ public class AtlasRelationshipType extends AtlasStructType {
                 attributeDef.addConstraint(constraint);
             }
 
-            attribute = new AtlasAttribute(entityType, attributeDef,
-                                           typeRegistry.getType(attrTypeName), getTypeName(), relationshipLabel);
+            AtlasType attrType = typeRegistry.getType(attrTypeName);
+
+            if (attrType instanceof AtlasArrayType) {
+                AtlasArrayType arrayType = (AtlasArrayType) attrType;
+
+                arrayType.setCardinality(attributeDef.getCardinality());
+            }
+
+            attribute = new AtlasAttribute(entityType, attributeDef, attrType, getTypeName(), relationshipLabel);
 
             attribute.setLegacyAttribute(endDef.getIsLegacyAttribute());
         } else {

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

@@ -99,6 +99,7 @@ public class AtlasStructType extends AtlasType {
 
                 arrayType.setMinCount(attributeDef.getValuesMinCount());
                 arrayType.setMaxCount(attributeDef.getValuesMaxCount());
+                arrayType.setCardinality(cardinality);
             }
 
             //check if attribute type is not classification

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

@@ -38,6 +38,7 @@ import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef.Cardinali
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasConstraintDef;
 import org.apache.atlas.model.typedef.AtlasTypeDefHeader;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
+import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
 import org.apache.atlas.v1.model.typedef.AttributeDefinition;
 import org.apache.atlas.v1.model.typedef.ClassTypeDefinition;
 import org.apache.atlas.v1.model.typedef.Multiplicity;
@@ -413,10 +414,36 @@ public class AtlasTypeUtil {
         return new AtlasRelatedObjectId(getAtlasObjectId(entity));
     }
 
+    public static AtlasRelatedObjectId toAtlasRelatedObjectId(AtlasEntity entity, AtlasTypeRegistry typeRegistry) {
+        return new AtlasRelatedObjectId(getAtlasObjectId(entity, typeRegistry));
+    }
+
     public static AtlasObjectId getAtlasObjectId(AtlasEntity entity) {
         return new AtlasObjectId(entity.getGuid(), entity.getTypeName());
     }
 
+    public static AtlasObjectId getAtlasObjectId(AtlasEntity entity, AtlasTypeRegistry typeRegistry) {
+        String              typeName       = entity.getTypeName();
+        AtlasEntityType     entityType     = typeRegistry.getEntityTypeByName(typeName);
+        Map<String, Object> uniqAttributes = null;
+
+        if (entityType != null && MapUtils.isNotEmpty(entityType.getUniqAttributes())) {
+            for (AtlasAttribute attribute : entityType.getUniqAttributes().values()) {
+                Object attrValue = entity.getAttribute(attribute.getName());
+
+                if (attrValue != null) {
+                    if (uniqAttributes == null) {
+                        uniqAttributes = new HashMap<>();
+                    }
+
+                    uniqAttributes.put(attribute.getName(), attrValue);
+                }
+            }
+        }
+
+        return new AtlasObjectId(entity.getGuid(), typeName, uniqAttributes);
+    }
+
     public static AtlasObjectId getAtlasObjectId(AtlasEntityHeader header) {
         return new AtlasObjectId(header.getGuid(), header.getTypeName());
     }

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

@@ -63,12 +63,12 @@ public class TestAtlasStructType {
 
         multiValuedAttribMin.setName(MULTI_VAL_ATTR_NAME_MIN);
         multiValuedAttribMin.setTypeName(AtlasBaseTypeDef.getArrayTypeName(ATLAS_TYPE_INT));
-        multiValuedAttribMin.setCardinality(Cardinality.SET);
+        multiValuedAttribMin.setCardinality(Cardinality.LIST);
         multiValuedAttribMin.setValuesMinCount(MULTI_VAL_ATTR_MIN_COUNT);
 
         multiValuedAttribMax.setName(MULTI_VAL_ATTR_NAME_MAX);
         multiValuedAttribMax.setTypeName(AtlasBaseTypeDef.getArrayTypeName(ATLAS_TYPE_INT));
-        multiValuedAttribMax.setCardinality(Cardinality.LIST);
+        multiValuedAttribMax.setCardinality(Cardinality.SET);
         multiValuedAttribMax.setValuesMaxCount(MULTI_VAL_ATTR_MAX_COUNT);
 
         AtlasStructDef structDef = ModelTestUtil.newStructDef();

repository/src/main/java/org/apache/atlas/repository/converters/AtlasStructFormatConverter.java

@@ -190,7 +190,7 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
 
                             if (entities != null) {
                                 v2Value  = entities;
-                                attrType = new AtlasArrayType(entityType);
+                                attrType = new AtlasArrayType(entityType, arrayType.getMinCount(), arrayType.getMaxCount(), arrayType.getCardinality());
 
                                 if (LOG.isDebugEnabled()) {
                                     LOG.debug("{}: replaced objIdList with entityList", attr.getQualifiedName());

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

@@ -22,6 +22,7 @@ import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 
 import java.util.Iterator;
+import java.util.List;
 
 public class AtlasEntityStream implements EntityStream {
     protected final AtlasEntitiesWithExtInfo entitiesWithExtInfo;
@@ -33,6 +34,10 @@ public class AtlasEntityStream implements EntityStream {
         this(new AtlasEntitiesWithExtInfo(entity), null);
     }
 
+    public AtlasEntityStream(List<AtlasEntity> entities) {
+        this(new AtlasEntitiesWithExtInfo(entities), null);
+    }
+
     public AtlasEntityStream(AtlasEntityWithExtInfo entityWithExtInfo) {
         this(new AtlasEntitiesWithExtInfo(entityWithExtInfo), null);
     }

webapp/src/main/java/org/apache/atlas/notification/preprocessor/EntityPreprocessor.java

@@ -53,6 +53,7 @@ public abstract class EntityPreprocessor {
     public static final String ATTRIBUTE_TABLES         = "tables";
     public static final String ATTRIBUTE_INDEXES        = "indexes";
     public static final String ATTRIBUTE_FOREIGN_KEYS   = "foreign_keys";
+    public static final String ATTRIBUTE_INSTANCE       = "instance";
 
     public static final char   QNAME_SEP_CLUSTER_NAME = '@';
     public static final char   QNAME_SEP_ENTITY_NAME  = '.';

webapp/src/main/java/org/apache/atlas/notification/preprocessor/HivePreprocessor.java

@@ -18,8 +18,6 @@
 package org.apache.atlas.notification.preprocessor;
 
 import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.model.instance.AtlasObjectId;
-import org.apache.atlas.model.instance.AtlasRelatedObjectId;
 import org.apache.atlas.notification.preprocessor.PreprocessorContext.PreprocessAction;
 import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
@@ -27,16 +25,14 @@ import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
-import java.util.Set;
 
 public class HivePreprocessor {
     private static final Logger LOG = LoggerFactory.getLogger(HivePreprocessor.class);
 
     private static final String RELATIONSHIP_TYPE_HIVE_TABLE_COLUMNS        = "hive_table_columns";
     private static final String RELATIONSHIP_TYPE_HIVE_TABLE_PARTITION_KEYS = "hive_table_partitionkeys";
+    private static final String RELATIONSHIP_TYPE_HIVE_TABLE_STORAGEDESC    = "hive_table_storagedesc";
 
     static class HiveTablePreprocessor extends EntityPreprocessor {
         public HiveTablePreprocessor() {
@@ -67,63 +63,12 @@ public class HivePreprocessor {
                     entity.setAttribute(ATTRIBUTE_COLUMNS, null);
                     entity.setAttribute(ATTRIBUTE_PARTITION_KEYS, null);
                 } else if (context.getHiveTypesRemoveOwnedRefAttrs()) {
-                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_SD);
-
-                    removeColumnsAttributeAndRegisterToMove(entity, ATTRIBUTE_COLUMNS, RELATIONSHIP_TYPE_HIVE_TABLE_COLUMNS, context);
-                    removeColumnsAttributeAndRegisterToMove(entity, ATTRIBUTE_PARTITION_KEYS, RELATIONSHIP_TYPE_HIVE_TABLE_PARTITION_KEYS, context);
+                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_SD, RELATIONSHIP_TYPE_HIVE_TABLE_STORAGEDESC, ATTRIBUTE_TABLE);
+                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_COLUMNS, RELATIONSHIP_TYPE_HIVE_TABLE_COLUMNS, ATTRIBUTE_TABLE);
+                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_PARTITION_KEYS, RELATIONSHIP_TYPE_HIVE_TABLE_PARTITION_KEYS, ATTRIBUTE_TABLE);
                 }
             }
         }
-
-        private void removeColumnsAttributeAndRegisterToMove(AtlasEntity entity, String attrName, String relationshipType, PreprocessorContext context) {
-            Object attrVal = entity.removeAttribute(attrName);
-
-            if (attrVal != null) {
-                Set<String> guids = new HashSet<>();
-
-                context.collectGuids(attrVal, guids);
-
-                for (String guid : guids) {
-                    AtlasEntity colEntity = context.getEntity(guid);
-
-                    if (colEntity != null) {
-                        Object attrTable = null;
-
-                        if (colEntity.hasRelationshipAttribute(ATTRIBUTE_TABLE)) {
-                            attrTable = colEntity.getRelationshipAttribute(ATTRIBUTE_TABLE);
-                        } else if (colEntity.hasAttribute(ATTRIBUTE_TABLE)) {
-                            attrTable = colEntity.getAttribute(ATTRIBUTE_TABLE);
-                        }
-
-                        attrTable = setRelationshipType(attrTable, relationshipType);
-
-                        if (attrTable != null) {
-                            colEntity.setRelationshipAttribute(ATTRIBUTE_TABLE, attrTable);
-                        }
-
-                        context.addToReferredEntitiesToMove(guid);
-                    }
-                }
-            }
-        }
-
-        private AtlasRelatedObjectId setRelationshipType(Object attr, String relationshipType) {
-            AtlasRelatedObjectId ret = null;
-
-            if (attr instanceof AtlasRelatedObjectId) {
-                ret = (AtlasRelatedObjectId) attr;
-            } else if (attr instanceof AtlasObjectId) {
-                ret = new AtlasRelatedObjectId((AtlasObjectId) attr);
-            } else if (attr instanceof Map) {
-                ret = new AtlasRelatedObjectId((Map) attr);
-            }
-
-            if (ret != null) {
-                ret.setRelationshipType(relationshipType);
-            }
-
-            return ret;
-        }
     }
 
 

webapp/src/main/java/org/apache/atlas/notification/preprocessor/RdbmsPreprocessor.java

@@ -32,6 +32,12 @@ import java.util.Set;
 public class RdbmsPreprocessor {
     private static final Logger LOG = LoggerFactory.getLogger(RdbmsPreprocessor.class);
 
+    private static final String RELATIONSHIP_TYPE_RDBMS_INSTANCE_DATABASES = "rdbms_instance_databases";
+    private static final String RELATIONSHIP_TYPE_RDBMS_DB_TABLES          = "rdbms_db_tables";
+    private static final String RELATIONSHIP_TYPE_RDBMS_TABLE_COLUMNS      = "rdbms_table_columns";
+    private static final String RELATIONSHIP_TYPE_RDBMS_TABLE_INDEXES      = "rdbms_table_indexes";
+    private static final String RELATIONSHIP_TYPE_RDBMS_TABLE_FOREIGN_KEYS = "rdbms_table_foreign_key";
+
     static class RdbmsInstancePreprocessor extends RdbmsTypePreprocessor {
         public RdbmsInstancePreprocessor() {
             super(TYPE_RDBMS_INSTANCE);
@@ -121,17 +127,17 @@ public class RdbmsPreprocessor {
         private void clearRefAttributesAndMove(AtlasEntity entity, PreprocessorContext context) {
             switch (entity.getTypeName()) {
                 case TYPE_RDBMS_INSTANCE:
-                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_DATABASES);
+                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_DATABASES, RELATIONSHIP_TYPE_RDBMS_INSTANCE_DATABASES, ATTRIBUTE_INSTANCE);
                 break;
 
                 case TYPE_RDBMS_DB:
-                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_TABLES);
+                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_TABLES, RELATIONSHIP_TYPE_RDBMS_DB_TABLES, ATTRIBUTE_DB);
                 break;
 
                 case TYPE_RDBMS_TABLE:
-                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_COLUMNS);
-                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_INDEXES);
-                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_FOREIGN_KEYS);
+                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_COLUMNS, RELATIONSHIP_TYPE_RDBMS_TABLE_COLUMNS, ATTRIBUTE_TABLE);
+                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_INDEXES, RELATIONSHIP_TYPE_RDBMS_TABLE_INDEXES, ATTRIBUTE_TABLE);
+                    context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_FOREIGN_KEYS, RELATIONSHIP_TYPE_RDBMS_TABLE_FOREIGN_KEYS, ATTRIBUTE_TABLE);
                 break;
             }
         }