ATLAS-2882: refactored import transformer to set context in constructor; fixed…

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

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() {

intg/src/main/java/org/apache/atlas/entitytransform/EntityAttribute.java

+/**
+ * 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);
+    }
+}

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());

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

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());

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

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

intg/src/main/java/org/apache/atlas/entitytransform/NeedsContext.java

-/**
- * 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);
-}

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