ATLAS-2477: search with tag and type doesn't list the entities where the tag is propagated

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

@@ -45,6 +45,7 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -149,13 +150,24 @@ public class ClassificationSearchProcessor extends SearchProcessor {
             }
         } else {
             tagGraphQueryWithAttributes = null;
+            List<AtlasGraphQuery> orConditions = new LinkedList<>();
 
             if (classificationType != SearchContext.MATCH_ALL_CLASSIFICATION) {
-                entityGraphQueryTraitNames = graph.query().in(Constants.TRAIT_NAMES_PROPERTY_KEY, typeAndSubTypes);
-                entityPredicateTraitNames  = SearchPredicateUtil.getContainsAnyPredicateGenerator().generatePredicate(Constants.TRAIT_NAMES_PROPERTY_KEY, classificationType.getTypeAndAllSubTypes(), List.class);
+                orConditions.add(graph.query().createChildQuery().in(Constants.TRAIT_NAMES_PROPERTY_KEY, typeAndSubTypes));
+                orConditions.add(graph.query().createChildQuery().in(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, typeAndSubTypes));
+
+                entityGraphQueryTraitNames = graph.query().or(orConditions);
+                entityPredicateTraitNames  = PredicateUtils.orPredicate(
+                        SearchPredicateUtil.getContainsAnyPredicateGenerator().generatePredicate(Constants.TRAIT_NAMES_PROPERTY_KEY, classificationType.getTypeAndAllSubTypes(), List.class),
+                        SearchPredicateUtil.getContainsAnyPredicateGenerator().generatePredicate(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, classificationType.getTypeAndAllSubTypes(), List.class));
             } else {
-                entityGraphQueryTraitNames = graph.query().has(Constants.TRAIT_NAMES_PROPERTY_KEY, NOT_EQUAL, null);
-                entityPredicateTraitNames = SearchPredicateUtil.getNotEmptyPredicateGenerator().generatePredicate(Constants.TRAIT_NAMES_PROPERTY_KEY, null, List.class);
+                orConditions.add(graph.query().createChildQuery().has(Constants.TRAIT_NAMES_PROPERTY_KEY, NOT_EQUAL, null));
+                orConditions.add(graph.query().createChildQuery().has(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, NOT_EQUAL, null));
+
+                entityGraphQueryTraitNames = graph.query().or(orConditions);
+                entityPredicateTraitNames  = PredicateUtils.orPredicate(
+                        SearchPredicateUtil.getNotEmptyPredicateGenerator().generatePredicate(Constants.TRAIT_NAMES_PROPERTY_KEY, null, List.class),
+                        SearchPredicateUtil.getNotEmptyPredicateGenerator().generatePredicate(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, null, List.class));
             }
 
             if (context.getSearchParameters().getExcludeDeletedEntities()) {

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

@@ -36,6 +36,7 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Iterator;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Set;
 
@@ -93,9 +94,11 @@ public class EntitySearchProcessor extends SearchProcessor {
         final Predicate traitPredicate;
 
         if (classificationType == SearchContext.MATCH_ALL_CLASSIFICATION) {
-            traitPredicate = SearchPredicateUtil.getNotEmptyPredicateGenerator().generatePredicate(Constants.TRAIT_NAMES_PROPERTY_KEY, null, List.class);
+            traitPredicate = PredicateUtils.orPredicate(SearchPredicateUtil.getNotEmptyPredicateGenerator().generatePredicate(Constants.TRAIT_NAMES_PROPERTY_KEY, null, List.class),
+                                                        SearchPredicateUtil.getNotEmptyPredicateGenerator().generatePredicate(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, null, List.class));
         } else {
-            traitPredicate = SearchPredicateUtil.getContainsAnyPredicateGenerator().generatePredicate(Constants.TRAIT_NAMES_PROPERTY_KEY, classificationTypeAndSubTypes, List.class);
+            traitPredicate = PredicateUtils.orPredicate(SearchPredicateUtil.getContainsAnyPredicateGenerator().generatePredicate(Constants.TRAIT_NAMES_PROPERTY_KEY, classificationTypeAndSubTypes, List.class),
+                                                        SearchPredicateUtil.getContainsAnyPredicateGenerator().generatePredicate(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, classificationTypeAndSubTypes, List.class));
         }
 
         processSearchAttributes(entityType, filterCriteria, indexAttributes, graphAttributes, allAttributes);
@@ -149,12 +152,18 @@ public class EntitySearchProcessor extends SearchProcessor {
 
             // If we need to filter on the trait names then we need to build the query and equivalent in-memory predicate
             if (filterClassification) {
+                List<AtlasGraphQuery> orConditions = new LinkedList<>();
+
                 if (classificationType == SearchContext.MATCH_ALL_CLASSIFICATION) {
-                    query.has(Constants.TRAIT_NAMES_PROPERTY_KEY, NOT_EQUAL, null);
+                    orConditions.add(query.createChildQuery().has(Constants.TRAIT_NAMES_PROPERTY_KEY, NOT_EQUAL, null));
+                    orConditions.add(query.createChildQuery().has(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, NOT_EQUAL, null));
                 } else {
-                    query.in(Constants.TRAIT_NAMES_PROPERTY_KEY, classificationTypeAndSubTypes);
+                    orConditions.add(query.createChildQuery().in(Constants.TRAIT_NAMES_PROPERTY_KEY, classificationTypeAndSubTypes));
+                    orConditions.add(query.createChildQuery().in(Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY, classificationTypeAndSubTypes));
                 }
 
+                query.or(orConditions);
+
                 // Construct a parallel in-memory predicate
                 if (graphQueryPredicate != null) {
                     graphQueryPredicate = PredicateUtils.andPredicate(graphQueryPredicate, traitPredicate);

repository/src/main/java/org/apache/atlas/query/GremlinClause.java

@@ -46,7 +46,7 @@ enum GremlinClause {
     TEXT_CONTAINS("has('%s', org.janusgraph.core.attribute.Text.textRegex(%s))"),
     TEXT_PREFIX("has('%s', org.janusgraph.core.attribute.Text.textPrefix(%s))"),
     TEXT_SUFFIX("has('%s', org.janusgraph.core.attribute.Text.textRegex(\".*\" + %s))"),
-    TRAIT("has('__traitNames', within('%s'))"),
+    TRAIT("or(has('__traitNames', within('%s')), has('__propagatedTraitNames', within('%s')))"),
     SELECT_NOOP_FN("def f(r){ r }; "),
     SELECT_FN("def f(r){ t=[[%s]]; %s r.each({t.add([%s])}); t.unique(); }; "),
     SELECT_ONLY_AGG_FN("def f(r){ t=[[%s]]; %s t.add([%s]); t;}; "),

repository/src/main/java/org/apache/atlas/query/GremlinQueryComposer.java

@@ -112,7 +112,7 @@ public class GremlinQueryComposer {
             IdentifierHelper.Info ia = createInfo(typeInfo.get());
 
             if (ia.isTrait()) {
-                add(GremlinClause.TRAIT, ia);
+                addTrait(GremlinClause.TRAIT, ia);
             } else {
                 if (ia.hasSubtypes()) {
                     add(GremlinClause.HAS_TYPE_WITHIN, ia.getSubTypes());
@@ -144,7 +144,7 @@ public class GremlinQueryComposer {
         }
 
         IdentifierHelper.Info traitInfo = createInfo(traitName);
-        add(GremlinClause.TRAIT, traitInfo);
+        addTrait(GremlinClause.TRAIT, traitInfo);
     }
 
     public void addWhere(String lhs, String operator, String rhs) {
@@ -591,6 +591,14 @@ public class GremlinQueryComposer {
         queryClauses.add(idx, new GremlinClauseValue(clause, clause.get(args)));
     }
 
+    private void addTrait(GremlinClause clause, IdentifierHelper.Info idInfo) {
+        if (context != null && !context.validator.isValid(context, clause, idInfo)) {
+            return;
+        }
+
+        add(clause, idInfo.get(), idInfo.get());
+    }
+
     static class GremlinClauseValue {
         private final GremlinClause clause;
         private final String        value;

repository/src/main/java/org/apache/atlas/util/AtlasGremlin2QueryProvider.java

@@ -68,7 +68,7 @@ public class AtlasGremlin2QueryProvider extends AtlasGremlinQueryProvider {
             case BASIC_SEARCH_TYPE_FILTER:
                 return ".has('__typeName', T.in, typeNames)";
             case BASIC_SEARCH_CLASSIFICATION_FILTER:
-                return ".has('__traitNames', T.in, traitNames)";
+                return ".or(has('__traitNames', T.in, traitNames), has('__propagatedTraitNames', T.in, traitNames))";
             case BASIC_SEARCH_STATE_FILTER:
                 return ".has('__state', state)";
             case TO_RANGE_LIST:

repository/src/test/java/org/apache/atlas/query/DSLQueriesTest.java

@@ -203,7 +203,7 @@ public class DSLQueriesTest extends BasicTestSetup {
                 {"hive_db has name", 3},
                 {"from hive_table", 10},
                 {"hive_table", 10},
-                {"hive_table isa Dimension", 3},
+                {"hive_table isa Dimension", 5},
                 {"hive_column where hive_column isa PII", 4},
                 {"hive_column where hive_column isa PII select hive_column.qualifiedName", 4},
                 {"hive_column select hive_column.qualifiedName", 17},
@@ -218,12 +218,12 @@ public class DSLQueriesTest extends BasicTestSetup {
                 {"hive_table where (name = \"sales_fact\" and createTime > \"2014-01-01\" ) select name as _col_0, createTime as _col_1", 1},
                 {"hive_table where (name = \"sales_fact\" and createTime >= \"2014-12-11T02:35:58.440Z\" ) select name as _col_0, createTime as _col_1", 1},
                 {"hive_table where (name = \"sales_fact\" and db.name = \"Sales\") select name, createTime", 1},
-                {"Dimension", 5},
+                {"Dimension", 9},
                 {"JdbcAccess", 2},
-                {"ETL", 5},
-                {"Metric", 5},
+                {"ETL", 10},
+                {"Metric", 8},
                 {"PII", 4},
-                {"`Log Data`", 3},
+                {"`Log Data`", 4},
                 {"DataSet where name='sales_fact'", 1},
                 {"Asset where name='sales_fact'", 1}
         };
@@ -305,11 +305,11 @@ public class DSLQueriesTest extends BasicTestSetup {
                 {"hive_table limit 5", 5},
                 {"hive_table limit 5 offset 5", 5},
 
-                {"hive_table isa Dimension", 3},
+                {"hive_table isa Dimension", 5},
                 {"hive_table isa Dimension limit 2", 2},
                 {"hive_table isa Dimension limit 2 offset 0", 2},
                 {"hive_table isa Dimension limit 2 offset 1", 2},
-                {"hive_table isa Dimension limit 3 offset 1", 2},
+                {"hive_table isa Dimension limit 3 offset 1", 3},
                 {"hive_table where db.name='Sales' and db.clusterName='cl1'", 4},
 
                 {"hive_column where hive_column isa PII", 4},

repository/src/test/java/org/apache/atlas/query/GremlinQueryComposerTest.java

@@ -36,13 +36,13 @@ import static org.testng.Assert.fail;
 public class GremlinQueryComposerTest {
     @Test
     public void classification() {
-        String expected = "g.V().has('__traitNames', within('PII')).dedup().limit(25).toList()";
+        String expected = "g.V().or(has('__traitNames', within('PII')), has('__propagatedTraitNames', within('PII'))).dedup().limit(25).toList()";
         verify("PII", expected);
     }
 
     @Test()
     public void dimension() {
-        String expected = "g.V().has('__typeName', 'Table').has('__traitNames', within('Dimension')).dedup().limit(25).toList()";
+        String expected = "g.V().has('__typeName', 'Table').or(has('__traitNames', within('Dimension')), has('__propagatedTraitNames', within('Dimension'))).dedup().limit(25).toList()";
         verify("Table isa Dimension", expected);
         verify("Table is Dimension", expected);
         verify("Table where Table is Dimension", expected);
@@ -295,14 +295,14 @@ public class GremlinQueryComposerTest {
     @Test
     public void keywordsInWhereClause() {
         verify("Table as t where t has name and t isa Dimension",
-                "g.V().has('__typeName', 'Table').as('t').and(__.has('Table.name'),__.has('__traitNames', within('Dimension'))).dedup().limit(25).toList()");
+                "g.V().has('__typeName', 'Table').as('t').and(__.has('Table.name'),__.or(has('__traitNames', within('Dimension')), has('__propagatedTraitNames', within('Dimension')))).dedup().limit(25).toList()");
         verify("Table as t where t has name and t.name = 'sales_fact'",
                 "g.V().has('__typeName', 'Table').as('t').and(__.has('Table.name'),__.has('Table.name', eq('sales_fact'))).dedup().limit(25).toList()");
         verify("Table as t where t is Dimension and t.name = 'sales_fact'",
-                "g.V().has('__typeName', 'Table').as('t').and(__.has('__traitNames', within('Dimension')),__.has('Table.name', eq('sales_fact'))).dedup().limit(25).toList()");
-        verify("Table isa 'Dimension' and name = 'sales_fact'", "g.V().has('__typeName', 'Table').and(__.has('__traitNames', within('Dimension')),__.has('Table.name', eq('sales_fact'))).dedup().limit(25).toList()");
+                "g.V().has('__typeName', 'Table').as('t').and(__.or(has('__traitNames', within('Dimension')), has('__propagatedTraitNames', within('Dimension'))),__.has('Table.name', eq('sales_fact'))).dedup().limit(25).toList()");
+        verify("Table isa 'Dimension' and name = 'sales_fact'", "g.V().has('__typeName', 'Table').and(__.or(has('__traitNames', within('Dimension')), has('__propagatedTraitNames', within('Dimension'))),__.has('Table.name', eq('sales_fact'))).dedup().limit(25).toList()");
         verify("Table has name and name = 'sales_fact'", "g.V().has('__typeName', 'Table').and(__.has('Table.name'),__.has('Table.name', eq('sales_fact'))).dedup().limit(25).toList()");
-        verify("Table is 'Dimension' and Table has owner and name = 'sales_fact'", "g.V().has('__typeName', 'Table').and(__.has('__traitNames', within('Dimension')),__.has('Table.owner'),__.has('Table.name', eq('sales_fact'))).dedup().limit(25).toList()");
+        verify("Table is 'Dimension' and Table has owner and name = 'sales_fact'", "g.V().has('__typeName', 'Table').and(__.or(has('__traitNames', within('Dimension')), has('__propagatedTraitNames', within('Dimension'))),__.has('Table.owner'),__.has('Table.name', eq('sales_fact'))).dedup().limit(25).toList()");
         verify("Table has name and Table has owner and name = 'sales_fact'", "g.V().has('__typeName', 'Table').and(__.has('Table.name'),__.has('Table.owner'),__.has('Table.name', eq('sales_fact'))).dedup().limit(25).toList()");
     }