ATLAS-2296: basic search enhancement to optionally exclude sub-type entities and…

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

@@ -45,6 +45,8 @@ public class SearchParameters implements Serializable {
     private String  classification;
     private boolean excludeDeletedEntities;
     private boolean includeClassificationAttributes;
+    private boolean includeSubTypes                 = true;
+    private boolean includeSubClassifications       = true;
     private int     limit;
     private int     offset;
 
@@ -129,6 +131,36 @@ public class SearchParameters implements Serializable {
     }
 
     /**
+     * @return True iff sub-type entities are to be included
+     */
+    public boolean getIncludeSubTypes() {
+        return includeSubTypes;
+    }
+
+    /**
+     * Include sub-type entities in search
+     * @param includeSubTypes boolean flag
+     */
+    public void setIncludeSubTypes(boolean includeSubTypes) {
+        this.includeSubTypes = includeSubTypes;
+    }
+
+    /**
+     * @return True iff sub-classifications are to be included
+     */
+    public boolean getIncludeSubClassifications() {
+        return includeSubClassifications;
+    }
+
+    /**
+     * Include sub-classifications in search
+     * @param includeSubClassifications boolean flag
+     */
+    public void setIncludeSubClassifications(boolean includeSubClassifications) {
+        this.includeSubClassifications = includeSubClassifications;
+    }
+
+    /**
      * @return Max number of results to be returned
      */
     public int getLimit() {

intg/src/main/java/org/apache/atlas/type/AtlasClassificationType.java

@@ -38,6 +38,7 @@ public class AtlasClassificationType extends AtlasStructType {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasClassificationType.class);
 
     private final AtlasClassificationDef classificationDef;
+    private final String                 typeQryStr;
 
     private List<AtlasClassificationType> superTypes               = Collections.emptyList();
     private Set<String>                   allSuperTypes            = Collections.emptySet();
@@ -58,6 +59,7 @@ public class AtlasClassificationType extends AtlasStructType {
         super(classificationDef);
 
         this.classificationDef = classificationDef;
+        this.typeQryStr        = AtlasAttribute.escapeIndexQueryValue(Collections.singleton(getTypeName()));
     }
 
     /**
@@ -73,6 +75,7 @@ public class AtlasClassificationType extends AtlasStructType {
         super(classificationDef);
 
         this.classificationDef = classificationDef;
+        this.typeQryStr        = AtlasAttribute.escapeIndexQueryValue(Collections.singleton(getTypeName()));
 
         resolveReferences(typeRegistry);
     }
@@ -238,6 +241,8 @@ public class AtlasClassificationType extends AtlasStructType {
 
     public Set<String> getTypeAndAllSubTypes() { return typeAndAllSubTypes; }
 
+    public String getTypeQryStr() { return typeQryStr; }
+
     public String getTypeAndAllSubTypesQryStr() {
         if (StringUtils.isEmpty(typeAndAllSubTypesQryStr)) {
             typeAndAllSubTypesQryStr = AtlasAttribute.escapeIndexQueryValue(typeAndAllSubTypes);

intg/src/main/java/org/apache/atlas/type/AtlasEntityType.java

@@ -48,6 +48,7 @@ public class AtlasEntityType extends AtlasStructType {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityType.class);
 
     private final AtlasEntityDef entityDef;
+    private final String         typeQryStr;
 
     private List<AtlasEntityType>                    superTypes                 = Collections.emptyList();
     private Set<String>                              allSuperTypes              = Collections.emptySet();
@@ -64,12 +65,14 @@ public class AtlasEntityType extends AtlasStructType {
         super(entityDef);
 
         this.entityDef  = entityDef;
+        this.typeQryStr = AtlasAttribute.escapeIndexQueryValue(Collections.singleton(getTypeName()));
     }
 
     public AtlasEntityType(AtlasEntityDef entityDef, AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
         super(entityDef);
 
         this.entityDef  = entityDef;
+        this.typeQryStr = AtlasAttribute.escapeIndexQueryValue(Collections.singleton(getTypeName()));
 
         resolveReferences(typeRegistry);
     }
@@ -240,6 +243,8 @@ public class AtlasEntityType extends AtlasStructType {
         return typeAndAllSubTypesQryStr;
     }
 
+    public String getTypeQryStr() { return typeQryStr; }
+
     public boolean hasRelationshipAttribute(String attributeName) {
         return relationshipAttributes.containsKey(attributeName);
     }

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

@@ -41,6 +41,7 @@ import org.slf4j.LoggerFactory;
 import javax.script.ScriptEngine;
 import javax.script.ScriptException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -69,12 +70,19 @@ public class ClassificationSearchProcessor extends SearchProcessor {
 
         final AtlasClassificationType classificationType = context.getClassificationType();
         final FilterCriteria          filterCriteria     = context.getSearchParameters().getTagFilters();
-        final Set<String>             typeAndSubTypes       = classificationType.getTypeAndAllSubTypes();
-        final String                  typeAndSubTypesQryStr = classificationType.getTypeAndAllSubTypesQryStr();
         final Set<String>             indexAttributes    = new HashSet<>();
         final Set<String>             graphAttributes    = new HashSet<>();
         final Set<String>             allAttributes      = new HashSet<>();
+        final Set<String>             typeAndSubTypes;
+        final String                  typeAndSubTypesQryStr;
 
+        if (context.getSearchParameters().getIncludeSubClassifications()) {
+            typeAndSubTypes       = classificationType.getTypeAndAllSubTypes();
+            typeAndSubTypesQryStr = classificationType.getTypeAndAllSubTypesQryStr();
+        } else {
+            typeAndSubTypes       = Collections.singleton(classificationType.getTypeName());
+            typeAndSubTypesQryStr = classificationType.getTypeQryStr();
+        }
 
         processSearchAttributes(classificationType, filterCriteria, indexAttributes, graphAttributes, allAttributes);
 

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

@@ -55,16 +55,36 @@ public class EntitySearchProcessor extends SearchProcessor {
 
         final AtlasEntityType entityType      = context.getEntityType();
         final FilterCriteria  filterCriteria  = context.getSearchParameters().getEntityFilters();
-        final Set<String>     typeAndSubTypes       = entityType.getTypeAndAllSubTypes();
-        final String          typeAndSubTypesQryStr = entityType.getTypeAndAllSubTypesQryStr();
         final Set<String>     indexAttributes = new HashSet<>();
         final Set<String>     graphAttributes = new HashSet<>();
         final Set<String>     allAttributes   = new HashSet<>();
+        final Set<String>     typeAndSubTypes;
+        final String          typeAndSubTypesQryStr;
+
+        if (context.getSearchParameters().getIncludeSubTypes()) {
+            typeAndSubTypes       = entityType.getTypeAndAllSubTypes();
+            typeAndSubTypesQryStr = entityType.getTypeAndAllSubTypesQryStr();
+        } else {
+            typeAndSubTypes       = Collections.singleton(entityType.getTypeName());
+            typeAndSubTypesQryStr = entityType.getTypeQryStr();
+        }
 
         final AtlasClassificationType classificationType = context.getClassificationType();
-        final boolean                 filterClassification          = classificationType != null && !context.needClassificationProcessor();
-        final Set<String>             classificationTypeAndSubTypes = classificationType != null ? classificationType.getTypeAndAllSubTypes() : Collections.EMPTY_SET;
+        final boolean                 filterClassification;
+        final Set<String>             classificationTypeAndSubTypes;
+
+        if (classificationType != null) {
+            filterClassification = !context.needClassificationProcessor();
 
+            if (context.getSearchParameters().getIncludeSubClassifications()) {
+                classificationTypeAndSubTypes = classificationType.getTypeAndAllSubTypes();
+            } else {
+                classificationTypeAndSubTypes = Collections.singleton(classificationType.getTypeName());
+            }
+        } else {
+            filterClassification          = false;
+            classificationTypeAndSubTypes = Collections.emptySet();
+        }
 
         final Predicate typeNamePredicate = SearchPredicateUtil.getINPredicateGenerator()
                                                                .generatePredicate(Constants.TYPE_NAME_PROPERTY_KEY, typeAndSubTypes, String.class);