Merge pull request #2909 from atlanhq/dg-980-public-daap

[Master]- Dg-980: Data Product Visibility Persona indexsearch change.
atlanhq · Apr 2, 2024 · 9bf6522 · 9bf6522
2 parents b014b40 + 6651392
commit 9bf6522
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 0 deletions.
diff --git a/intg/src/main/java/org/apache/atlas/model/discovery/IndexSearchParams.java b/intg/src/main/java/org/apache/atlas/model/discovery/IndexSearchParams.java
@@ -27,6 +27,7 @@ public class IndexSearchParams extends SearchParams {
     * (this will include related attributes which has relationshipStatus as DELETED along with ACTIVE ones)
     * */
     private boolean allowDeletedRelations;
+    private boolean accessControlExclusive;
 
     @Override
     public String getQuery() {
@@ -46,6 +47,14 @@ public boolean isAllowDeletedRelations() {
         return allowDeletedRelations;
     }
 
+    public boolean isAccessControlExclusive() {
+        return accessControlExclusive;
+    }
+
+    public void setAccessControlExclusive(boolean accessControlExclusive) {
+        this.accessControlExclusive = accessControlExclusive;
+    }
+
     public void setAllowDeletedRelations(boolean allowDeletedRelations) {
         this.allowDeletedRelations = allowDeletedRelations;
     }
@@ -78,6 +87,7 @@ public String toString() {
                 ", persona='" + persona + '\'' +
                 ", queryString='" + queryString + '\'' +
                 ", allowDeletedRelations=" + allowDeletedRelations +
+                ", accessControlExclusive=" + accessControlExclusive +
                 ", utmTags="+ getUtmTags() +
                 '}';
     }

diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
@@ -1149,9 +1149,57 @@ private String getIndexName(IndexSearchParams params) throws AtlasBaseException
         String aliasName = parts[parts.length - 1];
 
         if (StringUtils.isNotEmpty(aliasName)) {
+            if(params.isAccessControlExclusive()) {
+                accessControlExclusiveDsl(params, aliasName);
+                aliasName = aliasName+","+VERTEX_INDEX_NAME;
+            }
             return aliasName;
         } else {
             throw new AtlasBaseException("ES alias not found for purpose/persona " + params.getPurpose());
         }
     }
+
+    private void accessControlExclusiveDsl(IndexSearchParams params, String aliasName) {
+
+        List<Map<String, Object>> mustClauses = new ArrayList<>();
+        Map<String, Object> clientQuery = (Map<String, Object>) params.getDsl().get("query");
+
+        mustClauses.add(clientQuery);
+
+        List<Map<String, Object>>filterClauses = new ArrayList<>();
+        filterClauses.add(getMap("terms", getMap("_index", Collections.singletonList(aliasName))));
+
+        Map<String, Object> boolQuery = new HashMap<>();
+        boolQuery.put("must", mustClauses);
+        boolQuery.put("filter",filterClauses);
+
+        List<Map<String, Object>> shouldClauses = new ArrayList<>();
+        shouldClauses.add(getMap("bool", boolQuery));
+        shouldClauses.add(getStaticBoolQuery());
+
+        Map<String, Object> topBoolQuery = getMap("bool", getMap("should", shouldClauses));
+
+        Map copyOfDsl = new HashMap(params.getDsl());
+        copyOfDsl.put("query", topBoolQuery);
+
+        params.setDsl(copyOfDsl);
+    }
+
+    private Map<String, Object> getStaticBoolQuery() {
+        List<Map<String, Object>> mustClauses = new ArrayList<>();
+        Map<String, Object> mustClause = getMap("bool", getMap("should", Arrays.asList(
+                getMap("term", getMap("daapVisibility", "Public")),
+                getMap("term", getMap("daapVisibility", "Protected"))
+        )));
+        mustClauses.add(mustClause);
+
+        List<Map<String, Object>>filterClauses = new ArrayList<>();
+        filterClauses.add(getMap("terms", getMap("_index", Collections.singletonList(VERTEX_INDEX_NAME))));
+
+        Map<String, Object> boolQuery = new HashMap<>();
+        boolQuery.put("must", mustClauses);
+        boolQuery.put("filter", filterClauses);
+
+        return getMap("bool", boolQuery);
+    }
 }