corrected indexsearch for policy fetch

atlanhq · Apr 17, 2024 · 3115e48 · 3115e48
1 parent b094875
commit 3115e48
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 9 deletions.
diff --git a/...che/atlas/repository/store/graph/v2/preprocessor/datamesh/AbstractDomainPreProcessor.java b/...che/atlas/repository/store/graph/v2/preprocessor/datamesh/AbstractDomainPreProcessor.java
@@ -39,11 +39,13 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
-import static org.apache.atlas.repository.Constants.NAME;
+import static org.apache.atlas.repository.Constants.*;
+import static org.apache.atlas.repository.store.graph.v2.preprocessor.PreProcessorUtils.PARENT_DOMAIN_QN;
+import static org.apache.atlas.repository.store.graph.v2.preprocessor.PreProcessorUtils.SUPER_DOMAIN_QN;
+import static org.apache.atlas.repository.util.AccessControlUtils.ATTR_POLICY_CATEGORY;
+import static org.apache.atlas.repository.util.AccessControlUtils.ATTR_POLICY_RESOURCES;
 import static org.apache.atlas.repository.util.AtlasEntityUtils.mapOf;
 
 public abstract class AbstractDomainPreProcessor implements PreProcessor {
@@ -66,7 +68,7 @@ public abstract class AbstractDomainPreProcessor implements PreProcessor {
         }
     }
 
-    public List<AtlasEntityHeader> indexSearchPaginated(Map<String, Object> dsl) throws AtlasBaseException {
+    public List<AtlasEntityHeader> indexSearchPaginated(Map<String, Object> dsl, String entityType) throws AtlasBaseException {
         IndexSearchParams searchParams = new IndexSearchParams();
         List<AtlasEntityHeader> ret = new ArrayList<>();
 
@@ -83,6 +85,11 @@ public List<AtlasEntityHeader> indexSearchPaginated(Map<String, Object> dsl) thr
             dsl.put("size", size);
             searchParams.setDsl(dsl);
 
+            if (entityType.equals(POLICY_ENTITY_TYPE)) {
+                Set<String> attributes = new HashSet<>(Arrays.asList(ATTR_POLICY_RESOURCES, ATTR_POLICY_CATEGORY));
+                searchParams.setAttributes(attributes);
+            }
+
             List<AtlasEntityHeader> headers = discovery.directIndexSearch(searchParams).getEntities();
 
             if (CollectionUtils.isNotEmpty(headers)) {

diff --git a/...apache/atlas/repository/store/graph/v2/preprocessor/datamesh/DataProductPreProcessor.java b/...apache/atlas/repository/store/graph/v2/preprocessor/datamesh/DataProductPreProcessor.java
@@ -206,7 +206,7 @@ private List<AtlasEntityHeader> getPolicy(String resource) throws AtlasBaseExcep
 
             Map<String, Object> dsl = mapOf("query", mapOf("bool", bool));
 
-            List<AtlasEntityHeader> policies = indexSearchPaginated(dsl);
+            List<AtlasEntityHeader> policies = indexSearchPaginated(dsl, POLICY_ENTITY_TYPE);
 
             return policies;
         } finally {
@@ -289,7 +289,7 @@ private void productExists(String productName, String parentDomainQualifiedName)
 
             Map<String, Object> dsl = mapOf("query", mapOf("bool", bool));
 
-            List<AtlasEntityHeader> products = indexSearchPaginated(dsl);
+            List<AtlasEntityHeader> products = indexSearchPaginated(dsl, DATA_PRODUCT_ENTITY_TYPE);
 
             if (CollectionUtils.isNotEmpty(products)) {
                 for (AtlasEntityHeader product : products) {

diff --git a/.../org/apache/atlas/repository/store/graph/v2/preprocessor/datamesh/DomainPreProcessor.java b/.../org/apache/atlas/repository/store/graph/v2/preprocessor/datamesh/DomainPreProcessor.java
@@ -339,7 +339,7 @@ private List<AtlasEntityHeader> getPolicy(String resource) throws AtlasBaseExcep
 
             Map<String, Object> dsl = mapOf("query", mapOf("bool", bool));
 
-            List<AtlasEntityHeader> policies = indexSearchPaginated(dsl);
+            List<AtlasEntityHeader> policies = indexSearchPaginated(dsl, POLICY_ENTITY_TYPE);
 
             return policies;
         } finally {
@@ -396,7 +396,7 @@ private void domainExists(String domainName, String parentDomainQualifiedName) t
 
             Map<String, Object> dsl = mapOf("query", mapOf("bool", bool));
 
-            List<AtlasEntityHeader> domains = indexSearchPaginated(dsl);
+            List<AtlasEntityHeader> domains = indexSearchPaginated(dsl, DATA_DOMAIN_ENTITY_TYPE);
 
             if (CollectionUtils.isNotEmpty(domains)) {
                 for (AtlasEntityHeader domain : domains) {