Merge pull request #3979 from atlanhq/dg1999

DG-1999 Optimize authoriseRemoveRelation Latency
atlanhq · Jan 16, 2025 · 92718c4 · 92718c4
2 parents 396fb19 + 7a8add0
commit 92718c4
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 11 deletions.
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
@@ -514,21 +514,26 @@ public List<AtlasVertex> addTagPropagation(AtlasVertex classificationVertex, Lis
 
     public void authorizeRemoveRelation(AtlasEdge edge) throws AtlasBaseException {
         AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("authoriseRemoveRelation");
-        AtlasEntityHeader end1Entity, end2Entity;
-        String relationShipType = getTypeName(edge);
-        AtlasRelationshipDef relationshipDef = typeRegistry.getRelationshipDefByName(relationShipType);
-        if (relationshipDef == null) {
-            return;
-        }
-
-        end1Entity = entityRetriever.toAtlasEntityHeaderWithClassifications(edge.getOutVertex());
-        end2Entity = entityRetriever.toAtlasEntityHeaderWithClassifications(edge.getInVertex());
+        if(!RequestContext.get().isAuthorisedRemoveRelation()) {
+            if (isRequestFromWorkFlow()) {
+                RequestContext.get().setAuthorisedRemoveRelation(true);
+            }
+            AtlasEntityHeader end1Entity, end2Entity;
+            String relationShipType = getTypeName(edge);
+            AtlasRelationshipDef relationshipDef = typeRegistry.getRelationshipDefByName(relationShipType);
+            if (relationshipDef == null) {
+                return;
+            }
 
-        AtlasAuthorizationUtils.verifyAccess(new AtlasRelationshipAccessRequest(typeRegistry, AtlasPrivilege.RELATIONSHIP_REMOVE, relationShipType, end1Entity, end2Entity ));
+            end1Entity = entityRetriever.toAtlasEntityHeaderWithClassifications(edge.getOutVertex());
+            end2Entity = entityRetriever.toAtlasEntityHeaderWithClassifications(edge.getInVertex());
 
+            AtlasAuthorizationUtils.verifyAccess(new AtlasRelationshipAccessRequest(typeRegistry, AtlasPrivilege.RELATIONSHIP_REMOVE, relationShipType, end1Entity, end2Entity));
+        }
         RequestContext.get().endMetricRecord(metric);
     }
 
+
     public Map<AtlasVertex, List<AtlasVertex>> removeTagPropagation(AtlasEdge edge) throws AtlasBaseException {
         AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("removeTagPropagationEdge");
 
@@ -1568,6 +1573,14 @@ public void resetHasLineageOnInputOutputDelete(Collection<AtlasEdge> removedEdge
         }
         RequestContext.get().endMetricRecord(metricRecorder);
     }
+    private boolean isRequestFromWorkFlow() {
+        String workflowID = RequestContext.get().getRequestContextHeaders().getOrDefault("x-atlan-agent-workflow-id", "");
+        boolean isWorkFlowRequest = !workflowID.isEmpty();
+        if(isWorkFlowRequest){
+            LOG.info("Authorised one time request for workflow with id : {} ", workflowID);
+        }
+        return isWorkFlowRequest;
+    }
 
     private String getLabel(String guid, String label){
        return  guid + ":" + label;

diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/HardDeleteHandlerV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/HardDeleteHandlerV1.java
@@ -18,6 +18,7 @@
 
 package org.apache.atlas.repository.store.graph.v1;
 
+import org.apache.atlas.RequestContext;
 import org.apache.atlas.annotation.ConditionalOnAtlasProperty;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.repository.graph.GraphHelper;

diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/SoftDeleteHandlerV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/SoftDeleteHandlerV1.java
@@ -73,9 +73,9 @@ protected void deleteEdge(AtlasEdge edge, boolean force) throws AtlasBaseExcepti
                 LOG.debug("==> SoftDeleteHandlerV1.deleteEdge({}, {})", GraphHelper.string(edge), force);
             }
             boolean isRelationshipEdge = isRelationshipEdge(edge);
-
             authorizeRemoveRelation(edge);
 
+
             if (DEFERRED_ACTION_ENABLED && RequestContext.get().getCurrentTask() == null) {
                 if (CollectionUtils.isNotEmpty(getPropagatableClassifications(edge))) {
                     RequestContext.get().addToDeletedEdgesIds(edge.getIdForDisplay());

diff --git a/server-api/src/main/java/org/apache/atlas/RequestContext.java b/server-api/src/main/java/org/apache/atlas/RequestContext.java
@@ -88,6 +88,9 @@ public class RequestContext {
     private int attemptCount = 1;
     private boolean isImportInProgress = false;
     private boolean     isInNotificationProcessing = false;
+
+
+    private boolean     authorisedRemoveRelation = false;
     private boolean     isInTypePatching           = false;
     private boolean     createShellEntityForNonExistingReference = false;
     private boolean     skipFailedEntities = false;
@@ -205,6 +208,13 @@ public void clearEntityCache() {
         this.entityCache.clear();
     }
 
+    public boolean isAuthorisedRemoveRelation() {
+        return authorisedRemoveRelation;
+    }
+
+    public void setAuthorisedRemoveRelation(boolean authorisedRemoveRelation) {
+        this.authorisedRemoveRelation = authorisedRemoveRelation;
+    }
     public Set<String> getRelationAttrsForSearch() {
         return relationAttrsForSearch;
     }