ATLAS-2882: refactored import transformer to set context in constructor; fixed 
incorrect objId match


Project: http://git-wip-us.apache.org/repos/asf/atlas/repo
Commit: http://git-wip-us.apache.org/repos/asf/atlas/commit/c1b42468
Tree: http://git-wip-us.apache.org/repos/asf/atlas/tree/c1b42468
Diff: http://git-wip-us.apache.org/repos/asf/atlas/diff/c1b42468

Branch: refs/heads/branch-1.0
Commit: c1b4246829e0f1b56ca0fd848ea4e5aeba380183
Parents: 2111b95
Author: Madhan Neethiraj <mad...@apache.org>
Authored: Thu Sep 20 17:27:28 2018 -0700
Committer: Ashutosh Mestry <ames...@hortonworks.com>
Committed: Thu Nov 1 15:42:56 2018 -0700

----------------------------------------------------------------------
 .../apache/atlas/entitytransform/Action.java    | 150 +++++++++--------
 .../entitytransform/AtlasEntityTransformer.java |  27 ++--
 .../entitytransform/BaseEntityHandler.java      | 138 +++++-----------
 .../apache/atlas/entitytransform/Condition.java | 141 ++++++++--------
 .../atlas/entitytransform/EntityAttribute.java  |  69 ++++++++
 .../entitytransform/HdfsPathEntityHandler.java  |  28 ++--
 .../HiveColumnEntityHandler.java                |  28 ++--
 .../HiveDatabaseEntityHandler.java              |  24 +--
 .../HiveStorageDescriptorEntityHandler.java     |  26 +--
 .../entitytransform/HiveTableEntityHandler.java |  26 +--
 .../atlas/entitytransform/NeedsContext.java     |  23 ---
 .../entitytransform/TransformerContext.java     |   8 +-
 .../TransformationHandlerTest.java              | 160 +++++++++++++++----
 13 files changed, 470 insertions(+), 378 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/Action.java
----------------------------------------------------------------------
diff --git a/intg/src/main/java/org/apache/atlas/entitytransform/Action.java 
b/intg/src/main/java/org/apache/atlas/entitytransform/Action.java
index fa18558..0c6102d 100644
--- a/intg/src/main/java/org/apache/atlas/entitytransform/Action.java
+++ b/intg/src/main/java/org/apache/atlas/entitytransform/Action.java
@@ -22,6 +22,8 @@ import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.typedef.AtlasClassificationDef;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
+import org.apache.atlas.store.AtlasTypeDefStore;
+import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.commons.lang.StringUtils;
 import 
org.apache.atlas.entitytransform.BaseEntityHandler.AtlasTransformableEntity;
 import org.slf4j.Logger;
@@ -31,35 +33,35 @@ import java.util.ArrayList;
 import java.util.Collections;
 
 
+
 public abstract class Action {
     private static final Logger LOG = LoggerFactory.getLogger(Action.class);
 
-    private static final String ENTITY_KEY                  = "__entity";
     private static final String ACTION_DELIMITER           = ":";
-    private static final String ACTION_ADD_CLASSIFICATION  = 
"ADDCLASSIFICATION";
+    private static final String ACTION_ADD_CLASSIFICATION  = 
"ADD_CLASSIFICATION";
     private static final String ACTION_NAME_SET            = "SET";
     private static final String ACTION_NAME_REPLACE_PREFIX = "REPLACE_PREFIX";
     private static final String ACTION_NAME_TO_LOWER       = "TO_LOWER";
     private static final String ACTION_NAME_TO_UPPER       = "TO_UPPER";
     private static final String ACTION_NAME_CLEAR          = "CLEAR";
 
-    protected final String attributeName;
+    protected final EntityAttribute attribute;
 
 
-    protected Action(String attributeName) {
-        this.attributeName = attributeName;
+    protected Action(EntityAttribute attribute) {
+        this.attribute = attribute;
     }
 
-    public String getAttributeName() { return attributeName; }
+    public EntityAttribute getAttribute() { return attribute; }
 
     public boolean isValid() {
-        return StringUtils.isNotEmpty(attributeName);
+        return true;
     }
 
     public abstract void apply(AtlasTransformableEntity entity);
 
 
-    public static Action createAction(String key, String value) {
+    public static Action createAction(String key, String value, 
TransformerContext context) {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> Action.createAction(key={}, value={})", key, value);
         }
@@ -74,33 +76,35 @@ public abstract class Action {
         actionValue = StringUtils.trim(actionValue);
         value       = StringUtils.trim(value);
 
+        EntityAttribute attribute = new EntityAttribute(StringUtils.trim(key), 
context);
+
         switch (actionName.toUpperCase()) {
             case ACTION_ADD_CLASSIFICATION:
-                ret = new AddClassificationAction(actionValue);
-                break;
+                ret = new AddClassificationAction(attribute, actionValue, 
context);
+            break;
 
             case ACTION_NAME_REPLACE_PREFIX:
-                ret = new PrefixReplaceAction(key, actionValue);
+                ret = new PrefixReplaceAction(attribute, actionValue);
             break;
 
             case ACTION_NAME_TO_LOWER:
-                ret = new ToLowerCaseAction(key);
+                ret = new ToLowerCaseAction(attribute);
             break;
 
             case ACTION_NAME_TO_UPPER:
-                ret = new ToUpperCaseAction(key);
+                ret = new ToUpperCaseAction(attribute);
             break;
 
             case ACTION_NAME_SET:
-                ret = new SetAction(key, actionValue);
+                ret = new SetAction(attribute, actionValue);
             break;
 
             case ACTION_NAME_CLEAR:
-                ret = new ClearAction(key);
-                break;
+                ret = new ClearAction(attribute);
+            break;
 
             default:
-                ret = new SetAction(key, value); // treat unspecified/unknown 
action as 'SET'
+                ret = new SetAction(attribute, value); // treat 
unspecified/unknown action as 'SET'
             break;
         }
 
@@ -115,71 +119,79 @@ public abstract class Action {
     public static class SetAction extends Action {
         private final String attributeValue;
 
-        public SetAction(String attributeName, String attributeValue) {
-            super(attributeName);
+        public SetAction(EntityAttribute attribute, String attributeValue) {
+            super(attribute);
 
             this.attributeValue = attributeValue;
         }
 
         @Override
         public void apply(AtlasTransformableEntity entity) {
-            if (isValid()) {
-                entity.setAttribute(attributeName, attributeValue);
-            }
+            entity.setAttribute(attribute, attributeValue);
         }
     }
 
-    public static class AddClassificationAction extends Action implements 
NeedsContext {
-
+    public static class AddClassificationAction extends Action {
         private final String classificationName;
-        private TransformerContext transformerContext;
 
-        public AddClassificationAction(String classificationName) {
-            super(ENTITY_KEY);
+        public AddClassificationAction(EntityAttribute attribute, String 
classificationName, TransformerContext context) {
+            super(attribute);
 
             this.classificationName = classificationName;
+
+            createClassificationDefIfNotExists(classificationName, context);
         }
 
         @Override
         public void apply(AtlasTransformableEntity transformableEntity) {
-            AtlasEntity entity = transformableEntity.entity;
+            AtlasEntity entity = transformableEntity.getEntity();
+
             if (entity.getClassifications() == null) {
-                entity.setClassifications(new 
ArrayList<AtlasClassification>());
+                entity.setClassifications(new ArrayList<>());
             }
 
+            boolean hasClassification = false;
+
             for (AtlasClassification c : entity.getClassifications()) {
-                if (c.getTypeName().equals(classificationName)) {
-                    return;
+                hasClassification = c.getTypeName().equals(classificationName);
+
+                if (hasClassification) {
+                    break;
                 }
             }
 
-            entity.getClassifications().add(new 
AtlasClassification(classificationName));
+            if (!hasClassification) {
+                entity.getClassifications().add(new 
AtlasClassification(classificationName));
+            }
         }
 
-        @Override
-        public void setContext(TransformerContext transformerContext) {
-            this.transformerContext = transformerContext;
-            getCreateTag(classificationName);
-        }
+        private void createClassificationDefIfNotExists(String 
classificationName, TransformerContext context) {
+            AtlasTypeRegistry typeRegistry = context != null ? 
context.getTypeRegistry() : null;
 
-        private void getCreateTag(String classificationName) {
-            if (transformerContext == null) {
-                return;
-            }
+            if (typeRegistry != null) {
+                try {
+                    AtlasClassificationDef classificationDef = 
typeRegistry.getClassificationDefByName(classificationName);
 
-            try {
-                AtlasClassificationDef classificationDef = 
transformerContext.getTypeRegistry().getClassificationDefByName(classificationName);
-                if (classificationDef != null) {
-                    return;
-                }
+                    if (classificationDef == null) {
+                        AtlasTypeDefStore typeDefStore = 
context.getTypeDefStore();
+
+                        if (typeDefStore != null) {
+                            classificationDef = new 
AtlasClassificationDef(classificationName);
 
-                classificationDef = new 
AtlasClassificationDef(classificationName);
-                AtlasTypesDef typesDef = new AtlasTypesDef();
-                
typesDef.setClassificationDefs(Collections.singletonList(classificationDef));
-                transformerContext.getTypeDefStore().createTypesDef(typesDef);
-                LOG.info("created classification: {}", classificationName);
-            } catch (AtlasBaseException e) {
-                LOG.error("Error creating classification: {}", 
classificationName, e);
+                            AtlasTypesDef typesDef = new AtlasTypesDef();
+
+                            
typesDef.setClassificationDefs(Collections.singletonList(classificationDef));
+
+                            typeDefStore.createTypesDef(typesDef);
+
+                            LOG.info("created classification: {}", 
classificationName);
+                        } else {
+                            LOG.warn("skipped creation of classification {}. 
typeDefStore is null", classificationName);
+                        }
+                    }
+                } catch (AtlasBaseException e) {
+                    LOG.error("Error creating classification: {}", 
classificationName, e);
+                }
             }
         }
     }
@@ -188,8 +200,8 @@ public abstract class Action {
         private final String fromPrefix;
         private final String toPrefix;
 
-        public PrefixReplaceAction(String attributeName, String actionValue) {
-            super(attributeName);
+        public PrefixReplaceAction(EntityAttribute attribute, String 
actionValue) {
+            super(attribute);
 
             // actionValue => =:prefixToReplace=replacedValue
             if (actionValue != null) {
@@ -224,61 +236,61 @@ public abstract class Action {
         @Override
         public void apply(AtlasTransformableEntity entity) {
             if (isValid()) {
-                Object currValue = entity.getAttribute(attributeName);
+                Object currValue = entity.getAttribute(attribute);
                 String strValue  = currValue != null ? currValue.toString() : 
null;
 
                 if (strValue != null && strValue.startsWith(fromPrefix)) {
-                    entity.setAttribute(attributeName, 
StringUtils.replace(strValue, fromPrefix, toPrefix, 1));
+                    entity.setAttribute(attribute, 
StringUtils.replace(strValue, fromPrefix, toPrefix, 1));
                 }
             }
         }
     }
 
     public static class ToLowerCaseAction extends Action {
-        public ToLowerCaseAction(String attributeName) {
-            super(attributeName);
+        public ToLowerCaseAction(EntityAttribute attribute) {
+            super(attribute);
         }
 
         @Override
         public void apply(AtlasTransformableEntity entity) {
             if (isValid()) {
-                Object currValue = entity.getAttribute(attributeName);
+                Object currValue = entity.getAttribute(attribute);
                 String strValue  = currValue instanceof String ? (String) 
currValue : null;
 
                 if (strValue != null) {
-                    entity.setAttribute(attributeName, strValue.toLowerCase());
+                    entity.setAttribute(attribute, strValue.toLowerCase());
                 }
             }
         }
     }
 
     public static class ToUpperCaseAction extends Action {
-        public ToUpperCaseAction(String attributeName) {
-            super(attributeName);
+        public ToUpperCaseAction(EntityAttribute attribute) {
+            super(attribute);
         }
 
         @Override
         public void apply(AtlasTransformableEntity entity) {
             if (isValid()) {
-                Object currValue = entity.getAttribute(attributeName);
+                Object currValue = entity.getAttribute(attribute);
                 String strValue  = currValue instanceof String ? (String) 
currValue : null;
 
                 if (strValue != null) {
-                    entity.setAttribute(attributeName, strValue.toUpperCase());
+                    entity.setAttribute(attribute, strValue.toUpperCase());
                 }
             }
         }
     }
 
     public static class ClearAction extends Action {
-        public ClearAction(String attributeName) {
-            super(attributeName);
+        public ClearAction(EntityAttribute attribute) {
+            super(attribute);
         }
 
         @Override
         public void apply(AtlasTransformableEntity entity) {
-            if (isValid() && entity.hasAttribute(attributeName)) {
-                entity.setAttribute(attributeName, null);
+            if (isValid() && entity.hasAttribute(attribute)) {
+                entity.setAttribute(attribute, null);
             }
         }
     }

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/AtlasEntityTransformer.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/AtlasEntityTransformer.java
 
b/intg/src/main/java/org/apache/atlas/entitytransform/AtlasEntityTransformer.java
index e9b2afd..27a57a6 100644
--- 
a/intg/src/main/java/org/apache/atlas/entitytransform/AtlasEntityTransformer.java
+++ 
b/intg/src/main/java/org/apache/atlas/entitytransform/AtlasEntityTransformer.java
@@ -19,32 +19,25 @@ package org.apache.atlas.entitytransform;
 
 import 
org.apache.atlas.entitytransform.BaseEntityHandler.AtlasTransformableEntity;
 import org.apache.atlas.model.impexp.AttributeTransform;
-import org.apache.atlas.model.instance.AtlasObjectId;
-import org.apache.atlas.type.AtlasType;
-import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.collections.MapUtils;
-import org.apache.commons.lang.StringUtils;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
+
 public class AtlasEntityTransformer {
     private final List<Condition> conditions;
     private final List<Action>    actions;
 
-    public AtlasEntityTransformer(AttributeTransform attributeTransform) {
-        this(attributeTransform.getConditions(), 
attributeTransform.getAction());
-    }
 
-    public AtlasEntityTransformer(AtlasObjectId objectId, Map<String, String> 
actions) {
-        this(Collections.singletonMap("__entity", AtlasType.toJson(objectId)), 
actions);
+    public AtlasEntityTransformer(AttributeTransform attributeTransform, 
TransformerContext context) {
+        this(attributeTransform.getConditions(), 
attributeTransform.getAction(), context);
     }
 
-    public AtlasEntityTransformer(Map<String, String> conditions, Map<String, 
String> actions) {
-        this.conditions = createConditions(conditions);
-        this.actions    = createActions(actions);
+    public AtlasEntityTransformer(Map<String, String> conditions, Map<String, 
String> actions, TransformerContext context) {
+        this.conditions = createConditions(conditions, context);
+        this.actions    = createActions(actions, context);
     }
 
     public List<Condition> getConditions() {
@@ -71,12 +64,12 @@ public class AtlasEntityTransformer {
         }
     }
 
-    private List<Condition> createConditions(Map<String, String> conditions) {
+    private List<Condition> createConditions(Map<String, String> conditions, 
TransformerContext context) {
         List<Condition> ret = new ArrayList<>();
 
         if (MapUtils.isNotEmpty(conditions)) {
             for (Map.Entry<String, String> entry : conditions.entrySet()) {
-                Condition condition = 
Condition.createCondition(entry.getKey(), entry.getValue());
+                Condition condition = 
Condition.createCondition(entry.getKey(), entry.getValue(), context);
 
                 ret.add(condition);
             }
@@ -85,12 +78,12 @@ public class AtlasEntityTransformer {
         return ret;
     }
 
-    private List<Action> createActions(Map<String, String> actions) {
+    private List<Action> createActions(Map<String, String> actions, 
TransformerContext context) {
         List<Action> ret = new ArrayList<>();
 
         if (MapUtils.isNotEmpty(actions)) {
             for (Map.Entry<String, String> entry : actions.entrySet()) {
-                Action action = Action.createAction(entry.getKey(), 
entry.getValue());
+                Action action = Action.createAction(entry.getKey(), 
entry.getValue(), context);
 
                 ret.add(action);
             }

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/BaseEntityHandler.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/BaseEntityHandler.java 
b/intg/src/main/java/org/apache/atlas/entitytransform/BaseEntityHandler.java
index dd6c665..975e4dd 100644
--- a/intg/src/main/java/org/apache/atlas/entitytransform/BaseEntityHandler.java
+++ b/intg/src/main/java/org/apache/atlas/entitytransform/BaseEntityHandler.java
@@ -17,12 +17,9 @@
  */
 package org.apache.atlas.entitytransform;
 
-import org.apache.atlas.model.impexp.AtlasExportRequest;
 import org.apache.atlas.model.impexp.AttributeTransform;
 import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.store.AtlasTypeDefStore;
 import org.apache.atlas.type.AtlasType;
-import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
@@ -31,34 +28,24 @@ import org.slf4j.LoggerFactory;
 import java.util.ArrayList;
 import java.util.List;
 
-import static 
org.apache.atlas.entitytransform.TransformationConstants.TYPE_NAME_ATTRIBUTE_NAME_SEP;
 
 public class BaseEntityHandler {
     private static final Logger LOG = 
LoggerFactory.getLogger(BaseEntityHandler.class);
 
     protected final List<AtlasEntityTransformer> transformers;
-    protected final boolean                      hasCustomAttributeTransformer;
-    private TransformerContext                   transformerContext;
 
-    public BaseEntityHandler(List<AtlasEntityTransformer> transformers) {
-        this(transformers, null);
-    }
 
-    public BaseEntityHandler(List<AtlasEntityTransformer> transformers, 
List<String> customTransformAttributes) {
-        this.transformers                  = transformers;
-        this.hasCustomAttributeTransformer = 
hasTransformerForAnyAttribute(customTransformAttributes);
-    }
-
-    public boolean hasCustomAttributeTransformer() {
-        return hasCustomAttributeTransformer;
+    public BaseEntityHandler(List<AtlasEntityTransformer> transformers) {
+        this.transformers = transformers;
     }
 
     public AtlasEntity transform(AtlasEntity entity) {
-        if (!CollectionUtils.isNotEmpty(transformers)) {
+        if (CollectionUtils.isEmpty(transformers)) {
             return entity;
         }
 
         AtlasTransformableEntity transformableEntity = 
getTransformableEntity(entity);
+
         if (transformableEntity == null) {
             return entity;
         }
@@ -72,22 +59,6 @@ public class BaseEntityHandler {
         return entity;
     }
 
-    private void setContextForActions(List<Action> actions) {
-        for(Action action : actions) {
-            if (action instanceof NeedsContext) {
-                ((NeedsContext) action).setContext(transformerContext);
-            }
-        }
-    }
-
-    private void setContextForConditions(List<Condition> conditions) {
-        for(Condition condition : conditions) {
-            if (condition instanceof NeedsContext) {
-                ((NeedsContext) condition).setContext(transformerContext);
-            }
-        }
-    }
-
     public AtlasTransformableEntity getTransformableEntity(AtlasEntity entity) 
{
         return new AtlasTransformableEntity(entity);
     }
@@ -97,39 +68,38 @@ public class BaseEntityHandler {
             LOG.debug("==> 
BaseEntityHandler.createEntityHandlers(transforms={})", transforms);
         }
 
-        List<AtlasEntityTransformer> transformers = new ArrayList<>();
+        List<BaseEntityHandler> ret = new ArrayList<>();
 
-        for (AttributeTransform transform : transforms) {
-            transformers.add(new AtlasEntityTransformer(transform));
-        }
+        if (CollectionUtils.isNotEmpty(transforms)) {
+            List<AtlasEntityTransformer> transformers = new ArrayList<>();
 
-        BaseEntityHandler[] handlers = new BaseEntityHandler[] {
-                new HdfsPathEntityHandler(transformers),
-                new HiveDatabaseEntityHandler(transformers),
-                new HiveTableEntityHandler(transformers),
-                new HiveColumnEntityHandler(transformers),
-                new HiveStorageDescriptorEntityHandler(transformers)
-        };
+            for (AttributeTransform transform : transforms) {
+                transformers.add(new AtlasEntityTransformer(transform, 
context));
+            }
 
-        List<BaseEntityHandler> ret = new ArrayList<>();
+            if (hasTransformerForAnyAttribute(transformers, 
HdfsPathEntityHandler.CUSTOM_TRANSFORM_ATTRIBUTES)) {
+                ret.add(new HdfsPathEntityHandler(transformers));
+            }
 
-        // include customer handlers, only if its customer attribute is 
transformed
-        for (BaseEntityHandler handler : handlers) {
-            if (handler.hasCustomAttributeTransformer()) {
-                ret.add(handler);
-                handler.setContext(context);
+            if (hasTransformerForAnyAttribute(transformers, 
HiveDatabaseEntityHandler.CUSTOM_TRANSFORM_ATTRIBUTES)) {
+                ret.add(new HiveDatabaseEntityHandler(transformers));
             }
-        }
 
-        if (CollectionUtils.isEmpty(ret)) {
-            BaseEntityHandler be = new BaseEntityHandler(transformers);
-            be.setContext(context);
+            if (hasTransformerForAnyAttribute(transformers, 
HiveTableEntityHandler.CUSTOM_TRANSFORM_ATTRIBUTES)) {
+                ret.add(new HiveTableEntityHandler(transformers));
+            }
 
-            ret.add(be);
-        }
+            if (hasTransformerForAnyAttribute(transformers, 
HiveColumnEntityHandler.CUSTOM_TRANSFORM_ATTRIBUTES)) {
+                ret.add(new HiveColumnEntityHandler(transformers));
+            }
+
+            if (hasTransformerForAnyAttribute(transformers, 
HiveStorageDescriptorEntityHandler.CUSTOM_TRANSFORM_ATTRIBUTES)) {
+                ret.add(new HiveStorageDescriptorEntityHandler(transformers));
+            }
 
-        if (CollectionUtils.isEmpty(ret)) {
-            ret.add(new BaseEntityHandler(transformers));
+            if (CollectionUtils.isEmpty(ret)) {
+                ret.add(new BaseEntityHandler(transformers));
+            }
         }
 
         if (LOG.isDebugEnabled()) {
@@ -139,11 +109,11 @@ public class BaseEntityHandler {
         return ret;
     }
 
-    private boolean hasTransformerForAnyAttribute(List<String> attributes) {
+    private static boolean 
hasTransformerForAnyAttribute(List<AtlasEntityTransformer> transformers, 
List<String> attributes) {
         if (CollectionUtils.isNotEmpty(transformers) && 
CollectionUtils.isNotEmpty(attributes)) {
             for (AtlasEntityTransformer transformer : transformers) {
                 for (Action action : transformer.getActions()) {
-                    if (attributes.contains(action.getAttributeName())) {
+                    if 
(attributes.contains(action.getAttribute().getAttributeKey())) {
                         return true;
                     }
                 }
@@ -152,20 +122,6 @@ public class BaseEntityHandler {
 
         return false;
     }
-    public void setContext(AtlasTypeRegistry typeRegistry, AtlasTypeDefStore 
typeDefStore, AtlasExportRequest request) {
-        setContext(new TransformerContext(typeRegistry, typeDefStore, 
request));
-    }
-
-    public void setContext(TransformerContext context) {
-        this.transformerContext = context;
-
-        for (AtlasEntityTransformer transformer : transformers) {
-            if (transformerContext != null) {
-                setContextForActions(transformer.getActions());
-                setContextForConditions(transformer.getConditions());
-            }
-        }
-    }
 
     public static class AtlasTransformableEntity {
         protected final AtlasEntity entity;
@@ -178,38 +134,26 @@ public class BaseEntityHandler {
             return entity;
         }
 
-        public Object getAttribute(String attributeName) {
-            Object ret = null;
-
-            if (entity != null && attributeName != null) {
-                ret = entity.getAttribute(attributeName);
+        public Object getAttribute(EntityAttribute attribute) {
+            final Object ret;
 
-                if (ret == null) { // try after dropping typeName prefix, if 
attributeName contains it
-                    int idxSep = 
attributeName.indexOf(TYPE_NAME_ATTRIBUTE_NAME_SEP);
-
-                    if (idxSep != -1) {
-                        ret = 
entity.getAttribute(attributeName.substring(idxSep + 1));
-                    }
-                }
+            if (attribute.appliesToEntityType(entity.getTypeName())) {
+                ret = entity.getAttribute(attribute.getAttributeName());
+            } else {
+                ret = null;
             }
 
             return ret;
         }
 
-        public void setAttribute(String attributeName, String attributeValue) {
-            if (entity != null && attributeName != null) {
-                int idxSep = 
attributeName.indexOf(TYPE_NAME_ATTRIBUTE_NAME_SEP); // drop typeName prefix, 
if attributeName contains it
-
-                if (idxSep != -1) {
-                    entity.setAttribute(attributeName.substring(idxSep + 1), 
attributeValue);
-                } else {
-                    entity.setAttribute(attributeName, attributeValue);
-                }
+        public void setAttribute(EntityAttribute attribute, String 
attributeValue) {
+            if (attribute.appliesToEntityType(entity.getTypeName())) {
+                entity.setAttribute(attribute.getAttributeName(), 
attributeValue);
             }
         }
 
-        public boolean hasAttribute(String attributeName) {
-            return getAttribute(attributeName) != null;
+        public boolean hasAttribute(EntityAttribute attribute) {
+            return getAttribute(attribute) != null;
         }
 
         public void transformComplete() {

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/Condition.java
----------------------------------------------------------------------
diff --git a/intg/src/main/java/org/apache/atlas/entitytransform/Condition.java 
b/intg/src/main/java/org/apache/atlas/entitytransform/Condition.java
index 174b9b4..3bf49f0 100644
--- a/intg/src/main/java/org/apache/atlas/entitytransform/Condition.java
+++ b/intg/src/main/java/org/apache/atlas/entitytransform/Condition.java
@@ -17,7 +17,9 @@
  */
 package org.apache.atlas.entitytransform;
 
+import com.google.common.annotations.VisibleForTesting;
 import 
org.apache.atlas.entitytransform.BaseEntityHandler.AtlasTransformableEntity;
+import org.apache.atlas.model.impexp.AtlasExportRequest;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.commons.lang.StringUtils;
@@ -30,6 +32,7 @@ import java.util.Map;
 import java.util.Objects;
 
 
+
 public abstract class Condition {
     private static final Logger LOG = LoggerFactory.getLogger(Condition.class);
 
@@ -43,18 +46,19 @@ public abstract class Condition {
     private static final String CONDITION_NAME_STARTS_WITH_IGNORE_CASE = 
"STARTS_WITH_IGNORE_CASE";
     private static final String CONDITION_NAME_HAS_VALUE               = 
"HAS_VALUE";
 
-    protected final String attributeName;
+    protected final EntityAttribute attribute;
+
 
-    protected Condition(String attributeName) {
-        this.attributeName = attributeName;
+    protected Condition(EntityAttribute attribute) {
+        this.attribute = attribute;
     }
 
-    public String getAttributeName() { return attributeName; }
+    public EntityAttribute getAttribute() { return attribute; }
 
     public abstract boolean matches(AtlasTransformableEntity entity);
 
 
-    public static Condition createCondition(String key, String value) {
+    public static Condition createCondition(String key, String value, 
TransformerContext context) {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> Condition.createCondition(key={}, value={})", key, 
value);
         }
@@ -69,41 +73,43 @@ public abstract class Condition {
         conditionValue = StringUtils.trim(conditionValue);
         value          = StringUtils.trim(value);
 
+        EntityAttribute attribute = new EntityAttribute(StringUtils.trim(key), 
context);
+
         switch (conditionName.toUpperCase()) {
             case CONDITION_ENTITY_ALL:
-                ret = new ObjectIdEquals(key, CONDITION_ENTITY_ALL);
+                ret = new ObjectIdEquals(attribute, CONDITION_ENTITY_ALL, 
context);
                 break;
 
             case CONDITION_ENTITY_TOP_LEVEL:
-                ret = new ObjectIdEquals(key, CONDITION_ENTITY_TOP_LEVEL);
+                ret = new ObjectIdEquals(attribute, 
CONDITION_ENTITY_TOP_LEVEL, context);
                 break;
 
             case CONDITION_ENTITY_OBJECT_ID:
-                ret = new ObjectIdEquals(key, conditionValue);
+                ret = new ObjectIdEquals(attribute, conditionValue, context);
                 break;
 
             case CONDITION_NAME_EQUALS:
-                ret = new EqualsCondition(key, conditionValue);
+                ret = new EqualsCondition(attribute, conditionValue);
             break;
 
             case CONDITION_NAME_EQUALS_IGNORE_CASE:
-                ret = new EqualsIgnoreCaseCondition(key, conditionValue);
+                ret = new EqualsIgnoreCaseCondition(attribute, conditionValue);
             break;
 
             case CONDITION_NAME_STARTS_WITH:
-                ret = new StartsWithCondition(key, conditionValue);
+                ret = new StartsWithCondition(attribute, conditionValue);
             break;
 
             case CONDITION_NAME_STARTS_WITH_IGNORE_CASE:
-                ret = new StartsWithIgnoreCaseCondition(key, conditionValue);
+                ret = new StartsWithIgnoreCaseCondition(attribute, 
conditionValue);
             break;
 
             case CONDITION_NAME_HAS_VALUE:
-                ret = new HasValueCondition(key, conditionValue);
+                ret = new HasValueCondition(attribute);
                 break;
 
             default:
-                ret = new EqualsCondition(key, value); // treat 
unspecified/unknown condition as 'EQUALS'
+                ret = new EqualsCondition(attribute, value); // treat 
unspecified/unknown condition as 'EQUALS'
             break;
         }
 
@@ -118,15 +124,15 @@ public abstract class Condition {
     public static class EqualsCondition extends Condition {
         protected final String attributeValue;
 
-        public EqualsCondition(String attributeName, String attributeValue) {
-            super(attributeName);
+        public EqualsCondition(EntityAttribute attribute, String 
attributeValue) {
+            super(attribute);
 
             this.attributeValue = attributeValue;
         }
 
         @Override
         public boolean matches(AtlasTransformableEntity entity) {
-            Object attributeValue = entity != null ? 
entity.getAttribute(attributeName) : null;
+            Object attributeValue = entity != null ? 
entity.getAttribute(attribute) : null;
 
             return attributeValue != null && 
StringUtils.equals(attributeValue.toString(), this.attributeValue);
         }
@@ -136,15 +142,15 @@ public abstract class Condition {
     public static class EqualsIgnoreCaseCondition extends Condition {
         protected final String attributeValue;
 
-        public EqualsIgnoreCaseCondition(String attributeName, String 
attributeValue) {
-            super(attributeName);
+        public EqualsIgnoreCaseCondition(EntityAttribute attribute, String 
attributeValue) {
+            super(attribute);
 
             this.attributeValue = attributeValue;
         }
 
         @Override
         public boolean matches(AtlasTransformableEntity entity) {
-            Object attributeValue = entity != null ? 
entity.getAttribute(attributeName) : null;
+            Object attributeValue = entity != null ? 
entity.getAttribute(attribute) : null;
 
             return attributeValue != null && 
StringUtils.equalsIgnoreCase(attributeValue.toString(), this.attributeValue);
         }
@@ -154,15 +160,15 @@ public abstract class Condition {
     public static class StartsWithCondition extends Condition {
         protected final String prefix;
 
-        public StartsWithCondition(String attributeName, String prefix) {
-            super(attributeName);
+        public StartsWithCondition(EntityAttribute attribute, String prefix) {
+            super(attribute);
 
             this.prefix = prefix;
         }
 
         @Override
         public boolean matches(AtlasTransformableEntity entity) {
-            Object attributeValue = entity != null ? 
entity.getAttribute(attributeName) : null;
+            Object attributeValue = entity != null ? 
entity.getAttribute(attribute) : null;
 
             return attributeValue != null && 
StringUtils.startsWith(attributeValue.toString(), this.prefix);
         }
@@ -172,96 +178,89 @@ public abstract class Condition {
     public static class StartsWithIgnoreCaseCondition extends Condition {
         protected final String prefix;
 
-        public StartsWithIgnoreCaseCondition(String attributeName, String 
prefix) {
-            super(attributeName);
+        public StartsWithIgnoreCaseCondition(EntityAttribute attribute, String 
prefix) {
+            super(attribute);
 
             this.prefix = prefix;
         }
 
         @Override
         public boolean matches(AtlasTransformableEntity entity) {
-            Object attributeValue = entity != null ? 
entity.getAttribute(attributeName) : null;
+            Object attributeValue = entity != null ? 
entity.getAttribute(attribute) : null;
 
             return attributeValue != null && 
StringUtils.startsWithIgnoreCase(attributeValue.toString(), this.prefix);
         }
     }
 
-    static class ObjectIdEquals extends Condition implements NeedsContext {
+    static class ObjectIdEquals extends Condition {
+        private final boolean             isMatchAll;
         private final List<AtlasObjectId> objectIds;
-        private String scope;
-        private TransformerContext transformerContext;
 
-        public ObjectIdEquals(String key, String conditionValue) {
-            super(key);
+        public ObjectIdEquals(EntityAttribute attribute, String scope, 
TransformerContext context) {
+            super(attribute);
 
-            objectIds = new ArrayList<>();
-            this.scope = conditionValue;
+            this.isMatchAll = StringUtils.equals(scope, CONDITION_ENTITY_ALL);
+            this.objectIds  = new ArrayList<>();
+
+            if (!isMatchAll && context != null && context.getExportRequest() 
!= null) {
+                AtlasExportRequest request = context.getExportRequest();
+
+                for(AtlasObjectId objectId : request.getItemsToExport()) {
+                    addObjectId(objectId);
+                }
+            }
         }
 
         @Override
         public boolean matches(AtlasTransformableEntity entity) {
-            for (AtlasObjectId objectId : objectIds) {
-                return isMatch(objectId, entity.entity);
-            }
+            if (isMatchAll) {
+                return true;
+            } else {
+                for (AtlasObjectId objectId : objectIds) {
+                    if (isMatch(objectId, entity.getEntity())) {
+                        return true;
+                    }
+                }
 
-            return objectIds.size() == 0;
+                return false;
+            }
         }
 
-        public void add(AtlasObjectId objectId) {
-            this.objectIds.add(objectId);
+        @VisibleForTesting
+        void addObjectId(AtlasObjectId objId) {
+            this.objectIds.add(objId);
         }
 
         private boolean isMatch(AtlasObjectId objectId, AtlasEntity entity) {
-            boolean ret = true;
             if (!StringUtils.isEmpty(objectId.getGuid())) {
                 return Objects.equals(objectId.getGuid(), entity.getGuid());
             }
 
-            ret = Objects.equals(objectId.getTypeName(), entity.getTypeName());
-            if (!ret) {
-                return ret;
-            }
+            boolean ret = Objects.equals(objectId.getTypeName(), 
entity.getTypeName());
+
+            if (ret) {
+                for (Map.Entry<String, Object> entry : 
objectId.getUniqueAttributes().entrySet()) {
+                    ret = ret && 
Objects.equals(entity.getAttribute(entry.getKey()), entry.getValue());
 
-            for (Map.Entry<String, Object> entry : 
objectId.getUniqueAttributes().entrySet()) {
-                ret = ret && 
Objects.equals(entity.getAttribute(entry.getKey()), entry.getValue());
-                if (!ret) {
-                    break;
+                    if (!ret) {
+                        break;
+                    }
                 }
             }
 
             return ret;
         }
-
-        @Override
-        public void setContext(TransformerContext transformerContext) {
-            this.transformerContext = transformerContext;
-            if(StringUtils.isEmpty(scope) || 
scope.equals(CONDITION_ENTITY_ALL)) {
-                return;
-            }
-
-            addObjectIdsFromExportRequest();
-        }
-
-        private void addObjectIdsFromExportRequest() {
-            for(AtlasObjectId objectId : 
this.transformerContext.getExportRequest().getItemsToExport()) {
-                add(objectId);
-            }
-        }
     }
 
 
     public static class HasValueCondition extends Condition {
-        protected final String attributeValue;
-
-        public HasValueCondition(String attributeName, String attributeValue) {
-            super(attributeName);
-
-            this.attributeValue = attributeValue;
+        public HasValueCondition(EntityAttribute attribute) {
+            super(attribute);
         }
 
         @Override
         public boolean matches(AtlasTransformableEntity entity) {
-            Object attributeValue = entity != null ? 
entity.getAttribute(attributeName) : null;
+            Object attributeValue = entity != null ? 
entity.getAttribute(attribute) : null;
 
             return attributeValue != null ? 
StringUtils.isNotEmpty(attributeValue.toString()) : false;
         }

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/EntityAttribute.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/EntityAttribute.java 
b/intg/src/main/java/org/apache/atlas/entitytransform/EntityAttribute.java
new file mode 100644
index 0000000..040c7cb
--- /dev/null
+++ b/intg/src/main/java/org/apache/atlas/entitytransform/EntityAttribute.java
@@ -0,0 +1,69 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.atlas.entitytransform;
+
+import org.apache.atlas.type.AtlasEntityType;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.commons.lang.StringUtils;
+
+import static 
org.apache.atlas.entitytransform.TransformationConstants.TYPE_NAME_ATTRIBUTE_NAME_SEP;
+
+public class EntityAttribute {
+    private final String          attributeKey;
+    private final AtlasEntityType entityType;
+    private final String          attributeName;
+
+    public EntityAttribute(String attributeKey, TransformerContext context) {
+        this.attributeKey = attributeKey;
+
+        int idx = attributeKey != null ? 
attributeKey.indexOf(TYPE_NAME_ATTRIBUTE_NAME_SEP) : -1;
+
+        if (idx != -1) {
+            this.attributeName = StringUtils.trim(attributeKey.substring(idx + 
1));
+
+            AtlasTypeRegistry typeRegistry = context != null ? 
context.getTypeRegistry() : null;
+
+            if (typeRegistry != null) {
+                String typeName = StringUtils.trim(attributeKey.substring(0, 
idx));
+
+                this.entityType = typeRegistry.getEntityTypeByName(typeName);
+            } else {
+                this.entityType = null;
+            }
+        } else {
+            this.entityType    = null;
+            this.attributeName = attributeKey;
+        }
+    }
+
+    public String getAttributeKey() {
+        return attributeKey;
+    }
+
+    public AtlasEntityType getEntityType() {
+        return entityType;
+    }
+
+    public String getAttributeName() {
+        return attributeName;
+    }
+
+    public boolean appliesToEntityType(String typeName) {
+        return entityType == null || StringUtils.isEmpty(typeName) || 
entityType.getTypeAndAllSubTypes().contains(typeName);
+    }
+}

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/HdfsPathEntityHandler.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/HdfsPathEntityHandler.java
 
b/intg/src/main/java/org/apache/atlas/entitytransform/HdfsPathEntityHandler.java
index 1a398ea..2df98a6 100644
--- 
a/intg/src/main/java/org/apache/atlas/entitytransform/HdfsPathEntityHandler.java
+++ 
b/intg/src/main/java/org/apache/atlas/entitytransform/HdfsPathEntityHandler.java
@@ -35,10 +35,10 @@ import static 
org.apache.atlas.entitytransform.TransformationConstants.QUALIFIED
 
 
 public class HdfsPathEntityHandler extends BaseEntityHandler {
-    private static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HDFS_PATH_NAME_ATTRIBUTE, HDFS_PATH_PATH_ATTRIBUTE, 
HDFS_CLUSTER_NAME_ATTRIBUTE);
+    static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HDFS_PATH_NAME_ATTRIBUTE, HDFS_PATH_PATH_ATTRIBUTE, 
HDFS_CLUSTER_NAME_ATTRIBUTE);
 
     public HdfsPathEntityHandler(List<AtlasEntityTransformer> transformers) {
-        super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES);
+        super(transformers);
     }
 
     @Override
@@ -56,8 +56,8 @@ public class HdfsPathEntityHandler extends BaseEntityHandler {
         private String  path;
         private String  name;
         private String  pathPrefix;
-        private boolean isPathUpdated              = false;
-        private boolean isCustomerAttributeUpdated = false;
+        private boolean isPathUpdated            = false;
+        private boolean isCustomAttributeUpdated = false;
 
 
         public HdfsPathEntity(AtlasEntity entity) {
@@ -93,8 +93,8 @@ public class HdfsPathEntityHandler extends BaseEntityHandler {
         }
 
         @Override
-        public Object getAttribute(String attributeName) {
-            switch (attributeName) {
+        public Object getAttribute(EntityAttribute attribute) {
+            switch (attribute.getAttributeKey()) {
                 case HDFS_CLUSTER_NAME_ATTRIBUTE:
                     return clusterName;
 
@@ -105,40 +105,40 @@ public class HdfsPathEntityHandler extends 
BaseEntityHandler {
                     return path;
             }
 
-            return super.getAttribute(attributeName);
+            return super.getAttribute(attribute);
         }
 
         @Override
-        public void setAttribute(String attributeName, String attributeValue) {
-            switch (attributeName) {
+        public void setAttribute(EntityAttribute attribute, String 
attributeValue) {
+            switch (attribute.getAttributeKey()) {
                 case HDFS_CLUSTER_NAME_ATTRIBUTE:
                     clusterName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HDFS_PATH_NAME_ATTRIBUTE:
                     name = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HDFS_PATH_PATH_ATTRIBUTE:
                     path = attributeValue;
 
                     isPathUpdated              = true;
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 default:
-                    super.setAttribute(attributeName, attributeValue);
+                    super.setAttribute(attribute, attributeValue);
                 break;
             }
         }
 
         @Override
         public void transformComplete() {
-            if (isCustomerAttributeUpdated) {
+            if (isCustomAttributeUpdated) {
                 entity.setAttribute(CLUSTER_NAME_ATTRIBUTE, clusterName);
                 entity.setAttribute(NAME_ATTRIBUTE, name);
                 entity.setAttribute(PATH_ATTRIBUTE, toPath());

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/HiveColumnEntityHandler.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/HiveColumnEntityHandler.java
 
b/intg/src/main/java/org/apache/atlas/entitytransform/HiveColumnEntityHandler.java
index fca94b6..686e11c 100644
--- 
a/intg/src/main/java/org/apache/atlas/entitytransform/HiveColumnEntityHandler.java
+++ 
b/intg/src/main/java/org/apache/atlas/entitytransform/HiveColumnEntityHandler.java
@@ -27,10 +27,10 @@ import static 
org.apache.atlas.entitytransform.TransformationConstants.*;
 
 
 public class HiveColumnEntityHandler extends BaseEntityHandler {
-    private static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HIVE_DB_NAME_ATTRIBUTE, HIVE_TABLE_NAME_ATTRIBUTE, 
HIVE_COLUMN_NAME_ATTRIBUTE, HIVE_DB_CLUSTER_NAME_ATTRIBUTE);
+    static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HIVE_DB_NAME_ATTRIBUTE, HIVE_TABLE_NAME_ATTRIBUTE, 
HIVE_COLUMN_NAME_ATTRIBUTE, HIVE_DB_CLUSTER_NAME_ATTRIBUTE);
 
     public HiveColumnEntityHandler(List<AtlasEntityTransformer> transformers) {
-        super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES);
+        super(transformers);
     }
 
     @Override
@@ -48,7 +48,7 @@ public class HiveColumnEntityHandler extends 
BaseEntityHandler {
         private String  tableName;
         private String  columnName;
         private String  clusterName;
-        private boolean isCustomerAttributeUpdated = false;
+        private boolean isCustomAttributeUpdated = false;
 
         public HiveColumnEntity(AtlasEntity entity) {
             super(entity);
@@ -73,8 +73,8 @@ public class HiveColumnEntityHandler extends 
BaseEntityHandler {
         }
 
         @Override
-        public Object getAttribute(String attributeName) {
-            switch (attributeName) {
+        public Object getAttribute(EntityAttribute attribute) {
+            switch (attribute.getAttributeKey()) {
                 case HIVE_DB_NAME_ATTRIBUTE:
                     return databaseName;
 
@@ -88,45 +88,45 @@ public class HiveColumnEntityHandler extends 
BaseEntityHandler {
                     return clusterName;
             }
 
-            return super.getAttribute(attributeName);
+            return super.getAttribute(attribute);
         }
 
         @Override
-        public void setAttribute(String attributeName, String attributeValue) {
-            switch (attributeName) {
+        public void setAttribute(EntityAttribute attribute, String 
attributeValue) {
+            switch (attribute.getAttributeKey()) {
                 case HIVE_DB_NAME_ATTRIBUTE:
                     databaseName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HIVE_TABLE_NAME_ATTRIBUTE:
                     tableName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HIVE_COLUMN_NAME_ATTRIBUTE:
                     columnName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HIVE_DB_CLUSTER_NAME_ATTRIBUTE:
                     clusterName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 default:
-                    super.setAttribute(attributeName, attributeValue);
+                    super.setAttribute(attribute, attributeValue);
                 break;
             }
         }
 
         @Override
         public void transformComplete() {
-            if (isCustomerAttributeUpdated) {
+            if (isCustomAttributeUpdated) {
                 entity.setAttribute(NAME_ATTRIBUTE, columnName);
                 entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, 
toQualifiedName());
             }

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/HiveDatabaseEntityHandler.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/HiveDatabaseEntityHandler.java
 
b/intg/src/main/java/org/apache/atlas/entitytransform/HiveDatabaseEntityHandler.java
index 8a2e813..b8032aa 100644
--- 
a/intg/src/main/java/org/apache/atlas/entitytransform/HiveDatabaseEntityHandler.java
+++ 
b/intg/src/main/java/org/apache/atlas/entitytransform/HiveDatabaseEntityHandler.java
@@ -26,10 +26,10 @@ import java.util.List;
 import static org.apache.atlas.entitytransform.TransformationConstants.*;
 
 public class HiveDatabaseEntityHandler extends BaseEntityHandler {
-    private static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HIVE_DB_NAME_ATTRIBUTE, HIVE_DB_CLUSTER_NAME_ATTRIBUTE);
+    static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HIVE_DB_NAME_ATTRIBUTE, HIVE_DB_CLUSTER_NAME_ATTRIBUTE);
 
     public HiveDatabaseEntityHandler(List<AtlasEntityTransformer> 
transformers) {
-        super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES);
+        super(transformers);
     }
 
     @Override
@@ -45,7 +45,7 @@ public class HiveDatabaseEntityHandler extends 
BaseEntityHandler {
     private static class HiveDatabaseEntity extends AtlasTransformableEntity {
         private String  databaseName;
         private String  clusterName;
-        private boolean isCustomerAttributeUpdated = false;
+        private boolean isCustomAttributeUpdated = false;
 
         public HiveDatabaseEntity(AtlasEntity entity) {
             super(entity);
@@ -64,8 +64,8 @@ public class HiveDatabaseEntityHandler extends 
BaseEntityHandler {
         }
 
         @Override
-        public Object getAttribute(String attributeName) {
-            switch (attributeName) {
+        public Object getAttribute(EntityAttribute attribute) {
+            switch (attribute.getAttributeKey()) {
                 case HIVE_DB_NAME_ATTRIBUTE:
                     return databaseName;
 
@@ -73,33 +73,33 @@ public class HiveDatabaseEntityHandler extends 
BaseEntityHandler {
                     return clusterName;
             }
 
-            return super.getAttribute(attributeName);
+            return super.getAttribute(attribute);
         }
 
         @Override
-        public void setAttribute(String attributeName, String attributeValue) {
-            switch (attributeName) {
+        public void setAttribute(EntityAttribute attribute, String 
attributeValue) {
+            switch (attribute.getAttributeKey()) {
                 case HIVE_DB_NAME_ATTRIBUTE:
                     databaseName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HIVE_DB_CLUSTER_NAME_ATTRIBUTE:
                     clusterName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 default:
-                    super.setAttribute(attributeName, attributeValue);
+                    super.setAttribute(attribute, attributeValue);
                 break;
             }
         }
 
         @Override
         public void transformComplete() {
-            if (isCustomerAttributeUpdated) {
+            if (isCustomAttributeUpdated) {
                 entity.setAttribute(NAME_ATTRIBUTE, databaseName);
                 entity.setAttribute(CLUSTER_NAME_ATTRIBUTE, clusterName);
                 entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, 
toQualifiedName());

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/HiveStorageDescriptorEntityHandler.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/HiveStorageDescriptorEntityHandler.java
 
b/intg/src/main/java/org/apache/atlas/entitytransform/HiveStorageDescriptorEntityHandler.java
index 6a7b17b..dc4edfb 100644
--- 
a/intg/src/main/java/org/apache/atlas/entitytransform/HiveStorageDescriptorEntityHandler.java
+++ 
b/intg/src/main/java/org/apache/atlas/entitytransform/HiveStorageDescriptorEntityHandler.java
@@ -26,11 +26,11 @@ import java.util.List;
 import static org.apache.atlas.entitytransform.TransformationConstants.*;
 
 public class HiveStorageDescriptorEntityHandler extends BaseEntityHandler {
-    private static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HIVE_DB_NAME_ATTRIBUTE, HIVE_TABLE_NAME_ATTRIBUTE, 
HIVE_DB_CLUSTER_NAME_ATTRIBUTE);
+    static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HIVE_DB_NAME_ATTRIBUTE, HIVE_TABLE_NAME_ATTRIBUTE, 
HIVE_DB_CLUSTER_NAME_ATTRIBUTE);
 
 
     public HiveStorageDescriptorEntityHandler(List<AtlasEntityTransformer> 
transformers) {
-        super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES);
+        super(transformers);
     }
 
     @Override
@@ -47,7 +47,7 @@ public class HiveStorageDescriptorEntityHandler extends 
BaseEntityHandler {
         private String  tableName;
         private String  clusterName;
         private String  location;
-        private boolean isCustomerAttributeUpdated = false;
+        private boolean isCustomAttributeUpdated = false;
 
 
         public HiveStorageDescriptorEntity(AtlasEntity entity) {
@@ -80,8 +80,8 @@ public class HiveStorageDescriptorEntityHandler extends 
BaseEntityHandler {
         }
 
         @Override
-        public Object getAttribute(String attributeName) {
-            switch (attributeName) {
+        public Object getAttribute(EntityAttribute attribute) {
+            switch (attribute.getAttributeKey()) {
                 case HIVE_DB_NAME_ATTRIBUTE:
                     return databaseName;
 
@@ -92,39 +92,39 @@ public class HiveStorageDescriptorEntityHandler extends 
BaseEntityHandler {
                     return clusterName;
             }
 
-            return super.getAttribute(attributeName);
+            return super.getAttribute(attribute);
         }
 
         @Override
-        public void setAttribute(String attributeName, String attributeValue) {
-            switch (attributeName) {
+        public void setAttribute(EntityAttribute attribute, String 
attributeValue) {
+            switch (attribute.getAttributeKey()) {
                 case HIVE_DB_NAME_ATTRIBUTE:
                     databaseName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HIVE_TABLE_NAME_ATTRIBUTE:
                     tableName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HIVE_DB_CLUSTER_NAME_ATTRIBUTE:
                     clusterName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 default:
-                    super.setAttribute(attributeName, attributeValue);
+                    super.setAttribute(attribute, attributeValue);
                 break;
             }
         }
 
         @Override
         public void transformComplete() {
-            if (isCustomerAttributeUpdated) {
+            if (isCustomAttributeUpdated) {
                 entity.setAttribute(LOCATION_ATTRIBUTE, toLocation());
                 entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, 
toQualifiedName());
             }

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/HiveTableEntityHandler.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/HiveTableEntityHandler.java
 
b/intg/src/main/java/org/apache/atlas/entitytransform/HiveTableEntityHandler.java
index b008e6c..9eb44d7 100644
--- 
a/intg/src/main/java/org/apache/atlas/entitytransform/HiveTableEntityHandler.java
+++ 
b/intg/src/main/java/org/apache/atlas/entitytransform/HiveTableEntityHandler.java
@@ -26,11 +26,11 @@ import java.util.List;
 import static org.apache.atlas.entitytransform.TransformationConstants.*;
 
 public class HiveTableEntityHandler extends BaseEntityHandler {
-    private static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HIVE_DB_NAME_ATTRIBUTE, HIVE_TABLE_NAME_ATTRIBUTE, 
HIVE_DB_CLUSTER_NAME_ATTRIBUTE);
+    static final List<String> CUSTOM_TRANSFORM_ATTRIBUTES = 
Arrays.asList(HIVE_DB_NAME_ATTRIBUTE, HIVE_TABLE_NAME_ATTRIBUTE, 
HIVE_DB_CLUSTER_NAME_ATTRIBUTE);
 
 
     public HiveTableEntityHandler(List<AtlasEntityTransformer> transformers) {
-        super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES);
+        super(transformers);
     }
 
     @Override
@@ -46,7 +46,7 @@ public class HiveTableEntityHandler extends BaseEntityHandler 
{
         private String  databaseName;
         private String  tableName;
         private String  clusterName;
-        private boolean isCustomerAttributeUpdated = false;
+        private boolean isCustomAttributeUpdated = false;
 
 
         public HiveTableEntity(AtlasEntity entity) {
@@ -69,8 +69,8 @@ public class HiveTableEntityHandler extends BaseEntityHandler 
{
         }
 
         @Override
-        public Object getAttribute(String attributeName) {
-            switch (attributeName) {
+        public Object getAttribute(EntityAttribute attribute) {
+            switch (attribute.getAttributeKey()) {
                 case HIVE_TABLE_NAME_ATTRIBUTE:
                     return tableName;
 
@@ -81,39 +81,39 @@ public class HiveTableEntityHandler extends 
BaseEntityHandler {
                     return clusterName;
             }
 
-            return super.getAttribute(attributeName);
+            return super.getAttribute(attribute);
         }
 
         @Override
-        public void setAttribute(String attributeName, String attributeValue) {
-            switch (attributeName) {
+        public void setAttribute(EntityAttribute attribute, String 
attributeValue) {
+            switch (attribute.getAttributeKey()) {
                 case HIVE_TABLE_NAME_ATTRIBUTE:
                     tableName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HIVE_DB_NAME_ATTRIBUTE:
                     databaseName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 case HIVE_DB_CLUSTER_NAME_ATTRIBUTE:
                     clusterName = attributeValue;
 
-                    isCustomerAttributeUpdated = true;
+                    isCustomAttributeUpdated = true;
                 break;
 
                 default:
-                    super.setAttribute(attributeName, attributeValue);
+                    super.setAttribute(attribute, attributeValue);
                 break;
             }
         }
 
         @Override
         public void transformComplete() {
-            if (isCustomerAttributeUpdated) {
+            if (isCustomAttributeUpdated) {
                 entity.setAttribute(NAME_ATTRIBUTE, tableName);
                 entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, 
toQualifiedName());
             }

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/NeedsContext.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/NeedsContext.java 
b/intg/src/main/java/org/apache/atlas/entitytransform/NeedsContext.java
deleted file mode 100644
index 5c16bcf..0000000
--- a/intg/src/main/java/org/apache/atlas/entitytransform/NeedsContext.java
+++ /dev/null
@@ -1,23 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.atlas.entitytransform;
-
-public interface NeedsContext {
-    void setContext(TransformerContext transformerContext);
-}

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/main/java/org/apache/atlas/entitytransform/TransformerContext.java
----------------------------------------------------------------------
diff --git 
a/intg/src/main/java/org/apache/atlas/entitytransform/TransformerContext.java 
b/intg/src/main/java/org/apache/atlas/entitytransform/TransformerContext.java
index a7a77b5..4b2ece6 100644
--- 
a/intg/src/main/java/org/apache/atlas/entitytransform/TransformerContext.java
+++ 
b/intg/src/main/java/org/apache/atlas/entitytransform/TransformerContext.java
@@ -23,13 +23,13 @@ import org.apache.atlas.store.AtlasTypeDefStore;
 import org.apache.atlas.type.AtlasTypeRegistry;
 
 public class TransformerContext {
-    private final AtlasTypeRegistry typeRegistry;
-    private final AtlasTypeDefStore typeDefStore;
+    private final AtlasTypeRegistry  typeRegistry;
+    private final AtlasTypeDefStore  typeDefStore;
     private final AtlasExportRequest exportRequest;
 
     public TransformerContext(AtlasTypeRegistry typeRegistry, 
AtlasTypeDefStore typeDefStore, AtlasExportRequest exportRequest) {
-        this.typeRegistry = typeRegistry;
-        this.typeDefStore = typeDefStore;
+        this.typeRegistry  = typeRegistry;
+        this.typeDefStore  = typeDefStore;
         this.exportRequest = exportRequest;
     }
 

http://git-wip-us.apache.org/repos/asf/atlas/blob/c1b42468/intg/src/test/java/org/apache/atlas/entitytransform/TransformationHandlerTest.java
----------------------------------------------------------------------
diff --git 
a/intg/src/test/java/org/apache/atlas/entitytransform/TransformationHandlerTest.java
 
b/intg/src/test/java/org/apache/atlas/entitytransform/TransformationHandlerTest.java
index c76f959..d6b0ede 100644
--- 
a/intg/src/test/java/org/apache/atlas/entitytransform/TransformationHandlerTest.java
+++ 
b/intg/src/test/java/org/apache/atlas/entitytransform/TransformationHandlerTest.java
@@ -17,20 +17,31 @@
  */
 package org.apache.atlas.entitytransform;
 
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.model.impexp.AtlasExportRequest;
 import org.apache.atlas.model.impexp.AttributeTransform;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasObjectId;
-import org.apache.atlas.type.AtlasType;
+import org.apache.atlas.model.typedef.AtlasEntityDef;
+import org.apache.atlas.model.typedef.AtlasTypesDef;
+import org.apache.atlas.store.AtlasTypeDefStore;
+import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.commons.lang.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.testng.annotations.Test;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 import static 
org.apache.atlas.entitytransform.TransformationConstants.HDFS_PATH;
+import static 
org.apache.atlas.entitytransform.TransformationConstants.HIVE_COLUMN;
+import static 
org.apache.atlas.entitytransform.TransformationConstants.HIVE_DATABASE;
+import static 
org.apache.atlas.entitytransform.TransformationConstants.HIVE_STORAGE_DESCRIPTOR;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.assertNull;
@@ -38,6 +49,17 @@ import static org.testng.Assert.assertTrue;
 import static 
org.apache.atlas.entitytransform.TransformationConstants.HIVE_TABLE;
 
 public class TransformationHandlerTest {
+    private static final Logger LOG = 
LoggerFactory.getLogger(TransformationHandlerTest.class);
+
+    private static final String TYPENAME_REFERENCEABLE = "Referenceable";
+    private static final String TYPENAME_ASSET         = "Asset";
+    private static final String TYPENAME_NON_ASSET     = "non_asset";
+
+    private static final String[] CLUSTER_NAMES  = new String[] { "cl1", 
"prod" };
+    private static final String[] DATABASE_NAMES = new String[] { "hr", 
"sales", "engg" };
+    private static final String[] TABLE_NAMES    = new String[] { "employees", 
"products", "invoice" };
+    private static final String[] COLUMN_NAMES   = new String[] { "name", 
"age", "dob" };
+
     @Test
     public void testHdfsClusterRenameHandler() {
         // Rename clusterName from cl1 to cl2
@@ -139,8 +161,8 @@ public class TransformationHandlerTest {
     @Test
     public void testEntityClearAttributesActionWithNoCondition() {
         // clear replicatedFrom attribute for hive_table entities without any 
condition
-        Map<String, String> actions = new HashMap<String, String>() {{  
put("__entity.replicatedTo", "CLEAR:");
-                                                                        
put("__entity.replicatedFrom", "CLEAR:"); }};
+        Map<String, String> actions = new HashMap<String, String>() {{  
put("Referenceable.replicatedTo", "CLEAR:");
+                                                                        
put("Referenceable.replicatedFrom", "CLEAR:"); }};
 
         AttributeTransform transform = new AttributeTransform(null, actions);
 
@@ -334,38 +356,75 @@ public class TransformationHandlerTest {
 
     @Test
     public void verifyAddClassification() {
-        AtlasEntityTransformer entityTransformer = new AtlasEntityTransformer(
-                Collections.singletonMap("hdfs_path.qualifiedName", "EQUALS: 
hr@cl1"),
-                Collections.singletonMap("__entity", "addClassification: 
replicated")
-        );
+        AtlasEntityTransformer transformer = new 
AtlasEntityTransformer(Collections.singletonMap("hive_db.qualifiedName", 
"EQUALS: hr@cl1"),
+                                                                        
Collections.singletonMap("Referenceable.", "ADD_CLASSIFICATION: replicated"),
+                                                                        
getTransformerContext());
+
+        List<BaseEntityHandler> handlers = Collections.singletonList(new 
BaseEntityHandler(Collections.singletonList(transformer)));
 
-        List<BaseEntityHandler> handlers = new ArrayList<>();
-        handlers.add(new 
BaseEntityHandler(Collections.singletonList(entityTransformer)));
         assertApplyTransform(handlers);
     }
 
     @Test
     public void verifyAddClassificationUsingScope() {
-        AtlasObjectId objectId = new AtlasObjectId("hive_db", 
Collections.singletonMap("qualifiedName", "hr@cl1"));
-        AtlasEntityTransformer entityTransformer = new AtlasEntityTransformer(
-                Collections.singletonMap("__entity", "topLevel: "),
-                Collections.singletonMap("__entity", "addClassification: 
replicated")
-        );
-
-        List<BaseEntityHandler> handlers = new ArrayList<>();
-        handlers.add(new 
BaseEntityHandler(Collections.singletonList(entityTransformer)));
-        Condition condition = 
handlers.get(0).transformers.get(0).getConditions().get(0);
-        Condition.ObjectIdEquals objectIdEquals = (Condition.ObjectIdEquals) 
condition;
-        objectIdEquals.add(objectId);
+        AtlasExportRequest exportRequest = new AtlasExportRequest();
+
+        exportRequest.setItemsToExport(Collections.singletonList(new 
AtlasObjectId("hive_db", Collections.singletonMap("qualifiedName", "hr@cl1"))));
+
+        AtlasEntityTransformer transformer = new 
AtlasEntityTransformer(Collections.singletonMap("Referenceable.", "topLevel: "),
+                                                                        
Collections.singletonMap("Referenceable", "ADD_CLASSIFICATION: replicated"),
+                                                                        new 
TransformerContext(getTypeRegistry(), null, exportRequest));
+
+        List<BaseEntityHandler> handlers = Collections.singletonList(new 
BaseEntityHandler(Collections.singletonList(transformer)));
 
         assertApplyTransform(handlers);
     }
 
+    @Test
+    public void verifyEntityTypeInAttributeName() {
+        AttributeTransform p = new AttributeTransform();
+        p.addAction("Asset.name", "SET: renamed");
+
+        List<BaseEntityHandler> handlers = 
initializeHandlers(Collections.singletonList(p));
+
+        AtlasEntity assetEntity    = new AtlasEntity(TYPENAME_ASSET, "name", 
"originalName");
+        AtlasEntity assetSubEntity = new AtlasEntity(HIVE_DATABASE, "name", 
"originalName");
+        AtlasEntity nonAssetEntity = new AtlasEntity(TYPENAME_NON_ASSET, 
"name", "originalName");
+
+        applyTransforms(assetEntity, handlers);
+        applyTransforms(assetSubEntity, handlers);
+        applyTransforms(nonAssetEntity, handlers);
+
+        assertEquals((String) assetEntity.getAttribute("name"), "renamed", 
"Asset.name expected to be updated for Asset entity");
+        assertEquals((String) assetSubEntity.getAttribute("name"), "renamed", 
"Asset.name expected to be updated for Asset sub-type entity");
+        assertEquals((String) nonAssetEntity.getAttribute("name"), 
"originalName", "Asset.name expected to be not updated for non-Asset type 
entity");
+    }
+
+    @Test
+    public void verifyNoEntityTypeInAttributeName() {
+        AttributeTransform p = new AttributeTransform();
+        p.addAction("name", "SET: renamed");
+
+        List<BaseEntityHandler> handlers = 
initializeHandlers(Collections.singletonList(p));
+
+        AtlasEntity assetEntity    = new AtlasEntity(TYPENAME_ASSET, "name", 
"originalName");
+        AtlasEntity assetSubEntity = new AtlasEntity(HIVE_DATABASE, "name", 
"originalName");
+        AtlasEntity nonAssetEntity = new AtlasEntity(TYPENAME_NON_ASSET, 
"name", "originalName");
+
+        applyTransforms(assetEntity, handlers);
+        applyTransforms(assetSubEntity, handlers);
+        applyTransforms(nonAssetEntity, handlers);
+
+        assertEquals((String) assetEntity.getAttribute("name"), "renamed", 
"name expected to be updated for Asset entity");
+        assertEquals((String) assetSubEntity.getAttribute("name"), "renamed", 
"name expected to be updated for Asset sub-type entity");
+        assertEquals((String) nonAssetEntity.getAttribute("name"), "renamed", 
"name expected to be not updated for non-Asset type entity");
+    }
+
     private void assertApplyTransform(List<BaseEntityHandler> handlers) {
         for (AtlasEntity entity : getAllEntities()) {
             applyTransforms(entity, handlers);
 
-            if(entity.getAttribute("qualifiedName").equals("hr@cl1")) {
+            if(entity.getTypeName().equals("hive_db") && 
entity.getAttribute("qualifiedName").equals("hr@cl1")) {
                 assertNotNull(entity.getClassifications());
             } else{
                 assertNull(entity.getClassifications());
@@ -374,7 +433,7 @@ public class TransformationHandlerTest {
     }
 
     private List<BaseEntityHandler> 
initializeHandlers(List<AttributeTransform> params) {
-        return BaseEntityHandler.createEntityHandlers(params, null);
+        return BaseEntityHandler.createEntityHandlers(params, 
getTransformerContext());
     }
 
     private void applyTransforms(AtlasEntity entity, List<BaseEntityHandler> 
handlers) {
@@ -383,15 +442,50 @@ public class TransformationHandlerTest {
         }
     }
 
-    final String[] clusterNames  = new String[] { "cl1", "prod" };
-    final String[] databaseNames = new String[] { "hr", "sales", "engg" };
-    final String[] tableNames    = new String[] { "employees", "products", 
"invoice" };
-    final String[] columnNames   = new String[] { "name", "age", "dob" };
+    private TransformerContext getTransformerContext() {
+        return new TransformerContext(getTypeRegistry(), null, null);
+    }
+
+    private AtlasTypeRegistry getTypeRegistry() {
+        AtlasTypeRegistry ret = new AtlasTypeRegistry();
+
+        AtlasEntityDef defReferenceable = new 
AtlasEntityDef(TYPENAME_REFERENCEABLE);
+        AtlasEntityDef defAsset         = new AtlasEntityDef(TYPENAME_ASSET);
+        AtlasEntityDef defHdfsPath      = new AtlasEntityDef(HDFS_PATH);
+        AtlasEntityDef defHiveDb        = new AtlasEntityDef(HIVE_DATABASE);
+        AtlasEntityDef defHiveTable     = new AtlasEntityDef(HIVE_TABLE);
+        AtlasEntityDef defHiveColumn    = new AtlasEntityDef(HIVE_COLUMN);
+        AtlasEntityDef defHiveStorDesc  = new 
AtlasEntityDef(HIVE_STORAGE_DESCRIPTOR);
+        AtlasEntityDef defNonAsset      = new 
AtlasEntityDef(TYPENAME_NON_ASSET);
+
+        defAsset.addSuperType(TYPENAME_REFERENCEABLE);
+        defHdfsPath.addSuperType(TYPENAME_ASSET);
+        defHiveDb.addSuperType(TYPENAME_ASSET);
+        defHiveTable.addSuperType(TYPENAME_ASSET);
+        defHiveColumn.addSuperType(TYPENAME_ASSET);
+        defNonAsset.addSuperType(TYPENAME_REFERENCEABLE);
+
+        AtlasTypesDef typesDef = new AtlasTypesDef();
+
+        typesDef.setEntityDefs(Arrays.asList(defReferenceable, defAsset, 
defHdfsPath, defHiveDb, defHiveTable, defHiveColumn, defHiveStorDesc, 
defNonAsset));
+
+        try {
+            AtlasTypeRegistry.AtlasTransientTypeRegistry ttr = 
ret.lockTypeRegistryForUpdate();
+
+            ttr.addTypes(typesDef);
+
+            ret.releaseTypeRegistryForUpdate(ttr, true);
+        } catch (AtlasBaseException excp) {
+            LOG.warn("failed to initialize type-registry", excp);
+        }
+
+        return ret;
+    }
 
     private List<AtlasEntity> getHdfsPathEntities() {
         List<AtlasEntity> ret = new ArrayList<>();
 
-        for (String clusterName : clusterNames) {
+        for (String clusterName : CLUSTER_NAMES) {
             ret.add(getHdfsPathEntity1(clusterName));
             ret.add(getHdfsPathEntity2(clusterName));
         }
@@ -402,18 +496,18 @@ public class TransformationHandlerTest {
     private List<AtlasEntity> getAllEntities() {
         List<AtlasEntity> ret = new ArrayList<>();
 
-        for (String clusterName : clusterNames) {
+        for (String clusterName : CLUSTER_NAMES) {
             ret.add(getHdfsPathEntity1(clusterName));
             ret.add(getHdfsPathEntity2(clusterName));
 
-            for (String databaseName : databaseNames) {
+            for (String databaseName : DATABASE_NAMES) {
                 ret.add(getHiveDbEntity(clusterName, databaseName));
 
-                for (String tableName : tableNames) {
+                for (String tableName : TABLE_NAMES) {
                     ret.add(getHiveTableEntity(clusterName, databaseName, 
tableName));
                     ret.add(getHiveStorageDescriptorEntity(clusterName, 
databaseName, tableName));
 
-                    for (String columnName : columnNames) {
+                    for (String columnName : COLUMN_NAMES) {
                         ret.add(getHiveColumnEntity(clusterName, databaseName, 
tableName, columnName));
                     }
                 }
@@ -518,4 +612,8 @@ public class TransformationHandlerTest {
 
         return entity;
     }
+
+    private AtlasEntity getNonAssetEntity() {
+        return new AtlasEntity(TYPENAME_NON_ASSET);
+    }
 }
\ No newline at end of file

Reply via email to