Split to avoid max_clause_limit issue in ES DSL query

atlanhq · Jan 18, 2024 · 70d13d5 · 70d13d5
1 parent 425c307
commit 70d13d5
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 24 deletions.
diff --git a/repository/src/main/java/org/apache/atlas/authorizer/AccessorsExtractor.java b/repository/src/main/java/org/apache/atlas/authorizer/AccessorsExtractor.java
@@ -64,25 +64,8 @@ public class AccessorsExtractor {
 
     public static AtlasAccessorResponse getAccessors(AtlasAccessorRequest request) throws AtlasBaseException {
         return getAccessorsInMemory(request);
-        //return getAccessorsES(request);
     }
 
-    public static AtlasAccessorResponse getAccessorsES(AtlasAccessorRequest request) throws AtlasBaseException {
-        AtlasAccessorResponse response = new AtlasAccessorResponse();
-
-        String action = AtlasPrivilege.valueOf(request.getAction()).getType();
-
-        List<RangerPolicy> resourcePolicies = PoliciesStore.getRelevantPolicies(null, null, "atlas", Arrays.asList(action), POLICY_TYPE_ALLOW, true);
-        resourcePolicies.addAll(PoliciesStore.getRelevantPolicies(null, null, "atlas_tag", Arrays.asList(action), POLICY_TYPE_ALLOW, true));
-
-
-        List<RangerPolicy> abacPolicies = PoliciesStore.getRelevantPolicies(null, null, "atlas_abac", Arrays.asList(action), POLICY_TYPE_ALLOW, true);
-
-
-        return response;
-    }
-
-
     private static void collectSubjects(AtlasAccessorResponse response, List<RangerPolicy> matchedPolicies) {
         for (RangerPolicy policy: matchedPolicies) {
             List<RangerPolicy.RangerPolicyItem> policyItems = null;

diff --git a/repository/src/main/java/org/apache/atlas/authorizer/AuthorizerUtils.java b/repository/src/main/java/org/apache/atlas/authorizer/AuthorizerUtils.java
@@ -42,6 +42,7 @@ public class AuthorizerUtils {
 
     public static final String POLICY_TYPE_ALLOW = "allow";
     public static final String POLICY_TYPE_DENY = "deny";
+    public static final int MAX_CLAUSE_LIMIT = 1024;
 
     private static AtlasTypeRegistry typeRegistry;
     private static EntityGraphRetriever entityRetriever;
@@ -237,7 +238,7 @@ public static Map<String, Object>  getPreFilterDsl(String persona, String purpos
         return ListAuthorizer.getElasticsearchDSL(persona, purpose, actions);
     }
 
-    private   <T> T getResourceAsObject(String resourceName, Class<T> clazz) throws IOException {
+    private <T> T getResourceAsObject(String resourceName, Class<T> clazz) throws IOException {
         InputStream stream = getClass().getResourceAsStream(resourceName);
         return AtlasType.fromJson(stream, clazz);
     }

diff --git a/repository/src/main/java/org/apache/atlas/authorizer/authorizers/EntityAuthorizer.java b/repository/src/main/java/org/apache/atlas/authorizer/authorizers/EntityAuthorizer.java
@@ -3,6 +3,7 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Lists;
 import org.apache.atlas.RequestContext;
 import org.apache.atlas.authorize.AtlasAuthorizationUtils;
 import org.apache.atlas.authorizer.AccessResult;
@@ -15,6 +16,7 @@
 import org.apache.atlas.repository.graphdb.janus.AtlasElasticsearchQuery;
 import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
 import org.apache.atlas.utils.AtlasPerfMetrics;
+import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang.StringUtils;
 import org.elasticsearch.client.RestClient;
 import org.slf4j.Logger;
@@ -23,6 +25,7 @@
 import java.util.*;
 import java.util.stream.Collectors;
 
+import static org.apache.atlas.authorizer.AuthorizerUtils.MAX_CLAUSE_LIMIT;
 import static org.apache.atlas.authorizer.AuthorizerUtils.POLICY_TYPE_ALLOW;
 import static org.apache.atlas.authorizer.AuthorizerUtils.POLICY_TYPE_DENY;
 import static org.apache.atlas.authorizer.authorizers.AuthorizerCommon.getMap;
@@ -400,8 +403,8 @@ public static AccessResult isAccessAllowedEvaluator(String entityTypeName, Strin
 
     public static Map<String, Object> getElasticsearchDSL(String persona, String purpose, List<String> actions) {
         AtlasPerfMetrics.MetricRecorder recorder = RequestContext.get().startMetricRecord("EntityAuthorizer.getElasticsearchDSL");
-        Map<String, Object> allowDsl = getElasticsearchDSLForPolicyType(persona, purpose, actions, POLICY_TYPE_ALLOW);
-        Map<String, Object> denyDsl = getElasticsearchDSLForPolicyType(persona, purpose, actions, POLICY_TYPE_DENY);
+        Map<String, Object> allowDsl = ListAuthorizer.getElasticsearchDSLForPolicyType(persona, purpose, actions, POLICY_TYPE_ALLOW);
+        Map<String, Object> denyDsl = ListAuthorizer.getElasticsearchDSLForPolicyType(persona, purpose, actions, POLICY_TYPE_DENY);
         Map<String, Object> finaDsl = new HashMap<>();
         if (allowDsl != null) {
             finaDsl.put("filter", allowDsl);
@@ -427,7 +430,7 @@ private static Integer getCountFromElasticsearch(String query) throws AtlasBaseE
         return count;
     }
 
-    public static Map<String, Object> getElasticsearchDSLForPolicyType(String persona, String purpose, List<String> actions, String policyType) {
+    /*public static Map<String, Object> getElasticsearchDSLForPolicyType(String persona, String purpose, List<String> actions, String policyType) {
         List<RangerPolicy> resourcePolicies = PoliciesStore.getRelevantPolicies(persona, purpose, "atlas", actions, policyType);
         List<Map<String, Object>> resourcePoliciesClauses = ListAuthorizer.getDSLForResourcePolicies(resourcePolicies);
 
@@ -453,11 +456,24 @@ public static Map<String, Object> getElasticsearchDSLForPolicyType(String person
             }
 
         } else {
-            boolClause.put("should", shouldClauses);
+            //boolClause.put("should", shouldClauses);
+            if (shouldClauses.size() > MAX_CLAUSE_LIMIT) {
+                List<Map<String, Object>> splittedShould = new ArrayList<>();
+                List<List<Map<String, Object>>> partitionedShouldClause = Lists.partition(shouldClauses, MAX_CLAUSE_LIMIT);
+
+                for (List<Map<String, Object>> chunk : partitionedShouldClause) {
+                    splittedShould.add(getMap("bool", getMap("should", chunk)));
+                }
+                boolClause.put("should", splittedShould);
+
+            } else {
+                boolClause.put("should", shouldClauses);
+            }
+
             boolClause.put("minimum_should_match", 1);
         }
 
         return getMap("bool", boolClause);
 
-    }
+    }*/
 }
diff --git a/repository/src/main/java/org/apache/atlas/authorizer/authorizers/ListAuthorizer.java b/repository/src/main/java/org/apache/atlas/authorizer/authorizers/ListAuthorizer.java
@@ -3,6 +3,7 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Lists;
 import org.apache.atlas.RequestContext;
 import org.apache.atlas.authorize.AtlasAuthorizationUtils;
 import org.apache.atlas.authorizer.JsonToElasticsearchQuery;
@@ -17,6 +18,7 @@
 import java.util.*;
 import java.util.stream.Collectors;
 
+import static org.apache.atlas.authorizer.AuthorizerUtils.MAX_CLAUSE_LIMIT;
 import static org.apache.atlas.authorizer.AuthorizerUtils.POLICY_TYPE_ALLOW;
 import static org.apache.atlas.authorizer.AuthorizerUtils.POLICY_TYPE_DENY;
 import static org.apache.atlas.authorizer.authorizers.AuthorizerCommon.*;
@@ -70,7 +72,20 @@ public static Map<String, Object> getElasticsearchDSLForPolicyType(String person
             }
 
         } else {
-            boolClause.put("should", shouldClauses);
+            //boolClause.put("should", shouldClauses);
+            if (shouldClauses.size() > MAX_CLAUSE_LIMIT) {
+                List<Map<String, Object>> splittedShould = new ArrayList<>();
+                List<List<Map<String, Object>>> partitionedShouldClause = Lists.partition(shouldClauses, MAX_CLAUSE_LIMIT);
+
+                for (List<Map<String, Object>> chunk : partitionedShouldClause) {
+                    splittedShould.add(getMap("bool", getMap("should", chunk)));
+                }
+                boolClause.put("should", splittedShould);
+
+            } else {
+                boolClause.put("should", shouldClauses);
+            }
+
             boolClause.put("minimum_should_match", 1);
         }