ATLAS-3048: Remove use of hard coded qualifiedName.

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

@@ -65,7 +65,10 @@ public class ClassificationAssociator {
 
             Map<String, AtlasEntityHeader> guidEntityHeaderMap = new HashMap<>();
             for (String guid : guids) {
-                AtlasEntityHeader entityHeader = entityRetriever.toAtlasEntityHeaderWithClassifications(guid);
+                AtlasEntityHeader entityHeader = getEntityHeaderByGuid(guid);
+                if (entityHeader == null) {
+                    continue;
+                }
 
                 guidEntityHeaderMap.put(guid, entityHeader);
             }
@@ -74,6 +77,16 @@ public class ClassificationAssociator {
             return new AtlasEntityHeaders(guidEntityHeaderMap);
         }
 
+        private AtlasEntityHeader getEntityHeaderByGuid(String guid) {
+            try {
+                return entityRetriever.toAtlasEntityHeaderWithClassifications(guid);
+            } catch (AtlasBaseException e) {
+                LOG.error("Error fetching entity: {}", guid, e);
+            }
+
+            return null;
+        }
+
         private long incrementTimestamp(long t) {
             return t + 1;
         }
@@ -95,6 +108,7 @@ public class ClassificationAssociator {
         private final AtlasEntityStore entitiesStore;
         private final EntityGraphRetriever entityRetriever;
         private StringBuilder actionSummary = new StringBuilder();
+        private Map<String, String> typeNameUniqueAttributeNameMap = new HashMap<>();
 
         public Updater(AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore) {
             this.typeRegistry = typeRegistry;
@@ -105,6 +119,7 @@ public class ClassificationAssociator {
         public String setClassifications(Map<String, AtlasEntityHeader> map) {
             for (AtlasEntityHeader incomingEntityHeader : map.values()) {
                 String typeName = incomingEntityHeader.getTypeName();
+
                 AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
                 if (entityType == null) {
                     LOG.warn("Entity type: {}: Not found: {}!", typeName, STATUS_SKIPPED);
@@ -112,7 +127,11 @@ public class ClassificationAssociator {
                     continue;
                 }
 
-                String qualifiedName = getQualifiedName(incomingEntityHeader);
+                String qualifiedName = getUniqueAttributeName(entityType, incomingEntityHeader);
+                if (StringUtils.isEmpty(qualifiedName)) {
+                    qualifiedName = "<no unique name>";
+                }
+
                 AtlasEntityHeader entityToBeChanged = getByUniqueAttributes(entityType, qualifiedName, incomingEntityHeader.getAttributes());
                 if (entityToBeChanged == null) {
                     summarizeFormat("Entity:%s:%s:[Not found]:%s", entityType.getTypeName(), qualifiedName, STATUS_SKIPPED);
@@ -246,8 +265,21 @@ public class ClassificationAssociator {
             return list.stream().map(AtlasClassification::getTypeName).collect(Collectors.joining(", "));
         }
 
-        private String getQualifiedName(AtlasEntityHeader entityHeader) {
-            return (String) entityHeader.getAttribute(ATTR_NAME_QUALIFIED_NAME);
+        private String getUniqueAttributeName(AtlasEntityType entityType, AtlasEntityHeader entityHeader) {
+            String uniqueAttrName = ATTR_NAME_QUALIFIED_NAME;
+            if (!entityHeader.getAttributes().containsKey(uniqueAttrName)) {
+                uniqueAttrName = getUniqueAttributeName(entityType);
+            }
+
+            return uniqueAttrName;
+        }
+
+        private String getUniqueAttributeName(AtlasEntityType entityType) {
+            return entityType.getUniqAttributes()
+                    .entrySet()
+                    .stream()
+                    .findFirst()
+                    .get().getKey();
         }
 
         private void summarize(String... s) {

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

@@ -19,18 +19,25 @@
 package org.apache.atlas.repository.store.graph.v2;
 
 import com.fasterxml.jackson.core.type.TypeReference;
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasEntityHeaders;
+import org.apache.atlas.model.typedef.AtlasEntityDef;
+import org.apache.atlas.model.typedef.AtlasStructDef;
 import org.apache.atlas.repository.audit.EntityAuditRepository;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.type.AtlasEntityType;
+import org.apache.atlas.type.AtlasStructType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.utils.AtlasJson;
 import org.apache.atlas.utils.TestResourceFileUtils;
 import org.apache.commons.lang.StringUtils;
+import org.elasticsearch.common.util.CollectionUtils;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
@@ -161,6 +168,26 @@ public class ClassificationAssociatorTest {
         assertSummaryElement(summaryArray[0], "Entity", STATUS_SKIPPED, "");
     }
 
+
+    @Test
+    public void updaterEntityWithUniqueName() throws IOException, AtlasBaseException {
+        AtlasEntityDef ed = getAtlasEntityDefFromFile("col-entity-def-unique-name");
+
+        AtlasEntityHeaders entityHeaderMap = getEntityHeaderMapFromFile("header-PII-no-qualifiedName");
+        AtlasEntityStore entitiesStore = mock(AtlasEntityStore.class);
+        AtlasTypeRegistry typeRegistry = new AtlasTypeRegistry();
+        AtlasTypeRegistry.AtlasTransientTypeRegistry ttr = typeRegistry.lockTypeRegistryForUpdate();
+        ttr.addTypes(CollectionUtils.newSingletonArrayList(ed));
+
+        ClassificationAssociatorUpdaterForSpy updater = new ClassificationAssociatorUpdaterForSpy(ttr, entitiesStore, "col-entity-PII");
+        String summary = updater.setClassifications(entityHeaderMap.getGuidHeaderMap());
+
+        TypeReference<String[]> typeReference = new TypeReference<String[]>() {};
+        String[] summaryArray = AtlasJson.fromJson(summary, typeReference);
+        assertEquals(summaryArray.length, 1);
+        assertSummaryElement(summaryArray[0], "Update", STATUS_DONE, "PII");
+    }
+
     @Test
     public void updaterTests() throws IOException {
         updaterAssert("header-None", "col-entity-None");
@@ -227,6 +254,9 @@ public class ClassificationAssociatorTest {
     private AtlasEntityHeaders getEntityHeaderMapFromFile(String filename) throws IOException {
         return TestResourceFileUtils.readObjectFromJson(TEST_FILES_SUBDIR, filename, AtlasEntityHeaders.class);
     }
+    private AtlasEntityDef getAtlasEntityDefFromFile(String filename) throws IOException {
+        return TestResourceFileUtils.readObjectFromJson(TEST_FILES_SUBDIR, filename, AtlasEntityDef.class);
+    }
 
     private void updaterAssert(String incoming, String entity, String... opNamePair) throws IOException {
         String[] summary = setupUpdater(incoming, entity, opNamePair.length);

repository/src/test/resources/json/classification-association/col-entity-def-unique-name.json

+{
+  "attributeDefs": [
+    {
+      "cardinality": "SINGLE",
+      "includeInNotification": false,
+      "isIndexable": true,
+      "isOptional": false,
+      "isUnique": true,
+      "name": "name",
+      "typeName": "string",
+      "valuesMaxCount": 1,
+      "valuesMinCount": 1
+    },
+    {
+      "cardinality": "SINGLE",
+      "includeInNotification": false,
+      "isIndexable": false,
+      "isOptional": true,
+      "isUnique": false,
+      "name": "optional",
+      "typeName": "string",
+      "valuesMaxCount": 1,
+      "valuesMinCount": 0
+    },
+    {
+      "cardinality": "LIST",
+      "includeInNotification": false,
+      "isIndexable": false,
+      "isOptional": false,
+      "isUnique": false,
+      "name": "collection",
+      "typeName": "array<string>",
+      "valuesMaxCount": 2147483647,
+      "valuesMinCount": 1
+    },
+    {
+      "cardinality": "SET",
+      "includeInNotification": false,
+      "isIndexable": false,
+      "isOptional": false,
+      "isUnique": false,
+      "name": "set",
+      "typeName": "array<string>",
+      "valuesMaxCount": 2147483647,
+      "valuesMinCount": 1
+    },
+    {
+      "cardinality": "SINGLE",
+      "includeInNotification": false,
+      "isIndexable": false,
+      "isOptional": true,
+      "isUnique": false,
+      "name": "short",
+      "typeName": "short",
+      "valuesMaxCount": 1,
+      "valuesMinCount": 0
+    }
+  ],
+  "category": "ENTITY",
+  "createTime": 1549271328342,
+  "createdBy": "atlas",
+  "name": "hive_column",
+  "subTypes": [],
+  "superTypes": [],
+  "typeVersion": "1.0",
+  "updateTime": 1549271328342,
+  "updatedBy": "atlas",
+  "version": 1
+}
\ No newline at end of file

repository/src/test/resources/json/classification-association/header-PII-no-qualifiedName.json

+{
+  "guidHeaderMap": {
+    "0ce68113-77fe-4ed1-9585-69371202bd74": {
+      "typeName": "hive_column",
+      "attributes": {
+        "owner": "hive",
+        "uniqueName": "hortoniabank.us_customers.nationalid@cl1",
+        "name": "nationalid"
+      },
+      "guid": "0ce68113-77fe-4ed1-9585-69371202bd74",
+      "status": "ACTIVE",
+      "displayText": "nationalid",
+      "classificationNames": [
+        "PII"
+      ],
+      "classifications": [
+        {
+          "typeName": "PII",
+          "attributes": {
+            "type": "ssn"
+          },
+          "entityGuid": "0ce68113-77fe-4ed1-9585-69371202bd74",
+          "entityStatus": "ACTIVE",
+          "propagate": true,
+          "removePropagationsOnEntityDelete": false
+        }
+      ],
+      "meaningNames": [],
+      "meanings": []
+    }
+  }
+}