MESH-645 : Repair API for daapOutputPortGuids

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

@@ -0,0 +1,70 @@
+package org.apache.atlas.repository.store.graph.v2.repair;
+
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.AtlasErrorCode;
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
+
+import javax.inject.Inject;
+import java.util.Set;
+
+@Component
+public class AtlasRepairAttributeService {
+
+    private static final Logger LOG = LoggerFactory.getLogger(AtlasRepairAttributeService.class);
+
+    private final RepairAttributeFactory repairAttributeFactory;
+
+    @Inject
+    public AtlasRepairAttributeService(RepairAttributeFactory repairAttributeFactory) {
+        this.repairAttributeFactory = repairAttributeFactory;
+    }
+
+    public void repairAttributes(String attributeName, String repairType, Set<String> entityGuids)
+            throws AtlasBaseException {
+
+        validateRequest(attributeName, repairType, entityGuids);
+
+        LOG.info("Starting attribute repair - Type: {}, Attribute: {}, Entities: {}",
+                repairType, attributeName, entityGuids.size());
+
+        try {
+            AtlasRepairAttributeStrategy strategy = repairAttributeFactory.getStrategy(repairType, entityGuids);
+
+            strategy.validate(entityGuids, attributeName);
+            strategy.repair(entityGuids, attributeName);
+
+            LOG.info("Successfully completed attribute repair - Type: {}, Attribute: {}, Entities: {}",
+                    repairType, attributeName, entityGuids.size());
+
+        } catch (Exception e) {
+            LOG.error("Error during attribute repair - Type: {}, Attribute: {}, Entities: {}",
+                    repairType, attributeName, entityGuids.size(), e);
+            throw e;
+        }
+    }
+
+    private void validateRequest(String attributeName, String repairType, Set<String> entityGuids)
+            throws AtlasBaseException {
+
+        if (StringUtils.isEmpty(attributeName)) {
+            throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST, "Attribute name cannot be empty");
+        }
+
+        if (StringUtils.isEmpty(repairType)) {
+            throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST, "Repair type cannot be empty");
+        }
+
+        if (CollectionUtils.isEmpty(entityGuids)) {
+            throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST, "Entity GUIDs cannot be empty");
+        }
+
+        if (entityGuids.size() > 1000) {
+            throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST,
+                    "Too many entities. Maximum allowed: 1000, provided: " + entityGuids.size());
+        }
+    }
+}

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

@@ -0,0 +1,34 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.atlas.repository.store.graph.v2.repair;
+
+
+import org.apache.atlas.exception.AtlasBaseException;
+
+import java.util.Set;
+
+public interface AtlasRepairAttributeStrategy {
+
+    String getRepairType();
+    void repair(Set<String> entityGuids, String attributeName) throws AtlasBaseException;
+    void validate (Set<String> entityGuids, String attributeName) throws AtlasBaseException;
+
+}
+
+

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

@@ -0,0 +1,149 @@
+package org.apache.atlas.repository.store.graph.v2.repair;
+
+import org.apache.atlas.AtlasErrorCode;
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
+import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
+import org.apache.atlas.repository.store.graph.v2.TransactionInterceptHelper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.*;
+
+import static org.apache.atlas.repository.store.graph.v2.preprocessor.PreProcessorUtils.OUTPUT_PORT_GUIDS_ATTR;
+
+public class RemoveInvalidGuidsRepairStrategy implements AtlasRepairAttributeStrategy {
+
+    private static final Logger LOG = LoggerFactory.getLogger(RemoveInvalidGuidsRepairStrategy.class);
+
+    private final EntityGraphRetriever entityRetriever;
+
+
+    private final TransactionInterceptHelper transactionInterceptHelper;
+
+    private static final String REPAIR_TYPE = "REMOVE_INVALID_GUIDS";
+
+    public RemoveInvalidGuidsRepairStrategy(EntityGraphRetriever entityRetriever, TransactionInterceptHelper transactionInterceptHelper) {
+        this.entityRetriever = entityRetriever;
+        this.transactionInterceptHelper = transactionInterceptHelper;
+    }
+
+    @Override
+    public String getRepairType() {
+        return REPAIR_TYPE;
+    }
+
+    @Override
+    public void validate(Set<String> entityGuids, String attributeName) throws AtlasBaseException {
+        for (String entityGuid : entityGuids) {
+            AtlasVertex entityVertex = entityRetriever.getEntityVertex(entityGuid);
+
+            if (entityVertex == null) {
+                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, "Entity vertex not found for guid: " + entityGuid);
+            }
+        }
+    }
+
+    @Override
+    public void repair(Set<String> entityGuids, String attributeName) throws AtlasBaseException {
+        try {
+            int count = 0;
+            int totalUpdatedCount = 0;
+
+            for (String entityGuid : entityGuids) {
+                AtlasVertex entityVertex = entityRetriever.getEntityVertex(entityGuid);
+
+                if (entityVertex == null) {
+                    LOG.error("Entity vertex not found for guid: {}", entityGuid);
+                    continue;
+                }
+
+                if (!entityVertex.getPropertyKeys().contains(attributeName)) {
+                    LOG.info("Attribute: {} not found for entity: {}. Skipping repair for this entity.", attributeName, entityGuid);
+                    continue;
+                }
+
+                if (OUTPUT_PORT_GUIDS_ATTR.equals(attributeName)) {
+                    boolean isCommitRequired = repairAttr(entityVertex);
+                    if (isCommitRequired){
+                        count++;
+                        totalUpdatedCount++;
+                    } else {
+                        LOG.info("No changes to commit for entity: {}", entityGuid);
+                    }
+
+                    if (count == 50) {
+                        LOG.info("Committing batch of 50 entities...");
+                        commitChanges();
+                        count = 0;
+                    }
+                }
+            }
+
+            if (count > 0) {
+                LOG.info("Committing remaining {} entities...", count);
+                commitChanges();
+            }
+
+            LOG.info("Total Vertex updated: {}", totalUpdatedCount);
+
+        } catch(Exception e){
+            LOG.error("Error while performing repair: {}", entityGuids, e);
+            throw e;
+        }
+    }
+
+    private boolean repairAttr(AtlasVertex vertex) throws AtlasBaseException {
+        try{
+            boolean isCommitRequired = false;
+
+            List<String> outputPortGuids = vertex.getMultiValuedProperty(OUTPUT_PORT_GUIDS_ATTR, String.class);
+            if (outputPortGuids == null || outputPortGuids.isEmpty()) {
+                LOG.info("No guids found in attribute: {} for entity: {}. Skipping repair for this entity.", OUTPUT_PORT_GUIDS_ATTR, vertex.getProperty("guid", String.class));
+                return false;
+            }
+
+            List<String> validGuids = new ArrayList<>();
+            List<String> invalidGuids = new ArrayList<>();
+
+            for (String guid : outputPortGuids) {
+                AtlasVertex portVertex = entityRetriever.getEntityVertex(guid);
+                if (portVertex != null) {
+                    validGuids.add(guid);
+                } else {
+                    invalidGuids.add(guid);
+                }
+            }
+
+            if (!invalidGuids.isEmpty()) {
+                LOG.info("Removing invalid guids: {} from attribute: {} for entity: {}", invalidGuids, OUTPUT_PORT_GUIDS_ATTR, vertex.getProperty("guid", String.class));
+
+                vertex.removeProperty(OUTPUT_PORT_GUIDS_ATTR);
+
+                for (String validGuid : validGuids) {
+                    AtlasGraphUtilsV2.addEncodedProperty(vertex, OUTPUT_PORT_GUIDS_ATTR, validGuid);
+                }
+
+                isCommitRequired = true;
+            } else {
+                LOG.info("All guids in attribute: {} for entity: {} are valid. No repair needed.", OUTPUT_PORT_GUIDS_ATTR, vertex.getProperty("guid", String.class));
+            }
+
+            return isCommitRequired;
+        } catch (Exception e) {
+            LOG.error("Failed to repair attribute for entity: ", e);
+            throw e;
+        }
+    }
+
+    public void commitChanges() throws AtlasBaseException {
+        try {
+            transactionInterceptHelper.intercept();
+            LOG.info("Committed a entity to the graph");
+        } catch (Exception e){
+            LOG.error("Failed to commit asset: ", e);
+            throw e;
+        }
+    }
+}

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

@@ -0,0 +1,38 @@
+package org.apache.atlas.repository.store.graph.v2.repair;
+
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.AtlasErrorCode;
+import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
+import org.apache.atlas.repository.store.graph.v2.TransactionInterceptHelper;
+import org.springframework.stereotype.Component;
+
+import javax.inject.Inject;
+import java.util.Set;
+
+@Component
+public class RepairAttributeFactory {
+
+    private final EntityGraphRetriever entityRetriever;
+    private final TransactionInterceptHelper transactionInterceptHelper;
+
+    @Inject
+    public RepairAttributeFactory(EntityGraphRetriever entityRetriever,
+                                  TransactionInterceptHelper transactionInterceptHelper) {
+        this.entityRetriever = entityRetriever;
+        this.transactionInterceptHelper = transactionInterceptHelper;
+    }
+
+    public AtlasRepairAttributeStrategy getStrategy(String repairType, Set<String> entityGuids) throws AtlasBaseException {
+        switch (repairType) {
+            case "REMOVE_INVALID_GUIDS":
+                return new RemoveInvalidGuidsRepairStrategy(entityRetriever, transactionInterceptHelper);
+            default:
+                throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST,
+                        "Unsupported repair type: " + repairType + ". Supported types: [REMOVE_INVALID_GUIDS]");
+        }
+    }
+
+    public boolean isValidRepairType(String repairType) {
+        return "REMOVE_INVALID_GUIDS".equals(repairType);
+    }
+}

webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java

@@ -43,6 +43,7 @@
 import org.apache.atlas.repository.converters.AtlasInstanceConverter;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.repository.store.graph.v2.*;
+import org.apache.atlas.repository.store.graph.v2.repair.AtlasRepairAttributeService;
 import org.apache.atlas.service.FeatureFlagStore;
 import org.apache.atlas.type.AtlasClassificationType;
 import org.apache.atlas.type.AtlasEntityType;
@@ -109,14 +110,16 @@ public class EntityREST {
     private final ESBasedAuditRepository  esBasedAuditRepository;
     private final EntityGraphRetriever entityGraphRetriever;
     private final EntityMutationService entityMutationService;
+    private final AtlasRepairAttributeService repairAttributeService;
 
     @Inject
-    public EntityREST(AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore, ESBasedAuditRepository  esBasedAuditRepository, EntityGraphRetriever retriever, EntityMutationService entityMutationService) {
+    public EntityREST(AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore, ESBasedAuditRepository  esBasedAuditRepository, EntityGraphRetriever retriever, EntityMutationService entityMutationService, AtlasRepairAttributeService repairAttributeService) {
         this.typeRegistry      = typeRegistry;
         this.entitiesStore     = entitiesStore;
         this.esBasedAuditRepository = esBasedAuditRepository;
         this.entityGraphRetriever = retriever;
         this.entityMutationService = entityMutationService;
+        this.repairAttributeService = repairAttributeService;
     }
 
     /**
@@ -1856,6 +1859,38 @@ public void repairEntityIndexBulk(Set<String> guids) throws AtlasBaseException {
         }
     }
 
+    /**
+     * Repair attributes for the entity GUID.
+     * @param guids
+     * @throws AtlasBaseException
+     */
+
+    @POST
+    @Path("/guid/bulk/repairattributes")
+    public void repairEntityAttributesBulk(Set<String> guids, @QueryParam("repairType") String repairType, @QueryParam("repairAttributeName") String repairAttributeName) throws AtlasBaseException {
+
+        Servlets.validateQueryParamLength("repairType", repairType);
+        Servlets.validateQueryParamLength("repairAttributeName", repairAttributeName);
+
+        AtlasAuthorizationUtils.verifyAccess(new AtlasAdminAccessRequest(AtlasPrivilege.ADMIN_REPAIR_INDEX), "Admin Repair Attributes");
+
+        AtlasPerfTracer perf = null;
+
+        try {
+            if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
+                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.repairEntityAttributesBulk(" + guids.size() + ")");
+            }
+
+            repairAttributeService.repairAttributes(repairAttributeName, repairType, guids);
+
+        } catch (Exception e) {
+            LOG.error("Exception while repairEntityAttributesBulk ", e);
+            throw new AtlasBaseException(e);
+        } finally {
+            AtlasPerfTracer.log(perf);
+        }
+    }
+
     @POST
     @Path("/repairindex/{typename}")
     public void repairIndexByTypeName(@PathParam("typename") String typename, @QueryParam("delay") @DefaultValue("0") int delay, @QueryParam("limit") @DefaultValue("1000") int limit, @QueryParam("offset") @DefaultValue("0") int offset, @QueryParam("batchSize") @DefaultValue("1000") int batchSize) throws AtlasBaseException {