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

.github/workflows/maven.yml

@@ -28,6 +28,7 @@ on:
       - master
       - lineageondemand
       - makerlogic
+      - tagpropv1master
 
 jobs:
   build:
@@ -67,7 +68,7 @@ jobs:
       - name: Build with Maven
         run: |
           branch_name=${{ env.BRANCH_NAME }}
-          if [[ $branch_name == 'main' || $branch_name == 'master' || $branch_name == 'lineageondemand' ]]
+          if [[ $branch_name == 'main' || $branch_name == 'master' || $branch_name == 'tagpropv1master' ]]
           then
               echo "build without dashboard"
               chmod +x ./build.sh && ./build.sh build_without_dashboard

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

@@ -359,6 +359,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/patches/AtlasPatch.java

@@ -49,13 +49,15 @@ public class AtlasPatch implements Serializable {
     private long        createdTime;
     private long        updatedTime;
     private PatchStatus status;
+    private long                assetsCountToPropagate;
+    private long                assetsCountPropagated;
 
     public enum PatchStatus { UNKNOWN, APPLIED, SKIPPED, FAILED }
 
     public AtlasPatch() { }
 
     public AtlasPatch(String id, String patchName, String type, String action, PatchStatus status,
-                      String updatedBy, String createdBy, long createdTime, long updatedTime) {
+                      String updatedBy, String createdBy, long createdTime, long updatedTime, long assetsCountToPropagate, long assetsCountPropagated) {
         this.id          = id;
         this.description = patchName;
         this.type        = type;
@@ -65,6 +67,8 @@ public AtlasPatch(String id, String patchName, String type, String action, Patch
         this.createdBy   = createdBy;
         this.createdTime = createdTime;
         this.updatedTime = updatedTime;
+        this.assetsCountToPropagate = assetsCountToPropagate;
+        this.assetsCountPropagated = assetsCountPropagated;
     }
 
     public String getId() {
@@ -139,6 +143,18 @@ public void setUpdatedTime(long updatedTime) {
         this.updatedTime = updatedTime;
     }
 
+    public void setAssetsCountToPropagate(Long assetsCount) {
+        this.assetsCountToPropagate = assetsCount;
+    }
+
+    public Long getAssetsCountToPropagate() {
+        return assetsCountToPropagate;
+    }
+
+    public Long getAssetsCountPropagated(){
+        return  assetsCountPropagated;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -157,7 +173,7 @@ public boolean equals(Object o) {
 
     @Override
     public int hashCode() {
-        return Objects.hash(id, description, type, action, updatedBy, createdBy, createdTime, updatedTime, status);
+        return Objects.hash(id, description, type, action, updatedBy, createdBy, createdTime, updatedTime, status, assetsCountToPropagate, assetsCountPropagated);
     }
 
     @Override
@@ -173,6 +189,8 @@ public String toString() {
         sb.append(", createdTime=").append(createdTime);
         sb.append(", updatedTime=").append(updatedTime);
         sb.append(", status=").append(status);
+        sb.append(", assetsCountToPropagate=").append(assetsCountToPropagate);
+        sb.append(", assetsCountPropagated=").append(assetsCountPropagated);
         sb.append('}');
 
         return sb.toString();

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/patches/AtlasPatchRegistry.java

@@ -210,6 +210,7 @@ private static AtlasPatch toAtlasPatch(AtlasVertex vertex) {
         ret.setUpdatedBy(getEncodedProperty(vertex, MODIFIED_BY_KEY, String.class));
         ret.setCreatedTime(getEncodedProperty(vertex, TIMESTAMP_PROPERTY_KEY, Long.class));
         ret.setUpdatedTime(getEncodedProperty(vertex, MODIFICATION_TIMESTAMP_PROPERTY_KEY, Long.class));
+        ret.setAssetsCountToPropagate(getEncodedProperty(vertex, TASK_ASSET_COUNT_TO_PROPAGATE, Long.class));
         ret.setStatus(getPatchStatus(vertex));
 
         return ret;

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

@@ -60,8 +60,8 @@
 import org.slf4j.LoggerFactory;
 import static org.apache.atlas.repository.graph.GraphHelper.getTypeName;
 
-
 import java.util.*;
+import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
 import static org.apache.atlas.AtlasClient.DATA_SET_SUPER_TYPE;
@@ -96,7 +96,22 @@ 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 void updateTaskVertexProperty(String propertyKey, long value, boolean isIncremental, BiConsumer<AtlasTask, Long> taskSetter) {
+        AtlasTask currentTask = RequestContext.get().getCurrentTask();
+        AtlasVertex currentTaskVertex = (AtlasVertex) graph.query()
+                .has(TASK_GUID, currentTask.getGuid())
+                .vertices().iterator().next();
+
+        Long valueFromTaskVertex = currentTaskVertex.getProperty(propertyKey, Long.class);
+        long valueToPushToTaskVertex = isIncremental ? (valueFromTaskVertex != null ? valueFromTaskVertex : 0L) + value : value;
+        if (taskSetter != null) {
+            taskSetter.accept(currentTask, valueToPushToTaskVertex);
+        }
 
+        currentTaskVertex.setProperty(propertyKey, valueToPushToTaskVertex);
+    }
 
     public DeleteHandlerV1(AtlasGraph graph, AtlasTypeRegistry typeRegistry, boolean shouldUpdateInverseReference, boolean softDelete, TaskManagement taskManagement) {
         this.typeRegistry                  = typeRegistry;
@@ -1220,16 +1235,32 @@ public void updateTagPropagations(AtlasEdge edge, AtlasRelationship relationship
                 }
             }
 
+            taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_TO_PROPAGATE, graph, addPropagationsMap.size() + removePropagationsMap.size(), false, AtlasTask::setAssetsCountToPropagate);
+
+            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, true, AtlasTask::setAssetsCountPropagated);
+                    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, true, AtlasTask::setAssetsCountPropagated);
+                    propagatedCount = 0;
+                }
+            }
+            if (propagatedCount != 0){
+                taskManagement.updateTaskVertexProperty(TASK_ASSET_COUNT_PROPAGATED, graph, propagatedCount, true, AtlasTask::setAssetsCountPropagated);
             }
         } else {
             // update blocked propagated classifications only if there is no change is tag propagation (don't update both)