Merge pull request #2680 from atlanhq/PLT-408

check to restrict moving categories
atlanhq · Dec 20, 2023 · 699bfef · 699bfef
2 parents 145ec39 + 8e6d9c5
commit 699bfef
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 5 deletions.
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
@@ -21,6 +21,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
+import com.google.common.collect.Iterators;
 import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.AtlasException;
@@ -389,6 +390,17 @@ public static AtlasEdge getClassificationEdge(AtlasVertex entityVertex, AtlasVer
         return ret;
     }
 
+    public static Integer getCountOfCategoryEdges(AtlasVertex entityVertex) {
+
+        Iterator<AtlasEdge> edges = getOutGoingEdgesByLabel(entityVertex, CATEGORY_TERMS_EDGE_LABEL);
+
+        if (edges!=null) {
+            return Iterators.size(edges);
+        }
+
+        return 0;
+    }
+
     public static boolean isClassificationAttached(AtlasVertex entityVertex, AtlasVertex classificationVertex) {
         AtlasPerfMetrics.MetricRecorder isClassificationAttachedMetricRecorder  = RequestContext.get().startMetricRecord("isClassificationAttached");
         String                          classificationId                        = classificationVertex.getIdForDisplay();

diff --git a/...rg/apache/atlas/repository/store/graph/v2/preprocessor/glossary/CategoryPreProcessor.java b/...rg/apache/atlas/repository/store/graph/v2/preprocessor/glossary/CategoryPreProcessor.java
@@ -54,6 +54,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.Objects;
 import java.util.stream.Collectors;
 
 import static org.apache.atlas.AtlasErrorCode.BAD_REQUEST;
@@ -273,6 +274,7 @@ public void moveChildTermToAnotherGlossary(AtlasVertex termVertex,
 
             //check duplicate term name
             termExists(termName, targetGlossaryQualifiedName);
+            ensureOnlyOneCategoryIsAssociated(termVertex);
 
             String currentTermQualifiedName = termVertex.getProperty(QUALIFIED_NAME, String.class);
             String updatedTermQualifiedName = currentTermQualifiedName.replace(sourceGlossaryQualifiedName, targetGlossaryQualifiedName);
@@ -309,6 +311,15 @@ public void moveChildTermToAnotherGlossary(AtlasVertex termVertex,
         }
     }
 
+    private void ensureOnlyOneCategoryIsAssociated(AtlasVertex vertex) throws AtlasBaseException {
+        final Integer numOfCategoryEdges = GraphHelper.getCountOfCategoryEdges(vertex);
+
+        if(numOfCategoryEdges>1) {
+                throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST, "Cannot move term with multiple " +
+                        "categories associated to another glossary");
+        }
+    }
+
     private void validateParentForGlossaryChange(AtlasEntity category,
                                                  AtlasVertex categoryVertex,
                                                  String targetGlossaryQualifiedName) throws AtlasBaseException {

diff --git a/...va/org/apache/atlas/repository/store/graph/v2/preprocessor/glossary/TermPreProcessor.java b/...va/org/apache/atlas/repository/store/graph/v2/preprocessor/glossary/TermPreProcessor.java
@@ -115,8 +115,6 @@ private void processUpdateTerm(AtlasEntity entity, AtlasVertex vertex) throws At
             throw new AtlasBaseException(AtlasErrorCode.INVALID_DISPLAY_NAME);
         }
 
-        validateCategory(entity);
-
         AtlasEntity storedTerm = entityRetriever.toAtlasEntity(vertex);
         AtlasRelatedObjectId currentGlossary = (AtlasRelatedObjectId) storedTerm.getRelationshipAttribute(ANCHOR);
         AtlasEntityHeader currentGlossaryHeader = entityRetriever.toAtlasEntityHeader(currentGlossary.getGuid());
@@ -126,12 +124,16 @@ private void processUpdateTerm(AtlasEntity entity, AtlasVertex vertex) throws At
 
         String newGlossaryQualifiedName = (String) anchor.getAttribute(QUALIFIED_NAME);
 
+        if(!currentGlossaryQualifiedName.equals(newGlossaryQualifiedName)) {
+            ensureOnlyOneCategoryIsAssociated(entity);
+        }
+
+        validateCategory(entity);
+
         if (!currentGlossaryQualifiedName.equals(newGlossaryQualifiedName)){
             //Auth check
             isAuthorized(currentGlossaryHeader, anchor);
 
-            ensureOnlyOneCategoryIsAssociated(entity);
-
             String updatedTermQualifiedName = moveTermToAnotherGlossary(entity, vertex, currentGlossaryQualifiedName, newGlossaryQualifiedName, termQualifiedName);
 
             if (checkEntityTermAssociation(termQualifiedName)) {
@@ -177,7 +179,6 @@ private static void ensureOnlyOneCategoryIsAssociated(AtlasEntity entity) throws
 
     private void validateCategory(AtlasEntity entity) throws AtlasBaseException {
         String glossaryQualifiedName = (String) anchor.getAttribute(QUALIFIED_NAME);
-
         if (entity.hasRelationshipAttribute(ATTR_CATEGORIES) && entity.getRelationshipAttribute(ATTR_CATEGORIES) != null) {
             List<AtlasObjectId> categories = (List<AtlasObjectId>) entity.getRelationshipAttribute(ATTR_CATEGORIES);