Merge branch 'exclude-classification-names' into beta

atlanhq · Apr 9, 2024 · 3ad0a77 · 3ad0a77
2 parents 9a9db66 + 67a1895
commit 3ad0a77
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 3 deletions.
diff --git a/intg/src/main/java/org/apache/atlas/model/discovery/SearchParams.java b/intg/src/main/java/org/apache/atlas/model/discovery/SearchParams.java
@@ -20,6 +20,8 @@ public class SearchParams {
     boolean excludeClassifications;
     boolean enableFullRestriction;
 
+    boolean includeClassificationNames = false;
+
     RequestMetadata requestMetadata = new RequestMetadata();
 
     Async async = new Async();
@@ -109,6 +111,14 @@ public void setExcludeMeanings(boolean excludeMeanings) {
         this.excludeMeanings = excludeMeanings;
     }
 
+    public boolean isIncludeClassificationNames() {
+        return includeClassificationNames;
+    }
+
+    public void setIncludeClassificationNames(boolean includeClassificationNames) {
+        this.includeClassificationNames = includeClassificationNames;
+    }
+
     public boolean isSaveSearchLog() {
         return requestMetadata.saveSearchLog;
     }

diff --git a/repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java b/repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java
@@ -797,7 +797,18 @@ public static List<String> getTraitNames(AtlasVertex entityVertex) {
     public static List<String> getPropagatedTraitNames(AtlasVertex entityVertex) {
         return getTraitNames(entityVertex, true);
     }
-
+    public static List<String> getAllTraitNamesFromAttribute(AtlasVertex entityVertex) {
+        List<String>     ret   = new ArrayList<>();
+        List<String>    traitNames = entityVertex.getMultiValuedProperty(TRAIT_NAMES_PROPERTY_KEY, String.class);
+        if (traitNames != null) {
+            ret.addAll(traitNames);
+        }
+        List<String>    propagatedTraitNames = entityVertex.getMultiValuedProperty(PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, String.class);
+        if (propagatedTraitNames != null) {
+            ret.addAll(propagatedTraitNames);
+        }
+        return ret;
+    }
     public static List<String> getAllTraitNames(AtlasVertex entityVertex) {
         return getTraitNames(entityVertex, null);
     }

diff --git a/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -1019,8 +1019,13 @@ private AtlasEntityHeader mapVertexToAtlasEntityHeader(AtlasVertex entityVertex,
             ret.setTypeName(typeName);
             ret.setGuid(guid);
             ret.setStatus(GraphHelper.getStatus(entityVertex));
-            if(RequestContext.get().includeClassifications()){
-                ret.setClassificationNames(getAllTraitNames(entityVertex));
+            RequestContext context = RequestContext.get();
+            boolean includeClassifications = context.includeClassifications();
+            boolean includeClassificationNames = context.isIncludeClassificationNames();
+            if(includeClassifications){
+                ret.setClassificationNames(getAllTraitNamesFromAttribute(entityVertex));
+            } else if (!includeClassifications && includeClassificationNames) {
+                ret.setClassifications(getAllClassifications(entityVertex));
             }
             ret.setIsIncomplete(isIncomplete);
             ret.setLabels(getLabels(entityVertex));

diff --git a/server-api/src/main/java/org/apache/atlas/RequestContext.java b/server-api/src/main/java/org/apache/atlas/RequestContext.java
@@ -88,6 +88,8 @@ public class RequestContext {
     private boolean     allowDeletedRelationsIndexsearch = false;
     private boolean     includeMeanings = true;
     private boolean     includeClassifications = true;
+
+    private boolean     includeClassificationNames = false;
     private String      currentTypePatchAction = "";
     private AtlasTask   currentTask;
     private String traceId;
@@ -710,6 +712,14 @@ public String getRequestUri() {
         return this.requestUri;
     }
 
+    public boolean isIncludeClassificationNames() {
+        return includeClassificationNames;
+    }
+
+    public void setIncludeClassificationNames(boolean includeClassificationNames) {
+        this.includeClassificationNames = includeClassificationNames;
+    }
+
     public class EntityGuidPair {
         private final Object entity;
         private final String guid;

diff --git a/webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java b/webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java
@@ -393,6 +393,7 @@ public AtlasSearchResult indexSearch(@Context HttpServletRequest servletRequest,
 
         RequestContext.get().setIncludeMeanings(!parameters.isExcludeMeanings());
         RequestContext.get().setIncludeClassifications(!parameters.isExcludeClassifications());
+        RequestContext.get().setIncludeClassificationNames(parameters.isIncludeClassificationNames());
         try     {
             if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
                 perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.indexSearch(" + parameters + ")");