Merge pull request #2888 from atlanhq/PLT-1036

Handled the edge case(propagate:True,restrictLineage:True,restrictHie…
atlanhq · Mar 22, 2024 · b190e15 · b190e15
2 parents 219a123 + e0f74ab
commit b190e15
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 11 deletions.
diff --git a/...itory/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityComparator.java b/...itory/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityComparator.java
@@ -60,6 +60,12 @@ public AtlasEntityDiffResult getDiffResult(AtlasEntity updatedEntity, AtlasVerte
         return getDiffResult(updatedEntity, null, storedVertex, findOnlyFirstDiff);
     }
 
+    private void verifyClassificationsPropagationMode(List<AtlasClassification> incomingClassifications) throws AtlasBaseException {
+        for(AtlasClassification incomingClassification : incomingClassifications){
+            entityRetriever.determinePropagationMode(incomingClassification.getRestrictPropagationThroughLineage(),incomingClassification.getRestrictPropagationThroughHierachy());
+        }
+    }
+
     private AtlasEntityDiffResult getDiffResult(AtlasEntity updatedEntity, AtlasEntity storedEntity, AtlasVertex storedVertex, boolean findOnlyFirstDiff) throws AtlasBaseException {
         AtlasEntity                              diffEntity                       = new AtlasEntity(updatedEntity.getTypeName());
         AtlasEntityType                          entityType                       = typeRegistry.getEntityTypeByName(updatedEntity.getTypeName());

diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2.java
@@ -1632,7 +1632,7 @@ private EntityMutationContext preCreateOrUpdate(EntityStream entityStream, Entit
         for (Map.Entry<String, String> element : referencedGuids.entrySet()) {
             String guid = element.getKey();
             AtlasEntity entity = entityStream.getByGuid(guid);
-
+            entityRetriever.verifyClassificationsPropagationMode(entity.getClassifications());
             if (entity != null) { // entity would be null if guid is not in the stream but referenced by an entity in the stream
                 AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
 

diff --git a/...ry/src/main/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociator.java b/...ry/src/main/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociator.java
@@ -233,7 +233,7 @@ private Map<String, List<AtlasClassification>> computeChanges(AtlasEntityHeader
             List<AtlasClassification> incomingClassifications = listOps.filter(incomingEntityHeader.getGuid(), incomingEntityHeader.getClassifications());
             List<AtlasClassification> entityClassifications = listOps.filter(entityToBeUpdated.getGuid(), entityToBeUpdated.getClassifications());
 
-            verifyClassificationsPropagationMode(incomingClassifications);
+            entityRetriever.verifyClassificationsPropagationMode(incomingClassifications);
 
             if (CollectionUtils.isEmpty(incomingClassifications) && CollectionUtils.isEmpty(entityClassifications)) {
                 return null;
@@ -248,15 +248,7 @@ private Map<String, List<AtlasClassification>> computeChanges(AtlasEntityHeader
             return operationListMap;
         }
 
-        /**
-         * Checks for if the AtlasClassification has valid config of restrict flags
-         * Both Restrict flags can't be true with propagate flag allowed
-         */
-        private void verifyClassificationsPropagationMode(List<AtlasClassification> incomingClassifications) throws AtlasBaseException {
-            for(AtlasClassification incomingClassification : incomingClassifications){
-                entityRetriever.determinePropagationMode(incomingClassification.getRestrictPropagationThroughLineage(),incomingClassification.getRestrictPropagationThroughHierachy());
-            }
-        }
+
 
         private void bucket(String op, Map<String, List<AtlasClassification>> operationListMap, List<AtlasClassification> results) {
             if (CollectionUtils.isEmpty(results)) {

diff --git a/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -546,6 +546,16 @@ public Map<AtlasVertex, List<AtlasVertex>> getClassificationPropagatedEntitiesMa
         RequestContext.get().endMetricRecord(metricRecorder);
         return ret;
     }
+    /**
+     * Checks for if the AtlasClassification has valid config of restrict flags
+     * Both Restrict flags can't be true with propagate flag allowed
+     */
+    public void verifyClassificationsPropagationMode(List<AtlasClassification> incomingClassifications) throws AtlasBaseException {
+        for(AtlasClassification incomingClassification : incomingClassifications){
+            if(Boolean.TRUE.equals(incomingClassification.isPropagate()))
+                determinePropagationMode(incomingClassification.getRestrictPropagationThroughLineage(),incomingClassification.getRestrictPropagationThroughHierachy());
+        }
+    }
 
     public String determinePropagationMode(Boolean currentRestrictPropagationThroughLineage, Boolean currentRestrictPropagationThroughHierarchy) throws AtlasBaseException {
         String propagationMode;