This is an automated email from the ASF dual-hosted git repository.
chaitalicod pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/atlas.git
The following commit(s) were added to refs/heads/master by this push:
new df4d824878 ATLAS-5375: Fix stale PUT response for
/api/atlas/v2/relationship block/unblock propagation (#729)
df4d824878 is described below
commit df4d8248780a2d1d516400603166805fa645d7fb
Author: sheetalshah1007 <[email protected]>
AuthorDate: Wed Aug 26 12:12:41 2026 +0530
ATLAS-5375: Fix stale PUT response for /api/atlas/v2/relationship
block/unblock propagation (#729)
* ATLAS-5375: Fix stale PUT response for /api/atlas/v2/relationship
block/unblock propagation
* ATLAS-5375: Review fixes — Step 4 comment and null guard for task queue
---
.../atlas/model/instance/AtlasRelationship.java | 14 +
.../apache/atlas/repository/graph/GraphHelper.java | 10 +-
.../repository/store/graph/v1/DeleteHandlerV1.java | 142 +++++---
.../store/graph/v2/AtlasGraphUtilsV2.java | 20 +-
.../store/graph/v2/AtlasRelationshipStoreV2.java | 13 +-
.../store/graph/v2/EntityGraphMapper.java | 4 +-
.../store/graph/v2/EntityGraphRetriever.java | 9 +-
.../v2/tasks/ClassificationPropagationTasks.java | 8 +-
.../store/graph/v2/tasks/ClassificationTask.java | 13 +
.../tasks/ClassificationPropagationTasksTest.java | 10 +-
.../graph/v2/tasks/ClassificationTaskTest.java | 2 +-
.../ClassificationPropagationWithTasksTest.java | 367 ++++++++++++++++++++-
12 files changed, 532 insertions(+), 80 deletions(-)
diff --git
a/intg/src/main/java/org/apache/atlas/model/instance/AtlasRelationship.java
b/intg/src/main/java/org/apache/atlas/model/instance/AtlasRelationship.java
index 87ebbb871e..2bb8ade1b8 100644
--- a/intg/src/main/java/org/apache/atlas/model/instance/AtlasRelationship.java
+++ b/intg/src/main/java/org/apache/atlas/model/instance/AtlasRelationship.java
@@ -87,6 +87,7 @@ public class AtlasRelationship extends AtlasStruct implements
Serializable {
private Long version;
private Set<AtlasClassification> propagatedClassifications;
private Set<AtlasClassification> blockedPropagatedClassifications;
+ private Set<String> pendingTasks; // read-only field i.e.
value provided is ignored during relationship create/update
public AtlasRelationship() {
super();
@@ -239,6 +240,7 @@ public class AtlasRelationship extends AtlasStruct
implements Serializable {
if (other != null) {
init(other.guid, other.homeId, other.provenanceType, other.end1,
other.end2, other.label, other.propagateTags, other.status, other.createdBy,
other.updatedBy,
other.createTime, other.updateTime, other.version,
other.propagatedClassifications, other.blockedPropagatedClassifications);
+ setPendingTasks(other.getPendingTasks());
}
}
@@ -362,6 +364,14 @@ public class AtlasRelationship extends AtlasStruct
implements Serializable {
this.blockedPropagatedClassifications =
blockedPropagatedClassifications;
}
+ public Set<String> getPendingTasks() {
+ return pendingTasks;
+ }
+
+ public void setPendingTasks(Set<String> pendingTasks) {
+ this.pendingTasks = pendingTasks;
+ }
+
@Override
public StringBuilder toString(StringBuilder sb) {
if (sb == null) {
@@ -389,6 +399,9 @@ public class AtlasRelationship extends AtlasStruct
implements Serializable {
sb.append(", blockedPropagatedClassifications=[");
dumpObjects(blockedPropagatedClassifications, sb);
sb.append("]");
+ sb.append(", pendingTasks=[");
+ dumpObjects(pendingTasks, sb);
+ sb.append("]");
sb.append('}');
return sb;
@@ -460,6 +473,7 @@ public class AtlasRelationship extends AtlasStruct
implements Serializable {
setVersion(version);
setPropagatedClassifications(propagatedClassifications);
setBlockedPropagatedClassifications(blockedPropagatedClassifications);
+ setPendingTasks(null);
}
public enum Status { ACTIVE, DELETED }
diff --git
a/repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java
b/repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java
index 303a9f520f..0e5f71c0c1 100755
---
a/repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java
+++
b/repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java
@@ -428,10 +428,16 @@ public final class GraphHelper {
}
public static List<AtlasVertex> getPropagatableClassifications(AtlasEdge
edge) {
+ if (edge != null && getStatus(edge) != DELETED) {
+ return getPropagatableClassifications(edge,
getPropagateTags(edge));
+ }
+ return new ArrayList<>();
+ }
+
+ public static List<AtlasVertex> getPropagatableClassifications(AtlasEdge
edge, PropagateTags propagateTags) {
List<AtlasVertex> ret = new ArrayList<>();
- if (edge != null && getStatus(edge) != DELETED) {
- PropagateTags propagateTags = getPropagateTags(edge);
+ if (edge != null && getStatus(edge) != DELETED && propagateTags !=
null) {
AtlasVertex outVertex = edge.getOutVertex();
AtlasVertex inVertex = edge.getInVertex();
diff --git
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java
index fece9705b5..3844b24dd7 100644
---
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java
+++
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java
@@ -125,6 +125,7 @@ import static
org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.getSt
import static
org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.isReference;
import static
org.apache.atlas.repository.store.graph.v2.tasks.ClassificationPropagateTaskFactory.CLASSIFICATION_PROPAGATION_ADD;
import static
org.apache.atlas.repository.store.graph.v2.tasks.ClassificationPropagateTaskFactory.CLASSIFICATION_PROPAGATION_DELETE;
+import static
org.apache.atlas.repository.store.graph.v2.tasks.ClassificationPropagateTaskFactory.CLASSIFICATION_PROPAGATION_RELATIONSHIP_UPDATE;
import static
org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection.OUT;
import static org.apache.atlas.type.Constants.PENDING_TASKS_PROPERTY_KEY;
@@ -950,59 +951,107 @@ public abstract class DeleteHandlerV1 {
}
public void updateTagPropagations(AtlasEdge edge, AtlasRelationship
relationship) throws AtlasBaseException {
- PropagateTags oldTagPropagation = getPropagateTags(edge);
- PropagateTags newTagPropagation = relationship.getPropagateTags();
+ updateTagPropagations(edge, relationship, false, null, null);
+ }
- if (newTagPropagation != oldTagPropagation) {
- List<AtlasVertex> currentClassificationVertices
= getPropagatableClassifications(edge);
- Map<AtlasVertex, List<AtlasVertex>> currentClassificationsMap
=
entityRetriever.getClassificationPropagatedEntitiesMapping(currentClassificationVertices);
+ public void updateTagPropagations(AtlasEdge edge, AtlasRelationship
relationship, boolean isAsyncExecution, String oldTagPropStr,
java.util.List<String> oldBlockedList) throws AtlasBaseException {
+ // isAsyncExecution=false: relationship update — write edge, queue
background task (or update entities inline if tasks disabled).
+ // isAsyncExecution=true: background task — edge already updated;
update entity tags using old state from task-params (oldTagPropStr /
oldBlockedList).
+ boolean isLegacyTask = isAsyncExecution && oldTagPropStr == null;
- // Update propagation edge
- AtlasGraphUtilsV2.setEncodedProperty(edge,
RELATIONSHIPTYPE_TAG_PROPAGATION_KEY, newTagPropagation.name());
+ // Step 1: Capture the edge's propagation settings before the update
(propagateTags and blocked IDs).
+ // Read from the graph on a relationship update, or from task params
on a background task.
+ PropagateTags oldTagPropagation = isAsyncExecution && !isLegacyTask ?
PropagateTags.valueOf(oldTagPropStr) : getPropagateTags(edge);
+ List<String> oldBlocked = isAsyncExecution && !isLegacyTask ?
oldBlockedList : getBlockedClassificationIds(edge);
- List<AtlasVertex> updatedClassificationVertices
= getPropagatableClassifications(edge);
- List<AtlasVertex> classificationVerticesUnion
= (List<AtlasVertex>) CollectionUtils.union(currentClassificationVertices,
updatedClassificationVertices);
- Map<AtlasVertex, List<AtlasVertex>> updatedClassificationsMap
=
entityRetriever.getClassificationPropagatedEntitiesMapping(classificationVerticesUnion);
+ if (oldBlocked == null) {
+ oldBlocked = Collections.emptyList();
+ }
- // compute add/remove propagations list
- Map<AtlasVertex, List<AtlasVertex>> addPropagationsMap = new
HashMap<>();
- Map<AtlasVertex, List<AtlasVertex>> removePropagationsMap = new
HashMap<>();
+ PropagateTags newTagPropagation = relationship.getPropagateTags();
+ boolean propagationChanged = oldTagPropagation !=
newTagPropagation;
- if (MapUtils.isEmpty(currentClassificationsMap) &&
MapUtils.isNotEmpty(updatedClassificationsMap)) {
- addPropagationsMap.putAll(updatedClassificationsMap);
- } else if (MapUtils.isNotEmpty(currentClassificationsMap) &&
MapUtils.isEmpty(updatedClassificationsMap)) {
- removePropagationsMap.putAll(currentClassificationsMap);
- } else {
- for (AtlasVertex classificationVertex :
updatedClassificationsMap.keySet()) {
- List<AtlasVertex> currentPropagatingEntities =
currentClassificationsMap.getOrDefault(classificationVertex,
Collections.emptyList());
- List<AtlasVertex> updatedPropagatingEntities =
updatedClassificationsMap.getOrDefault(classificationVertex,
Collections.emptyList());
- List<AtlasVertex> entitiesAdded =
(List<AtlasVertex>) CollectionUtils.subtract(updatedPropagatingEntities,
currentPropagatingEntities);
- List<AtlasVertex> entitiesRemoved =
(List<AtlasVertex>) CollectionUtils.subtract(currentPropagatingEntities,
updatedPropagatingEntities);
-
- if (CollectionUtils.isNotEmpty(entitiesAdded)) {
- addPropagationsMap.put(classificationVertex,
entitiesAdded);
- }
+ // Step 2: When the client sends blockedPropagatedClassifications;
validate those tags against
+ // propagatable classifications, compare to the pre-update blocked
list, and set blockedChanged plus the
+ // vertex lists needed for edge write (Step 4) and entity updates
(Step 6). Null field = omit, preserve existing blocks.
+ List<AtlasVertex> classificationsToBlock = new ArrayList<>();
+ List<String> classificationIdsToBlock = new ArrayList<>();
- if (CollectionUtils.isNotEmpty(entitiesRemoved)) {
- removePropagationsMap.put(classificationVertex,
entitiesRemoved);
- }
+ Set<AtlasClassification> newBlockedClassifications =
relationship.getBlockedPropagatedClassifications();
+ boolean blockedChanged = false;
+ List<AtlasVertex> propagatableClassifications = null;
+ List<AtlasVertex> currBlockedClassifications =
Collections.emptyList();
+
+ if (newBlockedClassifications != null) {
+ propagatableClassifications = getPropagatableClassifications(edge,
oldTagPropagation);
+
+ for (AtlasClassification blockedClassification :
newBlockedClassifications) {
+ AtlasVertex classificationVertex =
validateBlockedPropagatedClassification(propagatableClassifications,
blockedClassification);
+
+ if (classificationVertex != null) {
+ classificationsToBlock.add(classificationVertex);
+
classificationIdsToBlock.add(classificationVertex.getIdForDisplay());
}
}
- for (AtlasVertex classificationVertex :
addPropagationsMap.keySet()) {
- List<AtlasVertex> entitiesToAddPropagation =
addPropagationsMap.get(classificationVertex);
+ blockedChanged = !CollectionUtils.isEqualCollection(oldBlocked,
classificationIdsToBlock);
- addTagPropagation(classificationVertex,
entitiesToAddPropagation);
+ if (blockedChanged) {
+ currBlockedClassifications =
getVerticesForIds(propagatableClassifications, oldBlocked);
}
+ }
- for (AtlasVertex classificationVertex :
removePropagationsMap.keySet()) {
- List<AtlasVertex> entitiesToRemovePropagation =
removePropagationsMap.get(classificationVertex);
+ // Step 3: No-op — nothing changed; skip edge write, task queue, and
entity propagation.
+ if (!propagationChanged && !blockedChanged) {
+ return;
+ }
- removeTagPropagation(classificationVertex,
entitiesToRemovePropagation);
+ // Step 4: Write propagateTags or blocked list to the edge (not both;
propagateTags wins if both changed). Skipped on background task unless legacy.
+ if (!isAsyncExecution || isLegacyTask) {
+ if (propagationChanged) {
+ AtlasGraphUtilsV2.setEncodedProperty(edge,
RELATIONSHIPTYPE_TAG_PROPAGATION_KEY, newTagPropagation.name());
+ } else if (blockedChanged) {
+ setBlockedClassificationIds(edge, classificationIdsToBlock);
+ }
+
+ // Step 5: Tasks enabled — edge written above; defer entity tags
to background task with pre-change state.
+ if (!isAsyncExecution && DEFERRED_ACTION_ENABLED) {
+
createAndQueueTask(CLASSIFICATION_PROPAGATION_RELATIONSHIP_UPDATE, edge,
relationship,
+ oldTagPropagation != null ? oldTagPropagation.name() :
null, oldBlocked);
+ return;
+ }
+ }
+
+ // Step 6: Add/remove propagated tags on downstream entities. — inline
when tasks disabled, or in the background task.
+ List<AtlasVertex> affectedClassifications = new ArrayList<>();
+
+ if (propagationChanged) {
+ if (propagatableClassifications == null) {
+ propagatableClassifications =
getPropagatableClassifications(edge, oldTagPropagation);
+ }
+
+ affectedClassifications.addAll(propagatableClassifications);
+
affectedClassifications.addAll(getPropagatableClassifications(edge,
newTagPropagation));
+ } else if (blockedChanged) {
+ List<AtlasVertex> blockedDiff = (List<AtlasVertex>)
CollectionUtils.disjunction(currBlockedClassifications, classificationsToBlock);
+ affectedClassifications.addAll(blockedDiff);
+ }
+
+ Set<AtlasVertex> uniqueAffectedClassifications = new
HashSet<>(affectedClassifications);
+
+ for (AtlasVertex classificationVertex : uniqueAffectedClassifications)
{
+ List<AtlasVertex> propagationsToRemove = new ArrayList<>();
+ List<AtlasVertex> propagationsToAdd = new ArrayList<>();
+
+
entityRetriever.evaluateClassificationPropagation(classificationVertex,
propagationsToAdd, propagationsToRemove);
+
+ if (CollectionUtils.isNotEmpty(propagationsToAdd)) {
+ addTagPropagation(classificationVertex, propagationsToAdd);
+ }
+
+ if (CollectionUtils.isNotEmpty(propagationsToRemove)) {
+ removeTagPropagation(classificationVertex,
propagationsToRemove);
}
- } else {
- // update blocked propagated classifications only if there is no
change is tag propagation (don't update both)
- handleBlockedClassifications(edge,
relationship.getBlockedPropagatedClassifications());
}
}
@@ -1056,9 +1105,13 @@ public abstract class DeleteHandlerV1 {
}
public void createAndQueueTask(String taskType, AtlasEdge
relationshipEdge, AtlasRelationship relationship) {
+ createAndQueueTask(taskType, relationshipEdge, relationship, null,
null);
+ }
+
+ public void createAndQueueTask(String taskType, AtlasEdge
relationshipEdge, AtlasRelationship relationship, String oldTagPropagation,
java.util.List<String> oldBlockedClassifications) {
String currentUser =
RequestContext.getCurrentUser();
String relationshipEdgeId =
relationshipEdge.getIdForDisplay();
- Map<String, Object> taskParams =
ClassificationTask.toParameters(relationshipEdgeId, relationship);
+ Map<String, Object> taskParams =
ClassificationTask.toParameters(relationshipEdgeId, relationship,
oldTagPropagation, oldBlockedClassifications);
AtlasTask task =
taskManagement.createTask(taskType, currentUser, taskParams);
AtlasGraphUtilsV2.addItemToListProperty(relationshipEdge,
EDGE_PENDING_TASKS_PROPERTY_KEY, task.getGuid());
@@ -1577,19 +1630,20 @@ public abstract class DeleteHandlerV1 {
// propagated classifications should contain blocked propagated
classification
private AtlasVertex
validateBlockedPropagatedClassification(List<AtlasVertex>
classificationVertices, AtlasClassification classification) {
- AtlasVertex ret = null;
+ if (classification == null ||
CollectionUtils.isEmpty(classificationVertices)) {
+ return null;
+ }
for (AtlasVertex vertex : classificationVertices) {
String classificationName = getClassificationName(vertex);
String entityGuid = getClassificationEntityGuid(vertex);
if (classificationName.equals(classification.getTypeName()) &&
entityGuid.equals(classification.getEntityGuid())) {
- ret = vertex;
- break;
+ return vertex;
}
}
- return ret;
+ return null;
}
private void setBlockedClassificationIds(AtlasEdge edge, List<String>
classificationIds) {
diff --git
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasGraphUtilsV2.java
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasGraphUtilsV2.java
index 2bf3e2dff0..b75a47d0b3 100644
---
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasGraphUtilsV2.java
+++
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasGraphUtilsV2.java
@@ -911,20 +911,26 @@ public class AtlasGraphUtilsV2 {
public static void addItemToListProperty(AtlasEdge edge, String property,
String value) {
List<String> list = (List<String>) getListFromProperty(edge, property);
- list.add(value);
+ if (!list.contains(value)) {
+ List<String> newList = new ArrayList<>(list);
+ newList.add(value);
- edge.setListProperty(property, list);
+ edge.setListProperty(property, newList);
+ }
}
public static void removeItemFromListProperty(AtlasEdge edge, String
property, String value) {
List<String> list = (List<String>) getListFromProperty(edge, property);
- list.remove(value);
+ if (list.contains(value)) {
+ List<String> newList = new ArrayList<>(list);
+ newList.remove(value);
- if (CollectionUtils.isEmpty(list)) {
- edge.removeProperty(property);
- } else {
- edge.setListProperty(property, list);
+ if (CollectionUtils.isEmpty(newList)) {
+ edge.removeProperty(property);
+ } else {
+ edge.setListProperty(property, newList);
+ }
}
}
diff --git
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasRelationshipStoreV2.java
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasRelationshipStoreV2.java
index de8b3b2a61..1fb68ca063 100644
---
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasRelationshipStoreV2.java
+++
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasRelationshipStoreV2.java
@@ -17,7 +17,6 @@
*/
package org.apache.atlas.repository.store.graph.v2;
-import org.apache.atlas.AtlasConfiguration;
import org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.RequestContext;
import org.apache.atlas.annotation.GraphTransaction;
@@ -83,7 +82,6 @@ import static
org.apache.atlas.repository.Constants.RELATIONSHIP_TYPE_PROPERTY_K
import static org.apache.atlas.repository.Constants.VERSION_PROPERTY_KEY;
import static
org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.getState;
import static
org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.getTypeName;
-import static
org.apache.atlas.repository.store.graph.v2.tasks.ClassificationPropagateTaskFactory.CLASSIFICATION_PROPAGATION_RELATIONSHIP_UPDATE;
@Component
public class AtlasRelationshipStoreV2 implements AtlasRelationshipStore {
@@ -91,7 +89,6 @@ public class AtlasRelationshipStoreV2 implements
AtlasRelationshipStore {
private static final Long DEFAULT_RELATIONSHIP_VERSION = 0L;
private static final boolean NOTIFICATIONS_ENABLED =
NOTIFICATION_RELATIONSHIPS_ENABLED.getBoolean();
- private static final boolean DEFERRED_ACTION_ENABLED =
AtlasConfiguration.TASKS_USE_ENABLED.getBoolean();
private final AtlasGraph graph;
private final AtlasTypeRegistry typeRegistry;
@@ -449,11 +446,7 @@ public class AtlasRelationshipStoreV2 implements
AtlasRelationshipStore {
}
private void updateTagPropagations(AtlasEdge relationshipEdge,
AtlasRelationship relationship) throws AtlasBaseException {
- if (DEFERRED_ACTION_ENABLED) {
- createAndQueueTask(CLASSIFICATION_PROPAGATION_RELATIONSHIP_UPDATE,
relationshipEdge, relationship);
- } else {
-
deleteDelegate.getHandler().updateTagPropagations(relationshipEdge,
relationship);
- }
+ deleteDelegate.getHandler().updateTagPropagations(relationshipEdge,
relationship);
}
private void validateRelationship(AtlasRelationship relationship) throws
AtlasBaseException {
@@ -771,10 +764,6 @@ public class AtlasRelationshipStoreV2 implements
AtlasRelationshipStore {
}
}
- private void createAndQueueTask(String taskType, AtlasEdge
relationshipEdge, AtlasRelationship relationship) {
- deleteDelegate.getHandler().createAndQueueTask(taskType,
relationshipEdge, relationship);
- }
-
private void verifyRelationshipReadAccess(AtlasEdge edge) throws
AtlasBaseException {
AtlasEntityHeader end1Entity =
entityRetriever.toAtlasEntityHeaderWithClassifications(edge.getOutVertex());
AtlasEntityHeader end2Entity =
entityRetriever.toAtlasEntityHeaderWithClassifications(edge.getInVertex());
diff --git
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java
index caa7e8f5d9..c65330f42e 100644
---
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java
+++
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java
@@ -1364,10 +1364,10 @@ public class EntityGraphMapper {
}
@GraphTransaction
- public void updateTagPropagations(String relationshipEdgeId,
AtlasRelationship relationship) throws AtlasBaseException {
+ public void updateTagPropagations(String relationshipEdgeId,
AtlasRelationship relationship, String oldTagPropagation,
java.util.List<String> oldBlockedClassifications) throws AtlasBaseException {
AtlasEdge relationshipEdge = graph.getEdge(relationshipEdgeId);
- deleteDelegate.getHandler().updateTagPropagations(relationshipEdge,
relationship);
+ deleteDelegate.getHandler().updateTagPropagations(relationshipEdge,
relationship, true, oldTagPropagation, oldBlockedClassifications);
entityChangeNotifier.notifyPropagatedEntities();
}
diff --git
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
index ffd5296b07..6063477933 100644
---
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
+++
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -1864,6 +1864,11 @@ public class EntityGraphRetriever {
// set propagated and blocked propagated classifications
readClassificationsFromEdge(edge, relationshipWithExtInfo,
extendedInfo);
+ List<String> pendingTasks =
edge.getListProperty(org.apache.atlas.repository.Constants.EDGE_PENDING_TASKS_PROPERTY_KEY);
+ if (CollectionUtils.isNotEmpty(pendingTasks)) {
+ relationship.setPendingTasks(new HashSet<>(pendingTasks));
+ }
+
return relationshipWithExtInfo;
}
@@ -1873,6 +1878,8 @@ public class EntityGraphRetriever {
AtlasRelationship relationship =
relationshipWithExtInfo.getRelationship();
Set<AtlasClassification> propagatedClassifications = new HashSet<>();
Set<AtlasClassification> blockedClassifications = new HashSet<>();
+ Set<String> blockedIdSet =
CollectionUtils.isEmpty(blockedClassificationIds)
+ ? Collections.emptySet() : new
HashSet<>(blockedClassificationIds);
for (AtlasVertex classificationVertex : classificationVertices) {
String classificationId =
classificationVertex.getIdForDisplay();
@@ -1882,7 +1889,7 @@ public class EntityGraphRetriever {
continue;
}
- if (blockedClassificationIds.contains(classificationId)) {
+ if (blockedIdSet.contains(classificationId)) {
blockedClassifications.add(classification);
} else {
propagatedClassifications.add(classification);
diff --git
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationPropagationTasks.java
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationPropagationTasks.java
index 2118bc42e0..96ebc39e73 100644
---
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationPropagationTasks.java
+++
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationPropagationTasks.java
@@ -67,8 +67,14 @@ public class ClassificationPropagationTasks {
protected void run(Map<String, Object> parameters) throws
AtlasBaseException {
String relationshipEdgeId = (String)
parameters.get(PARAM_RELATIONSHIP_EDGE_ID);
AtlasRelationship relationship = AtlasType.fromJson((String)
parameters.get(PARAM_RELATIONSHIP_OBJECT), AtlasRelationship.class);
+ String oldTagPropagation = (String)
parameters.get(PARAM_OLD_TAG_PROPAGATION);
- entityGraphMapper.updateTagPropagations(relationshipEdgeId,
relationship);
+ java.util.List<String> oldBlockedClassifications = null;
+ if (parameters.containsKey(PARAM_OLD_BLOCKED_CLASSIFICATIONS)) {
+ oldBlockedClassifications = AtlasType.fromJson((String)
parameters.get(PARAM_OLD_BLOCKED_CLASSIFICATIONS), java.util.List.class);
+ }
+
+ entityGraphMapper.updateTagPropagations(relationshipEdgeId,
relationship, oldTagPropagation, oldBlockedClassifications);
}
}
}
diff --git
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationTask.java
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationTask.java
index 5a4981c294..098a798633 100644
---
a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationTask.java
+++
b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationTask.java
@@ -50,6 +50,8 @@ public abstract class ClassificationTask extends AbstractTask
{
public static final String PARAM_RELATIONSHIP_GUID =
"relationshipGuid";
public static final String PARAM_RELATIONSHIP_OBJECT =
"relationshipObject";
public static final String PARAM_RELATIONSHIP_EDGE_ID =
"relationshipEdgeId";
+ public static final String PARAM_OLD_TAG_PROPAGATION =
"oldTagPropagation";
+ public static final String PARAM_OLD_BLOCKED_CLASSIFICATIONS =
"oldBlockedClassifications";
public static final String PARAM_IMPORT_IN_PROGRESS =
"isImportInProgress";
public static final boolean DEFAULT_IMPORT_IN_PROGRESS = false;
@@ -84,11 +86,22 @@ public abstract class ClassificationTask extends
AbstractTask {
}
public static Map<String, Object> toParameters(String relationshipEdgeId,
AtlasRelationship relationship) {
+ return toParameters(relationshipEdgeId, relationship, null, null);
+ }
+
+ public static Map<String, Object> toParameters(String relationshipEdgeId,
AtlasRelationship relationship, String oldTagPropagation,
java.util.List<String> oldBlockedClassifications) {
Map<String, Object> ret = new HashMap<>();
ret.put(PARAM_RELATIONSHIP_EDGE_ID, relationshipEdgeId);
ret.put(PARAM_RELATIONSHIP_OBJECT, AtlasType.toJson(relationship));
+ if (oldTagPropagation != null) {
+ ret.put(PARAM_OLD_TAG_PROPAGATION, oldTagPropagation);
+ }
+ if (oldBlockedClassifications != null) {
+ ret.put(PARAM_OLD_BLOCKED_CLASSIFICATIONS,
AtlasType.toJson(oldBlockedClassifications));
+ }
+
return ret;
}
diff --git
a/repository/src/test/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationPropagationTasksTest.java
b/repository/src/test/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationPropagationTasksTest.java
index 601a224930..86dfb1cc11 100644
---
a/repository/src/test/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationPropagationTasksTest.java
+++
b/repository/src/test/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationPropagationTasksTest.java
@@ -184,11 +184,11 @@ public class ClassificationPropagationTasksTest {
parameters.put(PARAM_RELATIONSHIP_EDGE_ID, "relationship-edge-id");
parameters.put(PARAM_RELATIONSHIP_OBJECT, relationshipJson);
- doNothing().when(entityGraphMapper).updateTagPropagations(anyString(),
any(AtlasRelationship.class));
+ doNothing().when(entityGraphMapper).updateTagPropagations(anyString(),
any(AtlasRelationship.class), any(), any());
invokeRunMethod(updateTask, parameters);
-
verify(entityGraphMapper).updateTagPropagations(eq("relationship-edge-id"),
eq(relationship));
+
verify(entityGraphMapper).updateTagPropagations(eq("relationship-edge-id"),
eq(relationship), any(), any());
}
@Test
@@ -205,7 +205,7 @@ public class ClassificationPropagationTasksTest {
parameters.put(PARAM_RELATIONSHIP_EDGE_ID, "relationship-edge-id");
parameters.put(PARAM_RELATIONSHIP_OBJECT, relationshipJson);
- doThrow(new AtlasBaseException("Test
exception")).when(entityGraphMapper).updateTagPropagations(anyString(),
any(AtlasRelationship.class));
+ doThrow(new AtlasBaseException("Test
exception")).when(entityGraphMapper).updateTagPropagations(anyString(),
any(AtlasRelationship.class), any(), any());
expectThrows(AtlasBaseException.class, () ->
invokeRunMethod(updateTask, parameters));
}
@@ -222,7 +222,7 @@ public class ClassificationPropagationTasksTest {
invokeRunMethod(updateTask, parameters);
- verify(entityGraphMapper).updateTagPropagations(null, null);
+ verify(entityGraphMapper).updateTagPropagations(null, null, null,
null);
}
@Test
@@ -293,7 +293,7 @@ public class ClassificationPropagationTasksTest {
invokeRunMethod(updateTask, parameters);
- verify(entityGraphMapper).updateTagPropagations(null, null);
+ verify(entityGraphMapper).updateTagPropagations(null, null, null,
null);
}
// Helper methods
diff --git
a/repository/src/test/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationTaskTest.java
b/repository/src/test/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationTaskTest.java
index 23c79dd318..3f38c70abe 100644
---
a/repository/src/test/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationTaskTest.java
+++
b/repository/src/test/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationTaskTest.java
@@ -278,7 +278,7 @@ public class ClassificationTaskTest {
@Test
public void testToParametersWithNullValues() {
- Map<String, Object> result = ClassificationTask.toParameters(null,
null, null, null);
+ Map<String, Object> result = ClassificationTask.toParameters((String)
null, (String) null, null, null);
assertNotNull(result);
assertEquals(result.get(PARAM_ENTITY_GUID), null);
diff --git
a/repository/src/test/java/org/apache/atlas/repository/tagpropagation/ClassificationPropagationWithTasksTest.java
b/repository/src/test/java/org/apache/atlas/repository/tagpropagation/ClassificationPropagationWithTasksTest.java
index 1ecb42bd8b..3eba0446e4 100644
---
a/repository/src/test/java/org/apache/atlas/repository/tagpropagation/ClassificationPropagationWithTasksTest.java
+++
b/repository/src/test/java/org/apache/atlas/repository/tagpropagation/ClassificationPropagationWithTasksTest.java
@@ -20,23 +20,33 @@ package org.apache.atlas.repository.tagpropagation;
import org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.RequestContext;
import org.apache.atlas.TestModules;
+import org.apache.atlas.discovery.AtlasLineageService;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.exception.EntityNotFoundException;
import org.apache.atlas.model.instance.AtlasClassification;
import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasRelationship;
+import org.apache.atlas.model.lineage.AtlasLineageInfo;
+import org.apache.atlas.model.lineage.AtlasLineageInfo.LineageRelation;
import org.apache.atlas.model.typedef.AtlasClassificationDef;
+import org.apache.atlas.model.typedef.AtlasRelationshipDef.PropagateTags;
+import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.apache.atlas.repository.AtlasTestBase;
import org.apache.atlas.repository.graph.GraphHelper;
+import org.apache.atlas.repository.graphdb.AtlasEdge;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.repository.impexp.ImportService;
import org.apache.atlas.repository.impexp.ZipFileResourceTestUtils;
import org.apache.atlas.repository.store.graph.AtlasEntityStore;
+import org.apache.atlas.repository.store.graph.AtlasRelationshipStore;
import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
import org.apache.atlas.repository.store.graph.v2.EntityGraphMapper;
import org.apache.atlas.store.AtlasTypeDefStore;
import org.apache.atlas.tasks.TaskManagement;
import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.testng.SkipException;
import org.testng.annotations.BeforeClass;
@@ -49,8 +59,13 @@ import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.Collections;
+import java.util.HashSet;
import java.util.List;
+import java.util.Set;
+import static org.apache.atlas.model.lineage.AtlasLineageInfo.LineageDirection;
+import static
org.apache.atlas.model.typedef.AtlasRelationshipDef.PropagateTags.NONE;
+import static
org.apache.atlas.model.typedef.AtlasRelationshipDef.PropagateTags.ONE_TO_TWO;
import static
org.apache.atlas.repository.graph.GraphHelper.getClassificationVertex;
import static
org.apache.atlas.repository.graph.GraphHelper.getPropagatedTraitNames;
import static
org.apache.atlas.repository.impexp.ZipFileResourceTestUtils.runImportWithNoParameters;
@@ -60,6 +75,7 @@ import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
@Guice(modules = TestModules.TestOnlyModule.class)
public class ClassificationPropagationWithTasksTest extends AtlasTestBase {
@@ -67,12 +83,17 @@ public class ClassificationPropagationWithTasksTest extends
AtlasTestBase {
private static final String IMPORT_DELETE_FILE =
"deleted_tab_propagation.zip";
- private static final String HDFS_PATH_EMPLOYEES =
"a3955120-ac17-426f-a4af-972ec8690e5f";
-
- private static final String HIVE_TABLE =
"089c1ad4-9dde-4f9e-80c8-12a3046be337";
+ private static final String HDFS_PATH_EMPLOYEES =
"a3955120-ac17-426f-a4af-972ec8690e5f";
+ private static final String EMPLOYEES1_TABLE =
"cdf0040e-739e-4590-a137-964d10e73573";
+ private static final String EMPLOYEES2_TABLE =
"0a3e66b6-472c-48b3-8453-abdd24f9494f";
+ private static final String EMPLOYEES_UNION_TABLE =
"1ceac963-1a2b-476a-a269-10396187d406";
+ private static final String EMPLOYEES_UNION_PROCESS =
"470a2d1e-b1fd-47de-8f2d-8dfd0a0275a7";
+ private static final String HIVE_TABLE =
"089c1ad4-9dde-4f9e-80c8-12a3046be337";
private static final String HIVE_TABLE_CTAS =
"e83551b7-bbef-45aa-99d5-3d98c0ac737b";
+ private AtlasLineageInfo lineageInfo;
+
@Inject
private AtlasTypeDefStore typeDefStore;
@@ -91,6 +112,15 @@ public class ClassificationPropagationWithTasksTest extends
AtlasTestBase {
@Inject
private TaskManagement tasksManagement;
+ @Inject
+ private AtlasRelationshipStore relationshipStore;
+
+ @Inject
+ private AtlasLineageService lineageService;
+
+ @Inject
+ private AtlasGraph graph;
+
public static InputStream getZipSource(String fileName) throws IOException
{
return ZipFileResourceTestUtils.getFileInputStream(fileName);
}
@@ -219,7 +249,7 @@ public class ClassificationPropagationWithTasksTest extends
AtlasTestBase {
assertNotNull(impactedEntities);
}
- @Test
+ @Test(priority = 100)
public void runImportForDeletedEntityLineage() throws Exception {
runImportWithNoParameters(importService,
getZipSource(IMPORT_DELETE_FILE));
final String tagName = "classification1";
@@ -241,6 +271,258 @@ public class ClassificationPropagationWithTasksTest
extends AtlasTestBase {
assertNotNull(propagatedTraitNames);
}
+ @Test(dependsOnMethods =
"updateRelationship_afterTaskCompletion_entitiesAndPendingTasksCleared")
+ public void
updateRelationship_propagateTagsChange_ignoresBlockedInSamePut() throws
Exception {
+ setupRelationshipBlockScenario();
+
+ AtlasRelationship relationship =
getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);
+
+ assertEquals(relationship.getPropagateTags(), ONE_TO_TWO);
+
assertFalse(relationship.getBlockedPropagatedClassifications().isEmpty());
+
+ List<String> blockedIdsBefore =
GraphHelper.getBlockedClassificationIds(
+ new GraphHelper(graph).getEdgeForGUID(relationship.getGuid()));
+
+ assertFalse(blockedIdsBefore.isEmpty());
+
+ relationship.setPropagateTags(NONE);
+
relationship.setBlockedPropagatedClassifications(Collections.emptySet());
+
+ AtlasRelationship putReturn = relationshipStore.update(relationship);
+
+ assertEquals(putReturn.getPropagateTags(), NONE);
+
+ List<String> blockedIdsAfter = GraphHelper.getBlockedClassificationIds(
+ new GraphHelper(graph).getEdgeForGUID(putReturn.getGuid()));
+
+ assertEquals(blockedIdsAfter, blockedIdsBefore,
+ "Blocked list must not change when propagateTags also changed
in same PUT");
+ }
+
+ @Test
+ public void updateRelationship_unblock_putAndGetBeforeTaskCompletion()
throws Exception {
+ setupRelationshipBlockScenario();
+
+ AtlasEntity employees1 = getEntity(EMPLOYEES1_TABLE);
+ AtlasEntity employees2 = getEntity(EMPLOYEES2_TABLE);
+
+ AtlasClassification piiTag2 = new AtlasClassification("PII");
+ piiTag2.setEntityGuid(employees1.getGuid());
+
+ AtlasClassification piiTag3 = new AtlasClassification("PII");
+ piiTag3.setEntityGuid(employees2.getGuid());
+
+ AtlasRelationship relationship =
getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);
+
+
relationship.setBlockedPropagatedClassifications(Collections.emptySet());
+
+ AtlasRelationship putReturn = relationshipStore.update(relationship);
+
+ assertNotNull(putReturn);
+ assertTrue(putReturn.getBlockedPropagatedClassifications().isEmpty());
+
assertClassificationExistInList(putReturn.getPropagatedClassifications(),
piiTag2);
+
assertClassificationExistInList(putReturn.getPropagatedClassifications(),
piiTag3);
+
+ AtlasRelationship getReturn =
relationshipStore.getById(putReturn.getGuid());
+
+ assertTrue(getReturn.getBlockedPropagatedClassifications().isEmpty());
+
+ if (putReturn.getPendingTasks() != null &&
!putReturn.getPendingTasks().isEmpty()) {
+ assertNotNull(getReturn.getPendingTasks());
+ assertFalse(getReturn.getPendingTasks().isEmpty());
+ assertClassificationNotExistInEntity(EMPLOYEES_UNION_TABLE,
piiTag2);
+ assertClassificationNotExistInEntity(EMPLOYEES_UNION_TABLE,
piiTag3);
+ } else {
+ assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, piiTag2);
+ assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, piiTag3);
+ }
+ }
+
+ @Test(dependsOnMethods =
"updateRelationship_unblock_putAndGetBeforeTaskCompletion")
+ public void
updateRelationship_afterTaskCompletion_entitiesAndPendingTasksCleared() throws
Exception {
+ setupRelationshipBlockScenario();
+
+ AtlasEntity employees1 = getEntity(EMPLOYEES1_TABLE);
+ AtlasEntity employees2 = getEntity(EMPLOYEES2_TABLE);
+
+ AtlasClassification piiTag2 = new AtlasClassification("PII");
+ piiTag2.setEntityGuid(employees1.getGuid());
+
+ AtlasClassification piiTag3 = new AtlasClassification("PII");
+ piiTag3.setEntityGuid(employees2.getGuid());
+
+ AtlasRelationship relationship =
getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);
+ PropagateTags propagateTags = relationship.getPropagateTags();
+ List<String> oldBlockedIds = GraphHelper.getBlockedClassificationIds(
+ new GraphHelper(graph).getEdgeForGUID(relationship.getGuid()));
+
+
relationship.setBlockedPropagatedClassifications(Collections.emptySet());
+
+ AtlasRelationship putReturn = relationshipStore.update(relationship);
+
+ AtlasRelationship unblockRequest =
relationshipStore.getById(putReturn.getGuid());
+
unblockRequest.setBlockedPropagatedClassifications(Collections.emptySet());
+
+ waitForRelationshipTasksToComplete(putReturn.getGuid(),
unblockRequest, propagateTags, oldBlockedIds);
+
+ AtlasRelationship finalRelationship =
relationshipStore.getById(putReturn.getGuid());
+
+ assertNull(finalRelationship.getPendingTasks());
+ assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, piiTag2);
+ assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, piiTag3);
+ }
+
+ private void setupRelationshipBlockScenario() throws AtlasBaseException {
+ AtlasEntity hdfsPath = getEntity(HDFS_PATH_EMPLOYEES);
+ AtlasEntity employees1 = getEntity(EMPLOYEES1_TABLE);
+ AtlasEntity employees2 = getEntity(EMPLOYEES2_TABLE);
+
+ AtlasClassification piiTag1 = new AtlasClassification("PII");
+ piiTag1.setPropagate(true);
+ piiTag1.setEntityGuid(hdfsPath.getGuid());
+ piiTag1.setAttribute("type", "from hdfs_path entity");
+ piiTag1.setAttribute("valid", true);
+
+ AtlasClassification piiTag2 = new AtlasClassification("PII");
+ piiTag2.setPropagate(true);
+ piiTag2.setEntityGuid(employees1.getGuid());
+ piiTag2.setAttribute("type", "from employees1 entity");
+ piiTag2.setAttribute("valid", true);
+
+ AtlasClassification piiTag3 = new AtlasClassification("PII");
+ piiTag3.setPropagate(true);
+ piiTag3.setEntityGuid(employees2.getGuid());
+ piiTag3.setAttribute("type", "from employees2 entity");
+ piiTag3.setAttribute("valid", true);
+
+ addClassificationIfMissing(hdfsPath, piiTag1);
+ addClassificationIfMissing(employees1, piiTag2);
+ addClassificationIfMissing(employees2, piiTag3);
+
+ propagateClassificationIfPresent(hdfsPath.getGuid(),
piiTag1.getTypeName());
+ propagateClassificationIfPresent(employees1.getGuid(),
piiTag2.getTypeName());
+ propagateClassificationIfPresent(employees2.getGuid(),
piiTag3.getTypeName());
+
+ assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, piiTag2);
+ assertClassificationExistInEntity(EMPLOYEES_UNION_TABLE, piiTag3);
+
+ AtlasRelationship relationship =
getRelationship(EMPLOYEES_UNION_PROCESS, EMPLOYEES_UNION_TABLE);
+ PropagateTags oldPropagateTags = relationship.getPropagateTags();
+ List<String> oldBlockedIds = Collections.emptyList();
+
+ Set<AtlasClassification> blockTags =
selectPropagatedClassifications(relationship, employees1.getGuid(),
employees2.getGuid());
+
+ assertFalse(blockTags.isEmpty(), "Expected propagated PII on
process->union relationship before block");
+
+ relationship.setBlockedPropagatedClassifications(blockTags);
+
+ AtlasRelationship blockReturn = relationshipStore.update(relationship);
+
+ if (blockReturn.getPendingTasks() != null &&
!blockReturn.getPendingTasks().isEmpty()) {
+ applyRelationshipTaskIfPending(blockReturn, oldPropagateTags,
oldBlockedIds);
+ }
+
+ assertClassificationNotExistInEntity(EMPLOYEES_UNION_TABLE, piiTag2);
+ assertClassificationNotExistInEntity(EMPLOYEES_UNION_TABLE, piiTag3);
+ }
+
+ private void addClassificationIfMissing(AtlasEntity entity,
AtlasClassification classification) throws AtlasBaseException {
+ List<AtlasClassification> classifications =
entity.getClassifications();
+
+ if (CollectionUtils.isNotEmpty(classifications)) {
+ for (AtlasClassification existing : classifications) {
+ if
(existing.getTypeName().equals(classification.getTypeName())) {
+ return;
+ }
+ }
+ }
+
+ entityStore.addClassifications(entity.getGuid(),
Collections.singletonList(classification));
+ }
+
+ private void propagateClassificationIfPresent(String entityGuid, String
classificationName) throws AtlasBaseException {
+ AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(entityGuid);
+
+ if (entityVertex == null) {
+ return;
+ }
+
+ AtlasVertex classificationVertex =
getClassificationVertex(entityVertex, classificationName);
+
+ if (classificationVertex != null) {
+ entityGraphMapper.propagateClassification(entityGuid,
classificationVertex.getId().toString(), StringUtils.EMPTY);
+ }
+ }
+
+ private Set<AtlasClassification>
selectPropagatedClassifications(AtlasRelationship relationship, String...
sourceEntityGuids) {
+ Set<AtlasClassification> result = new HashSet<>();
+ Set<AtlasClassification> propagated =
relationship.getPropagatedClassifications();
+
+ if (CollectionUtils.isEmpty(propagated)) {
+ return result;
+ }
+
+ Set<String> sourceGuids = new
HashSet<>(Arrays.asList(sourceEntityGuids));
+
+ for (AtlasClassification classification : propagated) {
+ if ("PII".equals(classification.getTypeName()) &&
sourceGuids.contains(classification.getEntityGuid())) {
+ result.add(classification);
+ }
+ }
+
+ return result;
+ }
+
+ private void applyRelationshipTaskIfPending(AtlasRelationship
updateRequest, PropagateTags oldPropagateTags,
+ List<String> oldBlockedIds)
throws AtlasBaseException {
+ AtlasEdge edge = new
GraphHelper(graph).getEdgeForGUID(updateRequest.getGuid());
+
+ entityGraphMapper.updateTagPropagations(edge.getIdForDisplay(),
updateRequest,
+ oldPropagateTags != null ? oldPropagateTags.name() : null,
oldBlockedIds);
+ }
+
+ private void waitForRelationshipTasksToComplete(String relationshipGuid,
AtlasRelationship updateRequest,
+ PropagateTags
oldPropagateTags, List<String> oldBlockedIds)
+ throws AtlasBaseException, InterruptedException {
+ AtlasRelationship relationship =
relationshipStore.getById(relationshipGuid);
+
+ if (relationship.getPendingTasks() != null &&
!relationship.getPendingTasks().isEmpty()) {
+ applyRelationshipTaskIfPending(updateRequest, oldPropagateTags,
oldBlockedIds);
+ relationship = relationshipStore.getById(relationshipGuid);
+ }
+
+ if (relationship.getPendingTasks() == null ||
relationship.getPendingTasks().isEmpty()) {
+ return;
+ }
+
+ for (int i = 0; i < 30; i++) {
+ relationship = relationshipStore.getById(relationshipGuid);
+
+ if (relationship.getPendingTasks() == null ||
relationship.getPendingTasks().isEmpty()) {
+ return;
+ }
+
+ Thread.sleep(100);
+ }
+
+ fail("Relationship tasks did not complete in time for guid: " +
relationshipGuid);
+ }
+
+ private void assertClassificationExistInList(Set<AtlasClassification>
classifications, AtlasClassification expected) {
+ if (classifications == null) {
+ fail("Propagated classifications list is null");
+ }
+
+ for (AtlasClassification classification : classifications) {
+ if (classification.getTypeName().equals(expected.getTypeName())
+ &&
classification.getEntityGuid().equals(expected.getEntityGuid())) {
+ return;
+ }
+ }
+
+ fail("Propagated classification not found in relationship response");
+ }
+
private void loadModelFilesAndImportTestData() {
try {
loadModelFromJson("0000-Area0/0010-base_model.json", typeDefStore,
typeRegistry);
@@ -250,6 +532,8 @@ public class ClassificationPropagationWithTasksTest extends
AtlasTestBase {
loadSampleClassificationDefs();
runImportWithNoParameters(importService,
getZipSource(IMPORT_FILE));
+
+ initializeLineageInfo();
} catch (AtlasBaseException | IOException e) {
throw new SkipException("Model loading failed!");
}
@@ -258,8 +542,17 @@ public class ClassificationPropagationWithTasksTest
extends AtlasTestBase {
private void loadSampleClassificationDefs() throws AtlasBaseException {
AtlasClassificationDef tagX = new AtlasClassificationDef("tagX");
AtlasClassificationDef tagY = new AtlasClassificationDef("tagY");
+ AtlasClassificationDef pii = new AtlasClassificationDef("PII");
- typeDefStore.createTypesDef(new AtlasTypesDef(Collections.emptyList(),
Collections.emptyList(), Arrays.asList(tagX, tagY), Collections.emptyList(),
Collections.emptyList()));
+ pii.addAttribute(new AtlasAttributeDef("type", "string"));
+ pii.addAttribute(new AtlasAttributeDef("valid", "boolean"));
+
+ typeDefStore.createTypesDef(new AtlasTypesDef(Collections.emptyList(),
Collections.emptyList(),
+ Arrays.asList(tagX, tagY, pii), Collections.emptyList(),
Collections.emptyList()));
+ }
+
+ private void initializeLineageInfo() throws AtlasBaseException {
+ lineageInfo = lineageService.getAtlasLineageInfo(HDFS_PATH_EMPLOYEES,
LineageDirection.BOTH, 3);
}
private AtlasEntity getEntity(String entityGuid) throws AtlasBaseException
{
@@ -267,4 +560,68 @@ public class ClassificationPropagationWithTasksTest
extends AtlasTestBase {
return entityWithExtInfo.getEntity();
}
+
+ private AtlasRelationship getRelationship(String fromEntityGuid, String
toEntityGuid) throws AtlasBaseException {
+ String relationshipGuid =
findRelationshipGuid(lineageInfo.getRelations(), fromEntityGuid, toEntityGuid);
+
+ if (relationshipGuid == null) {
+ relationshipGuid =
findRelationshipGuid(lineageInfo.getRelations(), toEntityGuid, fromEntityGuid);
+ }
+
+ if (relationshipGuid == null) {
+ AtlasLineageInfo outputLineage =
lineageService.getAtlasLineageInfo(fromEntityGuid, LineageDirection.OUTPUT, 1);
+
+ relationshipGuid =
findRelationshipGuid(outputLineage.getRelations(), fromEntityGuid,
toEntityGuid);
+
+ if (relationshipGuid == null) {
+ relationshipGuid =
findRelationshipGuid(outputLineage.getRelations(), toEntityGuid,
fromEntityGuid);
+ }
+ }
+
+ assertNotNull(relationshipGuid, "Relationship not found between " +
fromEntityGuid + " and " + toEntityGuid);
+
+ return relationshipStore.getById(relationshipGuid);
+ }
+
+ private String findRelationshipGuid(Set<LineageRelation> relations, String
fromEntityGuid, String toEntityGuid) {
+ if (relations == null) {
+ return null;
+ }
+
+ for (LineageRelation relation : relations) {
+ if (relation.getFromEntityId().equals(fromEntityGuid) &&
relation.getToEntityId().equals(toEntityGuid)) {
+ return relation.getRelationshipId();
+ }
+ }
+
+ return null;
+ }
+
+ private void assertClassificationExistInEntity(String entityGuid,
AtlasClassification classification) throws AtlasBaseException {
+ List<AtlasClassification> classifications =
getEntity(entityGuid).getClassifications();
+
+ if (CollectionUtils.isNotEmpty(classifications)) {
+ for (AtlasClassification c : classifications) {
+ if (c.getTypeName().equals(classification.getTypeName())
+ &&
c.getEntityGuid().equals(classification.getEntityGuid())) {
+ return;
+ }
+ }
+ }
+
+ fail("Propagated classification is not present in entity!");
+ }
+
+ private void assertClassificationNotExistInEntity(String entityGuid,
AtlasClassification classification) throws AtlasBaseException {
+ List<AtlasClassification> classifications =
getEntity(entityGuid).getClassifications();
+
+ if (CollectionUtils.isNotEmpty(classifications)) {
+ for (AtlasClassification c : classifications) {
+ if (c.getTypeName().equals(classification.getTypeName())
+ &&
c.getEntityGuid().equals(classification.getEntityGuid())) {
+ fail("Propagated classification should not be present in
entity!");
+ }
+ }
+ }
+ }
}
