DG-1924 | Add 'assetsCountToPropagate' and 'assetsCountPropagated' in task vertex.

.github/workflows/maven.yml

@@ -28,6 +28,7 @@ on:
       - master
       - lineageondemand
       - makerlogic
+      - tagpropv1master
 
 jobs:
   build:
@@ -66,8 +67,15 @@ jobs:
 
       - name: Build with Maven
         run: |
-          echo "build without dashboard"
-          chmod +x ./build.sh && ./build.sh
+          branch_name=${{ env.BRANCH_NAME }}
+          if [[ $branch_name == 'main' || $branch_name == 'master' || $branch_name == 'tagpropv1master' ]]
+          then
+              echo "build without dashboard"
+              chmod +x ./build.sh && ./build.sh build_without_dashboard
+          else
+              echo "build with dashboard"
+              chmod +x ./build.sh && ./build.sh
+          fi
 
       - name: Get Repository Name
         run:   echo "REPOSITORY_NAME=`echo "$GITHUB_REPOSITORY" | awk -F / '{print $2}' | sed -e "s/:refs//"`" >> $GITHUB_ENV

common/src/main/java/org/apache/atlas/repository/Constants.java

@@ -364,6 +364,8 @@ public final class Constants {
     public static final String TASK_CLASSIFICATION_ID       = encodePropertyKey(TASK_PREFIX + "classificationId");
     public static final String TASK_ENTITY_GUID             = encodePropertyKey(TASK_PREFIX + "entityGuid");
     public static final String TASK_CLASSIFICATION_TYPENAME = encodePropertyKey(TASK_PREFIX + "classificationTypeName");
+    public static final String TASK_ASSET_COUNT_TO_PROPAGATE = encodePropertyKey(TASK_PREFIX + "assetsCountToPropagate");
+    public static final String TASK_ASSET_COUNT_PROPAGATED = encodePropertyKey(TASK_PREFIX + "assetsCountPropagated");
     public static final String ACTIVE_STATE_VALUE           = "ACTIVE";
     public static final String TASK_HEADER_ATLAN_AGENT      = "x-atlan-agent";
     public static final String TASK_HEADER_ATLAN_AGENT_ID   = "x-atlan-agent-id";

intg/src/main/java/org/apache/atlas/model/tasks/AtlasTask.java

@@ -95,6 +95,8 @@ public static Status from(String s) {
     private String              classificationId;
     private String              entityGuid;
     private String              classificationTypeName;
+    private Long                assetsCountToPropagate;
+    private Long                assetsCountPropagated;
 
     public AtlasTask() {
     }
@@ -111,6 +113,8 @@ public AtlasTask(String type, String createdBy, Map<String, Object> parameters,
         this.attemptCount       = 0;
         this.classificationId   = classificationId;
         this.entityGuid         = entityGuid;
+        this.assetsCountToPropagate = 0L;
+        this.assetsCountPropagated = 0L;
     }
 
     public String getGuid() {
@@ -239,6 +243,22 @@ public String getEntityGuid() {
         return entityGuid;
     }
 
+    public void setAssetsCountToPropagate(Long assetsCount) {
+        this.assetsCountToPropagate = assetsCount;
+    }
+
+    public Long getAssetsCountToPropagate() {
+        return assetsCountToPropagate;
+    }
+
+    public void setAssetsCountPropagated(Long assetsCountPropagated) {
+        this.assetsCountPropagated = assetsCountPropagated;
+    }
+
+    public Long getAssetsCountPropagated(){
+        return  assetsCountPropagated;
+    }
+
     @JsonIgnore
     public void start() {
         this.setStatus(Status.IN_PROGRESS);
@@ -270,6 +290,8 @@ public String toString() {
                 ", attemptCount=" + attemptCount +
                 ", errorMessage='" + errorMessage + '\'' +
                 ", status=" + status +
+                ", assetsCountToPropagate=" + assetsCountToPropagate +
+                ", assetsCountPropagated=" + assetsCountPropagated +
                 '}';
     }
 }

repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedSearchIndexer.java

@@ -407,6 +407,8 @@ private void initialize(AtlasGraph graph) throws RepositoryException, IndexExcep
             createCommonVertexIndex(management, TASK_ENTITY_GUID, UniqueKind.NONE, String.class, SINGLE, false, false, true);
             createCommonVertexIndex(management, TASK_ERROR_MESSAGE, UniqueKind.NONE, String.class, SINGLE, false, false);
             createCommonVertexIndex(management, TASK_ATTEMPT_COUNT, UniqueKind.NONE, Integer.class, SINGLE, false, false);
+            createCommonVertexIndex(management, TASK_ASSET_COUNT_TO_PROPAGATE, UniqueKind.NONE, Long.class, SINGLE, false, false);
+            createCommonVertexIndex(management, TASK_ASSET_COUNT_PROPAGATED, UniqueKind.NONE, Long.class, SINGLE, false, false);
 
             createCommonVertexIndex(management, TASK_UPDATED_TIME, UniqueKind.NONE, Long.class, SINGLE, false, false);
             createCommonVertexIndex(management, TASK_TIME_TAKEN_IN_SECONDS, UniqueKind.NONE, Long.class, SINGLE, false, false);

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java

@@ -58,6 +58,7 @@
 import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import static org.apache.atlas.repository.graph.GraphHelper.getTypeName;
 
 import java.util.*;
 import java.util.stream.Collectors;
@@ -95,6 +96,7 @@ public abstract class DeleteHandlerV1 {
     private   final TaskManagement       taskManagement;
     private   final AtlasGraph           graph;
     private   final TaskUtil             taskUtil;
+    private static final int CHUNK_SIZE            = AtlasConfiguration.TASKS_GRAPH_COMMIT_CHUNK_SIZE.getInt();
 
     public DeleteHandlerV1(AtlasGraph graph, AtlasTypeRegistry typeRegistry, boolean shouldUpdateInverseReference, boolean softDelete, TaskManagement taskManagement) {
         this.typeRegistry                  = typeRegistry;
@@ -1235,16 +1237,32 @@ public void updateTagPropagations(AtlasEdge edge, AtlasRelationship relationship
                 }
             }
 
+            taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_TO_PROPAGATE, graph, addPropagationsMap.size() + removePropagationsMap.size());
+
+            int propagatedCount = 0;
             for (AtlasVertex classificationVertex : addPropagationsMap.keySet()) {
                 List<AtlasVertex> entitiesToAddPropagation = addPropagationsMap.get(classificationVertex);
 
                 addTagPropagation(classificationVertex, entitiesToAddPropagation);
+                propagatedCount++;
+                if (propagatedCount == CHUNK_SIZE){
+                    taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_PROPAGATED, graph, propagatedCount - 1);
+                    propagatedCount = 0;
+                }
             }
 
             for (AtlasVertex classificationVertex : removePropagationsMap.keySet()) {
                 List<AtlasVertex> entitiesToRemovePropagation = removePropagationsMap.get(classificationVertex);
 
                 removeTagPropagation(classificationVertex, entitiesToRemovePropagation);
+                propagatedCount++;
+                if (propagatedCount == CHUNK_SIZE){
+                    taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_PROPAGATED, graph, propagatedCount);
+                    propagatedCount = 0;
+                }
+            }
+            if (propagatedCount != 0){
+                taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_PROPAGATED, graph, propagatedCount);
             }
         } else {
             // update blocked propagated classifications only if there is no change is tag propagation (don't update both)

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java

@@ -19,6 +19,7 @@
 
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Iterators;
 import org.apache.atlas.*;
 import org.apache.atlas.annotation.GraphTransaction;
 import org.apache.atlas.authorize.AtlasAuthorizationUtils;
@@ -78,6 +79,7 @@
 
 import javax.inject.Inject;
 import java.util.*;
+import java.util.function.BiConsumer;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -94,6 +96,7 @@
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.PARTIAL_UPDATE;
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.UPDATE;
 import static org.apache.atlas.model.tasks.AtlasTask.Status.IN_PROGRESS;
+import static org.apache.atlas.model.tasks.AtlasTask.Status.PENDING;
 import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef.Cardinality.SET;
 import static org.apache.atlas.repository.Constants.*;
 import static org.apache.atlas.repository.graph.GraphHelper.getClassificationEdge;
@@ -3158,9 +3161,10 @@ public void cleanUpClassificationPropagation(String classificationName, int batc
         long classificationEdgeCount = 0;
         long classificationEdgeInMemoryCount = 0;
         Iterator<AtlasVertex> tagVertices = GraphHelper.getClassificationVertices(graph, classificationName, CLEANUP_BATCH_SIZE);
+
         List<AtlasVertex> tagVerticesProcessed = new ArrayList<>(0);
         List<AtlasVertex> currentAssetVerticesBatch = new ArrayList<>(0);
-
+        int totalCount = 0;
         while (tagVertices != null && tagVertices.hasNext()) {
             if (cleanedUpCount >= CLEANUP_MAX){
                 return;
@@ -3179,6 +3183,8 @@ public void cleanUpClassificationPropagation(String classificationName, int batc
             }
 
             int currentAssetsBatchSize = currentAssetVerticesBatch.size();
+            totalCount += currentAssetsBatchSize;
+
             if (currentAssetsBatchSize > 0) {
                 LOG.info("To clean up tag {} from {} entities", classificationName, currentAssetsBatchSize);
                 int offset = 0;
@@ -3210,17 +3216,20 @@ public void cleanUpClassificationPropagation(String classificationName, int batc
                                     classificationEdgeInMemoryCount = 0;
                                 }
                             }
+
                             try {
                                 AtlasEntity entity = repairClassificationMappings(vertex);
                                 entityChangeNotifier.onClassificationDeletedFromEntity(entity, deletedClassifications);
                             } catch (IllegalStateException | AtlasBaseException e) {
                                 e.printStackTrace();
                             }
+
                         }
 
                         transactionInterceptHelper.intercept();
-
                         offset += CHUNK_SIZE;
+
+
                     } finally {
                         LOG.info("For offset {} , classificationEdge were : {}", offset, classificationEdgeCount);
                         classificationEdgeCount = 0;
@@ -3236,7 +3245,6 @@ public void cleanUpClassificationPropagation(String classificationName, int batc
                         e.printStackTrace();
                     }
                 }
-                transactionInterceptHelper.intercept();
 
                 cleanedUpCount += currentAssetsBatchSize;
                 currentAssetVerticesBatch.clear();
@@ -3245,6 +3253,10 @@ public void cleanUpClassificationPropagation(String classificationName, int batc
             tagVertices = GraphHelper.getClassificationVertices(graph, classificationName, CLEANUP_BATCH_SIZE);
         }
 
+        taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_TO_PROPAGATE, graph, totalCount);
+        taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_PROPAGATED, graph, totalCount);
+
+        transactionInterceptHelper.intercept();
         LOG.info("Completed cleaning up classification {}", classificationName);
     }
 
@@ -3438,6 +3450,7 @@ public void addClassifications(final EntityMutationContext context, String guid,
 
     public List<String> propagateClassification(String entityGuid, String classificationVertexId, String relationshipGuid, Boolean previousRestrictPropagationThroughLineage,Boolean previousRestrictPropagationThroughHierarchy) throws AtlasBaseException {
         try {
+
             if (StringUtils.isEmpty(entityGuid) || StringUtils.isEmpty(classificationVertexId)) {
                 LOG.error("propagateClassification(entityGuid={}, classificationVertexId={}): entityGuid and/or classification vertex id is empty", entityGuid, classificationVertexId);
 
@@ -3481,9 +3494,10 @@ public List<String> propagateClassification(String entityGuid, String classifica
             Boolean toExclude = propagationMode == CLASSIFICATION_PROPAGATION_MODE_RESTRICT_LINEAGE ? true:false;
             List<AtlasVertex> impactedVertices = entityRetriever.getIncludedImpactedVerticesV2(entityVertex, relationshipGuid, classificationVertexId, edgeLabelsToCheck,toExclude);
 
+            taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_TO_PROPAGATE, graph, impactedVertices.size() - 1);
+
             if (CollectionUtils.isEmpty(impactedVertices)) {
                 LOG.debug("propagateClassification(entityGuid={}, classificationVertexId={}): found no entities to propagate the classification", entityGuid, classificationVertexId);
-
                 return null;
             }
 
@@ -3526,9 +3540,15 @@ public List<String> processClassificationPropagationAddition(List<AtlasVertex> v
 
                 propagatedEntitiesGuids.addAll(chunkedPropagatedEntitiesGuids);
 
+                transactionInterceptHelper.intercept();
+
+                int propagatedAssetsCount = (offset + CHUNK_SIZE >= impactedVerticesSize && impactedVerticesSize == verticesToPropagate.size())
+                        ? toIndex - offset - 1 // Subtract 1 for the last chunk
+                        : toIndex - offset;
+
                 offset += CHUNK_SIZE;
 
-                transactionInterceptHelper.intercept();
+                taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_PROPAGATED, graph, propagatedAssetsCount);
 
             } while (offset < impactedVerticesSize);
         } catch (AtlasBaseException exception) {
@@ -3538,7 +3558,7 @@ public List<String> processClassificationPropagationAddition(List<AtlasVertex> v
             RequestContext.get().endMetricRecord(classificationPropagationMetricRecorder);
         }
 
-    return propagatedEntitiesGuids;
+        return propagatedEntitiesGuids;
 
     }
 
@@ -4072,6 +4092,9 @@ public void updateClassificationTextPropagation(String classificationVertexId) t
         AtlasClassification classification = entityRetriever.toAtlasClassification(classificationVertex);
         LOG.info("Fetched classification : {} ", classification.toString());
         List<AtlasVertex> impactedVertices = graphHelper.getAllPropagatedEntityVertices(classificationVertex);
+
+        taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_TO_PROPAGATE, graph, impactedVertices.size());
+
         LOG.info("impactedVertices : {}", impactedVertices.size());
         int batchSize = 100;
         for (int i = 0; i < impactedVertices.size(); i += batchSize) {
@@ -4086,6 +4109,9 @@ public void updateClassificationTextPropagation(String classificationVertexId) t
                     entityChangeNotifier.onClassificationUpdatedToEntity(entity, Collections.singletonList(classification));
                 }
             }
+
+            taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_PROPAGATED, graph, end);
+
             transactionInterceptHelper.intercept();
             LOG.info("Updated classificationText from {} for {}", i, batchSize);
         }
@@ -4276,6 +4302,8 @@ public void classificationRefreshPropagation(String classificationId) throws Atl
                 .filter(vertex -> vertex != null)
                 .collect(Collectors.toList());
 
+        taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_TO_PROPAGATE, graph, verticesToRemove.size() + verticesToAddClassification.size());
+
         //Remove classifications from unreachable vertices
         processPropagatedClassificationDeletionFromVertices(verticesToRemove, currentClassificationVertex, classification);
 
@@ -4353,8 +4381,9 @@ private void processPropagatedClassificationDeletionFromVertices(List<AtlasVerte
                 List<AtlasEntity> updatedEntities = updateClassificationText(classification, updatedVertices);
                 entityChangeNotifier.onClassificationsDeletedFromEntities(updatedEntities, Collections.singletonList(classification));
 
+                int propagatedAssetsCount = toIndex - offset;
                 offset += CHUNK_SIZE;
-
+                taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_PROPAGATED, graph, propagatedAssetsCount);
                 transactionInterceptHelper.intercept();
 
             } while (offset < propagatedVerticesSize);
@@ -4372,6 +4401,8 @@ List<String> processClassificationEdgeDeletionInChunk(AtlasClassification classi
         int toIndex;
         int offset = 0;
 
+        taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_TO_PROPAGATE, graph, propagatedEdgesSize);
+
         do {
             toIndex = ((offset + CHUNK_SIZE > propagatedEdgesSize) ? propagatedEdgesSize : (offset + CHUNK_SIZE));
 
@@ -4387,8 +4418,11 @@ List<String> processClassificationEdgeDeletionInChunk(AtlasClassification classi
                 deletedPropagationsGuid.addAll(propagatedEntities.stream().map(x -> x.getGuid()).collect(Collectors.toList()));
             }
 
+            int propagatedAssetsCount = toIndex - offset;
+
             offset += CHUNK_SIZE;
 
+            taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_PROPAGATED, graph, propagatedAssetsCount);
             transactionInterceptHelper.intercept();
 
         } while (offset < propagatedEdgesSize);

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/tasks/ClassificationTask.java

@@ -38,6 +38,7 @@
 import java.util.Map;
 
 import static org.apache.atlas.model.tasks.AtlasTask.Status.*;
+
 import static org.apache.atlas.repository.store.graph.v2.tasks.ClassificationPropagateTaskFactory.CLASSIFICATION_PROPAGATION_RELATIONSHIP_UPDATE;
 
 public abstract class ClassificationTask extends AbstractTask {
@@ -97,24 +98,18 @@ public AtlasTask.Status perform() throws AtlasBaseException {
         }
 
         RequestContext.get().setUser(userName, null);
-
         try {
             setStatus(IN_PROGRESS);
-
             run(params);
-
             setStatus(COMPLETE);
         } catch (AtlasBaseException e) {
             LOG.error("Task: {}: Error performing task!", getTaskGuid(), e);
-
             setStatus(FAILED);
-
             throw e;
         } finally {
             RequestContext.get().endMetricRecord(metricRecorder);
             graph.commit();
         }
-
         return getStatus();
     }
 

repository/src/main/java/org/apache/atlas/tasks/AbstractTask.java

@@ -71,5 +71,12 @@ protected AtlasTask getTaskDef() {
         return this.task;
     }
 
+    protected void setAssetsCountToPropagate(Long assetsCount) {
+        task.setAssetsCountToPropagate(assetsCount);;
+    }
+
+    public Long getAssetsCountToPropagate() {
+        return this.task.getAssetsCountToPropagate();
+    }
     public abstract Status perform() throws Exception;
 }