ATLAS-3211 :- Update Hive hook with Relationship Attributes.

addons/hive-bridge/src/main/java/org/apache/atlas/hive/hook/events/AlterTableRename.java

@@ -129,17 +129,13 @@ public class AlterTableRename extends BaseHiveEvent {
         // update qualifiedName for all columns, partitionKeys, storageDesc
         String renamedTableQualifiedName = (String) renamedTableEntity.getEntity().getAttribute(ATTRIBUTE_QUALIFIED_NAME);
 
-        renameColumns((List<AtlasObjectId>) oldTableEntity.getEntity().getAttribute(ATTRIBUTE_COLUMNS), oldTableEntity, renamedTableQualifiedName, ret);
-        renameColumns((List<AtlasObjectId>) oldTableEntity.getEntity().getAttribute(ATTRIBUTE_PARTITION_KEYS), oldTableEntity, renamedTableQualifiedName, ret);
+        renameColumns((List<AtlasObjectId>) oldTableEntity.getEntity().getRelationshipAttribute(ATTRIBUTE_COLUMNS), oldTableEntity, renamedTableQualifiedName, ret);
+        renameColumns((List<AtlasObjectId>) oldTableEntity.getEntity().getRelationshipAttribute(ATTRIBUTE_PARTITION_KEYS), oldTableEntity, renamedTableQualifiedName, ret);
         renameStorageDesc(oldTableEntity, renamedTableEntity, ret);
 
-        // remove columns, partitionKeys and storageDesc - as they have already been updated above
-        removeAttribute(renamedTableEntity, ATTRIBUTE_COLUMNS);
-        removeAttribute(renamedTableEntity, ATTRIBUTE_PARTITION_KEYS);
-        removeAttribute(renamedTableEntity, ATTRIBUTE_STORAGEDESC);
-
         // set previous name as the alias
         renamedTableEntity.getEntity().setAttribute(ATTRIBUTE_ALIASES, Collections.singletonList(oldTable.getTableName()));
+        renamedTableEntity.getEntity().setRelationshipAttributes(null);
 
         String        oldTableQualifiedName = (String) oldTableEntity.getEntity().getAttribute(ATTRIBUTE_QUALIFIED_NAME);
         AtlasObjectId oldTableId            = new AtlasObjectId(oldTableEntity.getEntity().getTypeName(), ATTRIBUTE_QUALIFIED_NAME, oldTableQualifiedName);
@@ -179,35 +175,17 @@ public class AlterTableRename extends BaseHiveEvent {
             AtlasObjectId oldSdId = new AtlasObjectId(oldSd.getTypeName(), ATTRIBUTE_QUALIFIED_NAME, oldSd.getAttribute(ATTRIBUTE_QUALIFIED_NAME));
 
             newSd.removeAttribute(ATTRIBUTE_TABLE);
+            newSd.setRelationshipAttributes(null);
 
             notifications.add(new EntityPartialUpdateRequestV2(getUserName(), oldSdId, new AtlasEntityWithExtInfo(newSd)));
         }
     }
 
-    private void removeAttribute(AtlasEntityWithExtInfo entity, String attributeName) {
-        Object attributeValue = entity.getEntity().getAttribute(attributeName);
-
-        entity.getEntity().getAttributes().remove(attributeName);
-
-        if (attributeValue instanceof AtlasObjectId) {
-            AtlasObjectId objectId = (AtlasObjectId) attributeValue;
-
-            entity.removeReferredEntity(objectId.getGuid());
-        } else if (attributeValue instanceof Collection) {
-            for (Object item : (Collection) attributeValue)
-                if (item instanceof AtlasObjectId) {
-                    AtlasObjectId objectId = (AtlasObjectId) item;
-
-                    entity.removeReferredEntity(objectId.getGuid());
-                }
-        }
-    }
-
     private AtlasEntity getStorageDescEntity(AtlasEntityWithExtInfo tableEntity) {
         AtlasEntity ret = null;
 
         if (tableEntity != null && tableEntity.getEntity() != null) {
-            Object attrSdId = tableEntity.getEntity().getAttribute(ATTRIBUTE_STORAGEDESC);
+            Object attrSdId = tableEntity.getEntity().getRelationshipAttribute(ATTRIBUTE_STORAGEDESC);
 
             if (attrSdId instanceof AtlasObjectId) {
                 ret = tableEntity.getReferredEntity(((AtlasObjectId) attrSdId).getGuid());

addons/hive-bridge/src/main/java/org/apache/atlas/hive/hook/events/CreateHiveProcess.java

@@ -207,9 +207,9 @@ public class CreateHiveProcess extends BaseHiveEvent {
 
             columnLineageProcess.setAttribute(ATTRIBUTE_NAME, hiveProcess.getAttribute(ATTRIBUTE_QUALIFIED_NAME) + ":" + outputColumn.getAttribute(ATTRIBUTE_NAME));
             columnLineageProcess.setAttribute(ATTRIBUTE_QUALIFIED_NAME, hiveProcess.getAttribute(ATTRIBUTE_QUALIFIED_NAME) + ":" + outputColumn.getAttribute(ATTRIBUTE_NAME));
-            columnLineageProcess.setAttribute(ATTRIBUTE_INPUTS, getObjectIds(inputColumns));
-            columnLineageProcess.setAttribute(ATTRIBUTE_OUTPUTS, Collections.singletonList(getObjectId(outputColumn)));
-            columnLineageProcess.setAttribute(ATTRIBUTE_QUERY, getObjectId(hiveProcess));
+            columnLineageProcess.setRelationshipAttribute(ATTRIBUTE_INPUTS, getObjectIdsWithRelationshipType(inputColumns, BaseHiveEvent.RELATIONSHIP_DATASET_PROCESS_INPUTS));
+            columnLineageProcess.setRelationshipAttribute(ATTRIBUTE_OUTPUTS, Collections.singletonList(getObjectIdWithRelationshipType(outputColumn, BaseHiveEvent.RELATIONSHIP_PROCESS_DATASET_OUTPUTS)));
+            columnLineageProcess.setRelationshipAttribute(ATTRIBUTE_QUERY, getObjectIdWithRelationshipType(hiveProcess, BaseHiveEvent.RELATIONSHIP_HIVE_PROCESS_COLUMN_LINEAGE));
             columnLineageProcess.setAttribute(ATTRIBUTE_DEPENDENCY_TYPE, entry.getValue().getType());
             columnLineageProcess.setAttribute(ATTRIBUTE_EXPRESSION, entry.getValue().getExpr());
 

intg/src/main/java/org/apache/atlas/model/instance/AtlasRelatedObjectId.java

@@ -83,6 +83,12 @@ public class AtlasRelatedObjectId extends AtlasObjectId implements Serializable
         super(other);
     }
 
+    public AtlasRelatedObjectId(AtlasObjectId objId, String relationshipType) {
+        this(objId);
+
+        setRelationshipType(relationshipType);
+    }
+
     public AtlasRelatedObjectId(Map objIdMap) {
         super(objIdMap);
 

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

@@ -414,6 +414,10 @@ public class AtlasTypeUtil {
         return new AtlasRelatedObjectId(getAtlasObjectId(entity));
     }
 
+    public static AtlasRelatedObjectId toAtlasRelatedObjectId(AtlasEntity entity, String relationshipType){
+        return new AtlasRelatedObjectId(getAtlasObjectId(entity), relationshipType);
+    }
+
     public static AtlasRelatedObjectId toAtlasRelatedObjectId(AtlasEntity entity, AtlasTypeRegistry typeRegistry) {
         return new AtlasRelatedObjectId(getAtlasObjectId(entity, typeRegistry));
     }