Commit 31c3bea1 by Madhan Neethiraj Committed by Ashutosh Mestry

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

ATLAS-2882: refactored import transformer to set context in constructor; fixed incorrect objId match
parent 9d4f9728
...@@ -22,6 +22,8 @@ import org.apache.atlas.model.instance.AtlasClassification; ...@@ -22,6 +22,8 @@ import org.apache.atlas.model.instance.AtlasClassification;
import org.apache.atlas.model.instance.AtlasEntity; import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.typedef.AtlasClassificationDef; import org.apache.atlas.model.typedef.AtlasClassificationDef;
import org.apache.atlas.model.typedef.AtlasTypesDef; 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.commons.lang.StringUtils;
import org.apache.atlas.entitytransform.BaseEntityHandler.AtlasTransformableEntity; import org.apache.atlas.entitytransform.BaseEntityHandler.AtlasTransformableEntity;
import org.slf4j.Logger; import org.slf4j.Logger;
...@@ -31,35 +33,35 @@ import java.util.ArrayList; ...@@ -31,35 +33,35 @@ import java.util.ArrayList;
import java.util.Collections; import java.util.Collections;
public abstract class Action { public abstract class Action {
private static final Logger LOG = LoggerFactory.getLogger(Action.class); 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_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_SET = "SET";
private static final String ACTION_NAME_REPLACE_PREFIX = "REPLACE_PREFIX"; 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_LOWER = "TO_LOWER";
private static final String ACTION_NAME_TO_UPPER = "TO_UPPER"; private static final String ACTION_NAME_TO_UPPER = "TO_UPPER";
private static final String ACTION_NAME_CLEAR = "CLEAR"; private static final String ACTION_NAME_CLEAR = "CLEAR";
protected final String attributeName; protected final EntityAttribute attribute;
protected Action(String attributeName) { protected Action(EntityAttribute attribute) {
this.attributeName = attributeName; this.attribute = attribute;
} }
public String getAttributeName() { return attributeName; } public EntityAttribute getAttribute() { return attribute; }
public boolean isValid() { public boolean isValid() {
return StringUtils.isNotEmpty(attributeName); return true;
} }
public abstract void apply(AtlasTransformableEntity entity); 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()) { if (LOG.isDebugEnabled()) {
LOG.debug("==> Action.createAction(key={}, value={})", key, value); LOG.debug("==> Action.createAction(key={}, value={})", key, value);
} }
...@@ -74,33 +76,35 @@ public abstract class Action { ...@@ -74,33 +76,35 @@ public abstract class Action {
actionValue = StringUtils.trim(actionValue); actionValue = StringUtils.trim(actionValue);
value = StringUtils.trim(value); value = StringUtils.trim(value);
EntityAttribute attribute = new EntityAttribute(StringUtils.trim(key), context);
switch (actionName.toUpperCase()) { switch (actionName.toUpperCase()) {
case ACTION_ADD_CLASSIFICATION: case ACTION_ADD_CLASSIFICATION:
ret = new AddClassificationAction(actionValue); ret = new AddClassificationAction(attribute, actionValue, context);
break; break;
case ACTION_NAME_REPLACE_PREFIX: case ACTION_NAME_REPLACE_PREFIX:
ret = new PrefixReplaceAction(key, actionValue); ret = new PrefixReplaceAction(attribute, actionValue);
break; break;
case ACTION_NAME_TO_LOWER: case ACTION_NAME_TO_LOWER:
ret = new ToLowerCaseAction(key); ret = new ToLowerCaseAction(attribute);
break; break;
case ACTION_NAME_TO_UPPER: case ACTION_NAME_TO_UPPER:
ret = new ToUpperCaseAction(key); ret = new ToUpperCaseAction(attribute);
break; break;
case ACTION_NAME_SET: case ACTION_NAME_SET:
ret = new SetAction(key, actionValue); ret = new SetAction(attribute, actionValue);
break; break;
case ACTION_NAME_CLEAR: case ACTION_NAME_CLEAR:
ret = new ClearAction(key); ret = new ClearAction(attribute);
break; break;
default: default:
ret = new SetAction(key, value); // treat unspecified/unknown action as 'SET' ret = new SetAction(attribute, value); // treat unspecified/unknown action as 'SET'
break; break;
} }
...@@ -115,71 +119,79 @@ public abstract class Action { ...@@ -115,71 +119,79 @@ public abstract class Action {
public static class SetAction extends Action { public static class SetAction extends Action {
private final String attributeValue; private final String attributeValue;
public SetAction(String attributeName, String attributeValue) { public SetAction(EntityAttribute attribute, String attributeValue) {
super(attributeName); super(attribute);
this.attributeValue = attributeValue; this.attributeValue = attributeValue;
} }
@Override @Override
public void apply(AtlasTransformableEntity entity) { public void apply(AtlasTransformableEntity entity) {
if (isValid()) { entity.setAttribute(attribute, attributeValue);
entity.setAttribute(attributeName, attributeValue);
}
} }
} }
public static class AddClassificationAction extends Action implements NeedsContext { public static class AddClassificationAction extends Action {
private final String classificationName; private final String classificationName;
private TransformerContext transformerContext;
public AddClassificationAction(String classificationName) { public AddClassificationAction(EntityAttribute attribute, String classificationName, TransformerContext context) {
super(ENTITY_KEY); super(attribute);
this.classificationName = classificationName; this.classificationName = classificationName;
createClassificationDefIfNotExists(classificationName, context);
} }
@Override @Override
public void apply(AtlasTransformableEntity transformableEntity) { public void apply(AtlasTransformableEntity transformableEntity) {
AtlasEntity entity = transformableEntity.entity; AtlasEntity entity = transformableEntity.getEntity();
if (entity.getClassifications() == null) { if (entity.getClassifications() == null) {
entity.setClassifications(new ArrayList<AtlasClassification>()); entity.setClassifications(new ArrayList<>());
} }
boolean hasClassification = false;
for (AtlasClassification c : entity.getClassifications()) { for (AtlasClassification c : entity.getClassifications()) {
if (c.getTypeName().equals(classificationName)) { hasClassification = c.getTypeName().equals(classificationName);
return;
if (hasClassification) {
break;
} }
} }
entity.getClassifications().add(new AtlasClassification(classificationName)); if (!hasClassification) {
entity.getClassifications().add(new AtlasClassification(classificationName));
}
} }
@Override private void createClassificationDefIfNotExists(String classificationName, TransformerContext context) {
public void setContext(TransformerContext transformerContext) { AtlasTypeRegistry typeRegistry = context != null ? context.getTypeRegistry() : null;
this.transformerContext = transformerContext;
getCreateTag(classificationName);
}
private void getCreateTag(String classificationName) { if (typeRegistry != null) {
if (transformerContext == null) { try {
return; AtlasClassificationDef classificationDef = typeRegistry.getClassificationDefByName(classificationName);
}
try { if (classificationDef == null) {
AtlasClassificationDef classificationDef = transformerContext.getTypeRegistry().getClassificationDefByName(classificationName); AtlasTypeDefStore typeDefStore = context.getTypeDefStore();
if (classificationDef != null) {
return; if (typeDefStore != null) {
} classificationDef = new AtlasClassificationDef(classificationName);
classificationDef = new AtlasClassificationDef(classificationName); AtlasTypesDef typesDef = new AtlasTypesDef();
AtlasTypesDef typesDef = new AtlasTypesDef();
typesDef.setClassificationDefs(Collections.singletonList(classificationDef)); typesDef.setClassificationDefs(Collections.singletonList(classificationDef));
transformerContext.getTypeDefStore().createTypesDef(typesDef);
LOG.info("created classification: {}", classificationName); typeDefStore.createTypesDef(typesDef);
} catch (AtlasBaseException e) {
LOG.error("Error creating classification: {}", classificationName, e); 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 { ...@@ -188,8 +200,8 @@ public abstract class Action {
private final String fromPrefix; private final String fromPrefix;
private final String toPrefix; private final String toPrefix;
public PrefixReplaceAction(String attributeName, String actionValue) { public PrefixReplaceAction(EntityAttribute attribute, String actionValue) {
super(attributeName); super(attribute);
// actionValue => =:prefixToReplace=replacedValue // actionValue => =:prefixToReplace=replacedValue
if (actionValue != null) { if (actionValue != null) {
...@@ -224,61 +236,61 @@ public abstract class Action { ...@@ -224,61 +236,61 @@ public abstract class Action {
@Override @Override
public void apply(AtlasTransformableEntity entity) { public void apply(AtlasTransformableEntity entity) {
if (isValid()) { if (isValid()) {
Object currValue = entity.getAttribute(attributeName); Object currValue = entity.getAttribute(attribute);
String strValue = currValue != null ? currValue.toString() : null; String strValue = currValue != null ? currValue.toString() : null;
if (strValue != null && strValue.startsWith(fromPrefix)) { 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 static class ToLowerCaseAction extends Action {
public ToLowerCaseAction(String attributeName) { public ToLowerCaseAction(EntityAttribute attribute) {
super(attributeName); super(attribute);
} }
@Override @Override
public void apply(AtlasTransformableEntity entity) { public void apply(AtlasTransformableEntity entity) {
if (isValid()) { if (isValid()) {
Object currValue = entity.getAttribute(attributeName); Object currValue = entity.getAttribute(attribute);
String strValue = currValue instanceof String ? (String) currValue : null; String strValue = currValue instanceof String ? (String) currValue : null;
if (strValue != null) { if (strValue != null) {
entity.setAttribute(attributeName, strValue.toLowerCase()); entity.setAttribute(attribute, strValue.toLowerCase());
} }
} }
} }
} }
public static class ToUpperCaseAction extends Action { public static class ToUpperCaseAction extends Action {
public ToUpperCaseAction(String attributeName) { public ToUpperCaseAction(EntityAttribute attribute) {
super(attributeName); super(attribute);
} }
@Override @Override
public void apply(AtlasTransformableEntity entity) { public void apply(AtlasTransformableEntity entity) {
if (isValid()) { if (isValid()) {
Object currValue = entity.getAttribute(attributeName); Object currValue = entity.getAttribute(attribute);
String strValue = currValue instanceof String ? (String) currValue : null; String strValue = currValue instanceof String ? (String) currValue : null;
if (strValue != null) { if (strValue != null) {
entity.setAttribute(attributeName, strValue.toUpperCase()); entity.setAttribute(attribute, strValue.toUpperCase());
} }
} }
} }
} }
public static class ClearAction extends Action { public static class ClearAction extends Action {
public ClearAction(String attributeName) { public ClearAction(EntityAttribute attribute) {
super(attributeName); super(attribute);
} }
@Override @Override
public void apply(AtlasTransformableEntity entity) { public void apply(AtlasTransformableEntity entity) {
if (isValid() && entity.hasAttribute(attributeName)) { if (isValid() && entity.hasAttribute(attribute)) {
entity.setAttribute(attributeName, null); entity.setAttribute(attribute, null);
} }
} }
} }
......
...@@ -19,32 +19,25 @@ package org.apache.atlas.entitytransform; ...@@ -19,32 +19,25 @@ package org.apache.atlas.entitytransform;
import org.apache.atlas.entitytransform.BaseEntityHandler.AtlasTransformableEntity; import org.apache.atlas.entitytransform.BaseEntityHandler.AtlasTransformableEntity;
import org.apache.atlas.model.impexp.AttributeTransform; 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.collections.MapUtils;
import org.apache.commons.lang.StringUtils;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.Collections;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
public class AtlasEntityTransformer { public class AtlasEntityTransformer {
private final List<Condition> conditions; private final List<Condition> conditions;
private final List<Action> actions; private final List<Action> actions;
public AtlasEntityTransformer(AttributeTransform attributeTransform) {
this(attributeTransform.getConditions(), attributeTransform.getAction());
}
public AtlasEntityTransformer(AtlasObjectId objectId, Map<String, String> actions) { public AtlasEntityTransformer(AttributeTransform attributeTransform, TransformerContext context) {
this(Collections.singletonMap("__entity", AtlasType.toJson(objectId)), actions); this(attributeTransform.getConditions(), attributeTransform.getAction(), context);
} }
public AtlasEntityTransformer(Map<String, String> conditions, Map<String, String> actions) { public AtlasEntityTransformer(Map<String, String> conditions, Map<String, String> actions, TransformerContext context) {
this.conditions = createConditions(conditions); this.conditions = createConditions(conditions, context);
this.actions = createActions(actions); this.actions = createActions(actions, context);
} }
public List<Condition> getConditions() { public List<Condition> getConditions() {
...@@ -71,12 +64,12 @@ public class AtlasEntityTransformer { ...@@ -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<>(); List<Condition> ret = new ArrayList<>();
if (MapUtils.isNotEmpty(conditions)) { if (MapUtils.isNotEmpty(conditions)) {
for (Map.Entry<String, String> entry : conditions.entrySet()) { 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); ret.add(condition);
} }
...@@ -85,12 +78,12 @@ public class AtlasEntityTransformer { ...@@ -85,12 +78,12 @@ public class AtlasEntityTransformer {
return ret; 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<>(); List<Action> ret = new ArrayList<>();
if (MapUtils.isNotEmpty(actions)) { if (MapUtils.isNotEmpty(actions)) {
for (Map.Entry<String, String> entry : actions.entrySet()) { 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); ret.add(action);
} }
......
...@@ -17,12 +17,9 @@ ...@@ -17,12 +17,9 @@
*/ */
package org.apache.atlas.entitytransform; package org.apache.atlas.entitytransform;
import org.apache.atlas.model.impexp.AtlasExportRequest;
import org.apache.atlas.model.impexp.AttributeTransform; import org.apache.atlas.model.impexp.AttributeTransform;
import org.apache.atlas.model.instance.AtlasEntity; import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.store.AtlasTypeDefStore;
import org.apache.atlas.type.AtlasType; import org.apache.atlas.type.AtlasType;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.commons.collections.CollectionUtils; import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang.StringUtils; import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger; import org.slf4j.Logger;
...@@ -31,34 +28,24 @@ import org.slf4j.LoggerFactory; ...@@ -31,34 +28,24 @@ import org.slf4j.LoggerFactory;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import static org.apache.atlas.entitytransform.TransformationConstants.TYPE_NAME_ATTRIBUTE_NAME_SEP;
public class BaseEntityHandler { public class BaseEntityHandler {
private static final Logger LOG = LoggerFactory.getLogger(BaseEntityHandler.class); private static final Logger LOG = LoggerFactory.getLogger(BaseEntityHandler.class);
protected final List<AtlasEntityTransformer> transformers; 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) { public BaseEntityHandler(List<AtlasEntityTransformer> transformers) {
this.transformers = transformers; this.transformers = transformers;
this.hasCustomAttributeTransformer = hasTransformerForAnyAttribute(customTransformAttributes);
}
public boolean hasCustomAttributeTransformer() {
return hasCustomAttributeTransformer;
} }
public AtlasEntity transform(AtlasEntity entity) { public AtlasEntity transform(AtlasEntity entity) {
if (!CollectionUtils.isNotEmpty(transformers)) { if (CollectionUtils.isEmpty(transformers)) {
return entity; return entity;
} }
AtlasTransformableEntity transformableEntity = getTransformableEntity(entity); AtlasTransformableEntity transformableEntity = getTransformableEntity(entity);
if (transformableEntity == null) { if (transformableEntity == null) {
return entity; return entity;
} }
...@@ -72,22 +59,6 @@ public class BaseEntityHandler { ...@@ -72,22 +59,6 @@ public class BaseEntityHandler {
return entity; 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) { public AtlasTransformableEntity getTransformableEntity(AtlasEntity entity) {
return new AtlasTransformableEntity(entity); return new AtlasTransformableEntity(entity);
} }
...@@ -97,39 +68,38 @@ public class BaseEntityHandler { ...@@ -97,39 +68,38 @@ public class BaseEntityHandler {
LOG.debug("==> BaseEntityHandler.createEntityHandlers(transforms={})", transforms); LOG.debug("==> BaseEntityHandler.createEntityHandlers(transforms={})", transforms);
} }
List<AtlasEntityTransformer> transformers = new ArrayList<>(); List<BaseEntityHandler> ret = new ArrayList<>();
for (AttributeTransform transform : transforms) { if (CollectionUtils.isNotEmpty(transforms)) {
transformers.add(new AtlasEntityTransformer(transform)); List<AtlasEntityTransformer> transformers = new ArrayList<>();
}
BaseEntityHandler[] handlers = new BaseEntityHandler[] { for (AttributeTransform transform : transforms) {
new HdfsPathEntityHandler(transformers), transformers.add(new AtlasEntityTransformer(transform, context));
new HiveDatabaseEntityHandler(transformers), }
new HiveTableEntityHandler(transformers),
new HiveColumnEntityHandler(transformers),
new HiveStorageDescriptorEntityHandler(transformers)
};
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 if (hasTransformerForAnyAttribute(transformers, HiveDatabaseEntityHandler.CUSTOM_TRANSFORM_ATTRIBUTES)) {
for (BaseEntityHandler handler : handlers) { ret.add(new HiveDatabaseEntityHandler(transformers));
if (handler.hasCustomAttributeTransformer()) {
ret.add(handler);
handler.setContext(context);
} }
}
if (CollectionUtils.isEmpty(ret)) { if (hasTransformerForAnyAttribute(transformers, HiveTableEntityHandler.CUSTOM_TRANSFORM_ATTRIBUTES)) {
BaseEntityHandler be = new BaseEntityHandler(transformers); ret.add(new HiveTableEntityHandler(transformers));
be.setContext(context); }
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)) { if (CollectionUtils.isEmpty(ret)) {
ret.add(new BaseEntityHandler(transformers)); ret.add(new BaseEntityHandler(transformers));
}
} }
if (LOG.isDebugEnabled()) { if (LOG.isDebugEnabled()) {
...@@ -139,11 +109,11 @@ public class BaseEntityHandler { ...@@ -139,11 +109,11 @@ public class BaseEntityHandler {
return ret; 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)) { if (CollectionUtils.isNotEmpty(transformers) && CollectionUtils.isNotEmpty(attributes)) {
for (AtlasEntityTransformer transformer : transformers) { for (AtlasEntityTransformer transformer : transformers) {
for (Action action : transformer.getActions()) { for (Action action : transformer.getActions()) {
if (attributes.contains(action.getAttributeName())) { if (attributes.contains(action.getAttribute().getAttributeKey())) {
return true; return true;
} }
} }
...@@ -152,20 +122,6 @@ public class BaseEntityHandler { ...@@ -152,20 +122,6 @@ public class BaseEntityHandler {
return false; 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 { public static class AtlasTransformableEntity {
protected final AtlasEntity entity; protected final AtlasEntity entity;
...@@ -178,38 +134,26 @@ public class BaseEntityHandler { ...@@ -178,38 +134,26 @@ public class BaseEntityHandler {
return entity; return entity;
} }
public Object getAttribute(String attributeName) { public Object getAttribute(EntityAttribute attribute) {
Object ret = null; final Object ret;
if (entity != null && attributeName != null) {
ret = entity.getAttribute(attributeName);
if (ret == null) { // try after dropping typeName prefix, if attributeName contains it if (attribute.appliesToEntityType(entity.getTypeName())) {
int idxSep = attributeName.indexOf(TYPE_NAME_ATTRIBUTE_NAME_SEP); ret = entity.getAttribute(attribute.getAttributeName());
} else {
if (idxSep != -1) { ret = null;
ret = entity.getAttribute(attributeName.substring(idxSep + 1));
}
}
} }
return ret; return ret;
} }
public void setAttribute(String attributeName, String attributeValue) { public void setAttribute(EntityAttribute attribute, String attributeValue) {
if (entity != null && attributeName != null) { if (attribute.appliesToEntityType(entity.getTypeName())) {
int idxSep = attributeName.indexOf(TYPE_NAME_ATTRIBUTE_NAME_SEP); // drop typeName prefix, if attributeName contains it entity.setAttribute(attribute.getAttributeName(), attributeValue);
if (idxSep != -1) {
entity.setAttribute(attributeName.substring(idxSep + 1), attributeValue);
} else {
entity.setAttribute(attributeName, attributeValue);
}
} }
} }
public boolean hasAttribute(String attributeName) { public boolean hasAttribute(EntityAttribute attribute) {
return getAttribute(attributeName) != null; return getAttribute(attribute) != null;
} }
public void transformComplete() { public void transformComplete() {
......
...@@ -17,7 +17,9 @@ ...@@ -17,7 +17,9 @@
*/ */
package org.apache.atlas.entitytransform; package org.apache.atlas.entitytransform;
import com.google.common.annotations.VisibleForTesting;
import org.apache.atlas.entitytransform.BaseEntityHandler.AtlasTransformableEntity; 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.AtlasEntity;
import org.apache.atlas.model.instance.AtlasObjectId; import org.apache.atlas.model.instance.AtlasObjectId;
import org.apache.commons.lang.StringUtils; import org.apache.commons.lang.StringUtils;
...@@ -30,6 +32,7 @@ import java.util.Map; ...@@ -30,6 +32,7 @@ import java.util.Map;
import java.util.Objects; import java.util.Objects;
public abstract class Condition { public abstract class Condition {
private static final Logger LOG = LoggerFactory.getLogger(Condition.class); private static final Logger LOG = LoggerFactory.getLogger(Condition.class);
...@@ -43,18 +46,19 @@ public abstract class Condition { ...@@ -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_STARTS_WITH_IGNORE_CASE = "STARTS_WITH_IGNORE_CASE";
private static final String CONDITION_NAME_HAS_VALUE = "HAS_VALUE"; private static final String CONDITION_NAME_HAS_VALUE = "HAS_VALUE";
protected final String attributeName; protected final EntityAttribute attribute;
protected Condition(String attributeName) { protected Condition(EntityAttribute attribute) {
this.attributeName = attributeName; this.attribute = attribute;
} }
public String getAttributeName() { return attributeName; } public EntityAttribute getAttribute() { return attribute; }
public abstract boolean matches(AtlasTransformableEntity entity); 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()) { if (LOG.isDebugEnabled()) {
LOG.debug("==> Condition.createCondition(key={}, value={})", key, value); LOG.debug("==> Condition.createCondition(key={}, value={})", key, value);
} }
...@@ -69,41 +73,43 @@ public abstract class Condition { ...@@ -69,41 +73,43 @@ public abstract class Condition {
conditionValue = StringUtils.trim(conditionValue); conditionValue = StringUtils.trim(conditionValue);
value = StringUtils.trim(value); value = StringUtils.trim(value);
EntityAttribute attribute = new EntityAttribute(StringUtils.trim(key), context);
switch (conditionName.toUpperCase()) { switch (conditionName.toUpperCase()) {
case CONDITION_ENTITY_ALL: case CONDITION_ENTITY_ALL:
ret = new ObjectIdEquals(key, CONDITION_ENTITY_ALL); ret = new ObjectIdEquals(attribute, CONDITION_ENTITY_ALL, context);
break; break;
case CONDITION_ENTITY_TOP_LEVEL: case CONDITION_ENTITY_TOP_LEVEL:
ret = new ObjectIdEquals(key, CONDITION_ENTITY_TOP_LEVEL); ret = new ObjectIdEquals(attribute, CONDITION_ENTITY_TOP_LEVEL, context);
break; break;
case CONDITION_ENTITY_OBJECT_ID: case CONDITION_ENTITY_OBJECT_ID:
ret = new ObjectIdEquals(key, conditionValue); ret = new ObjectIdEquals(attribute, conditionValue, context);
break; break;
case CONDITION_NAME_EQUALS: case CONDITION_NAME_EQUALS:
ret = new EqualsCondition(key, conditionValue); ret = new EqualsCondition(attribute, conditionValue);
break; break;
case CONDITION_NAME_EQUALS_IGNORE_CASE: case CONDITION_NAME_EQUALS_IGNORE_CASE:
ret = new EqualsIgnoreCaseCondition(key, conditionValue); ret = new EqualsIgnoreCaseCondition(attribute, conditionValue);
break; break;
case CONDITION_NAME_STARTS_WITH: case CONDITION_NAME_STARTS_WITH:
ret = new StartsWithCondition(key, conditionValue); ret = new StartsWithCondition(attribute, conditionValue);
break; break;
case CONDITION_NAME_STARTS_WITH_IGNORE_CASE: case CONDITION_NAME_STARTS_WITH_IGNORE_CASE:
ret = new StartsWithIgnoreCaseCondition(key, conditionValue); ret = new StartsWithIgnoreCaseCondition(attribute, conditionValue);
break; break;
case CONDITION_NAME_HAS_VALUE: case CONDITION_NAME_HAS_VALUE:
ret = new HasValueCondition(key, conditionValue); ret = new HasValueCondition(attribute);
break; break;
default: default:
ret = new EqualsCondition(key, value); // treat unspecified/unknown condition as 'EQUALS' ret = new EqualsCondition(attribute, value); // treat unspecified/unknown condition as 'EQUALS'
break; break;
} }
...@@ -118,15 +124,15 @@ public abstract class Condition { ...@@ -118,15 +124,15 @@ public abstract class Condition {
public static class EqualsCondition extends Condition { public static class EqualsCondition extends Condition {
protected final String attributeValue; protected final String attributeValue;
public EqualsCondition(String attributeName, String attributeValue) { public EqualsCondition(EntityAttribute attribute, String attributeValue) {
super(attributeName); super(attribute);
this.attributeValue = attributeValue; this.attributeValue = attributeValue;
} }
@Override @Override
public boolean matches(AtlasTransformableEntity entity) { 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); return attributeValue != null && StringUtils.equals(attributeValue.toString(), this.attributeValue);
} }
...@@ -136,15 +142,15 @@ public abstract class Condition { ...@@ -136,15 +142,15 @@ public abstract class Condition {
public static class EqualsIgnoreCaseCondition extends Condition { public static class EqualsIgnoreCaseCondition extends Condition {
protected final String attributeValue; protected final String attributeValue;
public EqualsIgnoreCaseCondition(String attributeName, String attributeValue) { public EqualsIgnoreCaseCondition(EntityAttribute attribute, String attributeValue) {
super(attributeName); super(attribute);
this.attributeValue = attributeValue; this.attributeValue = attributeValue;
} }
@Override @Override
public boolean matches(AtlasTransformableEntity entity) { 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); return attributeValue != null && StringUtils.equalsIgnoreCase(attributeValue.toString(), this.attributeValue);
} }
...@@ -154,15 +160,15 @@ public abstract class Condition { ...@@ -154,15 +160,15 @@ public abstract class Condition {
public static class StartsWithCondition extends Condition { public static class StartsWithCondition extends Condition {
protected final String prefix; protected final String prefix;
public StartsWithCondition(String attributeName, String prefix) { public StartsWithCondition(EntityAttribute attribute, String prefix) {
super(attributeName); super(attribute);
this.prefix = prefix; this.prefix = prefix;
} }
@Override @Override
public boolean matches(AtlasTransformableEntity entity) { 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); return attributeValue != null && StringUtils.startsWith(attributeValue.toString(), this.prefix);
} }
...@@ -172,96 +178,89 @@ public abstract class Condition { ...@@ -172,96 +178,89 @@ public abstract class Condition {
public static class StartsWithIgnoreCaseCondition extends Condition { public static class StartsWithIgnoreCaseCondition extends Condition {
protected final String prefix; protected final String prefix;
public StartsWithIgnoreCaseCondition(String attributeName, String prefix) { public StartsWithIgnoreCaseCondition(EntityAttribute attribute, String prefix) {
super(attributeName); super(attribute);
this.prefix = prefix; this.prefix = prefix;
} }
@Override @Override
public boolean matches(AtlasTransformableEntity entity) { 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); 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 final List<AtlasObjectId> objectIds;
private String scope;
private TransformerContext transformerContext;
public ObjectIdEquals(String key, String conditionValue) { public ObjectIdEquals(EntityAttribute attribute, String scope, TransformerContext context) {
super(key); super(attribute);
objectIds = new ArrayList<>(); this.isMatchAll = StringUtils.equals(scope, CONDITION_ENTITY_ALL);
this.scope = conditionValue; this.objectIds = new ArrayList<>();
if (!isMatchAll && context != null && context.getExportRequest() != null) {
AtlasExportRequest request = context.getExportRequest();
for(AtlasObjectId objectId : request.getItemsToExport()) {
addObjectId(objectId);
}
}
} }
@Override @Override
public boolean matches(AtlasTransformableEntity entity) { public boolean matches(AtlasTransformableEntity entity) {
for (AtlasObjectId objectId : objectIds) { if (isMatchAll) {
return isMatch(objectId, entity.entity); 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) { @VisibleForTesting
this.objectIds.add(objectId); void addObjectId(AtlasObjectId objId) {
this.objectIds.add(objId);
} }
private boolean isMatch(AtlasObjectId objectId, AtlasEntity entity) { private boolean isMatch(AtlasObjectId objectId, AtlasEntity entity) {
boolean ret = true;
if (!StringUtils.isEmpty(objectId.getGuid())) { if (!StringUtils.isEmpty(objectId.getGuid())) {
return Objects.equals(objectId.getGuid(), entity.getGuid()); return Objects.equals(objectId.getGuid(), entity.getGuid());
} }
ret = Objects.equals(objectId.getTypeName(), entity.getTypeName()); boolean ret = Objects.equals(objectId.getTypeName(), entity.getTypeName());
if (!ret) {
return ret; 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()) { if (!ret) {
ret = ret && Objects.equals(entity.getAttribute(entry.getKey()), entry.getValue()); break;
if (!ret) { }
break;
} }
} }
return ret; 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 { public static class HasValueCondition extends Condition {
protected final String attributeValue; public HasValueCondition(EntityAttribute attribute) {
super(attribute);
public HasValueCondition(String attributeName, String attributeValue) {
super(attributeName);
this.attributeValue = attributeValue;
} }
@Override @Override
public boolean matches(AtlasTransformableEntity entity) { 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; return attributeValue != null ? StringUtils.isNotEmpty(attributeValue.toString()) : false;
} }
......
/**
* 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);
}
}
...@@ -35,10 +35,10 @@ import static org.apache.atlas.entitytransform.TransformationConstants.QUALIFIED ...@@ -35,10 +35,10 @@ import static org.apache.atlas.entitytransform.TransformationConstants.QUALIFIED
public class HdfsPathEntityHandler extends BaseEntityHandler { 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) { public HdfsPathEntityHandler(List<AtlasEntityTransformer> transformers) {
super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES); super(transformers);
} }
@Override @Override
...@@ -56,8 +56,8 @@ public class HdfsPathEntityHandler extends BaseEntityHandler { ...@@ -56,8 +56,8 @@ public class HdfsPathEntityHandler extends BaseEntityHandler {
private String path; private String path;
private String name; private String name;
private String pathPrefix; private String pathPrefix;
private boolean isPathUpdated = false; private boolean isPathUpdated = false;
private boolean isCustomerAttributeUpdated = false; private boolean isCustomAttributeUpdated = false;
public HdfsPathEntity(AtlasEntity entity) { public HdfsPathEntity(AtlasEntity entity) {
...@@ -93,8 +93,8 @@ public class HdfsPathEntityHandler extends BaseEntityHandler { ...@@ -93,8 +93,8 @@ public class HdfsPathEntityHandler extends BaseEntityHandler {
} }
@Override @Override
public Object getAttribute(String attributeName) { public Object getAttribute(EntityAttribute attribute) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HDFS_CLUSTER_NAME_ATTRIBUTE: case HDFS_CLUSTER_NAME_ATTRIBUTE:
return clusterName; return clusterName;
...@@ -105,40 +105,40 @@ public class HdfsPathEntityHandler extends BaseEntityHandler { ...@@ -105,40 +105,40 @@ public class HdfsPathEntityHandler extends BaseEntityHandler {
return path; return path;
} }
return super.getAttribute(attributeName); return super.getAttribute(attribute);
} }
@Override @Override
public void setAttribute(String attributeName, String attributeValue) { public void setAttribute(EntityAttribute attribute, String attributeValue) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HDFS_CLUSTER_NAME_ATTRIBUTE: case HDFS_CLUSTER_NAME_ATTRIBUTE:
clusterName = attributeValue; clusterName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HDFS_PATH_NAME_ATTRIBUTE: case HDFS_PATH_NAME_ATTRIBUTE:
name = attributeValue; name = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HDFS_PATH_PATH_ATTRIBUTE: case HDFS_PATH_PATH_ATTRIBUTE:
path = attributeValue; path = attributeValue;
isPathUpdated = true; isPathUpdated = true;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
default: default:
super.setAttribute(attributeName, attributeValue); super.setAttribute(attribute, attributeValue);
break; break;
} }
} }
@Override @Override
public void transformComplete() { public void transformComplete() {
if (isCustomerAttributeUpdated) { if (isCustomAttributeUpdated) {
entity.setAttribute(CLUSTER_NAME_ATTRIBUTE, clusterName); entity.setAttribute(CLUSTER_NAME_ATTRIBUTE, clusterName);
entity.setAttribute(NAME_ATTRIBUTE, name); entity.setAttribute(NAME_ATTRIBUTE, name);
entity.setAttribute(PATH_ATTRIBUTE, toPath()); entity.setAttribute(PATH_ATTRIBUTE, toPath());
......
...@@ -27,10 +27,10 @@ import static org.apache.atlas.entitytransform.TransformationConstants.*; ...@@ -27,10 +27,10 @@ import static org.apache.atlas.entitytransform.TransformationConstants.*;
public class HiveColumnEntityHandler extends BaseEntityHandler { 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) { public HiveColumnEntityHandler(List<AtlasEntityTransformer> transformers) {
super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES); super(transformers);
} }
@Override @Override
...@@ -48,7 +48,7 @@ public class HiveColumnEntityHandler extends BaseEntityHandler { ...@@ -48,7 +48,7 @@ public class HiveColumnEntityHandler extends BaseEntityHandler {
private String tableName; private String tableName;
private String columnName; private String columnName;
private String clusterName; private String clusterName;
private boolean isCustomerAttributeUpdated = false; private boolean isCustomAttributeUpdated = false;
public HiveColumnEntity(AtlasEntity entity) { public HiveColumnEntity(AtlasEntity entity) {
super(entity); super(entity);
...@@ -73,8 +73,8 @@ public class HiveColumnEntityHandler extends BaseEntityHandler { ...@@ -73,8 +73,8 @@ public class HiveColumnEntityHandler extends BaseEntityHandler {
} }
@Override @Override
public Object getAttribute(String attributeName) { public Object getAttribute(EntityAttribute attribute) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HIVE_DB_NAME_ATTRIBUTE: case HIVE_DB_NAME_ATTRIBUTE:
return databaseName; return databaseName;
...@@ -88,45 +88,45 @@ public class HiveColumnEntityHandler extends BaseEntityHandler { ...@@ -88,45 +88,45 @@ public class HiveColumnEntityHandler extends BaseEntityHandler {
return clusterName; return clusterName;
} }
return super.getAttribute(attributeName); return super.getAttribute(attribute);
} }
@Override @Override
public void setAttribute(String attributeName, String attributeValue) { public void setAttribute(EntityAttribute attribute, String attributeValue) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HIVE_DB_NAME_ATTRIBUTE: case HIVE_DB_NAME_ATTRIBUTE:
databaseName = attributeValue; databaseName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HIVE_TABLE_NAME_ATTRIBUTE: case HIVE_TABLE_NAME_ATTRIBUTE:
tableName = attributeValue; tableName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HIVE_COLUMN_NAME_ATTRIBUTE: case HIVE_COLUMN_NAME_ATTRIBUTE:
columnName = attributeValue; columnName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HIVE_DB_CLUSTER_NAME_ATTRIBUTE: case HIVE_DB_CLUSTER_NAME_ATTRIBUTE:
clusterName = attributeValue; clusterName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
default: default:
super.setAttribute(attributeName, attributeValue); super.setAttribute(attribute, attributeValue);
break; break;
} }
} }
@Override @Override
public void transformComplete() { public void transformComplete() {
if (isCustomerAttributeUpdated) { if (isCustomAttributeUpdated) {
entity.setAttribute(NAME_ATTRIBUTE, columnName); entity.setAttribute(NAME_ATTRIBUTE, columnName);
entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, toQualifiedName()); entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, toQualifiedName());
} }
......
...@@ -26,10 +26,10 @@ import java.util.List; ...@@ -26,10 +26,10 @@ import java.util.List;
import static org.apache.atlas.entitytransform.TransformationConstants.*; import static org.apache.atlas.entitytransform.TransformationConstants.*;
public class HiveDatabaseEntityHandler extends BaseEntityHandler { 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) { public HiveDatabaseEntityHandler(List<AtlasEntityTransformer> transformers) {
super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES); super(transformers);
} }
@Override @Override
...@@ -45,7 +45,7 @@ public class HiveDatabaseEntityHandler extends BaseEntityHandler { ...@@ -45,7 +45,7 @@ public class HiveDatabaseEntityHandler extends BaseEntityHandler {
private static class HiveDatabaseEntity extends AtlasTransformableEntity { private static class HiveDatabaseEntity extends AtlasTransformableEntity {
private String databaseName; private String databaseName;
private String clusterName; private String clusterName;
private boolean isCustomerAttributeUpdated = false; private boolean isCustomAttributeUpdated = false;
public HiveDatabaseEntity(AtlasEntity entity) { public HiveDatabaseEntity(AtlasEntity entity) {
super(entity); super(entity);
...@@ -64,8 +64,8 @@ public class HiveDatabaseEntityHandler extends BaseEntityHandler { ...@@ -64,8 +64,8 @@ public class HiveDatabaseEntityHandler extends BaseEntityHandler {
} }
@Override @Override
public Object getAttribute(String attributeName) { public Object getAttribute(EntityAttribute attribute) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HIVE_DB_NAME_ATTRIBUTE: case HIVE_DB_NAME_ATTRIBUTE:
return databaseName; return databaseName;
...@@ -73,33 +73,33 @@ public class HiveDatabaseEntityHandler extends BaseEntityHandler { ...@@ -73,33 +73,33 @@ public class HiveDatabaseEntityHandler extends BaseEntityHandler {
return clusterName; return clusterName;
} }
return super.getAttribute(attributeName); return super.getAttribute(attribute);
} }
@Override @Override
public void setAttribute(String attributeName, String attributeValue) { public void setAttribute(EntityAttribute attribute, String attributeValue) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HIVE_DB_NAME_ATTRIBUTE: case HIVE_DB_NAME_ATTRIBUTE:
databaseName = attributeValue; databaseName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HIVE_DB_CLUSTER_NAME_ATTRIBUTE: case HIVE_DB_CLUSTER_NAME_ATTRIBUTE:
clusterName = attributeValue; clusterName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
default: default:
super.setAttribute(attributeName, attributeValue); super.setAttribute(attribute, attributeValue);
break; break;
} }
} }
@Override @Override
public void transformComplete() { public void transformComplete() {
if (isCustomerAttributeUpdated) { if (isCustomAttributeUpdated) {
entity.setAttribute(NAME_ATTRIBUTE, databaseName); entity.setAttribute(NAME_ATTRIBUTE, databaseName);
entity.setAttribute(CLUSTER_NAME_ATTRIBUTE, clusterName); entity.setAttribute(CLUSTER_NAME_ATTRIBUTE, clusterName);
entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, toQualifiedName()); entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, toQualifiedName());
......
...@@ -26,11 +26,11 @@ import java.util.List; ...@@ -26,11 +26,11 @@ import java.util.List;
import static org.apache.atlas.entitytransform.TransformationConstants.*; import static org.apache.atlas.entitytransform.TransformationConstants.*;
public class HiveStorageDescriptorEntityHandler extends BaseEntityHandler { 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) { public HiveStorageDescriptorEntityHandler(List<AtlasEntityTransformer> transformers) {
super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES); super(transformers);
} }
@Override @Override
...@@ -47,7 +47,7 @@ public class HiveStorageDescriptorEntityHandler extends BaseEntityHandler { ...@@ -47,7 +47,7 @@ public class HiveStorageDescriptorEntityHandler extends BaseEntityHandler {
private String tableName; private String tableName;
private String clusterName; private String clusterName;
private String location; private String location;
private boolean isCustomerAttributeUpdated = false; private boolean isCustomAttributeUpdated = false;
public HiveStorageDescriptorEntity(AtlasEntity entity) { public HiveStorageDescriptorEntity(AtlasEntity entity) {
...@@ -80,8 +80,8 @@ public class HiveStorageDescriptorEntityHandler extends BaseEntityHandler { ...@@ -80,8 +80,8 @@ public class HiveStorageDescriptorEntityHandler extends BaseEntityHandler {
} }
@Override @Override
public Object getAttribute(String attributeName) { public Object getAttribute(EntityAttribute attribute) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HIVE_DB_NAME_ATTRIBUTE: case HIVE_DB_NAME_ATTRIBUTE:
return databaseName; return databaseName;
...@@ -92,39 +92,39 @@ public class HiveStorageDescriptorEntityHandler extends BaseEntityHandler { ...@@ -92,39 +92,39 @@ public class HiveStorageDescriptorEntityHandler extends BaseEntityHandler {
return clusterName; return clusterName;
} }
return super.getAttribute(attributeName); return super.getAttribute(attribute);
} }
@Override @Override
public void setAttribute(String attributeName, String attributeValue) { public void setAttribute(EntityAttribute attribute, String attributeValue) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HIVE_DB_NAME_ATTRIBUTE: case HIVE_DB_NAME_ATTRIBUTE:
databaseName = attributeValue; databaseName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HIVE_TABLE_NAME_ATTRIBUTE: case HIVE_TABLE_NAME_ATTRIBUTE:
tableName = attributeValue; tableName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HIVE_DB_CLUSTER_NAME_ATTRIBUTE: case HIVE_DB_CLUSTER_NAME_ATTRIBUTE:
clusterName = attributeValue; clusterName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
default: default:
super.setAttribute(attributeName, attributeValue); super.setAttribute(attribute, attributeValue);
break; break;
} }
} }
@Override @Override
public void transformComplete() { public void transformComplete() {
if (isCustomerAttributeUpdated) { if (isCustomAttributeUpdated) {
entity.setAttribute(LOCATION_ATTRIBUTE, toLocation()); entity.setAttribute(LOCATION_ATTRIBUTE, toLocation());
entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, toQualifiedName()); entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, toQualifiedName());
} }
......
...@@ -26,11 +26,11 @@ import java.util.List; ...@@ -26,11 +26,11 @@ import java.util.List;
import static org.apache.atlas.entitytransform.TransformationConstants.*; import static org.apache.atlas.entitytransform.TransformationConstants.*;
public class HiveTableEntityHandler extends BaseEntityHandler { 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) { public HiveTableEntityHandler(List<AtlasEntityTransformer> transformers) {
super(transformers, CUSTOM_TRANSFORM_ATTRIBUTES); super(transformers);
} }
@Override @Override
...@@ -46,7 +46,7 @@ public class HiveTableEntityHandler extends BaseEntityHandler { ...@@ -46,7 +46,7 @@ public class HiveTableEntityHandler extends BaseEntityHandler {
private String databaseName; private String databaseName;
private String tableName; private String tableName;
private String clusterName; private String clusterName;
private boolean isCustomerAttributeUpdated = false; private boolean isCustomAttributeUpdated = false;
public HiveTableEntity(AtlasEntity entity) { public HiveTableEntity(AtlasEntity entity) {
...@@ -69,8 +69,8 @@ public class HiveTableEntityHandler extends BaseEntityHandler { ...@@ -69,8 +69,8 @@ public class HiveTableEntityHandler extends BaseEntityHandler {
} }
@Override @Override
public Object getAttribute(String attributeName) { public Object getAttribute(EntityAttribute attribute) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HIVE_TABLE_NAME_ATTRIBUTE: case HIVE_TABLE_NAME_ATTRIBUTE:
return tableName; return tableName;
...@@ -81,39 +81,39 @@ public class HiveTableEntityHandler extends BaseEntityHandler { ...@@ -81,39 +81,39 @@ public class HiveTableEntityHandler extends BaseEntityHandler {
return clusterName; return clusterName;
} }
return super.getAttribute(attributeName); return super.getAttribute(attribute);
} }
@Override @Override
public void setAttribute(String attributeName, String attributeValue) { public void setAttribute(EntityAttribute attribute, String attributeValue) {
switch (attributeName) { switch (attribute.getAttributeKey()) {
case HIVE_TABLE_NAME_ATTRIBUTE: case HIVE_TABLE_NAME_ATTRIBUTE:
tableName = attributeValue; tableName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HIVE_DB_NAME_ATTRIBUTE: case HIVE_DB_NAME_ATTRIBUTE:
databaseName = attributeValue; databaseName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
case HIVE_DB_CLUSTER_NAME_ATTRIBUTE: case HIVE_DB_CLUSTER_NAME_ATTRIBUTE:
clusterName = attributeValue; clusterName = attributeValue;
isCustomerAttributeUpdated = true; isCustomAttributeUpdated = true;
break; break;
default: default:
super.setAttribute(attributeName, attributeValue); super.setAttribute(attribute, attributeValue);
break; break;
} }
} }
@Override @Override
public void transformComplete() { public void transformComplete() {
if (isCustomerAttributeUpdated) { if (isCustomAttributeUpdated) {
entity.setAttribute(NAME_ATTRIBUTE, tableName); entity.setAttribute(NAME_ATTRIBUTE, tableName);
entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, toQualifiedName()); entity.setAttribute(QUALIFIED_NAME_ATTRIBUTE, toQualifiedName());
} }
......
/**
* 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);
}
...@@ -23,13 +23,13 @@ import org.apache.atlas.store.AtlasTypeDefStore; ...@@ -23,13 +23,13 @@ import org.apache.atlas.store.AtlasTypeDefStore;
import org.apache.atlas.type.AtlasTypeRegistry; import org.apache.atlas.type.AtlasTypeRegistry;
public class TransformerContext { public class TransformerContext {
private final AtlasTypeRegistry typeRegistry; private final AtlasTypeRegistry typeRegistry;
private final AtlasTypeDefStore typeDefStore; private final AtlasTypeDefStore typeDefStore;
private final AtlasExportRequest exportRequest; private final AtlasExportRequest exportRequest;
public TransformerContext(AtlasTypeRegistry typeRegistry, AtlasTypeDefStore typeDefStore, AtlasExportRequest exportRequest) { public TransformerContext(AtlasTypeRegistry typeRegistry, AtlasTypeDefStore typeDefStore, AtlasExportRequest exportRequest) {
this.typeRegistry = typeRegistry; this.typeRegistry = typeRegistry;
this.typeDefStore = typeDefStore; this.typeDefStore = typeDefStore;
this.exportRequest = exportRequest; this.exportRequest = exportRequest;
} }
......
...@@ -17,20 +17,31 @@ ...@@ -17,20 +17,31 @@
*/ */
package org.apache.atlas.entitytransform; 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.impexp.AttributeTransform;
import org.apache.atlas.model.instance.AtlasEntity; import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.instance.AtlasObjectId; 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.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.Test; import org.testng.annotations.Test;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections; import java.util.Collections;
import java.util.HashMap; import java.util.HashMap;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
import static org.apache.atlas.entitytransform.TransformationConstants.HDFS_PATH; 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.assertEquals;
import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull; import static org.testng.Assert.assertNull;
...@@ -38,6 +49,17 @@ import static org.testng.Assert.assertTrue; ...@@ -38,6 +49,17 @@ import static org.testng.Assert.assertTrue;
import static org.apache.atlas.entitytransform.TransformationConstants.HIVE_TABLE; import static org.apache.atlas.entitytransform.TransformationConstants.HIVE_TABLE;
public class TransformationHandlerTest { 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 @Test
public void testHdfsClusterRenameHandler() { public void testHdfsClusterRenameHandler() {
// Rename clusterName from cl1 to cl2 // Rename clusterName from cl1 to cl2
...@@ -139,8 +161,8 @@ public class TransformationHandlerTest { ...@@ -139,8 +161,8 @@ public class TransformationHandlerTest {
@Test @Test
public void testEntityClearAttributesActionWithNoCondition() { public void testEntityClearAttributesActionWithNoCondition() {
// clear replicatedFrom attribute for hive_table entities without any condition // clear replicatedFrom attribute for hive_table entities without any condition
Map<String, String> actions = new HashMap<String, String>() {{ put("__entity.replicatedTo", "CLEAR:"); Map<String, String> actions = new HashMap<String, String>() {{ put("Referenceable.replicatedTo", "CLEAR:");
put("__entity.replicatedFrom", "CLEAR:"); }}; put("Referenceable.replicatedFrom", "CLEAR:"); }};
AttributeTransform transform = new AttributeTransform(null, actions); AttributeTransform transform = new AttributeTransform(null, actions);
...@@ -334,38 +356,75 @@ public class TransformationHandlerTest { ...@@ -334,38 +356,75 @@ public class TransformationHandlerTest {
@Test @Test
public void verifyAddClassification() { public void verifyAddClassification() {
AtlasEntityTransformer entityTransformer = new AtlasEntityTransformer( AtlasEntityTransformer transformer = new AtlasEntityTransformer(Collections.singletonMap("hive_db.qualifiedName", "EQUALS: hr@cl1"),
Collections.singletonMap("hdfs_path.qualifiedName", "EQUALS: hr@cl1"), Collections.singletonMap("Referenceable.", "ADD_CLASSIFICATION: replicated"),
Collections.singletonMap("__entity", "addClassification: 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); assertApplyTransform(handlers);
} }
@Test @Test
public void verifyAddClassificationUsingScope() { public void verifyAddClassificationUsingScope() {
AtlasObjectId objectId = new AtlasObjectId("hive_db", Collections.singletonMap("qualifiedName", "hr@cl1")); AtlasExportRequest exportRequest = new AtlasExportRequest();
AtlasEntityTransformer entityTransformer = new AtlasEntityTransformer(
Collections.singletonMap("__entity", "topLevel: "), exportRequest.setItemsToExport(Collections.singletonList(new AtlasObjectId("hive_db", Collections.singletonMap("qualifiedName", "hr@cl1"))));
Collections.singletonMap("__entity", "addClassification: replicated")
); AtlasEntityTransformer transformer = new AtlasEntityTransformer(Collections.singletonMap("Referenceable.", "topLevel: "),
Collections.singletonMap("Referenceable", "ADD_CLASSIFICATION: replicated"),
List<BaseEntityHandler> handlers = new ArrayList<>(); new TransformerContext(getTypeRegistry(), null, exportRequest));
handlers.add(new BaseEntityHandler(Collections.singletonList(entityTransformer)));
Condition condition = handlers.get(0).transformers.get(0).getConditions().get(0); List<BaseEntityHandler> handlers = Collections.singletonList(new BaseEntityHandler(Collections.singletonList(transformer)));
Condition.ObjectIdEquals objectIdEquals = (Condition.ObjectIdEquals) condition;
objectIdEquals.add(objectId);
assertApplyTransform(handlers); 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) { private void assertApplyTransform(List<BaseEntityHandler> handlers) {
for (AtlasEntity entity : getAllEntities()) { for (AtlasEntity entity : getAllEntities()) {
applyTransforms(entity, handlers); 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()); assertNotNull(entity.getClassifications());
} else{ } else{
assertNull(entity.getClassifications()); assertNull(entity.getClassifications());
...@@ -374,7 +433,7 @@ public class TransformationHandlerTest { ...@@ -374,7 +433,7 @@ public class TransformationHandlerTest {
} }
private List<BaseEntityHandler> initializeHandlers(List<AttributeTransform> params) { 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) { private void applyTransforms(AtlasEntity entity, List<BaseEntityHandler> handlers) {
...@@ -383,15 +442,50 @@ public class TransformationHandlerTest { ...@@ -383,15 +442,50 @@ public class TransformationHandlerTest {
} }
} }
final String[] clusterNames = new String[] { "cl1", "prod" }; private TransformerContext getTransformerContext() {
final String[] databaseNames = new String[] { "hr", "sales", "engg" }; return new TransformerContext(getTypeRegistry(), null, null);
final String[] tableNames = new String[] { "employees", "products", "invoice" }; }
final String[] columnNames = new String[] { "name", "age", "dob" };
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() { private List<AtlasEntity> getHdfsPathEntities() {
List<AtlasEntity> ret = new ArrayList<>(); List<AtlasEntity> ret = new ArrayList<>();
for (String clusterName : clusterNames) { for (String clusterName : CLUSTER_NAMES) {
ret.add(getHdfsPathEntity1(clusterName)); ret.add(getHdfsPathEntity1(clusterName));
ret.add(getHdfsPathEntity2(clusterName)); ret.add(getHdfsPathEntity2(clusterName));
} }
...@@ -402,18 +496,18 @@ public class TransformationHandlerTest { ...@@ -402,18 +496,18 @@ public class TransformationHandlerTest {
private List<AtlasEntity> getAllEntities() { private List<AtlasEntity> getAllEntities() {
List<AtlasEntity> ret = new ArrayList<>(); List<AtlasEntity> ret = new ArrayList<>();
for (String clusterName : clusterNames) { for (String clusterName : CLUSTER_NAMES) {
ret.add(getHdfsPathEntity1(clusterName)); ret.add(getHdfsPathEntity1(clusterName));
ret.add(getHdfsPathEntity2(clusterName)); ret.add(getHdfsPathEntity2(clusterName));
for (String databaseName : databaseNames) { for (String databaseName : DATABASE_NAMES) {
ret.add(getHiveDbEntity(clusterName, databaseName)); ret.add(getHiveDbEntity(clusterName, databaseName));
for (String tableName : tableNames) { for (String tableName : TABLE_NAMES) {
ret.add(getHiveTableEntity(clusterName, databaseName, tableName)); ret.add(getHiveTableEntity(clusterName, databaseName, tableName));
ret.add(getHiveStorageDescriptorEntity(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)); ret.add(getHiveColumnEntity(clusterName, databaseName, tableName, columnName));
} }
} }
...@@ -518,4 +612,8 @@ public class TransformationHandlerTest { ...@@ -518,4 +612,8 @@ public class TransformationHandlerTest {
return entity; return entity;
} }
private AtlasEntity getNonAssetEntity() {
return new AtlasEntity(TYPENAME_NON_ASSET);
}
} }
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment