ATLAS-3613: updated search to support namespace attributes

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

@@ -238,6 +238,11 @@ public class AtlasClassificationType extends AtlasStructType {
         classificationDef.setSubTypes(subTypes);
     }
 
+    @Override
+    public AtlasAttribute getSystemAttribute(String attributeName) {
+        return AtlasClassificationType.CLASSIFICATION_ROOT.allAttributes.get(attributeName);
+    }
+
     private void addSubType(AtlasClassificationType subType) {
         subTypes.add(subType.getTypeName());
     }

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

@@ -226,17 +226,19 @@ public class AtlasStructType extends AtlasType {
             ret = getSystemAttribute(attributeName);
         }
 
+        if (ret == null) {
+            ret = getNamespaceAttribute(attributeName);
+        }
+
         return ret;
     }
 
     public AtlasAttribute getSystemAttribute(String attributeName) {
-        AtlasAttribute ret = null;
-        if (this instanceof AtlasEntityType) {
-            ret = AtlasEntityType.ENTITY_ROOT.allAttributes.get(attributeName);
-        } else if (this instanceof AtlasClassificationType) {
-            ret = AtlasClassificationType.CLASSIFICATION_ROOT.allAttributes.get(attributeName);
-        }
-        return ret;
+        return null;
+    }
+
+    public AtlasAttribute getNamespaceAttribute(String attributeName) {
+        return null;
     }
 
     @Override

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

@@ -119,9 +119,9 @@ public class GremlinQueryComposer {
             if (ia.isTrait()) {
                 String traitName = ia.get();
 
-                if (traitName.equalsIgnoreCase(ALL_CLASSIFICATIONS)) {
+                if (traitName.equals(ALL_CLASSIFICATIONS)) {
                     addTrait(GremlinClause.ANY_TRAIT, ia);
-                } else if (traitName.equalsIgnoreCase(NO_CLASSIFICATIONS)) {
+                } else if (traitName.equals(NO_CLASSIFICATIONS)) {
                     addTrait(GremlinClause.NO_TRAIT, ia);
                 } else {
                     addTrait(GremlinClause.TRAIT, ia);

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

@@ -20,6 +20,8 @@ package org.apache.atlas.query;
 
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.type.AtlasNamespaceType;
+import org.apache.atlas.type.AtlasType;
 import org.apache.commons.lang.StringUtils;
 
 import java.util.regex.Matcher;
@@ -148,7 +150,6 @@ public class IdentifierHelper {
                         raw = context.getTypeNameFromAlias(this.raw);
                     }
 
-                    updateParts();
                     updateTypeInfo(lookup, context);
                     setIsTrait(context, lookup, attributeName);
                     updateEdgeInfo(lookup, context);
@@ -183,15 +184,29 @@ public class IdentifierHelper {
         }
 
         private void updateTypeInfo(org.apache.atlas.query.Lookup lookup, GremlinQueryComposer.Context context) {
+            parts = StringUtils.split(raw, ".");
+
+            // check if this is a namespace attribute
+            if (parts.length == 2) {
+                try {
+                    AtlasType type = lookup.getType(parts[0]);
+
+                    if (type instanceof AtlasNamespaceType) {
+                        parts = new String[1];
+                        parts[0] = raw;
+                    }
+                } catch (AtlasBaseException excp) {
+                    // ignore
+                }
+            }
+
             if (parts.length == 1) {
                 typeName = context.hasAlias(parts[0]) ?
                                    context.getTypeNameFromAlias(parts[0]) :
                                    context.getActiveTypeName();
                 qualifiedName = getDefaultQualifiedNameForSinglePartName(context, parts[0]);
                 attributeName = parts[0];
-            }
-
-            if (parts.length == 2) {
+            } else if (parts.length == 2) {
                 boolean isAttrOfActiveType = lookup.hasAttribute(context, parts[0]);
                 if (isAttrOfActiveType) {
                     attributeName = parts[0];
@@ -242,10 +257,6 @@ public class IdentifierHelper {
             }
         }
 
-        private void updateParts() {
-            parts = StringUtils.split(raw, ".");
-        }
-
         public String getQualifiedName() {
             return qualifiedName;
         }

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

@@ -21,7 +21,6 @@ package org.apache.atlas.query;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
-import org.apache.atlas.repository.Constants;
 import org.apache.atlas.type.*;
 import org.apache.commons.lang.StringUtils;
 
@@ -29,20 +28,12 @@ import java.util.*;
 
 import static org.apache.atlas.discovery.SearchContext.MATCH_ALL_CLASSIFIED;
 import static org.apache.atlas.discovery.SearchContext.MATCH_ALL_NOT_CLASSIFIED;
+import static org.apache.atlas.discovery.SearchContext.MATCH_ALL_ENTITY_TYPES;
 import static org.apache.atlas.model.discovery.SearchParameters.ALL_CLASSIFICATIONS;
 import static org.apache.atlas.model.discovery.SearchParameters.NO_CLASSIFICATIONS;
+import static org.apache.atlas.model.discovery.SearchParameters.ALL_ENTITY_TYPES;
 
 class RegistryBasedLookup implements Lookup {
-    private static final Set<String> SYSTEM_ATTRIBUTES = new HashSet<>(
-            Arrays.asList(Constants.GUID_PROPERTY_KEY,
-                    Constants.MODIFIED_BY_KEY,
-                    Constants.CREATED_BY_KEY,
-                    Constants.STATE_PROPERTY_KEY,
-                    Constants.TIMESTAMP_PROPERTY_KEY,
-                    Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY,
-                    Constants.HOME_ID_KEY
-            ));
-
     private static final Map<String, String> NUMERIC_ATTRIBUTES = new HashMap<String, String>() {{
             put(AtlasBaseTypeDef.ATLAS_TYPE_SHORT, "");
             put(AtlasBaseTypeDef.ATLAS_TYPE_INT, "");
@@ -63,10 +54,12 @@ class RegistryBasedLookup implements Lookup {
     public AtlasType getType(String typeName) throws AtlasBaseException {
         AtlasType ret;
 
-        if (typeName.equalsIgnoreCase(ALL_CLASSIFICATIONS)) {
+        if (typeName.equals(ALL_CLASSIFICATIONS)) {
             ret = MATCH_ALL_CLASSIFIED;
-        } else if (typeName.equalsIgnoreCase(NO_CLASSIFICATIONS)) {
+        } else if (typeName.equals(NO_CLASSIFICATIONS)) {
             ret = MATCH_ALL_NOT_CLASSIFIED;
+        } else if (typeName.equals(ALL_ENTITY_TYPES)) {
+            ret = MATCH_ALL_ENTITY_TYPES;
         } else {
             ret = typeRegistry.getType(typeName);
         }
@@ -81,15 +74,7 @@ class RegistryBasedLookup implements Lookup {
             return "";
         }
 
-        if(isSystemAttribute(name)) {
-            return name;
-        } else {
-            return et.getQualifiedAttributeName(name);
-        }
-    }
-
-    private boolean isSystemAttribute(String s) {
-        return SYSTEM_ATTRIBUTES.contains(s);
+        return et.getQualifiedAttributeName(name);
     }
 
     @Override
@@ -99,10 +84,6 @@ class RegistryBasedLookup implements Lookup {
             return false;
         }
 
-        if(isSystemAttribute(attributeName)) {
-            return true;
-        }
-
         AtlasType at = getAttributeType(et, attributeName);
         if(at == null) {
             return false;
@@ -144,8 +125,7 @@ class RegistryBasedLookup implements Lookup {
     public boolean hasAttribute(GremlinQueryComposer.Context context, String typeName) {
         AtlasEntityType entityType = context.getActiveEntityType();
 
-        return (entityType != null) &&
-                (isSystemAttribute(typeName) || entityType.hasAttribute(typeName) || entityType.hasRelationshipAttribute(typeName));
+        return getAttribute(entityType, typeName) != null;
     }
 
     @Override
@@ -178,6 +158,8 @@ class RegistryBasedLookup implements Lookup {
                 t = MATCH_ALL_CLASSIFIED;
             } else if (typeName.equalsIgnoreCase(NO_CLASSIFICATIONS)) {
                 t = MATCH_ALL_NOT_CLASSIFIED;
+            } else if (typeName.equalsIgnoreCase(ALL_ENTITY_TYPES)) {
+                t = MATCH_ALL_ENTITY_TYPES;
             } else {
                 t = typeRegistry.getType(typeName);
             }
@@ -257,6 +239,11 @@ class RegistryBasedLookup implements Lookup {
     @Override
     public String getVertexPropertyName(String typeName, String attrName) {
         AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
+
+        if (entityType == null && StringUtils.equals(typeName, ALL_ENTITY_TYPES)) {
+            entityType = MATCH_ALL_ENTITY_TYPES;
+        }
+
         AtlasStructType.AtlasAttribute attribute = getAttribute(entityType, attrName);
         if (attribute == null) {
             return null;

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2.java

@@ -1377,19 +1377,19 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
     private void validateNamespaceAttributes(AtlasVertex entityVertex, AtlasEntityType entityType, Map<String, Map<String, Object>> entityNamespaces, boolean isOverwrite) throws AtlasBaseException {
         List<String> messages = new ArrayList<>();
 
-        Map<String, List<AtlasNamespaceAttribute>> entityTypeNamespaces = entityType.getNamespaceAttributes();
+        Map<String, Map<String, AtlasNamespaceAttribute>> entityTypeNamespaces = entityType.getNamespaceAttributes();
 
         for (String nsName : entityNamespaces.keySet()) {
-            if (!entityNamespaces.containsKey(nsName)) {
+            if (!entityTypeNamespaces.containsKey(nsName)) {
                 messages.add(nsName + ": invalid namespace for entity type " + entityType.getTypeName());
 
                 continue;
             }
 
-            List<AtlasNamespaceAttribute> entityTypeNsAttributes = entityTypeNamespaces.get(nsName);
-            Map<String, Object>           entityNsAttributes     = entityNamespaces.get(nsName);
+            Map<String, AtlasNamespaceAttribute> entityTypeNsAttributes = entityTypeNamespaces.get(nsName);
+            Map<String, Object>                  entityNsAttributes     = entityNamespaces.get(nsName);
 
-            for (AtlasNamespaceAttribute nsAttribute : entityTypeNsAttributes) {
+            for (AtlasNamespaceAttribute nsAttribute : entityTypeNsAttributes.values()) {
                 AtlasType attrType  = nsAttribute.getAttributeType();
                 String    attrName  = nsAttribute.getName();
                 Object    attrValue = entityNsAttributes.get(attrName);

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java

@@ -427,14 +427,14 @@ public class EntityGraphMapper {
             LOG.debug("==> setNamespaceAttributes(entityVertex={}, entityType={}, entityNamespaces={}", entityVertex, entityType.getTypeName(), entityNamespaces);
         }
 
-        Map<String, List<AtlasNamespaceAttribute>> entityTypeNamespaces = entityType.getNamespaceAttributes();
+        Map<String, Map<String, AtlasNamespaceAttribute>> entityTypeNamespaces = entityType.getNamespaceAttributes();
 
-        for (Map.Entry<String, List<AtlasNamespaceAttribute>> entry : entityTypeNamespaces.entrySet()) {
-            String                        nsName                 = entry.getKey();
-            List<AtlasNamespaceAttribute> entityTypeNsAttributes = entry.getValue();
-            Map<String, Object>           entityNsAttributes     = MapUtils.isEmpty(entityNamespaces) ? null : entityNamespaces.get(nsName);
+        for (Map.Entry<String, Map<String, AtlasNamespaceAttribute>> entry : entityTypeNamespaces.entrySet()) {
+            String                               nsName                 = entry.getKey();
+            Map<String, AtlasNamespaceAttribute> entityTypeNsAttributes = entry.getValue();
+            Map<String, Object>                  entityNsAttributes     = MapUtils.isEmpty(entityNamespaces) ? null : entityNamespaces.get(nsName);
 
-            for (AtlasNamespaceAttribute nsAttribute : entityTypeNsAttributes) {
+            for (AtlasNamespaceAttribute nsAttribute : entityTypeNsAttributes.values()) {
                 String nsAttrName          = nsAttribute.getName();
                 Object nsAttrExistingValue = entityVertex.getProperty(nsAttribute.getVertexPropertyName(), Object.class);
                 Object nsAttrNewValue      = MapUtils.isEmpty(entityNsAttributes) ? null : entityNsAttributes.get(nsAttrName);
@@ -480,19 +480,19 @@ public class EntityGraphMapper {
             LOG.debug("==> addOrUpdateNamespaceAttributes(entityVertex={}, entityType={}, entityNamespaces={}", entityVertex, entityType.getTypeName(), entityNamespaces);
         }
 
-        Map<String, List<AtlasNamespaceAttribute>> entityTypeNamespaces = entityType.getNamespaceAttributes();
+        Map<String, Map<String, AtlasNamespaceAttribute>> entityTypeNamespaces = entityType.getNamespaceAttributes();
 
         if (MapUtils.isNotEmpty(entityTypeNamespaces) && MapUtils.isNotEmpty(entityNamespaces)) {
-            for (Map.Entry<String, List<AtlasNamespaceAttribute>> entry : entityTypeNamespaces.entrySet()) {
-                String                        nsName                 = entry.getKey();
-                List<AtlasNamespaceAttribute> entityTypeNsAttributes = entry.getValue();
-                Map<String, Object>           entityNsAttributes     = entityNamespaces.get(nsName);
+            for (Map.Entry<String, Map<String, AtlasNamespaceAttribute>> entry : entityTypeNamespaces.entrySet()) {
+                String                               nsName                 = entry.getKey();
+                Map<String, AtlasNamespaceAttribute> entityTypeNsAttributes = entry.getValue();
+                Map<String, Object>                  entityNsAttributes     = entityNamespaces.get(nsName);
 
                 if (MapUtils.isEmpty(entityNsAttributes)) {
                     continue;
                 }
 
-                for (AtlasNamespaceAttribute nsAttribute : entityTypeNsAttributes) {
+                for (AtlasNamespaceAttribute nsAttribute : entityTypeNsAttributes.values()) {
                     String nsAttrName = nsAttribute.getName();
 
                     if (!entityNsAttributes.containsKey(nsAttrName)) {
@@ -528,12 +528,12 @@ public class EntityGraphMapper {
             LOG.debug("==> removeNamespaceAttributes(entityVertex={}, entityType={}, entityNamespaces={}", entityVertex, entityType.getTypeName(), entityNamespaces);
         }
 
-        Map<String, List<AtlasNamespaceAttribute>> entityTypeNamespaces = entityType.getNamespaceAttributes();
+        Map<String, Map<String, AtlasNamespaceAttribute>> entityTypeNamespaces = entityType.getNamespaceAttributes();
 
         if (MapUtils.isNotEmpty(entityTypeNamespaces) && MapUtils.isNotEmpty(entityNamespaces)) {
-            for (Map.Entry<String, List<AtlasNamespaceAttribute>> entry : entityTypeNamespaces.entrySet()) {
-                String                        nsName                 = entry.getKey();
-                List<AtlasNamespaceAttribute> entityTypeNsAttributes = entry.getValue();
+            for (Map.Entry<String, Map<String, AtlasNamespaceAttribute>> entry : entityTypeNamespaces.entrySet()) {
+                String                               nsName                 = entry.getKey();
+                Map<String, AtlasNamespaceAttribute> entityTypeNsAttributes = entry.getValue();
 
                 if (!entityNamespaces.containsKey(nsName)) { // nothing to remove for this namespace
                     continue;
@@ -541,7 +541,7 @@ public class EntityGraphMapper {
 
                 Map<String, Object> entityNsAttributes = entityNamespaces.get(nsName);
 
-                for (AtlasNamespaceAttribute nsAttribute : entityTypeNsAttributes) {
+                for (AtlasNamespaceAttribute nsAttribute : entityTypeNsAttributes.values()) {
                     // if (entityNsAttributes is empty) remove all attributes in this namespace
                     // else remove the attribute only if its given in entityNsAttributes
                     if (MapUtils.isEmpty(entityNsAttributes) || entityNsAttributes.containsKey(nsAttribute.getName())) {

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

@@ -678,21 +678,26 @@ public class EntityGraphRetriever {
 
             if (CollectionUtils.isNotEmpty(attributes)) {
                 for (String attrName : attributes) {
-                    String nonQualifiedAttrName = toNonQualifiedName(attrName);
-                    if (ret.hasAttribute(attrName)) {
-                        continue;
-                    }
-
-                    AtlasAttribute attribute = entityType.getAttribute(nonQualifiedAttrName);
+                    AtlasAttribute attribute = entityType.getAttribute(attrName);
 
                     if (attribute == null) {
-                        attribute = entityType.getRelationshipAttribute(nonQualifiedAttrName, null);
+                        attrName = toNonQualifiedName(attrName);
+
+                        if (ret.hasAttribute(attrName)) {
+                            continue;
+                        }
+
+                        attribute = entityType.getAttribute(attrName);
+
+                        if (attribute == null) {
+                            attribute = entityType.getRelationshipAttribute(attrName, null);
+                        }
                     }
 
                     Object attrValue = getVertexAttribute(entityVertex, attribute);
 
                     if (attrValue != null) {
-                        ret.setAttribute(nonQualifiedAttrName, attrValue);
+                        ret.setAttribute(attrName, attrValue);
                     }
                 }
             }
@@ -765,15 +770,15 @@ public class EntityGraphRetriever {
     }
 
     private void mapNamespaceAttributes(AtlasVertex entityVertex, AtlasEntity entity) throws AtlasBaseException {
-        AtlasEntityType                            entityType           = typeRegistry.getEntityTypeByName(entity.getTypeName());
-        Map<String, List<AtlasNamespaceAttribute>> entityTypeNamespaces = entityType != null ? entityType.getNamespaceAttributes() : null;
+        AtlasEntityType                                   entityType           = typeRegistry.getEntityTypeByName(entity.getTypeName());
+        Map<String, Map<String, AtlasNamespaceAttribute>> entityTypeNamespaces = entityType != null ? entityType.getNamespaceAttributes() : null;
 
         if (MapUtils.isNotEmpty(entityTypeNamespaces)) {
-            for (Map.Entry<String, List<AtlasNamespaceAttribute>> entry : entityTypeNamespaces.entrySet()) {
-                String                        nsName       = entry.getKey();
-                List<AtlasNamespaceAttribute> nsAttributes = entry.getValue();
+            for (Map.Entry<String, Map<String, AtlasNamespaceAttribute>> entry : entityTypeNamespaces.entrySet()) {
+                String                               nsName       = entry.getKey();
+                Map<String, AtlasNamespaceAttribute> nsAttributes = entry.getValue();
 
-                for (AtlasNamespaceAttribute nsAttribute : nsAttributes) {
+                for (AtlasNamespaceAttribute nsAttribute : nsAttributes.values()) {
                     Object nsAttrValue = mapVertexToAttribute(entityVertex, nsAttribute, null, false, false);
 
                     if (nsAttrValue != null) {

repository/src/test/java/org/apache/atlas/repository/store/graph/v2/AtlasNamespaceDefStoreV2Test.java

@@ -25,7 +25,7 @@ import org.apache.atlas.model.typedef.AtlasNamespaceDef;
 import org.apache.atlas.model.typedef.AtlasStructDef;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.type.AtlasEntityType;
-import org.apache.atlas.type.AtlasNamespaceType;
+import org.apache.atlas.type.AtlasNamespaceType.AtlasNamespaceAttribute;
 import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.testng.Assert;
@@ -84,7 +84,7 @@ public class AtlasNamespaceDefStoreV2Test {
         createNamespaceTypes(namespaceName);
         Assert.assertEquals(typeRegistry.getAllNamespaceDefs().size(), 1);
         AtlasEntityType entityType = typeRegistry.getEntityTypeByName("hive_table");
-        Map<String, List<AtlasNamespaceType.AtlasNamespaceAttribute>> m1 = entityType.getNamespaceAttributes();
+        Map<String, Map<String, AtlasNamespaceAttribute>> m1 = entityType.getNamespaceAttributes();
         Assert.assertEquals(m1.get(namespaceName).size(), 2);
     }
 
@@ -117,7 +117,7 @@ public class AtlasNamespaceDefStoreV2Test {
         updateNamespaceTypeDefs(namespaceDef);
         typeDefStore.updateTypesDef(typesDefs);
         AtlasEntityType entityType = typeRegistry.getEntityTypeByName("hive_table");
-        Map<String, List<AtlasNamespaceType.AtlasNamespaceAttribute>> m1 = entityType.getNamespaceAttributes();
+        Map<String, Map<String, AtlasNamespaceAttribute>> m1 = entityType.getNamespaceAttributes();
         Assert.assertEquals(m1.get(namespaceName).size(), 3);
     }