ATLAS-2267: Update basic search to support classification attributes

intg/src/main/java/org/apache/atlas/model/discovery/SearchParameters.java

@@ -44,6 +44,7 @@ public class SearchParameters implements Serializable {
     private String  typeName;
     private String  classification;
     private boolean excludeDeletedEntities;
+    private boolean includeClassificationAttributes;
     private int     limit;
     private int     offset;
 
@@ -113,6 +114,21 @@ public class SearchParameters implements Serializable {
     }
 
     /**
+     * @return True if classification attributes are included in search result.
+     */
+    public boolean getIncludeClassificationAttributes() {
+        return includeClassificationAttributes;
+    }
+
+    /**
+     * Include classificatio attributes in search result.
+     * @param includeClassificationAttributes boolean flag
+     */
+    public void setIncludeClassificationAttributes(boolean includeClassificationAttributes) {
+        this.includeClassificationAttributes = includeClassificationAttributes;
+    }
+
+    /**
      * @return Max number of results to be returned
      */
     public int getLimit() {
@@ -195,6 +211,7 @@ public class SearchParameters implements Serializable {
         if (o == null || getClass() != o.getClass()) return false;
         SearchParameters that = (SearchParameters) o;
         return excludeDeletedEntities == that.excludeDeletedEntities &&
+                includeClassificationAttributes == that.includeClassificationAttributes &&
                 limit == that.limit &&
                 offset == that.offset &&
                 Objects.equals(query, that.query) &&
@@ -207,7 +224,8 @@ public class SearchParameters implements Serializable {
 
     @Override
     public int hashCode() {
-        return Objects.hash(query, typeName, classification, excludeDeletedEntities, limit, offset, entityFilters, tagFilters, attributes);
+        return Objects.hash(query, typeName, classification, excludeDeletedEntities, includeClassificationAttributes, 
+                            limit, offset, entityFilters, tagFilters, attributes);
     }
 
     public StringBuilder toString(StringBuilder sb) {
@@ -220,6 +238,7 @@ public class SearchParameters implements Serializable {
         sb.append(", typeName='").append(typeName).append('\'');
         sb.append(", classification='").append(classification).append('\'');
         sb.append(", excludeDeletedEntities=").append(excludeDeletedEntities);
+        sb.append(", includeClassificationAttributes=").append(includeClassificationAttributes);
         sb.append(", limit=").append(limit);
         sb.append(", offset=").append(offset);
         sb.append(", entityFilters=").append(entityFilters);

intg/src/main/java/org/apache/atlas/model/instance/AtlasEntityHeader.java

@@ -34,6 +34,7 @@ import javax.xml.bind.annotation.XmlSeeAlso;
 
 import org.apache.atlas.model.PList;
 import org.apache.atlas.model.SearchFilter.SortType;
+import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.model.typedef.AtlasEntityDef;
 
 import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.NONE;
@@ -51,10 +52,11 @@ import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.PUBLIC_
 public class AtlasEntityHeader extends AtlasStruct implements Serializable {
     private static final long serialVersionUID = 1L;
 
-    private String             guid                = null;
-    private AtlasEntity.Status status              = AtlasEntity.Status.ACTIVE;
-    private String             displayText         = null;
-    private List<String>       classificationNames = null;
+    private String                    guid                = null;
+    private AtlasEntity.Status        status              = AtlasEntity.Status.ACTIVE;
+    private String                    displayText         = null;
+    private List<String>              classificationNames = null;
+    private List<AtlasClassification> classifications     = null;
 
     public AtlasEntityHeader() {
         this(null, null);
@@ -72,6 +74,7 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
         super(typeName, attributes);
 
         setClassificationNames(null);
+        setClassifications(null);
     }
 
 
@@ -79,6 +82,7 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
         super(typeName, attributes);
         setGuid(guid);
         setClassificationNames(null);
+        setClassifications(null);
     }
 
 
@@ -90,6 +94,7 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
             setStatus(other.getStatus());
             setDisplayText(other.getDisplayText());
             setClassificationNames(other.getClassificationNames());
+            setClassifications(other.getClassifications());
         }
     }
 
@@ -125,6 +130,10 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
         this.classificationNames = classificationNames;
     }
 
+    public List<AtlasClassification> getClassifications() { return classifications; }
+
+    public void setClassifications(List<AtlasClassification> classifications) { this.classifications = classifications; }
+
     @Override
     public StringBuilder toString(StringBuilder sb) {
         if (sb == null) {
@@ -137,8 +146,10 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
         sb.append(", displayText=").append(displayText);
         sb.append(", classificationNames=[");
         dumpObjects(classificationNames, sb);
-        sb.append("],");
-        sb.append(", ");
+        sb.append("], ");
+        sb.append("classifications=[");
+        AtlasBaseTypeDef.dumpObjects(classifications, sb);
+        sb.append("], ");
         super.toString(sb);
         sb.append('}');
 
@@ -154,12 +165,13 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
         return Objects.equals(guid, that.guid) &&
                 status == that.status &&
                 Objects.equals(displayText, that.displayText) &&
-                Objects.equals(classificationNames, that.classificationNames);
+                Objects.equals(classificationNames, that.classificationNames) &&
+                Objects.equals(classifications, that.classifications);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(super.hashCode(), guid, status, displayText, classificationNames);
+        return Objects.hash(super.hashCode(), guid, status, displayText, classificationNames, classifications);
     }
 
     @Override

repository/src/main/java/org/apache/atlas/discovery/ClassificationSearchProcessor.java

@@ -48,6 +48,8 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
+import static org.apache.atlas.repository.graphdb.AtlasGraphQuery.ComparisionOperator.NOT_EQUAL;
+
 
 public class ClassificationSearchProcessor extends SearchProcessor {
     private static final Logger LOG      = LoggerFactory.getLogger(ClassificationSearchProcessor.class);
@@ -77,7 +79,10 @@ public class ClassificationSearchProcessor extends SearchProcessor {
         processSearchAttributes(classificationType, filterCriteria, indexAttributes, graphAttributes, allAttributes);
 
         // for classification search, if any attribute can't be handled by index query - switch to all filter by Graph query
-        boolean useIndexSearch = typeAndSubTypesQryStr.length() <= MAX_QUERY_STR_LENGTH_TAGS && CollectionUtils.isEmpty(graphAttributes) && canApplyIndexFilter(classificationType, filterCriteria, false);
+        boolean useIndexSearch = classificationType != SearchContext.MATCH_ALL_CLASSIFICATION &&
+                                 typeAndSubTypesQryStr.length() <= MAX_QUERY_STR_LENGTH_TAGS &&
+                                 CollectionUtils.isEmpty(graphAttributes) &&
+                                 canApplyIndexFilter(classificationType, filterCriteria, false);
 
         AtlasGraph graph = context.getGraph();
 
@@ -136,7 +141,13 @@ public class ClassificationSearchProcessor extends SearchProcessor {
             }
         } else {
             tagGraphQueryWithAttributes = null;
-            entityGraphQueryTraitNames  = graph.query().in(Constants.TRAIT_NAMES_PROPERTY_KEY, typeAndSubTypes);
+
+            if (classificationType != SearchContext.MATCH_ALL_CLASSIFICATION) {
+                entityGraphQueryTraitNames = graph.query().in(Constants.TRAIT_NAMES_PROPERTY_KEY, typeAndSubTypes);
+            } else {
+                entityGraphQueryTraitNames = graph.query().has(Constants.TRAIT_NAMES_PROPERTY_KEY, NOT_EQUAL, null);
+            }
+
             entityPredicateTraitNames   = SearchPredicateUtil.getContainsAnyPredicateGenerator()
                                                              .generatePredicate(Constants.TRAIT_NAMES_PROPERTY_KEY, classificationType.getTypeAndAllSubTypes(), List.class);
 

repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java

@@ -469,6 +469,10 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
             for (AtlasVertex atlasVertex : resultList) {
                 AtlasEntityHeader entity = entityRetriever.toAtlasEntityHeader(atlasVertex, resultAttributes);
 
+                if(searchParameters.getIncludeClassificationAttributes()) {
+                    entity.setClassifications(entityRetriever.getClassifications(atlasVertex));
+                }
+
                 ret.addEntity(entity);
 
                 // populate ret.referredEntities

repository/src/main/java/org/apache/atlas/discovery/SearchContext.java

@@ -22,6 +22,7 @@ import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.discovery.SearchParameters;
 import org.apache.atlas.model.discovery.SearchParameters.FilterCriteria;
+import org.apache.atlas.model.typedef.AtlasClassificationDef;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.type.AtlasClassificationType;
 import org.apache.atlas.type.AtlasEntityType;
@@ -49,14 +50,18 @@ public class SearchContext {
     private       SearchProcessor         searchProcessor;
     private       boolean                 terminateSearch = false;
 
+    public final static AtlasClassificationType MATCH_ALL_CLASSIFICATION = new AtlasClassificationType(new AtlasClassificationDef("*"));
+
     public SearchContext(SearchParameters searchParameters, AtlasTypeRegistry typeRegistry, AtlasGraph graph, Set<String> indexedKeys) throws AtlasBaseException {
+        String classificationName = searchParameters.getClassification();
+
         this.searchParameters   = searchParameters;
         this.typeRegistry       = typeRegistry;
         this.graph              = graph;
         this.indexedKeys        = indexedKeys;
-        entityAttributes        = new HashSet<>();
-        entityType              = typeRegistry.getEntityTypeByName(searchParameters.getTypeName());
-        classificationType      = typeRegistry.getClassificationTypeByName(searchParameters.getClassification());
+        this.entityAttributes   = new HashSet<>();
+        this.entityType         = typeRegistry.getEntityTypeByName(searchParameters.getTypeName());
+        this.classificationType = getClassificationType(classificationName);
 
         // Validate if the type name exists
         if (StringUtils.isNotEmpty(searchParameters.getTypeName()) && entityType == null) {
@@ -64,8 +69,8 @@ public class SearchContext {
         }
 
         // Validate if the classification exists
-        if (StringUtils.isNotEmpty(searchParameters.getClassification()) && classificationType == null) {
-            throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_CLASSIFICATION, searchParameters.getClassification());
+        if (StringUtils.isNotEmpty(classificationName) && classificationType == null) {
+            throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_CLASSIFICATION, classificationName);
         }
 
         // Invalid attributes will raise an exception with 400 error code
@@ -168,4 +173,16 @@ public class SearchContext {
             searchProcessor.addProcessor(processor);
         }
     }
+
+    private AtlasClassificationType getClassificationType(String classificationName) {
+        AtlasClassificationType ret;
+
+        if (StringUtils.equals(classificationName, MATCH_ALL_CLASSIFICATION.getTypeName())) {
+            ret = MATCH_ALL_CLASSIFICATION;
+        } else {
+            ret = typeRegistry.getClassificationTypeByName(classificationName);
+        }
+
+        return ret;
+    }
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java

@@ -323,6 +323,10 @@ public final class EntityGraphRetriever {
         return getClassifications(instanceVertex, null);
     }
 
+    public List<AtlasClassification> getClassifications(AtlasVertex instanceVertex) throws AtlasBaseException {
+        return getClassifications(instanceVertex, null);
+    }
+
     public AtlasClassification getClassification(String guid, String classificationName) throws AtlasBaseException {
 
         AtlasVertex instanceVertex = AtlasGraphUtilsV1.findByGuid(guid);