ATLAS-4495: Improve Glossary GET api performance when Glossaries have large number of Categories and Terms associated with them

Signed-off-by: Pinal Shah <pinal.shah@freestoneinfotech.com>
diff --git a/repository/src/main/java/org/apache/atlas/glossary/GlossaryService.java b/repository/src/main/java/org/apache/atlas/glossary/GlossaryService.java
index f81b538..110c447 100644
--- a/repository/src/main/java/org/apache/atlas/glossary/GlossaryService.java
+++ b/repository/src/main/java/org/apache/atlas/glossary/GlossaryService.java
@@ -18,7 +18,6 @@
 package org.apache.atlas.glossary;
 
 import org.apache.atlas.AtlasErrorCode;
-import org.apache.atlas.RequestContext;
 import org.apache.atlas.SortOrder;
 import org.apache.atlas.annotation.GraphTransaction;
 import org.apache.atlas.bulkimport.BulkImportResponse;
@@ -30,18 +29,19 @@
 import org.apache.atlas.model.glossary.relations.AtlasRelatedCategoryHeader;
 import org.apache.atlas.model.glossary.relations.AtlasRelatedTermHeader;
 import org.apache.atlas.model.glossary.relations.AtlasTermCategorizationHeader;
+import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasRelatedObjectId;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.ogm.DataAccess;
+import org.apache.atlas.repository.ogm.glossary.AtlasGlossaryDTO;
 import org.apache.atlas.repository.store.graph.AtlasRelationshipStore;
 import org.apache.atlas.repository.store.graph.v2.AtlasEntityChangeNotifier;
 import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.util.FileUtils;
 import org.apache.atlas.utils.AtlasJson;
-import org.apache.atlas.utils.AtlasPerfMetrics;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
@@ -85,6 +85,7 @@
     private final GlossaryCategoryUtils     glossaryCategoryUtils;
     private final AtlasTypeRegistry         atlasTypeRegistry;
     private final AtlasEntityChangeNotifier entityChangeNotifier;
+    private final AtlasGlossaryDTO          glossaryDTO;
 
     private static final char[] invalidNameChars = { '@', '.' };
 
@@ -93,12 +94,13 @@
 
     @Inject
     public GlossaryService(DataAccess dataAccess, final AtlasRelationshipStore relationshipStore,
-                           final AtlasTypeRegistry typeRegistry, AtlasEntityChangeNotifier entityChangeNotifier) {
+                           final AtlasTypeRegistry typeRegistry, AtlasEntityChangeNotifier entityChangeNotifier, final AtlasGlossaryDTO glossaryDTO) {
         this.dataAccess           = dataAccess;
         atlasTypeRegistry         = typeRegistry;
         glossaryTermUtils         = new GlossaryTermUtils(relationshipStore, typeRegistry, dataAccess);
         glossaryCategoryUtils     = new GlossaryCategoryUtils(relationshipStore, typeRegistry, dataAccess);
         this.entityChangeNotifier = entityChangeNotifier;
+        this.glossaryDTO          = glossaryDTO;
     }
 
     /**
@@ -116,19 +118,18 @@
             LOG.debug("==> GlossaryService.getGlossaries({}, {}, {})", limit, offset, sortOrder);
         }
 
-        List<String>     glossaryGuids    = AtlasGraphUtilsV2.findEntityGUIDsByType(GlossaryUtils.ATLAS_GLOSSARY_TYPENAME, sortOrder);
+        List<String> glossaryGuids = AtlasGraphUtilsV2.findEntityGUIDsByType(GlossaryUtils.ATLAS_GLOSSARY_TYPENAME, sortOrder);
         PaginationHelper paginationHelper = new PaginationHelper<>(glossaryGuids, offset, limit);
 
         List<AtlasGlossary> ret;
-        List<String>        guidsToLoad = paginationHelper.getPaginatedList();
-        if (CollectionUtils.isNotEmpty(guidsToLoad)) {
-            ret = guidsToLoad.stream().map(GlossaryUtils::getGlossarySkeleton).collect(Collectors.toList());
-            Iterable<AtlasGlossary> glossaries = dataAccess.load(ret);
-            ret.clear();
+        List<String> guidsToLoad = paginationHelper.getPaginatedList();
+        AtlasEntity.AtlasEntitiesWithExtInfo glossaryEntities;
 
-            // Set the displayText for all relations
-            for (AtlasGlossary glossary : glossaries) {
-                setInfoForRelations(glossary);
+        if (CollectionUtils.isNotEmpty(guidsToLoad)) {
+            glossaryEntities = dataAccess.getAtlasEntityStore().getByIds(guidsToLoad, true, false);
+            ret = new ArrayList<>();
+            for (AtlasEntity glossaryEntity : glossaryEntities.getEntities()) {
+                AtlasGlossary glossary = glossaryDTO.from(glossaryEntity);
                 ret.add(glossary);
             }
         } else {
diff --git a/repository/src/main/java/org/apache/atlas/repository/ogm/DataAccess.java b/repository/src/main/java/org/apache/atlas/repository/ogm/DataAccess.java
index f902b2a..e631524 100644
--- a/repository/src/main/java/org/apache/atlas/repository/ogm/DataAccess.java
+++ b/repository/src/main/java/org/apache/atlas/repository/ogm/DataAccess.java
@@ -46,6 +46,10 @@
     private final AtlasEntityStore entityStore;
     private final DTORegistry      dtoRegistry;
 
+    public AtlasEntityStore getAtlasEntityStore(){
+        return this.entityStore;
+    }
+
     @Inject
     public DataAccess(AtlasEntityStore entityStore, DTORegistry dtoRegistry) {
         this.entityStore = entityStore;
diff --git a/repository/src/main/java/org/apache/atlas/repository/ogm/glossary/AbstractGlossaryDTO.java b/repository/src/main/java/org/apache/atlas/repository/ogm/glossary/AbstractGlossaryDTO.java
index bee88c6..abefac2 100644
--- a/repository/src/main/java/org/apache/atlas/repository/ogm/glossary/AbstractGlossaryDTO.java
+++ b/repository/src/main/java/org/apache/atlas/repository/ogm/glossary/AbstractGlossaryDTO.java
@@ -46,6 +46,7 @@
 
         ret.setTermGuid(relatedObjectId.getGuid());
         ret.setRelationGuid(relatedObjectId.getRelationshipGuid());
+        ret.setDisplayText(relatedObjectId.getDisplayText());
 
         AtlasStruct relationshipAttributes = relatedObjectId.getRelationshipAttributes();
         if (relationshipAttributes != null) {
@@ -87,10 +88,13 @@
 
         ret.setCategoryGuid(relatedObjectId.getGuid());
         ret.setRelationGuid(relatedObjectId.getRelationshipGuid());
+        ret.setDisplayText(relatedObjectId.getDisplayText());
 
         AtlasStruct relationshipAttributes = relatedObjectId.getRelationshipAttributes();
         if (relationshipAttributes != null) {
             ret.setDescription((String) relationshipAttributes.getAttribute("description"));
+            ret.setParentCategoryGuid(relationshipAttributes.getAttribute("parentCategoryGuid") == null ? null :
+                    (String) relationshipAttributes.getAttribute("parentCategoryGuid"));
         }
 
         return ret;
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
index 9671df9..fcd8dfe 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -21,6 +21,7 @@
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TimeBoundary;
+import org.apache.atlas.model.glossary.AtlasGlossaryCategory;
 import org.apache.atlas.model.glossary.enums.AtlasTermAssignmentStatus;
 import org.apache.atlas.model.glossary.relations.AtlasTermAssignmentHeader;
 import org.apache.atlas.model.instance.AtlasClassification;
@@ -134,6 +135,10 @@
     public static final String CREATE_TIME    = "createTime";
     public static final String QUALIFIED_NAME = "qualifiedName";
 
+    private static final String GLOSSARY_CATEGORY_HIERARCHY_EDGE_LABEL = "r:AtlasGlossaryCategoryHierarchyLink";
+    private static final String GLOSSARY_CATEGORY_TYPE_NAME = AtlasGlossaryCategory.class.getSimpleName();
+    private static final String PARENT_GLOSSARY_CATEGORY_GUID = "parentCategoryGuid";
+
     private static final TypeReference<List<TimeBoundary>> TIME_BOUNDARIES_LIST_TYPE = new TypeReference<List<TimeBoundary>>() {};
     private final GraphHelper graphHelper;
 
@@ -1537,12 +1542,26 @@
                     }
                 }
 
+                populateGlossaryAttributesIfApplicable(ret, referenceVertex, entityTypeName);
             }
         }
 
         return ret;
     }
 
+    private void populateGlossaryAttributesIfApplicable(AtlasRelatedObjectId ret, AtlasVertex referenceVertex, String entityTypeName) {
+        if (!StringUtils.equals(entityTypeName, GLOSSARY_CATEGORY_TYPE_NAME)) {
+            return;
+        }
+        Iterator<AtlasEdge> edgeIterator = GraphHelper.getIncomingEdgesByLabel(referenceVertex, GLOSSARY_CATEGORY_HIERARCHY_EDGE_LABEL);
+        while (edgeIterator.hasNext()) {
+            AtlasEdge atlasEdge = edgeIterator.next();
+            AtlasVertex parentCategoryVertex = atlasEdge.getOutVertex();
+            String parentCategoryGuid = GraphHelper.getGuid(parentCategoryVertex);
+            ret.getRelationshipAttributes().setAttribute(PARENT_GLOSSARY_CATEGORY_GUID, parentCategoryGuid);
+        }
+    }
+
     private Object getDisplayText(AtlasVertex entityVertex, String entityTypeName) throws AtlasBaseException {
         return getDisplayText(entityVertex, typeRegistry.getEntityTypeByName(entityTypeName));
     }