Add check to restrict moving a category with terms linked to multiple…

… categories
atlanhq · Dec 20, 2023 · 3063d00 · 3063d00
1 parent 6b3a97b
commit 3063d00
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 15 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,8 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Iterators;
 import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.AtlasException;
@@ -389,6 +391,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
@@ -47,7 +47,14 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+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;
@@ -121,8 +128,8 @@ private void processCreateCategory(AtlasEntity entity, AtlasVertex vertex) throw
         validateParent(glossaryQualifiedName);
 
         entity.setAttribute(QUALIFIED_NAME, createQualifiedName(vertex));
-//        AtlasAuthorizationUtils.verifyAccess(new AtlasEntityAccessRequest(typeRegistry, AtlasPrivilege.ENTITY_CREATE, new AtlasEntityHeader(entity)),
-//                "create entity: type=", entity.getTypeName());
+        AtlasAuthorizationUtils.verifyAccess(new AtlasEntityAccessRequest(typeRegistry, AtlasPrivilege.ENTITY_CREATE, new AtlasEntityHeader(entity)),
+                "create entity: type=", entity.getTypeName());
 
         validateChildren(entity, null);
 
@@ -267,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);
@@ -303,17 +311,13 @@ public void moveChildTermToAnotherGlossary(AtlasVertex termVertex,
         }
     }
 
-    private void ensureOnlyOneCategoryIsAssociated(AtlasEntity entity) throws AtlasBaseException {
-        if(entity.hasRelationshipAttribute(ATTR_CATEGORIES) && Objects.nonNull(entity.getRelationshipAttribute(ATTR_CATEGORIES))) {
-            List<AtlasObjectId> categories = (List<AtlasObjectId>) entity.getRelationshipAttribute(ATTR_CATEGORIES);
+    private void ensureOnlyOneCategoryIsAssociated(AtlasVertex vertex) throws AtlasBaseException {
+        final Integer numOfCategoryEdges = GraphHelper.getCountOfCategoryEdges(vertex);
 
-            if(categories.size() > 1) {
+        if(numOfCategoryEdges>1) {
                 throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST, "Cannot move term with multiple " +
                         "categories associated to another glossary");
-            }
-
         }
-
     }
 
     private void validateParentForGlossaryChange(AtlasEntity category,

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