ATLAS-716 Entity update/delete notifications (shwethags)

addons/hive-bridge/src/main/java/org/apache/atlas/hive/bridge/HiveMetaStoreBridge.java

@@ -188,7 +188,7 @@ public class HiveMetaStoreBridge {
         List<String> guids = getAtlasClient().createEntity(entityJSON);
         LOG.debug("created instance for type " + typeName + ", guid: " + guids);
 
-        return new Referenceable(guids.get(0), referenceable.getTypeName(), null);
+        return new Referenceable(guids.get(guids.size() - 1), referenceable.getTypeName(), null);
     }
 
     /**

addons/hive-bridge/src/test/java/org/apache/atlas/hive/hook/HiveHookIT.java

@@ -29,7 +29,6 @@ import org.apache.atlas.hive.bridge.HiveMetaStoreBridge;
 import org.apache.atlas.hive.model.HiveDataModelGenerator;
 import org.apache.atlas.hive.model.HiveDataTypes;
 import org.apache.atlas.hive.rewrite.HiveASTRewriter;
-import org.apache.atlas.hive.rewrite.RewriteException;
 import org.apache.atlas.typesystem.Referenceable;
 import org.apache.atlas.typesystem.Struct;
 import org.apache.atlas.typesystem.persistence.Id;
@@ -66,6 +65,7 @@ import java.util.Map;
 
 import static org.apache.atlas.hive.hook.HiveHook.lower;
 import static org.apache.atlas.hive.hook.HiveHook.normalize;
+import static org.apache.atlas.hive.model.HiveDataModelGenerator.NAME;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.fail;
@@ -197,8 +197,8 @@ public class HiveHookIT {
         Assert.assertEquals(tableRef.get(HiveDataModelGenerator.TABLE_TYPE_ATTR), TableType.MANAGED_TABLE.name());
         Assert.assertEquals(tableRef.get(HiveDataModelGenerator.COMMENT), "table comment");
         String entityName = HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, tableName);
-        Assert.assertEquals(tableRef.get(HiveDataModelGenerator.NAME), entityName);
-        Assert.assertEquals(tableRef.get(HiveDataModelGenerator.NAME), "default." + tableName.toLowerCase() + "@" + CLUSTER_NAME);
+        Assert.assertEquals(tableRef.get(NAME), entityName);
+        Assert.assertEquals(tableRef.get(NAME), "default." + tableName.toLowerCase() + "@" + CLUSTER_NAME);
 
         Table t = hiveMetaStoreBridge.hiveClient.getTable(DEFAULT_DB, tableName);
         long createTime = Long.parseLong(t.getMetadata().getProperty(hive_metastoreConstants.DDL_TIME)) * HiveMetaStoreBridge.MILLIS_CONVERT_FACTOR;
@@ -631,7 +631,7 @@ public class HiveHookIT {
         final String newDBName = createDatabase();
 
         assertTableIsRegistered(DEFAULT_DB, tableName);
-        String columnGuid = assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, tableName), HiveDataModelGenerator.NAME));
+        String columnGuid = assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, tableName), NAME));
         String sdGuid = assertSDIsRegistered(HiveMetaStoreBridge.getStorageDescQFName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, tableName)), null);
         assertDatabaseIsRegistered(newDBName);
 
@@ -649,10 +649,10 @@ public class HiveHookIT {
         String query = String.format("alter table %s rename to %s", DEFAULT_DB + "." + tableName, newDBName + "." + newTableName);
         runCommand(query);
 
-        String newColGuid = assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, newDBName, newTableName), HiveDataModelGenerator.NAME));
+        String newColGuid = assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, newDBName, newTableName), NAME));
         Assert.assertEquals(newColGuid, columnGuid);
 
-        assertColumnIsNotRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, newDBName, tableName), HiveDataModelGenerator.NAME));
+        assertColumnIsNotRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, newDBName, tableName), NAME));
 
         assertTrait(columnGuid, colTraitDetails);
         String newSdGuid = assertSDIsRegistered(HiveMetaStoreBridge.getStorageDescQFName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, newDBName, newTableName)), null);
@@ -676,7 +676,16 @@ public class HiveHookIT {
     private List<Referenceable> getColumns(String dbName, String tableName) throws Exception {
         String tableId = assertTableIsRegistered(dbName, tableName);
         Referenceable tableRef = atlasClient.getEntity(tableId);
-        return ((List<Referenceable>)tableRef.get(HiveDataModelGenerator.COLUMNS));
+
+        //with soft delete, the deleted columns are returned as well. So, filter the deleted ones
+        List<Referenceable> columns = ((List<Referenceable>) tableRef.get(HiveDataModelGenerator.COLUMNS));
+        List<Referenceable> activeColumns = new ArrayList<>();
+        for (Referenceable col : columns) {
+            if (col.getId().getState() == Id.EntityState.ACTIVE) {
+                activeColumns.add(col);
+            }
+        }
+        return activeColumns;
     }
 
 
@@ -723,21 +732,15 @@ public class HiveHookIT {
                 colDropped));
 
         //Verify the number of columns present in the table
-        assertTableIsRegistered(DEFAULT_DB, tableName, new AssertPredicate() {
-            @Override
-            public void assertOnEntity(Referenceable tableRef) throws Exception {
-                List<Referenceable> columns = (List<Referenceable>) tableRef.get(HiveDataModelGenerator.COLUMNS);
-                Assert.assertEquals(columns.size(), 1);
-                Assert.assertEquals(columns.get(0).get(HiveDataModelGenerator.NAME), HiveDataModelGenerator.NAME);
-
-            }
-        });
+        List<Referenceable> columns = getColumns(DEFAULT_DB, tableName);
+        assertEquals(columns.size(), 1);
+        assertEquals(columns.get(0).get(NAME), "name");
     }
 
     @Test
     public void testAlterTableChangeColumn() throws Exception {
         //Change name
-        String oldColName = HiveDataModelGenerator.NAME;
+        String oldColName = NAME;
         String newColName = "name1";
         String tableName = createTable();
         String query = String.format("alter table %s change %s %s string", tableName, oldColName, newColName);
@@ -818,8 +821,8 @@ public class HiveHookIT {
                 @Override
                 public void assertOnEntity(Referenceable entity) throws Exception {
                     List<Referenceable> columns = (List<Referenceable>) entity.get(HiveDataModelGenerator.COLUMNS);
-                    assertEquals(columns.get(0).get(HiveDataModelGenerator.NAME), finalNewColName);
-                    assertEquals(columns.get(1).get(HiveDataModelGenerator.NAME), "id");
+                    assertEquals(columns.get(0).get(NAME), finalNewColName);
+                    assertEquals(columns.get(1).get(NAME), "id");
                 }
             }
         );
@@ -846,8 +849,8 @@ public class HiveHookIT {
                 @Override
                 public void assertOnEntity(Referenceable entity) throws Exception {
                     List<Referenceable> columns = (List<Referenceable>) entity.get(HiveDataModelGenerator.COLUMNS);
-                    assertEquals(columns.get(1).get(HiveDataModelGenerator.NAME), finalNewColName2);
-                    assertEquals(columns.get(0).get(HiveDataModelGenerator.NAME), "id");
+                    assertEquals(columns.get(1).get(NAME), finalNewColName2);
+                    assertEquals(columns.get(0).get(NAME), "id");
                 }
             }
         );
@@ -955,7 +958,7 @@ public class HiveHookIT {
 
         Referenceable hdfsPathRef = atlasClient.getEntity(hdfsPathId);
         Assert.assertEquals(hdfsPathRef.get("path"), testPathNormed);
-        Assert.assertEquals(hdfsPathRef.get(HiveDataModelGenerator.NAME), testPathNormed);
+        Assert.assertEquals(hdfsPathRef.get(NAME), testPathNormed);
 //        Assert.assertEquals(hdfsPathRef.get("name"), new Path(testPath).getName());
         Assert.assertEquals(hdfsPathRef.get(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME), testPathNormed);
 
@@ -964,7 +967,7 @@ public class HiveHookIT {
 
     private String assertHDFSPathIsRegistered(String path) throws Exception {
         LOG.debug("Searching for hdfs path {}", path);
-        return assertEntityIsRegistered(FSDataTypes.HDFS_PATH().toString(), HiveDataModelGenerator.NAME, path, null);
+        return assertEntityIsRegistered(FSDataTypes.HDFS_PATH().toString(), NAME, path, null);
     }
 
     @Test
@@ -1014,7 +1017,7 @@ public class HiveHookIT {
         ImmutableList<String> cols = ImmutableList.of("id");
         runBucketSortQuery(tableName, 5, cols, cols);
 
-        cols = ImmutableList.of("id", HiveDataModelGenerator.NAME);
+        cols = ImmutableList.of("id", NAME);
         runBucketSortQuery(tableName, 2, cols, cols);
     }
 
@@ -1077,7 +1080,7 @@ public class HiveHookIT {
 
         assertTableIsRegistered(DEFAULT_DB, tableName);
         assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, tableName), "id"));
-        assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, tableName), HiveDataModelGenerator.NAME));
+        assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, tableName), NAME));
 
         final String query = String.format("drop table %s ", tableName);
         runCommand(query);
@@ -1086,7 +1089,7 @@ public class HiveHookIT {
                 "id"));
         assertColumnIsNotRegistered(HiveMetaStoreBridge
             .getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, tableName),
-                HiveDataModelGenerator.NAME));
+                NAME));
         assertTableIsNotRegistered(DEFAULT_DB, tableName);
     }
 
@@ -1110,7 +1113,7 @@ public class HiveHookIT {
             HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, dbName, tableNames[0]), "id"));
         assertColumnIsNotRegistered(HiveMetaStoreBridge
             .getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, dbName, tableNames[0]),
-                HiveDataModelGenerator.NAME));
+                NAME));
 
         for(int i = 0; i < numTables; i++) {
             assertTableIsNotRegistered(dbName, tableNames[i]);
@@ -1175,7 +1178,7 @@ public class HiveHookIT {
 
         assertTableIsRegistered(DEFAULT_DB, viewName);
         assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, viewName), "id"));
-        assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, viewName), HiveDataModelGenerator.NAME));
+        assertColumnIsRegistered(HiveMetaStoreBridge.getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, viewName), NAME));
 
         query = String.format("drop view %s ", viewName);
 
@@ -1185,7 +1188,7 @@ public class HiveHookIT {
                     "id"));
         assertColumnIsNotRegistered(HiveMetaStoreBridge
             .getColumnQualifiedName(HiveMetaStoreBridge.getTableQualifiedName(CLUSTER_NAME, DEFAULT_DB, viewName),
-                HiveDataModelGenerator.NAME));
+                NAME));
         assertTableIsNotRegistered(DEFAULT_DB, viewName);
     }
 
@@ -1349,7 +1352,7 @@ public class HiveHookIT {
 
         if (inputTblName != null) {
             Referenceable inputTableRef = new Referenceable(HiveDataTypes.HIVE_TABLE.name(), new HashMap<String, Object>() {{
-                put(HiveDataModelGenerator.NAME, inputTblName);
+                put(NAME, inputTblName);
             }});
             inputs = new ArrayList<Referenceable>();
             inputs.add(inputTableRef);
@@ -1357,7 +1360,7 @@ public class HiveHookIT {
         List<Referenceable> outputs = null;
         if (outputTblName != null) {
             Referenceable outputTableRef = new Referenceable(HiveDataTypes.HIVE_TABLE.name(), new HashMap<String, Object>() {{
-                put(HiveDataModelGenerator.NAME, outputTblName);
+                put(NAME, outputTblName);
             }});
 
             outputs = new ArrayList<Referenceable>();

client/src/main/java/org/apache/atlas/AtlasClient.java

@@ -20,11 +20,14 @@ package org.apache.atlas;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
 import com.sun.jersey.api.client.Client;
 import com.sun.jersey.api.client.ClientHandlerException;
 import com.sun.jersey.api.client.ClientResponse;
 import com.sun.jersey.api.client.WebResource;
 import com.sun.jersey.api.client.config.DefaultClientConfig;
+import com.sun.jersey.api.client.filter.HTTPBasicAuthFilter;
 import com.sun.jersey.client.urlconnection.URLConnectionClientHandler;
 import org.apache.atlas.security.SecureClientUtils;
 import org.apache.atlas.typesystem.Referenceable;
@@ -45,6 +48,7 @@ import org.codehaus.jettison.json.JSONException;
 import org.codehaus.jettison.json.JSONObject;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+
 import javax.ws.rs.HttpMethod;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
@@ -54,8 +58,10 @@ import java.net.ConnectException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
-import com.sun.jersey.api.client.filter.HTTPBasicAuthFilter;
+import java.util.Map;
+
 import static org.apache.atlas.security.SecurityProperties.TLS_ENABLED;
 
 /**
@@ -65,9 +71,10 @@ public class AtlasClient {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasClient.class);
 
     public static final String NAME = "name";
-    public static final String GUID = "GUID";
     public static final String TYPE = "type";
     public static final String TYPENAME = "typeName";
+    public static final String GUID = "GUID";
+    public static final String ENTITIES = "entities";
 
     public static final String DEFINITION = "definition";
     public static final String ERROR = "error";
@@ -340,6 +347,61 @@ public class AtlasClient {
         return service;
     }
 
+    public static class EntityResult {
+        private static final Gson gson = new GsonBuilder().setPrettyPrinting().create();
+
+        public static final String OP_CREATED = "created";
+        public static final String OP_UPDATED = "updated";
+        public static final String OP_DELETED = "deleted";
+
+        Map<String, List<String>> entities = new HashMap<>();
+
+        public EntityResult() {
+            //For gson
+        }
+
+        public EntityResult(List<String> created, List<String> updated, List<String> deleted) {
+            add(OP_CREATED, created);
+            add(OP_UPDATED, updated);
+            add(OP_DELETED, deleted);
+        }
+
+        private void add(String type, List<String> list) {
+            if (list != null && list.size() > 0) {
+                entities.put(type, list);
+            }
+        }
+
+        private List<String> get(String type) {
+            List<String> list = entities.get(type);
+            if (list == null) {
+                list = new ArrayList<>();
+            }
+            return list;
+        }
+
+        public List<String> getCreatedEntities() {
+            return get(OP_CREATED);
+        }
+
+        public List<String> getUpdateEntities() {
+            return get(OP_UPDATED);
+        }
+
+        public List<String> getDeletedEntities() {
+            return get(OP_DELETED);
+        }
+
+        @Override
+        public String toString() {
+            return gson.toJson(this);
+        }
+
+        public static EntityResult fromString(String json) throws AtlasServiceException {
+            return gson.fromJson(json, EntityResult.class);
+        }
+    }
+
     /**
      * Return status of the service instance the client is pointing to.
      *
@@ -562,11 +624,15 @@ public class AtlasClient {
     protected List<String> createEntity(JSONArray entities) throws AtlasServiceException {
         LOG.debug("Creating entities: {}", entities);
         JSONObject response = callAPI(API.CREATE_ENTITY, entities.toString());
-        List<String> results = extractResults(response, GUID, new ExtractOperation<String, String>());
+        List<String> results = extractEntityResult(response).getCreatedEntities();
         LOG.debug("Create entities returned results: {}", results);
         return results;
     }
 
+    protected EntityResult extractEntityResult(JSONObject response) throws AtlasServiceException {
+        return EntityResult.fromString(response.toString());
+    }
+
     /**
      * Create the given entity
      * @param entitiesAsJson entity(type instance) as json
@@ -601,19 +667,19 @@ public class AtlasClient {
      * @return json array of guids which were updated/created
      * @throws AtlasServiceException
      */
-    public List<String> updateEntities(Referenceable... entities) throws AtlasServiceException {
+    public EntityResult updateEntities(Referenceable... entities) throws AtlasServiceException {
         return updateEntities(Arrays.asList(entities));
     }
 
-    protected List<String> updateEntities(JSONArray entities) throws AtlasServiceException {
+    protected EntityResult updateEntities(JSONArray entities) throws AtlasServiceException {
         LOG.debug("Updating entities: {}", entities);
         JSONObject response = callAPI(API.UPDATE_ENTITY, entities.toString());
-        List<String> results = extractResults(response, GUID, new ExtractOperation<String, String>());
+        EntityResult results = extractEntityResult(response);
         LOG.debug("Update entities returned results: {}", results);
         return results;
     }
 
-    public List<String> updateEntities(Collection<Referenceable> entities) throws AtlasServiceException {
+    public EntityResult updateEntities(Collection<Referenceable> entities) throws AtlasServiceException {
         JSONArray entitiesArray = getEntitiesArray(entities);
         return updateEntities(entitiesArray);
     }
@@ -625,9 +691,10 @@ public class AtlasClient {
      * @param attribute  property key
      * @param value     property value
      */
-    public void updateEntityAttribute(final String guid, final String attribute, String value) throws AtlasServiceException {
+    public EntityResult updateEntityAttribute(final String guid, final String attribute, String value)
+            throws AtlasServiceException {
         LOG.debug("Updating entity id: {}, attribute name: {}, attribute value: {}", guid, attribute, value);
-        callAPIWithRetries(API.UPDATE_ENTITY_PARTIAL, value, new ResourceCreator() {
+        JSONObject response = callAPIWithRetries(API.UPDATE_ENTITY_PARTIAL, value, new ResourceCreator() {
             @Override
             public WebResource createResource() {
                 API api = API.UPDATE_ENTITY_PARTIAL;
@@ -636,6 +703,7 @@ public class AtlasClient {
                 return resource;
             }
         });
+        return extractEntityResult(response);
     }
 
     @VisibleForTesting
@@ -665,10 +733,11 @@ public class AtlasClient {
      * @param guid      guid
      * @param entity entity definition
      */
-    public void updateEntity(String guid, Referenceable entity) throws AtlasServiceException {
+    public EntityResult updateEntity(String guid, Referenceable entity) throws AtlasServiceException {
         String entityJson = InstanceSerialization.toJson(entity, true);
         LOG.debug("Updating entity id {} with {}", guid, entityJson);
-        callAPI(API.UPDATE_ENTITY_PARTIAL, entityJson, guid);
+        JSONObject response = callAPI(API.UPDATE_ENTITY_PARTIAL, entityJson, guid);
+        return extractEntityResult(response);
     }
 
     /**
@@ -691,8 +760,9 @@ public class AtlasClient {
      * @param uniqueAttributeValue Attribute Value that uniquely identifies the entity
      * @param entity entity definition
      */
-    public String updateEntity(final String entityType, final String uniqueAttributeName, final String uniqueAttributeValue,
-                               Referenceable entity) throws AtlasServiceException {
+    public EntityResult updateEntity(final String entityType, final String uniqueAttributeName,
+                                     final String uniqueAttributeValue,
+                                     Referenceable entity) throws AtlasServiceException {
         final API api = API.UPDATE_ENTITY_PARTIAL;
         String entityJson = InstanceSerialization.toJson(entity, true);
         LOG.debug("Updating entity type: {}, attributeName: {}, attributeValue: {}, entity: {}", entityType,
@@ -707,7 +777,7 @@ public class AtlasClient {
                 return resource;
             }
         });
-        String result = getString(response, GUID);
+        EntityResult result = extractEntityResult(response);
         LOG.debug("Update entity returned result: {}", result);
         return result;
     }
@@ -724,10 +794,10 @@ public class AtlasClient {
      * Delete the specified entities from the repository
      * 
      * @param guids guids of entities to delete
-     * @return List of deleted entity guids
+     * @return List of entity ids updated/deleted
      * @throws AtlasServiceException
      */
-    public List<String> deleteEntities(final String ... guids) throws AtlasServiceException {
+    public EntityResult deleteEntities(final String ... guids) throws AtlasServiceException {
         LOG.debug("Deleting entities: {}", guids);
         JSONObject jsonResponse = callAPIWithRetries(API.DELETE_ENTITIES, null, new ResourceCreator() {
             @Override
@@ -740,7 +810,7 @@ public class AtlasClient {
                 return resource;
             }
         });
-        List<String> results = extractResults(jsonResponse, GUID, new ExtractOperation<String, String>());
+        EntityResult results = extractEntityResult(jsonResponse);
         LOG.debug("Delete entities returned results: {}", results);
         return results;
     }
@@ -750,9 +820,9 @@ public class AtlasClient {
      * @param entityType Type of the entity being deleted
      * @param uniqueAttributeName Attribute Name that uniquely identifies the entity
      * @param uniqueAttributeValue Attribute Value that uniquely identifies the entity
-     * @return List of deleted entity guids(including composite references from that entity)
+     * @return List of entity ids updated/deleted(including composite references from that entity)
      */
-    public List<String> deleteEntity(String entityType, String uniqueAttributeName, String uniqueAttributeValue)
+    public EntityResult deleteEntity(String entityType, String uniqueAttributeName, String uniqueAttributeValue)
             throws AtlasServiceException {
         LOG.debug("Deleting entity type: {}, attributeName: {}, attributeValue: {}", entityType, uniqueAttributeName,
                 uniqueAttributeValue);
@@ -762,7 +832,7 @@ public class AtlasClient {
         resource = resource.queryParam(ATTRIBUTE_NAME, uniqueAttributeName);
         resource = resource.queryParam(ATTRIBUTE_VALUE, uniqueAttributeValue);
         JSONObject jsonResponse = callAPIWithResource(API.DELETE_ENTITIES, resource, null);
-        List<String> results = extractResults(jsonResponse, GUID, new ExtractOperation<String, String>());
+        EntityResult results = extractEntityResult(jsonResponse);
         LOG.debug("Delete entities returned results: {}", results);
         return results;
     }
@@ -901,7 +971,7 @@ public class AtlasClient {
         return extractResults(jsonResponse, AtlasClient.EVENTS, new ExtractOperation<EntityAuditEvent, JSONObject>() {
             @Override
             EntityAuditEvent extractElement(JSONObject element) throws JSONException {
-                return EntityAuditEvent.GSON.fromJson(element.toString(), EntityAuditEvent.class);
+                return SerDe.GSON.fromJson(element.toString(), EntityAuditEvent.class);
             }
         });
 

client/src/main/java/org/apache/atlas/EntityAuditEvent.java

@@ -18,16 +18,14 @@
 
 package org.apache.atlas;
 
-import com.google.gson.Gson;
-import com.google.gson.GsonBuilder;
+import org.apache.atlas.typesystem.IReferenceableInstance;
+import org.apache.atlas.typesystem.json.InstanceSerialization;
 import org.apache.commons.lang.StringUtils;
 
 /**
  * Structure of entity audit event
  */
 public class EntityAuditEvent {
-    public static final Gson GSON = new GsonBuilder().create();
-
     public enum EntityAuditAction {
         ENTITY_CREATE, ENTITY_UPDATE, ENTITY_DELETE, TAG_ADD, TAG_DELETE
     }
@@ -38,16 +36,19 @@ public class EntityAuditEvent {
     private EntityAuditAction action;
     private String details;
     private String eventKey;
+    private IReferenceableInstance entityDefinition;
 
     public EntityAuditEvent() {
     }
 
-    public EntityAuditEvent(String entityId, Long ts, String user, EntityAuditAction action, String details) {
+    public EntityAuditEvent(String entityId, Long ts, String user, EntityAuditAction action, String details,
+                            IReferenceableInstance entityDefinition) throws AtlasException {
         this.entityId = entityId;
         this.timestamp = ts;
         this.user = user;
         this.action = action;
         this.details = details;
+        this.entityDefinition = entityDefinition;
     }
 
     @Override
@@ -62,10 +63,12 @@ public class EntityAuditEvent {
 
         EntityAuditEvent otherEvent = (EntityAuditEvent) other;
         return StringUtils.equals(entityId, otherEvent.entityId) &&
-                (timestamp == otherEvent.timestamp) &&
-                StringUtils.equals(user, otherEvent.user) && (action == otherEvent.action) &&
-                StringUtils.equals(details, otherEvent.details) &&
-                StringUtils.equals(eventKey, otherEvent.eventKey);
+                    (timestamp == otherEvent.timestamp) &&
+                    StringUtils.equals(user, otherEvent.user) &&
+                    (action == otherEvent.action) &&
+                    StringUtils.equals(details, otherEvent.details) &&
+                    StringUtils.equals(eventKey, otherEvent.eventKey) &&
+                    StringUtils.equals(getEntityDefinitionString(), otherEvent.getEntityDefinitionString());
     }
 
     @Override
@@ -75,11 +78,11 @@ public class EntityAuditEvent {
 
     @Override
     public String toString() {
-        return GSON.toJson(this);
+        return SerDe.GSON.toJson(this);
     }
 
     public static EntityAuditEvent fromString(String eventString) {
-        return GSON.fromJson(eventString, EntityAuditEvent.class);
+        return SerDe.GSON.fromJson(eventString, EntityAuditEvent.class);
     }
 
     public String getEntityId() {
@@ -129,4 +132,19 @@ public class EntityAuditEvent {
     public void setEventKey(String eventKey) {
         this.eventKey = eventKey;
     }
+
+    public IReferenceableInstance getEntityDefinition() {
+        return entityDefinition;
+    }
+
+    public String getEntityDefinitionString() {
+        if (entityDefinition != null) {
+            return InstanceSerialization.toJson(entityDefinition, true);
+        }
+        return null;
+    }
+
+    public void setEntityDefinition(String entityDefinition) {
+        this.entityDefinition = InstanceSerialization.fromJsonReferenceable(entityDefinition, true);
+    }
 }

client/src/main/java/org/apache/atlas/SerDe.java

+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p/>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p/>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.atlas;
+
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
+import com.google.gson.JsonDeserializationContext;
+import com.google.gson.JsonDeserializer;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonParser;
+import com.google.gson.JsonSerializationContext;
+import com.google.gson.JsonSerializer;
+import org.apache.atlas.typesystem.IReferenceableInstance;
+import org.apache.atlas.typesystem.IStruct;
+import org.apache.atlas.typesystem.Referenceable;
+import org.apache.atlas.typesystem.Struct;
+import org.apache.atlas.typesystem.json.InstanceSerialization;
+
+import java.lang.reflect.Type;
+
+public class SerDe {
+    public static final Gson GSON = new GsonBuilder().
+             registerTypeAdapter(IStruct.class, new StructDeserializer()).
+             registerTypeAdapter(IReferenceableInstance.class, new ReferenceableSerializerDeserializer()).
+             registerTypeAdapter(Referenceable.class, new ReferenceableSerializerDeserializer()).
+             create();
+
+    /**
+     * Serde for Struct used by AbstractNotificationConsumer.GSON.
+     */
+    public static final class StructDeserializer implements JsonDeserializer<IStruct>, JsonSerializer<IStruct> {
+        @Override
+        public IStruct deserialize(final JsonElement json, final Type type,
+                                   final JsonDeserializationContext context) {
+            return context.deserialize(json, Struct.class);
+        }
+
+        @Override
+        public JsonElement serialize(IStruct src, Type typeOfSrc, JsonSerializationContext context) {
+            String instanceJson = InstanceSerialization.toJson(src, true);
+            return new JsonParser().parse(instanceJson).getAsJsonObject();
+        }
+    }
+
+    /**
+     * Serde for Referenceable used by AbstractNotificationConsumer.GSON.
+     */
+    public static final class ReferenceableSerializerDeserializer implements JsonDeserializer<IStruct>,
+            JsonSerializer<IReferenceableInstance> {
+        @Override
+        public IReferenceableInstance deserialize(final JsonElement json, final Type type,
+                                                  final JsonDeserializationContext context) {
+
+            return InstanceSerialization.fromJsonReferenceable(json.toString(), true);
+        }
+
+        @Override
+        public JsonElement serialize(IReferenceableInstance src, Type typeOfSrc, JsonSerializationContext context) {
+            String instanceJson = InstanceSerialization.toJson(src, true);
+            return new JsonParser().parse(instanceJson).getAsJsonObject();
+        }
+    }
+}

client/src/test/java/org/apache/atlas/AtlasClientTest.java

@@ -21,8 +21,12 @@ import com.sun.jersey.api.client.Client;
 import com.sun.jersey.api.client.ClientHandlerException;
 import com.sun.jersey.api.client.ClientResponse;
 import com.sun.jersey.api.client.WebResource;
+import org.apache.atlas.typesystem.Referenceable;
+import org.apache.atlas.typesystem.json.InstanceSerialization;
 import org.apache.commons.configuration.Configuration;
 import org.apache.hadoop.security.UserGroupInformation;
+import org.codehaus.jettison.json.JSONObject;
+import org.mockito.Matchers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.testng.annotations.BeforeMethod;
@@ -33,9 +37,12 @@ import javax.ws.rs.core.UriBuilder;
 import java.net.ConnectException;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.util.Arrays;
+import java.util.List;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.anyString;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -76,6 +83,25 @@ public class AtlasClientTest {
         assertTrue(atlasClient.isServerReady());
     }
 
+    @Test
+    public void testCreateEntity() throws Exception {
+        setupRetryParams();
+        AtlasClient atlasClient = new AtlasClient(service, configuration);
+
+        WebResource.Builder builder = setupBuilder(AtlasClient.API.CREATE_ENTITY, service);
+        ClientResponse response = mock(ClientResponse.class);
+        when(response.getStatus()).thenReturn(Response.Status.CREATED.getStatusCode());
+
+        JSONObject jsonResponse = new JSONObject(new AtlasClient.EntityResult(Arrays.asList("id"), null, null).toString());
+        when(response.getEntity(String.class)).thenReturn(jsonResponse.toString());
+        String entityJson = InstanceSerialization.toJson(new Referenceable("type"), true);
+        when(builder.method(anyString(), Matchers.<Class>any(), anyString())).thenReturn(response);
+
+        List<String> ids = atlasClient.createEntity(entityJson);
+        assertEquals(ids.size(), 1);
+        assertEquals(ids.get(0), "id");
+    }
+
     private WebResource.Builder setupBuilder(AtlasClient.API api, WebResource webResource) {
         when(webResource.path(api.getPath())).thenReturn(service);
         WebResource.Builder builder = getBuilder(service);

notification/src/main/java/org/apache/atlas/notification/MessageVersion.java

@@ -97,6 +97,11 @@ public class MessageVersion implements Comparable<MessageVersion> {
     }
 
 
+    @Override
+    public String toString() {
+        return "MessageVersion[version=" + version + "]";
+    }
+
     // ----- helper methods --------------------------------------------------
 
     /**

notification/src/main/java/org/apache/atlas/notification/NotificationInterface.java

@@ -17,11 +17,11 @@
  */
 package org.apache.atlas.notification;
 
+import com.google.gson.reflect.TypeToken;
 import org.apache.atlas.notification.entity.EntityMessageDeserializer;
 import org.apache.atlas.notification.entity.EntityNotification;
 import org.apache.atlas.notification.hook.HookMessageDeserializer;
 import org.apache.atlas.notification.hook.HookNotification;
-import com.google.gson.reflect.TypeToken;
 
 import java.lang.reflect.Type;
 import java.util.List;

notification/src/main/java/org/apache/atlas/notification/VersionedMessageDeserializer.java

@@ -31,7 +31,7 @@ import java.lang.reflect.Type;
 public abstract class VersionedMessageDeserializer<T> implements MessageDeserializer<T> {
 
     public static final String VERSION_MISMATCH_MSG =
-        "Notification message version mismatch.  Expected %s but recieved %s";
+        "Notification message version mismatch. Expected %s but recieved %s. Message %s";
 
     private final Type versionedMessageType;
     private final MessageVersion expectedVersion;
@@ -90,18 +90,16 @@ public abstract class VersionedMessageDeserializer<T> implements MessageDeserial
 
         // message has newer version
         if (comp > 0) {
-            String msg = String.format(VERSION_MISMATCH_MSG, expectedVersion, versionedMessage.getVersion());
+            String msg =
+                    String.format(VERSION_MISMATCH_MSG, expectedVersion, versionedMessage.getVersion(), messageJson);
             notificationLogger.error(msg);
-            notificationLogger.info(messageJson);
             throw new IncompatibleVersionException(msg);
         }
 
         // message has older version
         if (comp < 0) {
-            notificationLogger.info(
-                String.format(VERSION_MISMATCH_MSG, expectedVersion, versionedMessage.getVersion()));
-
-            notificationLogger.info(messageJson);
+            notificationLogger.info(String.format(VERSION_MISMATCH_MSG, expectedVersion, versionedMessage.getVersion(),
+                    messageJson));
         }
     }
 }

notification/src/test/java/org/apache/atlas/notification/AbstractNotificationConsumerTest.java

@@ -27,9 +27,13 @@ import java.lang.reflect.Type;
 import java.util.LinkedList;
 import java.util.List;
 
+import static org.mockito.Matchers.endsWith;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
-import static org.testng.Assert.*;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
 
 /**
  * AbstractNotificationConsumer tests.
@@ -110,17 +114,17 @@ public class AbstractNotificationConsumerTest {
         assertTrue(consumer.hasNext());
 
         assertEquals(new TestMessage("sValue2", 98), consumer.next());
-        verify(logger).info(json2);
+        verify(logger).info(endsWith(json2));
 
         assertTrue(consumer.hasNext());
 
         assertEquals(new TestMessage("sValue3", 97), consumer.next());
-        verify(logger).info(json3);
+        verify(logger).info(endsWith(json3));
 
         assertTrue(consumer.hasNext());
 
         assertEquals(new TestMessage("sValue4", 96), consumer.next());
-        verify(logger).info(json4);
+        verify(logger).info(endsWith(json4));
 
         assertFalse(consumer.hasNext());
     }
@@ -154,7 +158,7 @@ public class AbstractNotificationConsumerTest {
             consumer.next();
             fail("Expected VersionMismatchException!");
         } catch (IncompatibleVersionException e) {
-            verify(logger).info(json2);
+            verify(logger).error(endsWith(json2));
         }
 
         assertFalse(consumer.hasNext());

release-log.txt

@@ -3,6 +3,7 @@ Apache Atlas Release Notes
 
 --trunk - unreleased
 INCOMPATIBLE CHANGES:
+ATLAS-716 Entity update/delete notifications (shwethags)
 ATLAS-619 Canonicalize hive queries (sumasai)
 ATLAS-497 Simple Authorization (saqeeb.s via yhemanth)
 ATLAS-661 REST API Authentication (nixonrodrigues via yhemanth)

repository/src/main/java/org/apache/atlas/repository/MetadataRepository.java

@@ -18,6 +18,7 @@
 
 package org.apache.atlas.repository;
 
+import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.ITypedStruct;
@@ -26,7 +27,6 @@ import org.apache.atlas.typesystem.exception.EntityNotFoundException;
 import org.apache.atlas.typesystem.exception.TraitNotFoundException;
 import org.apache.atlas.typesystem.types.AttributeInfo;
 import org.apache.atlas.typesystem.types.IDataType;
-import org.apache.atlas.typesystem.types.TypeUtils;
 
 import java.util.List;
 
@@ -111,7 +111,7 @@ public interface MetadataRepository {
      * @return guids of deleted entities
      * @throws RepositoryException
      */
-    TypeUtils.Pair<List<String>, List<ITypedReferenceableInstance>> deleteEntities(List<String> guids) throws RepositoryException;
+    AtlasClient.EntityResult deleteEntities(List<String> guids) throws RepositoryException;
     
     
     // Trait management functions
@@ -147,13 +147,13 @@ public interface MetadataRepository {
      * Adds/Updates the property to the entity that corresponds to the GUID
      * Supports only primitive attribute/Class Id updations.
      */
-    TypeUtils.Pair<List<String>, List<String>> updatePartial(ITypedReferenceableInstance entity) throws RepositoryException;
+    AtlasClient.EntityResult updatePartial(ITypedReferenceableInstance entity) throws RepositoryException;
 
     /**
      * Adds the property to the entity that corresponds to the GUID
      * @param entitiesToBeUpdated The entities to be updated
      */
-    TypeUtils.Pair<List<String>, List<String>> updateEntities(ITypedReferenceableInstance... entitiesToBeUpdated) throws RepositoryException;
+    AtlasClient.EntityResult updateEntities(ITypedReferenceableInstance... entitiesToBeUpdated) throws RepositoryException;
 
     /**
      * Returns the entity for the given type and qualified name

repository/src/main/java/org/apache/atlas/repository/audit/EntityAuditListener.java

@@ -55,8 +55,9 @@ public class EntityAuditListener implements EntityChangeListener {
     }
 
     private EntityAuditEvent createEvent(ITypedReferenceableInstance entity, long ts,
-                                         EntityAuditEvent.EntityAuditAction action, String details) {
-        return new EntityAuditEvent(entity.getId()._getId(), ts, RequestContext.get().getUser(), action, details);
+                                         EntityAuditEvent.EntityAuditAction action, String details)
+            throws AtlasException {
+        return new EntityAuditEvent(entity.getId()._getId(), ts, RequestContext.get().getUser(), action, details, entity);
     }
 
     @Override

repository/src/main/java/org/apache/atlas/repository/audit/HBaseBasedAuditRepository.java

@@ -78,6 +78,7 @@ public class HBaseBasedAuditRepository implements Service, EntityAuditRepository
     public static final byte[] COLUMN_ACTION = Bytes.toBytes("action");
     public static final byte[] COLUMN_DETAIL = Bytes.toBytes("detail");
     public static final byte[] COLUMN_USER = Bytes.toBytes("user");
+    public static final byte[] COLUMN_DEFINITION = Bytes.toBytes("def");
 
     private TableName tableName;
     private Connection connection;
@@ -110,6 +111,7 @@ public class HBaseBasedAuditRepository implements Service, EntityAuditRepository
                 addColumn(put, COLUMN_ACTION, event.getAction());
                 addColumn(put, COLUMN_USER, event.getUser());
                 addColumn(put, COLUMN_DETAIL, event.getDetails());
+                addColumn(put, COLUMN_DEFINITION, event.getEntityDefinitionString());
                 puts.add(put);
             }
             table.put(puts);
@@ -183,6 +185,7 @@ public class HBaseBasedAuditRepository implements Service, EntityAuditRepository
                 event.setUser(getResultString(result, COLUMN_USER));
                 event.setAction(EntityAuditEvent.EntityAuditAction.valueOf(getResultString(result, COLUMN_ACTION)));
                 event.setDetails(getResultString(result, COLUMN_DETAIL));
+                event.setEntityDefinition(getResultString(result, COLUMN_DEFINITION));
                 events.add(event);
             }
             LOG.info("Got events for entity id {}, starting timestamp {}, #records {}", entityId, startKey, events.size());

repository/src/main/java/org/apache/atlas/repository/graph/DeleteHandler.java

@@ -47,15 +47,16 @@ import static org.apache.atlas.repository.graph.GraphHelper.string;
 public abstract class DeleteHandler {
     public static final Logger LOG = LoggerFactory.getLogger(DeleteHandler.class);
 
-    private static final GraphHelper graphHelper = GraphHelper.getInstance();
+    protected static final GraphHelper graphHelper = GraphHelper.getInstance();
 
     protected TypeSystem typeSystem;
     private boolean shouldUpdateReverseAttribute;
+    private boolean softDelete;
 
-    public DeleteHandler(TypeSystem typeSystem, boolean shouldUpdateReverseAttribute) {
+    public DeleteHandler(TypeSystem typeSystem, boolean shouldUpdateReverseAttribute, boolean softDelete) {
         this.typeSystem = typeSystem;
         this.shouldUpdateReverseAttribute = shouldUpdateReverseAttribute;
-
+        this.softDelete = softDelete;
     }
 
     /**
@@ -64,16 +65,22 @@ public abstract class DeleteHandler {
      * @throws AtlasException
      */
     public void deleteEntity(Vertex instanceVertex) throws AtlasException {
+        RequestContext requestContext = RequestContext.get();
         String guid = GraphHelper.getIdFromVertex(instanceVertex);
+        Id.EntityState state = GraphHelper.getState(instanceVertex);
+        if (requestContext.getDeletedEntityIds().contains(guid) || state == Id.EntityState.DELETED) {
+            LOG.debug("Skipping deleting {} as its already deleted", guid);
+            return;
+        }
         String typeName = GraphHelper.getTypeName(instanceVertex);
-        RequestContext.get().recordDeletedEntity(guid, typeName);
+        requestContext.recordEntityDelete(guid, typeName);
 
         deleteAllTraits(instanceVertex);
 
-        deleteTypeVertex(instanceVertex);
+        deleteTypeVertex(instanceVertex, false);
     }
 
-    protected abstract void deleteEdge(Edge edge) throws AtlasException;
+    protected abstract void deleteEdge(Edge edge, boolean force) throws AtlasException;
 
     /**
      * Deletes a type vertex - can be entity(class type) or just vertex(struct/trait type)
@@ -81,11 +88,11 @@ public abstract class DeleteHandler {
      * @param typeCategory
      * @throws AtlasException
      */
-    protected void deleteTypeVertex(Vertex instanceVertex, DataTypes.TypeCategory typeCategory) throws AtlasException {
+    protected void deleteTypeVertex(Vertex instanceVertex, DataTypes.TypeCategory typeCategory, boolean force) throws AtlasException {
         switch (typeCategory) {
         case STRUCT:
         case TRAIT:
-            deleteTypeVertex(instanceVertex);
+            deleteTypeVertex(instanceVertex, force);
             break;
 
         case CLASS:
@@ -102,7 +109,7 @@ public abstract class DeleteHandler {
      * @param instanceVertex
      * @throws AtlasException
      */
-    protected void deleteTypeVertex(Vertex instanceVertex) throws AtlasException {
+    protected void deleteTypeVertex(Vertex instanceVertex, boolean force) throws AtlasException {
         LOG.debug("Deleting {}", string(instanceVertex));
         String typeName = GraphHelper.getTypeName(instanceVertex);
         IDataType type = typeSystem.getDataType(IDataType.class, typeName);
@@ -115,12 +122,12 @@ public abstract class DeleteHandler {
             switch (attributeInfo.dataType().getTypeCategory()) {
             case CLASS:
                 //If its class attribute, delete the reference
-                deleteReference(instanceVertex, edgeLabel, DataTypes.TypeCategory.CLASS, attributeInfo.isComposite);
+                deleteEdgeReference(instanceVertex, edgeLabel, DataTypes.TypeCategory.CLASS, attributeInfo.isComposite);
                 break;
 
             case STRUCT:
                 //If its struct attribute, delete the reference
-                deleteReference(instanceVertex, edgeLabel, DataTypes.TypeCategory.STRUCT);
+                deleteEdgeReference(instanceVertex, edgeLabel, DataTypes.TypeCategory.STRUCT, false);
                 break;
 
             case ARRAY:
@@ -133,7 +140,7 @@ public abstract class DeleteHandler {
                     if (edges != null) {
                         while (edges.hasNext()) {
                             Edge edge = edges.next();
-                            deleteReference(edge, elementType, attributeInfo);
+                            deleteEdgeReference(edge, elementType.getTypeCategory(), attributeInfo.isComposite, false);
                         }
                     }
                 }
@@ -151,22 +158,31 @@ public abstract class DeleteHandler {
                     if (keys != null) {
                         for (String key : keys) {
                             String mapEdgeLabel = GraphHelper.getQualifiedNameForMapKey(edgeLabel, key);
-                            deleteReference(instanceVertex, mapEdgeLabel, valueTypeCategory, attributeInfo.isComposite);
+                            deleteEdgeReference(instanceVertex, mapEdgeLabel, valueTypeCategory, attributeInfo.isComposite);
                         }
                     }
                 }
             }
         }
 
-        deleteVertex(instanceVertex, type.getTypeCategory());
+        deleteVertex(instanceVertex, force);
     }
 
-    public void deleteReference(Edge edge, IDataType dataType, AttributeInfo attributeInfo) throws AtlasException {
-        deleteReference(edge, dataType.getTypeCategory(), attributeInfo.isComposite);
-    }
-
-    public void deleteReference(Edge edge, DataTypes.TypeCategory typeCategory, boolean isComposite) throws AtlasException {
+    /**
+     * Force delete is used to remove struct/trait in case of entity updates
+     * @param edge
+     * @param typeCategory
+     * @param isComposite
+     * @param forceDeleteStructTrait
+     * @return returns true if the edge reference is hard deleted
+     * @throws AtlasException
+     */
+    public boolean deleteEdgeReference(Edge edge, DataTypes.TypeCategory typeCategory, boolean isComposite,
+                                    boolean forceDeleteStructTrait) throws AtlasException {
         LOG.debug("Deleting {}", string(edge));
+        boolean forceDelete =
+                (typeCategory == DataTypes.TypeCategory.STRUCT || typeCategory == DataTypes.TypeCategory.TRAIT)
+                        ? forceDeleteStructTrait : false;
         if (typeCategory == DataTypes.TypeCategory.STRUCT || typeCategory == DataTypes.TypeCategory.TRAIT
                 || (typeCategory == DataTypes.TypeCategory.CLASS && isComposite)) {
             //If the vertex is of type struct/trait, delete the edge and then the reference vertex as the vertex is not shared by any other entities.
@@ -175,32 +191,28 @@ public abstract class DeleteHandler {
             Vertex vertexForDelete = edge.getVertex(Direction.IN);
 
             //If deleting the edge and then the in vertex, reverse attribute shouldn't be updated
-            deleteEdge(edge, false);
-            deleteTypeVertex(vertexForDelete, typeCategory);
+            deleteEdge(edge, false, forceDelete);
+            deleteTypeVertex(vertexForDelete, typeCategory, forceDelete);
         } else {
             //If the vertex is of type class, and its not a composite attributes, the reference vertex' lifecycle is not controlled
             //through this delete. Hence just remove the reference edge. Leave the reference vertex as is
 
             //If deleting just the edge, reverse attribute should be updated for any references
             //For example, for the department type system, if the person's manager edge is deleted, subordinates of manager should be updated
-            deleteEdge(edge, true);
+            deleteEdge(edge, true, false);
         }
+        return !softDelete || forceDelete;
     }
 
-    public void deleteReference(Vertex instanceVertex, String edgeLabel, DataTypes.TypeCategory typeCategory)
-            throws AtlasException {
-        deleteReference(instanceVertex, edgeLabel, typeCategory, false);
-    }
-
-    public void deleteReference(Vertex instanceVertex, String edgeLabel, DataTypes.TypeCategory typeCategory,
-                                boolean isComposite) throws AtlasException {
-        Edge edge = GraphHelper.getEdgeForLabel(instanceVertex, edgeLabel);
+    public void deleteEdgeReference(Vertex outVertex, String edgeLabel, DataTypes.TypeCategory typeCategory,
+                                    boolean isComposite) throws AtlasException {
+        Edge edge = GraphHelper.getEdgeForLabel(outVertex, edgeLabel);
         if (edge != null) {
-            deleteReference(edge, typeCategory, isComposite);
+            deleteEdgeReference(edge, typeCategory, isComposite, false);
         }
     }
 
-    protected void deleteEdge(Edge edge, boolean updateReverseAttribute) throws AtlasException {
+    protected void deleteEdge(Edge edge, boolean updateReverseAttribute, boolean force) throws AtlasException {
         //update reverse attribute
         if (updateReverseAttribute) {
             AttributeInfo attributeInfo = getAttributeForEdge(edge.getLabel());
@@ -210,28 +222,28 @@ public abstract class DeleteHandler {
             }
         }
 
-        deleteEdge(edge);
+        deleteEdge(edge, force);
     }
 
-    protected void deleteVertex(Vertex instanceVertex, DataTypes.TypeCategory typeCategory) throws AtlasException {
+    protected void deleteVertex(Vertex instanceVertex, boolean force) throws AtlasException {
         //Update external references(incoming edges) to this vertex
         LOG.debug("Setting the external references to {} to null(removing edges)", string(instanceVertex));
         Iterator<Edge> edges = instanceVertex.getEdges(Direction.IN).iterator();
 
         while(edges.hasNext()) {
             Edge edge = edges.next();
-            String edgeState = edge.getProperty(Constants.STATE_PROPERTY_KEY);
-            if (Id.EntityState.ACTIVE.name().equals(edgeState)) {
+            Id.EntityState edgeState = GraphHelper.getState(edge);
+            if (edgeState == Id.EntityState.ACTIVE) {
                 //Delete only the active edge references
                 AttributeInfo attribute = getAttributeForEdge(edge.getLabel());
+                //TODO use delete edge instead??
                 deleteEdgeBetweenVertices(edge.getVertex(Direction.OUT), edge.getVertex(Direction.IN), attribute.name);
-                deleteEdge(edge);
             }
         }
-        _deleteVertex(instanceVertex);
+        _deleteVertex(instanceVertex, force);
     }
 
-    protected abstract void _deleteVertex(Vertex instanceVertex);
+    protected abstract void _deleteVertex(Vertex instanceVertex, boolean force);
 
     /**
      * Deletes the edge between outvertex and inVertex. The edge is for attribute attributeName of outVertex
@@ -245,7 +257,8 @@ public abstract class DeleteHandler {
                 attributeName);
         String typeName = GraphHelper.getTypeName(outVertex);
         String outId = GraphHelper.getIdFromVertex(outVertex);
-        if (outId != null && RequestContext.get().isDeletedEntity(outId)) {
+        Id.EntityState state = GraphHelper.getState(outVertex);
+        if ((outId != null && RequestContext.get().isDeletedEntity(outId)) || state == Id.EntityState.DELETED) {
             //If the reference vertex is marked for deletion, skip updating the reference
             return;
         }
@@ -261,8 +274,10 @@ public abstract class DeleteHandler {
             //If its class attribute, its the only edge between two vertices
             if (attributeInfo.multiplicity.nullAllowed()) {
                 edge = GraphHelper.getEdgeForLabel(outVertex, edgeLabel);
-            }
-            else {
+                if (shouldUpdateReverseAttribute) {
+                    GraphHelper.setProperty(outVertex, propertyName, null);
+                }
+            } else {
                 // Cannot unset a required attribute.
                 throw new NullRequiredAttributeException("Cannot unset required attribute " + GraphHelper.getQualifiedFieldName(type, attributeName) +
                     " on " + string(outVertex) + " edge = " + edgeLabel);
@@ -275,23 +290,26 @@ public abstract class DeleteHandler {
             if (elements != null) {
                 elements = new ArrayList<>(elements);   //Make a copy, else list.remove reflects on titan.getProperty()
                 for (String elementEdgeId : elements) {
-                    Edge elementEdge = graphHelper.getEdgeById(elementEdgeId);
+                    Edge elementEdge = graphHelper.getEdgeByEdgeId(outVertex, edgeLabel, elementEdgeId);
                     if (elementEdge == null) {
                         continue;
                     }
 
                     Vertex elementVertex = elementEdge.getVertex(Direction.IN);
                     if (elementVertex.getId().toString().equals(inVertex.getId().toString())) {
-                        if (attributeInfo.multiplicity.nullAllowed() || elements.size() > attributeInfo.multiplicity.lower) {
-                            edge = elementEdge;
-                        }
-                        else {
+                        edge = elementEdge;
+
+                        //TODO element.size includes deleted items as well. should exclude
+                        if (!attributeInfo.multiplicity.nullAllowed()
+                                && elements.size() <= attributeInfo.multiplicity.lower) {
                             // Deleting this edge would violate the attribute's lower bound.
                             throw new NullRequiredAttributeException(
-                                "Cannot remove array element from required attribute " +
-                                    GraphHelper.getQualifiedFieldName(type, attributeName) + " on " + string(outVertex) + " " + string(elementEdge));
+                                    "Cannot remove array element from required attribute " +
+                                            GraphHelper.getQualifiedFieldName(type, attributeName) + " on "
+                                            + string(outVertex) + " " + string(elementEdge));
                         }
-                        if (shouldUpdateReverseAttribute || attributeInfo.isComposite) {
+
+                        if (shouldUpdateReverseAttribute) {
                             //if composite attribute, remove the reference as well. else, just remove the edge
                             //for example, when table is deleted, process still references the table
                             //but when column is deleted, table will not reference the deleted column
@@ -299,8 +317,9 @@ public abstract class DeleteHandler {
                                     attributeName);
                             elements.remove(elementEdge.getId().toString());
                             GraphHelper.setProperty(outVertex, propertyName, elements);
+                            break;
+
                         }
-                        break;
                     }
                 }
             }
@@ -312,11 +331,12 @@ public abstract class DeleteHandler {
             if (keys != null) {
                 keys = new ArrayList<>(keys);   //Make a copy, else list.remove reflects on titan.getProperty()
                 for (String key : keys) {
-                    String keyPropertyName = propertyName + "." + key;
+                    String keyPropertyName = GraphHelper.getQualifiedNameForMapKey(propertyName, key);
                     String mapEdgeId = outVertex.getProperty(keyPropertyName);
-                    Edge mapEdge = graphHelper.getEdgeById(mapEdgeId);
+                    Edge mapEdge = graphHelper.getEdgeByEdgeId(outVertex, keyPropertyName, mapEdgeId);
                     Vertex mapVertex = mapEdge.getVertex(Direction.IN);
                     if (mapVertex.getId().toString().equals(inVertex.getId().toString())) {
+                        //TODO keys.size includes deleted items as well. should exclude
                         if (attributeInfo.multiplicity.nullAllowed() || keys.size() > attributeInfo.multiplicity.lower) {
                             edge = mapEdge;
                         }
@@ -327,7 +347,7 @@ public abstract class DeleteHandler {
                                     GraphHelper.getQualifiedFieldName(type, attributeName) + " on " + string(outVertex) + " " + string(mapEdge));
                         }
 
-                        if (shouldUpdateReverseAttribute || attributeInfo.isComposite) {
+                        if (shouldUpdateReverseAttribute) {
                             //remove this key
                             LOG.debug("Removing edge {}, key {} from the map attribute {}", string(mapEdge), key,
                                     attributeName);
@@ -351,9 +371,11 @@ public abstract class DeleteHandler {
         }
 
         if (edge != null) {
-            deleteEdge(edge);
+            deleteEdge(edge, false);
+            RequestContext requestContext = RequestContext.get();
             GraphHelper.setProperty(outVertex, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY,
-                    RequestContext.get().getRequestTime());
+                    requestContext.getRequestTime());
+            requestContext.recordEntityUpdate(outId);
         }
     }
 
@@ -389,7 +411,7 @@ public abstract class DeleteHandler {
 
         for (String traitNameToBeDeleted : traitNames) {
             String relationshipLabel = GraphHelper.getTraitLabel(typeName, traitNameToBeDeleted);
-            deleteReference(instanceVertex, relationshipLabel, DataTypes.TypeCategory.TRAIT);
+            deleteEdgeReference(instanceVertex, relationshipLabel, DataTypes.TypeCategory.TRAIT, false);
         }
     }
 }

repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepository.java

@@ -22,8 +22,10 @@ import com.google.common.base.Preconditions;
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
 import com.thinkaurelius.titan.core.TitanGraph;
+import com.tinkerpop.blueprints.Edge;
 import com.tinkerpop.blueprints.GraphQuery;
 import com.tinkerpop.blueprints.Vertex;
+import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.GraphTransaction;
 import org.apache.atlas.RequestContext;
@@ -40,7 +42,6 @@ import org.apache.atlas.typesystem.types.ClassType;
 import org.apache.atlas.typesystem.types.DataTypes;
 import org.apache.atlas.typesystem.types.IDataType;
 import org.apache.atlas.typesystem.types.TypeSystem;
-import org.apache.atlas.typesystem.types.TypeUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -258,8 +259,8 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
         try {
             final String entityTypeName = GraphHelper.getTypeName(instanceVertex);
             String relationshipLabel = GraphHelper.getTraitLabel(entityTypeName, traitNameToBeDeleted);
-
-            deleteHandler.deleteReference(instanceVertex, relationshipLabel, DataTypes.TypeCategory.TRAIT);
+            Edge edge = GraphHelper.getEdgeForLabel(instanceVertex, relationshipLabel);
+            deleteHandler.deleteEdgeReference(edge, DataTypes.TypeCategory.TRAIT, false, true);
 
             // update the traits in entity once trait removal is successful
             traitNames.remove(traitNameToBeDeleted);
@@ -284,14 +285,15 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
 
     @Override
     @GraphTransaction
-    public TypeUtils.Pair<List<String>, List<String>> updateEntities(ITypedReferenceableInstance... entitiesUpdated) throws RepositoryException {
+    public AtlasClient.EntityResult updateEntities(ITypedReferenceableInstance... entitiesUpdated) throws RepositoryException {
         LOG.info("updating entity {}", entitiesUpdated);
         try {
             TypedInstanceToGraphMapper instanceToGraphMapper = new TypedInstanceToGraphMapper(graphToInstanceMapper, deleteHandler);
             instanceToGraphMapper.mapTypedInstanceToGraph(TypedInstanceToGraphMapper.Operation.UPDATE_FULL,
                     entitiesUpdated);
             RequestContext requestContext = RequestContext.get();
-            return TypeUtils.Pair.of(requestContext.getCreatedEntityIds(), requestContext.getUpdatedEntityIds());
+            return new AtlasClient.EntityResult(requestContext.getCreatedEntityIds(),
+                    requestContext.getUpdatedEntityIds(), requestContext.getDeletedEntityIds());
         } catch (AtlasException e) {
             throw new RepositoryException(e);
         }
@@ -299,13 +301,14 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
 
     @Override
     @GraphTransaction
-    public TypeUtils.Pair<List<String>, List<String>> updatePartial(ITypedReferenceableInstance entity) throws RepositoryException {
+    public AtlasClient.EntityResult updatePartial(ITypedReferenceableInstance entity) throws RepositoryException {
         LOG.info("updating entity {}", entity);
         try {
             TypedInstanceToGraphMapper instanceToGraphMapper = new TypedInstanceToGraphMapper(graphToInstanceMapper, deleteHandler);
             instanceToGraphMapper.mapTypedInstanceToGraph(TypedInstanceToGraphMapper.Operation.UPDATE_PARTIAL, entity);
             RequestContext requestContext = RequestContext.get();
-            return TypeUtils.Pair.of(requestContext.getCreatedEntityIds(), requestContext.getUpdatedEntityIds());
+            return new AtlasClient.EntityResult(requestContext.getCreatedEntityIds(),
+                    requestContext.getUpdatedEntityIds(), requestContext.getDeletedEntityIds());
         } catch (AtlasException e) {
             throw new RepositoryException(e);
         }
@@ -313,7 +316,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
 
     @Override
     @GraphTransaction
-    public TypeUtils.Pair<List<String>, List<ITypedReferenceableInstance>> deleteEntities(List<String> guids) throws RepositoryException {
+    public AtlasClient.EntityResult deleteEntities(List<String> guids) throws RepositoryException {
 
         if (guids == null || guids.size() == 0) {
             throw new IllegalArgumentException("guids must be non-null and non-empty");
@@ -337,6 +340,7 @@ public class GraphBackedMetadataRepository implements MetadataRepository {
             }
         }
         RequestContext requestContext = RequestContext.get();
-        return new TypeUtils.Pair<>(requestContext.getDeletedEntityIds(), requestContext.getDeletedEntities());
+        return new AtlasClient.EntityResult(requestContext.getCreatedEntityIds(),
+                requestContext.getUpdatedEntityIds(), requestContext.getDeletedEntityIds());
     }
 }

repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java

@@ -107,11 +107,13 @@ public final class GraphHelper {
 
         // add timestamp information
         setProperty(vertexWithoutIdentity, Constants.TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
+        setProperty(vertexWithoutIdentity, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY,
+                RequestContext.get().getRequestTime());
 
         return vertexWithoutIdentity;
     }
 
-    public Edge addEdge(Vertex fromVertex, Vertex toVertex, String edgeLabel) {
+    private Edge addEdge(Vertex fromVertex, Vertex toVertex, String edgeLabel) {
         LOG.debug("Adding edge for {} -> label {} -> {}", string(fromVertex), edgeLabel, string(toVertex));
         Edge edge = titanGraph.addEdge(null, fromVertex, toVertex, edgeLabel);
 
@@ -127,12 +129,34 @@ public final class GraphHelper {
         Iterable<Edge> edges = inVertex.getEdges(Direction.IN, edgeLabel);
         for (Edge edge : edges) {
             if (edge.getVertex(Direction.OUT).getId().toString().equals(outVertex.getId().toString())) {
-                return edge;
+                Id.EntityState edgeState = getState(edge);
+                if (edgeState == null || edgeState == Id.EntityState.ACTIVE) {
+                    return edge;
+                }
             }
         }
         return addEdge(outVertex, inVertex, edgeLabel);
     }
 
+
+    public Edge getEdgeByEdgeId(Vertex outVertex, String edgeLabel, String edgeId) {
+        if (edgeId == null) {
+            return null;
+        }
+        return titanGraph.getEdge(edgeId);
+
+        //TODO get edge id is expensive. Use this logic. But doesn't work for now
+        /**
+        Iterable<Edge> edges = outVertex.getEdges(Direction.OUT, edgeLabel);
+        for (Edge edge : edges) {
+            if (edge.getId().toString().equals(edgeId)) {
+                return edge;
+            }
+        }
+        return null;
+         **/
+    }
+
     /**
      * Args of the format prop1, key1, prop2, key2...
      * Searches for a vertex with prop1=key1 && prop2=key2
@@ -180,15 +204,14 @@ public final class GraphHelper {
      * @return
      */
     public static Edge getEdgeForLabel(Vertex vertex, String edgeLabel) {
-        String vertexState = vertex.getProperty(Constants.STATE_PROPERTY_KEY);
-
         Iterator<Edge> iterator = GraphHelper.getOutGoingEdgesByLabel(vertex, edgeLabel);
         Edge latestDeletedEdge = null;
         long latestDeletedEdgeTime = Long.MIN_VALUE;
+
         while (iterator != null && iterator.hasNext()) {
             Edge edge = iterator.next();
-            String edgeState = edge.getProperty(Constants.STATE_PROPERTY_KEY);
-            if (edgeState == null || Id.EntityState.ACTIVE.name().equals(edgeState)) {
+            Id.EntityState edgeState = getState(edge);
+            if (edgeState == null || edgeState == Id.EntityState.ACTIVE) {
                 LOG.debug("Found {}", string(edge));
                 return edge;
             } else {
@@ -201,19 +224,8 @@ public final class GraphHelper {
         }
 
         //If the vertex is deleted, return latest deleted edge
-        if (Id.EntityState.DELETED.equals(vertexState)) {
-            LOG.debug("Found {}", string(latestDeletedEdge));
-            return latestDeletedEdge;
-        }
-
-        return null;
-    }
-
-    public Edge getEdgeById(String edgeId) {
-        if(edgeId != null) {
-            return titanGraph.getEdge(edgeId);
-        }
-        return null;
+        LOG.debug("Found {}", latestDeletedEdge == null ? "null" : string(latestDeletedEdge));
+        return latestDeletedEdge;
     }
 
     public static String vertexString(final Vertex vertex) {
@@ -343,6 +355,15 @@ public final class GraphHelper {
         return instanceVertex.getProperty(Constants.ENTITY_TYPE_PROPERTY_KEY);
     }
 
+    public static Id.EntityState getState(Element element) {
+        String state = getStateAsString(element);
+        return state == null ? null : Id.EntityState.valueOf(state);
+    }
+
+    public static String getStateAsString(Element element) {
+        return element.getProperty(Constants.STATE_PROPERTY_KEY);
+    }
+
     /**
      * For the given type, finds an unique attribute and checks if there is an existing instance with the same
      * unique value

repository/src/main/java/org/apache/atlas/repository/graph/GraphToTypedInstanceMapper.java

@@ -68,9 +68,9 @@ public final class GraphToTypedInstanceMapper {
         LOG.debug("Mapping graph root vertex {} to typed instance for guid {}", instanceVertex, guid);
         String typeName = instanceVertex.getProperty(Constants.ENTITY_TYPE_PROPERTY_KEY);
         List<String> traits = GraphHelper.getTraitNames(instanceVertex);
+        String state = GraphHelper.getStateAsString(instanceVertex);
 
-        Id id = new Id(guid, instanceVertex.<Integer>getProperty(Constants.VERSION_PROPERTY_KEY), typeName,
-                instanceVertex.<String>getProperty(Constants.STATE_PROPERTY_KEY));
+        Id id = new Id(guid, instanceVertex.<Integer>getProperty(Constants.VERSION_PROPERTY_KEY), typeName, state);
         LOG.debug("Created id {} for instance type {}", id, typeName);
 
         ClassType classType = typeSystem.getDataType(ClassType.class, typeName);
@@ -161,9 +161,9 @@ public final class GraphToTypedInstanceMapper {
 
         Edge edge;
         if (edgeId == null) {
-            edge = GraphHelper.getEdgeForLabel(instanceVertex, relationshipLabel);;
+            edge = GraphHelper.getEdgeForLabel(instanceVertex, relationshipLabel);
         } else {
-            edge = graphHelper.getEdgeById(edgeId);
+            edge = graphHelper.getEdgeByEdgeId(instanceVertex, relationshipLabel, edgeId);
         }
 
         if (edge != null) {
@@ -175,9 +175,10 @@ public final class GraphToTypedInstanceMapper {
                 LOG.debug("Found composite, mapping vertex to instance");
                 return mapGraphToTypedInstance(guid, referenceVertex);
             } else {
+                String state = GraphHelper.getStateAsString(referenceVertex);
                 Id referenceId =
                         new Id(guid, referenceVertex.<Integer>getProperty(Constants.VERSION_PROPERTY_KEY),
-                                dataType.getName());
+                                dataType.getName(), state);
                 LOG.debug("Found non-composite, adding id {} ", referenceId);
                 return referenceId;
             }
@@ -271,7 +272,7 @@ public final class GraphToTypedInstanceMapper {
         if (edgeId == null) {
             edge = GraphHelper.getEdgeForLabel(instanceVertex, relationshipLabel);
         } else {
-            edge = graphHelper.getEdgeById(edgeId);
+            edge = graphHelper.getEdgeByEdgeId(instanceVertex, relationshipLabel, edgeId);
         }
 
         if (edge != null) {

repository/src/main/java/org/apache/atlas/repository/graph/HardDeleteHandler.java

@@ -26,20 +26,18 @@ import org.apache.atlas.typesystem.types.TypeSystem;
 
 public class HardDeleteHandler extends DeleteHandler {
 
-    private static final GraphHelper graphHelper = GraphHelper.getInstance();
-
     @Inject
     public HardDeleteHandler(TypeSystem typeSystem) {
-        super(typeSystem, true);
+        super(typeSystem, true, false);
     }
 
     @Override
-    protected void _deleteVertex(Vertex instanceVertex) {
+    protected void _deleteVertex(Vertex instanceVertex, boolean force) {
         graphHelper.removeVertex(instanceVertex);
     }
 
     @Override
-    protected void deleteEdge(Edge edge) throws AtlasException {
+    protected void deleteEdge(Edge edge, boolean force) throws AtlasException {
         graphHelper.removeEdge(edge);
     }
 }

repository/src/main/java/org/apache/atlas/repository/graph/SoftDeleteHandler.java

@@ -32,24 +32,34 @@ import static org.apache.atlas.repository.Constants.STATE_PROPERTY_KEY;
 public class SoftDeleteHandler extends DeleteHandler {
     @Inject
     public SoftDeleteHandler(TypeSystem typeSystem) {
-        super(typeSystem, false);
+        super(typeSystem, false, true);
     }
 
     @Override
-    protected void _deleteVertex(Vertex instanceVertex) {
-        Id.EntityState state = Id.EntityState.valueOf((String) instanceVertex.getProperty(STATE_PROPERTY_KEY));
-        if (state != Id.EntityState.DELETED) {
-            GraphHelper.setProperty(instanceVertex, STATE_PROPERTY_KEY, Id.EntityState.DELETED.name());
-            GraphHelper.setProperty(instanceVertex, MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
+    protected void _deleteVertex(Vertex instanceVertex, boolean force) {
+        if (force) {
+            graphHelper.removeVertex(instanceVertex);
+        } else {
+            Id.EntityState state = GraphHelper.getState(instanceVertex);
+            if (state != Id.EntityState.DELETED) {
+                GraphHelper.setProperty(instanceVertex, STATE_PROPERTY_KEY, Id.EntityState.DELETED.name());
+                GraphHelper.setProperty(instanceVertex, MODIFICATION_TIMESTAMP_PROPERTY_KEY,
+                        RequestContext.get().getRequestTime());
+            }
         }
     }
 
     @Override
-    protected void deleteEdge(Edge edge) throws AtlasException {
-        Id.EntityState state = Id.EntityState.valueOf((String) edge.getProperty(STATE_PROPERTY_KEY));
-        if (state != Id.EntityState.DELETED) {
-            GraphHelper.setProperty(edge, STATE_PROPERTY_KEY, Id.EntityState.DELETED.name());
-            GraphHelper.setProperty(edge, MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
+    protected void deleteEdge(Edge edge, boolean force) throws AtlasException {
+        if (force) {
+            graphHelper.removeEdge(edge);
+        } else {
+            Id.EntityState state = GraphHelper.getState(edge);
+            if (state != Id.EntityState.DELETED) {
+                GraphHelper.setProperty(edge, STATE_PROPERTY_KEY, Id.EntityState.DELETED.name());
+                GraphHelper
+                        .setProperty(edge, MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
+            }
         }
     }
 }

repository/src/main/java/org/apache/atlas/repository/graph/TypedInstanceToGraphMapper.java

@@ -101,15 +101,15 @@ public final class TypedInstanceToGraphMapper {
             case CREATE:
                 List<String> ids = addOrUpdateAttributesAndTraits(operation, entitiesToCreate);
                 addFullTextProperty(entitiesToCreate, fulltextMapper);
-                requestContext.recordCreatedEntities(ids);
+                requestContext.recordEntityCreate(ids);
                 break;
 
             case UPDATE_FULL:
             case UPDATE_PARTIAL:
                 ids = addOrUpdateAttributesAndTraits(Operation.CREATE, entitiesToCreate);
-                requestContext.recordCreatedEntities(ids);
+                requestContext.recordEntityCreate(ids);
                 ids = addOrUpdateAttributesAndTraits(operation, entitiesToUpdate);
-                requestContext.recordUpdatedEntities(ids);
+                requestContext.recordEntityUpdate(ids);
 
                 addFullTextProperty(entitiesToCreate, fulltextMapper);
                 addFullTextProperty(entitiesToUpdate, fulltextMapper);
@@ -218,8 +218,8 @@ public final class TypedInstanceToGraphMapper {
                         attrValue, currentEdge, edgeLabel, operation);
 
                 if (currentEdge != null && !currentEdge.getId().toString().equals(newEdgeId)) {
-                    deleteHandler.deleteReference(currentEdge, attributeInfo.dataType().getTypeCategory(),
-                            attributeInfo.isComposite);
+                    deleteHandler.deleteEdgeReference(currentEdge, attributeInfo.dataType().getTypeCategory(),
+                            attributeInfo.isComposite, true);
                 }
                 break;
 
@@ -326,7 +326,7 @@ public final class TypedInstanceToGraphMapper {
         String propertyName = GraphHelper.getQualifiedFieldName(typedInstance, attributeInfo);
         List<String> currentElements = instanceVertex.getProperty(propertyName);
         IDataType elementType = ((DataTypes.ArrayType) attributeInfo.dataType()).getElemType();
-        List<String> newElementsCreated = new ArrayList<>();
+        List<Object> newElementsCreated = new ArrayList<>();
 
         if (!newAttributeEmpty) {
             if (newElements != null && !newElements.isEmpty()) {
@@ -336,75 +336,53 @@ public final class TypedInstanceToGraphMapper {
                             currentElements.get(index) : null;
                     LOG.debug("Adding/updating element at position {}, current element {}, new element {}", index,
                             currentElement, newElements.get(index));
-                    String newEntry = addOrUpdateCollectionEntry(instanceVertex, attributeInfo, elementType,
+                    Object newEntry = addOrUpdateCollectionEntry(instanceVertex, attributeInfo, elementType,
                             newElements.get(index), currentElement, propertyName, operation);
                     newElementsCreated.add(newEntry);
                 }
             }
         }
 
+        List<String> additionalEdges = removeUnusedEntries(instanceVertex, propertyName, currentElements,
+                newElementsCreated, elementType, attributeInfo);
+        newElementsCreated.addAll(additionalEdges);
+
         // for dereference on way out
         GraphHelper.setProperty(instanceVertex, propertyName, newElementsCreated);
-
-        removeUnusedEntries(currentElements, newElementsCreated, elementType, attributeInfo);
     }
 
-    private void removeUnusedEntries(List<String> currentEntries, List<String> newEntries, IDataType entryType,
-                                     AttributeInfo attributeInfo) throws AtlasException {
-        if (currentEntries == null || currentEntries.isEmpty()) {
-            return;
-        }
-
-        LOG.debug("Removing unused entries from the old collection");
-        if (entryType.getTypeCategory() == DataTypes.TypeCategory.STRUCT
-                || entryType.getTypeCategory() == DataTypes.TypeCategory.CLASS) {
-
-            //Get map of edge id to edge
-            Map<String, Edge> edgeMap = new HashMap<>();
-            getEdges(currentEntries, edgeMap);
-            getEdges(newEntries, edgeMap);
-
-            //Get final set of in vertices
-            Set<String> newInVertices = new HashSet<>();
-            for (String edgeId : newEntries) {
-                Vertex inVertex = edgeMap.get(edgeId).getVertex(Direction.IN);
-                newInVertices.add(inVertex.getId().toString());
-            }
-
-            //Remove the edges for (current edges - new edges)
-            List<String> cloneElements = new ArrayList<>(currentEntries);
-            cloneElements.removeAll(newEntries);
-            LOG.debug("Removing unused entries from the old collection - {}", cloneElements);
-
-            if (!cloneElements.isEmpty()) {
-                for (String edgeIdForDelete : cloneElements) {
-                    Edge edge = edgeMap.get(edgeIdForDelete);
-                    Vertex inVertex = edge.getVertex(Direction.IN);
-                    if (newInVertices.contains(inVertex.getId().toString())) {
-                        //If the edge.inVertex is in the new set of in vertices, just delete the edge
-                        deleteHandler.deleteEdge(edge, true);
-                    } else {
-                        //else delete the edge + vertex
-                        deleteHandler.deleteReference(edge, entryType.getTypeCategory(), attributeInfo.isComposite);
+    //Removes unused edges from the old collection, compared to the new collection
+    private List<String> removeUnusedEntries(Vertex instanceVertex, String edgeLabel,
+                                             Collection<String> currentEntries,
+                                             Collection<Object> newEntries,
+                                             IDataType entryType, AttributeInfo attributeInfo) throws AtlasException {
+        if (currentEntries != null && !currentEntries.isEmpty()) {
+            LOG.debug("Removing unused entries from the old collection");
+            if (entryType.getTypeCategory() == DataTypes.TypeCategory.STRUCT
+                    || entryType.getTypeCategory() == DataTypes.TypeCategory.CLASS) {
+
+                //Remove the edges for (current edges - new edges)
+                List<String> cloneElements = new ArrayList<>(currentEntries);
+                cloneElements.removeAll(newEntries);
+                List<String> additionalElements = new ArrayList<>();
+                LOG.debug("Removing unused entries from the old collection - {}", cloneElements);
+
+                if (!cloneElements.isEmpty()) {
+                    for (String edgeIdForDelete : cloneElements) {
+                        Edge edge = graphHelper.getEdgeByEdgeId(instanceVertex, edgeLabel, edgeIdForDelete);
+                        boolean deleted = deleteHandler.deleteEdgeReference(edge, entryType.getTypeCategory(),
+                                attributeInfo.isComposite, true);
+                        if (!deleted) {
+                            additionalElements.add(edgeIdForDelete);
+                        }
                     }
                 }
+                return additionalElements;
             }
         }
+        return new ArrayList<>();
     }
 
-    private void getEdges(List<String> edgeIds, Map<String, Edge> edgeMap) {
-        if (edgeIds == null) {
-            return;
-        }
-
-        for (String edgeId : edgeIds) {
-            if (!edgeMap.containsKey(edgeId)) {
-                edgeMap.put(edgeId, graphHelper.getEdgeById(edgeId));
-            }
-        }
-    }
-
-
     /******************************************** MAP **************************************************/
 
     private void mapMapCollectionToVertex(ITypedInstance typedInstance, Vertex instanceVertex,
@@ -421,38 +399,81 @@ public final class TypedInstanceToGraphMapper {
 
         IDataType elementType = ((DataTypes.MapType) attributeInfo.dataType()).getValueType();
         String propertyName = GraphHelper.getQualifiedFieldName(typedInstance, attributeInfo);
-        List<String> currentElements = new ArrayList<>();
-        List<String> newElementsCreated = new ArrayList<>();
-        List<String> newKeysCreated = new ArrayList<>();
+
+        Map<String, String> currentMap = new HashMap<>();
+        Map<String, Object> newMap = new HashMap<>();
+
+        List<String> currentKeys = instanceVertex.getProperty(propertyName);
+        if (currentKeys != null && !currentKeys.isEmpty()) {
+            for (String key : currentKeys) {
+                String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(propertyName, key);
+                String propertyValueForKey = instanceVertex.getProperty(propertyNameForKey).toString();
+                currentMap.put(key, propertyValueForKey);
+            }
+        }
 
         if (!newAttributeEmpty) {
             for (Map.Entry entry : newAttribute.entrySet()) {
-                String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(propertyName, entry.getKey().toString());
-                newKeysCreated.add(entry.getKey().toString());
+                String keyStr = entry.getKey().toString();
+                String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(propertyName, keyStr);
 
-                String currentEntry = instanceVertex.getProperty(propertyNameForKey);
-                if (currentEntry != null) {
-                    currentElements.add(currentEntry);
-                }
-
-                String newEntry = addOrUpdateCollectionEntry(instanceVertex, attributeInfo, elementType,
-                        entry.getValue(), currentEntry, propertyNameForKey, operation);
+                Object newEntry = addOrUpdateCollectionEntry(instanceVertex, attributeInfo, elementType,
+                        entry.getValue(), currentMap.get(keyStr), propertyNameForKey, operation);
 
                 //Add/Update/Remove property value
                 GraphHelper.setProperty(instanceVertex, propertyNameForKey, newEntry);
-                newElementsCreated.add(newEntry);
+                newMap.put(keyStr, newEntry);
             }
         }
 
+        Map<String, String> additionalMap =
+                removeUnusedMapEntries(instanceVertex, propertyName, currentMap, newMap, elementType, attributeInfo);
+
+        Set<String> newKeys = new HashSet<>(newMap.keySet());
+        newKeys.addAll(additionalMap.keySet());
+
         // for dereference on way out
-        GraphHelper.setProperty(instanceVertex, propertyName, newKeysCreated);
+        GraphHelper.setProperty(instanceVertex, propertyName, new ArrayList<>(newKeys));
+    }
+
+    //Remove unused entries from map
+    private Map<String, String> removeUnusedMapEntries(Vertex instanceVertex, String propertyName,
+                                                       Map<String, String> currentMap,
+                                                       Map<String, Object> newMap, IDataType elementType,
+                                                       AttributeInfo attributeInfo)
+            throws AtlasException {
+        boolean reference = (elementType.getTypeCategory() == DataTypes.TypeCategory.STRUCT
+                || elementType.getTypeCategory() == DataTypes.TypeCategory.CLASS);
+        Map<String, String> additionalMap = new HashMap<>();
+
+        for (String currentKey : currentMap.keySet()) {
+            boolean shouldDeleteKey = !newMap.containsKey(currentKey);
+            if (reference) {
+                String currentEdge = currentMap.get(currentKey);
+                //Delete the edge reference if its not part of new edges created/updated
+                if (!newMap.values().contains(currentEdge)) {
+                    String edgeLabel = GraphHelper.getQualifiedNameForMapKey(propertyName, currentKey);
+                    Edge edge = graphHelper.getEdgeByEdgeId(instanceVertex, edgeLabel, currentMap.get(currentKey));
+                    boolean deleted =
+                            deleteHandler.deleteEdgeReference(edge, elementType.getTypeCategory(), attributeInfo.isComposite, true);
+                    if (!deleted) {
+                        additionalMap.put(currentKey, currentEdge);
+                        shouldDeleteKey = false;
+                    }
+                }
+            }
 
-        removeUnusedEntries(currentElements, newElementsCreated, elementType, attributeInfo);
+            if (shouldDeleteKey) {
+                String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(propertyName, currentKey);
+                graphHelper.setProperty(instanceVertex, propertyNameForKey, null);
+            }
+        }
+        return additionalMap;
     }
 
     /******************************************** ARRAY & MAP **************************************************/
 
-    private String addOrUpdateCollectionEntry(Vertex instanceVertex, AttributeInfo attributeInfo,
+    private Object addOrUpdateCollectionEntry(Vertex instanceVertex, AttributeInfo attributeInfo,
                                               IDataType elementType, Object newAttributeValue, String currentValue,
                                               String propertyName, Operation operation)
             throws AtlasException {
@@ -460,7 +481,7 @@ public final class TypedInstanceToGraphMapper {
         switch (elementType.getTypeCategory()) {
         case PRIMITIVE:
         case ENUM:
-            return newAttributeValue != null ? newAttributeValue.toString() : null;
+            return newAttributeValue != null ? newAttributeValue : null;
 
         case ARRAY:
         case MAP:
@@ -471,7 +492,7 @@ public final class TypedInstanceToGraphMapper {
         case STRUCT:
         case CLASS:
             final String edgeLabel = GraphHelper.EDGE_LABEL_PREFIX + propertyName;
-            Edge currentEdge = graphHelper.getEdgeById(currentValue);
+            Edge currentEdge = graphHelper.getEdgeByEdgeId(instanceVertex, edgeLabel, currentValue);
             return addOrUpdateReference(instanceVertex, attributeInfo, elementType, newAttributeValue, currentEdge,
                     edgeLabel, operation);
 
@@ -526,7 +547,7 @@ public final class TypedInstanceToGraphMapper {
         mapInstanceToVertex(structInstance, structInstanceVertex, structInstance.fieldMapping().fields, false,
                 Operation.CREATE);
         // add an edge to the newly created vertex from the parent
-        Edge newEdge = graphHelper.addEdge(instanceVertex, structInstanceVertex, edgeLabel);
+        Edge newEdge = graphHelper.getOrCreateEdge(instanceVertex, structInstanceVertex, edgeLabel);
 
         return newEdge;
     }
@@ -575,7 +596,7 @@ public final class TypedInstanceToGraphMapper {
 
     private Edge addClassEdge(Vertex instanceVertex, Vertex toVertex, String edgeLabel) throws AtlasException {
         // add an edge to the class vertex from the instance
-        return graphHelper.addEdge(instanceVertex, toVertex, edgeLabel);
+        return graphHelper.getOrCreateEdge(instanceVertex, toVertex, edgeLabel);
     }
 
     private Vertex getClassVertex(ITypedReferenceableInstance typedReference) throws EntityNotFoundException {
@@ -644,7 +665,7 @@ public final class TypedInstanceToGraphMapper {
 
         // add an edge to the newly created vertex from the parent
         String relationshipLabel = GraphHelper.getTraitLabel(entityType.getName(), traitName);
-        graphHelper.addEdge(parentInstanceVertex, traitInstanceVertex, relationshipLabel);
+        graphHelper.getOrCreateEdge(parentInstanceVertex, traitInstanceVertex, relationshipLabel);
     }
 
     /******************************************** PRIMITIVES **************************************************/

repository/src/main/java/org/apache/atlas/repository/typestore/GraphBackedTypeStore.java

@@ -69,6 +69,8 @@ public class GraphBackedTypeStore implements ITypeStore {
 
     private final TitanGraph titanGraph;
 
+    private GraphHelper graphHelper = GraphHelper.getInstance();
+
     @Inject
     public GraphBackedTypeStore(GraphProvider<TitanGraph> graphProvider) {
         titanGraph = graphProvider.get();
@@ -155,7 +157,7 @@ public class GraphBackedTypeStore implements ITypeStore {
             for (String superTypeName : superTypes) {
                 HierarchicalType superType = typeSystem.getDataType(HierarchicalType.class, superTypeName);
                 Vertex superVertex = createVertex(superType.getTypeCategory(), superTypeName, superType.getDescription());
-                addEdge(vertex, superVertex, SUPERTYPE_EDGE_LABEL);
+                graphHelper.getOrCreateEdge(vertex, superVertex, SUPERTYPE_EDGE_LABEL);
             }
         }
     }
@@ -200,26 +202,11 @@ public class GraphBackedTypeStore implements ITypeStore {
             if (!coreTypes.contains(attrType.getName())) {
                 Vertex attrVertex = createVertex(attrType.getTypeCategory(), attrType.getName(), attrType.getDescription());
                 String label = getEdgeLabel(vertexTypeName, attribute.name);
-                addEdge(vertex, attrVertex, label);
+                graphHelper.getOrCreateEdge(vertex, attrVertex, label);
             }
         }
     }
 
-    private void addEdge(Vertex fromVertex, Vertex toVertex, String label) {
-        Iterator<Edge> edges = GraphHelper.getOutGoingEdgesByLabel(fromVertex, label);
-        // ATLAS-474: Check if this type system edge already exists, to avoid duplicates.
-        while (edges.hasNext()) {
-            Edge edge = edges.next();
-            if (edge.getVertex(Direction.IN).equals(toVertex)) {
-                LOG.debug("Edge from {} to {} with label {} already exists", 
-                    toString(fromVertex), toString(toVertex), label);
-                return;
-            }
-        }        
-        LOG.debug("Adding edge from {} to {} with label {}", toString(fromVertex), toString(toVertex), label);
-        titanGraph.addEdge(null, fromVertex, toVertex, label);
-    }
-
     @Override
     @GraphTransaction
     public TypesDef restore() throws AtlasException {

repository/src/main/java/org/apache/atlas/services/DefaultMetadataService.java

@@ -26,6 +26,7 @@ import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.EntityAuditEvent;
+import org.apache.atlas.RequestContext;
 import org.apache.atlas.classification.InterfaceAudience;
 import org.apache.atlas.ha.HAConfiguration;
 import org.apache.atlas.listener.ActiveStateChangeHandler;
@@ -58,8 +59,6 @@ import org.apache.atlas.typesystem.types.Multiplicity;
 import org.apache.atlas.typesystem.types.StructTypeDefinition;
 import org.apache.atlas.typesystem.types.TraitType;
 import org.apache.atlas.typesystem.types.TypeSystem;
-import org.apache.atlas.typesystem.types.TypeUtils;
-import org.apache.atlas.typesystem.types.TypeUtils.Pair;
 import org.apache.atlas.typesystem.types.ValueConversionException;
 import org.apache.atlas.typesystem.types.utils.TypesUtil;
 import org.apache.atlas.utils.ParamChecker;
@@ -305,16 +304,15 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
      * Creates an entity, instance of the type.
      *
      * @param entityInstanceDefinition json array of entity definitions
-     * @return guids - json array of guids
+     * @return guids - list of guids
      */
     @Override
-    public String createEntities(String entityInstanceDefinition) throws AtlasException {
+    public List<String> createEntities(String entityInstanceDefinition) throws AtlasException {
         ParamChecker.notEmpty(entityInstanceDefinition, "Entity instance definition");
 
         ITypedReferenceableInstance[] typedInstances = deserializeClassInstances(entityInstanceDefinition);
 
-        List<String> guids = createEntities(typedInstances);
-        return new JSONArray(guids).toString();
+        return createEntities(typedInstances);
     }
 
     public List<String> createEntities(ITypedReferenceableInstance[] typedInstances) throws AtlasException {
@@ -422,25 +420,26 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
      * @return guids - json array of guids
      */
     @Override
-    public String updateEntities(String entityInstanceDefinition) throws AtlasException {
+    public AtlasClient.EntityResult updateEntities(String entityInstanceDefinition) throws AtlasException {
 
         ParamChecker.notEmpty(entityInstanceDefinition, "Entity instance definition");
         ITypedReferenceableInstance[] typedInstances = deserializeClassInstances(entityInstanceDefinition);
 
-        TypeUtils.Pair<List<String>, List<String>> guids = repository.updateEntities(typedInstances);
-        return onEntitiesAddedUpdated(guids);
+        AtlasClient.EntityResult entityResult = repository.updateEntities(typedInstances);
+        onEntitiesAddedUpdated(entityResult);
+        return entityResult;
     }
 
-    private String onEntitiesAddedUpdated(TypeUtils.Pair<List<String>, List<String>> guids) throws AtlasException {
-        onEntitiesAdded(guids.left);
-        onEntitiesUpdated(guids.right);
-
-        guids.left.addAll(guids.right);
-        return new JSONArray(guids.left).toString();
+    private void onEntitiesAddedUpdated(AtlasClient.EntityResult entityResult) throws AtlasException {
+        onEntitiesAdded(entityResult.getCreatedEntities());
+        onEntitiesUpdated(entityResult.getUpdateEntities());
+        //Note: doesn't access deletedEntities from entityResult
+        onEntitiesDeleted(RequestContext.get().getDeletedEntities());
     }
 
     @Override
-    public String updateEntityAttributeByGuid(final String guid, String attributeName, String value) throws AtlasException {
+    public AtlasClient.EntityResult updateEntityAttributeByGuid(final String guid, String attributeName,
+                                                                String value) throws AtlasException {
         ParamChecker.notEmpty(guid, "entity id");
         ParamChecker.notEmpty(attributeName, "attribute name");
         ParamChecker.notEmpty(value, "attribute value");
@@ -469,8 +468,9 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
         }
 
         ((ReferenceableInstance)newInstance).replaceWithNewId(new Id(guid, 0, newInstance.getTypeName()));
-        TypeUtils.Pair<List<String>, List<String>> guids = repository.updatePartial(newInstance);
-        return onEntitiesAddedUpdated(guids);
+        AtlasClient.EntityResult entityResult = repository.updatePartial(newInstance);
+        onEntitiesAddedUpdated(entityResult);
+        return entityResult;
     }
 
     private ITypedReferenceableInstance validateEntityExists(String guid)
@@ -483,7 +483,8 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
     }
 
     @Override
-    public String updateEntityPartialByGuid(final String guid, Referenceable newEntity) throws AtlasException {
+    public AtlasClient.EntityResult updateEntityPartialByGuid(final String guid, Referenceable newEntity)
+            throws AtlasException {
         ParamChecker.notEmpty(guid, "guid cannot be null");
         ParamChecker.notNull(newEntity, "updatedEntity cannot be null");
         ITypedReferenceableInstance existInstance = validateEntityExists(guid);
@@ -491,11 +492,13 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
         ITypedReferenceableInstance newInstance = convertToTypedInstance(newEntity, existInstance.getTypeName());
         ((ReferenceableInstance)newInstance).replaceWithNewId(new Id(guid, 0, newInstance.getTypeName()));
 
-        TypeUtils.Pair<List<String>, List<String>> guids = repository.updatePartial(newInstance);
-        return onEntitiesAddedUpdated(guids);
+        AtlasClient.EntityResult entityResult = repository.updatePartial(newInstance);
+        onEntitiesAddedUpdated(entityResult);
+        return entityResult;
     }
 
-    private ITypedReferenceableInstance convertToTypedInstance(Referenceable updatedEntity, String typeName) throws AtlasException {
+    private ITypedReferenceableInstance convertToTypedInstance(Referenceable updatedEntity, String typeName)
+            throws AtlasException {
         ClassType type = typeSystem.getDataType(ClassType.class, typeName);
         ITypedReferenceableInstance newInstance = type.createInstance();
 
@@ -538,8 +541,9 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
     }
 
     @Override
-    public String updateEntityByUniqueAttribute(String typeName, String uniqueAttributeName, String attrValue,
-                                                Referenceable updatedEntity) throws AtlasException {
+    public AtlasClient.EntityResult updateEntityByUniqueAttribute(String typeName, String uniqueAttributeName,
+                                                                  String attrValue,
+                                                                  Referenceable updatedEntity) throws AtlasException {
         ParamChecker.notEmpty(typeName, "typeName");
         ParamChecker.notEmpty(uniqueAttributeName, "uniqueAttributeName");
         ParamChecker.notNull(attrValue, "unique attribute value");
@@ -550,8 +554,9 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
         final ITypedReferenceableInstance newInstance = convertToTypedInstance(updatedEntity, typeName);
         ((ReferenceableInstance)newInstance).replaceWithNewId(oldInstance.getId());
 
-        TypeUtils.Pair<List<String>, List<String>> guids = repository.updatePartial(newInstance);
-        return onEntitiesAddedUpdated(guids);
+        AtlasClient.EntityResult entityResult = repository.updatePartial(newInstance);
+        onEntitiesAddedUpdated(entityResult);
+        return entityResult;
     }
 
     private void validateTypeExists(String entityType) throws AtlasException {
@@ -726,13 +731,14 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
      * @see org.apache.atlas.services.MetadataService#deleteEntities(java.lang.String)
      */
     @Override
-    public List<String> deleteEntities(List<String> deleteCandidateGuids) throws AtlasException {
+    public AtlasClient.EntityResult deleteEntities(List<String> deleteCandidateGuids) throws AtlasException {
         ParamChecker.notEmpty(deleteCandidateGuids, "delete candidate guids");
         return deleteGuids(deleteCandidateGuids);
     }
 
     @Override
-    public List<String> deleteEntityByUniqueAttribute(String typeName, String uniqueAttributeName, String attrValue) throws AtlasException {
+    public AtlasClient.EntityResult deleteEntityByUniqueAttribute(String typeName, String uniqueAttributeName,
+                                                                  String attrValue) throws AtlasException {
         ParamChecker.notEmpty(typeName, "delete candidate typeName");
         ParamChecker.notEmpty(uniqueAttributeName, "delete candidate unique attribute name");
         ParamChecker.notEmpty(attrValue, "delete candidate unique attribute value");
@@ -745,12 +751,10 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
         return deleteGuids(deleteCandidateGuids);
     }
 
-    private List<String> deleteGuids(List<String> deleteCandidateGuids) throws AtlasException {
-        Pair<List<String>, List<ITypedReferenceableInstance>> deleteEntitiesResult = repository.deleteEntities(deleteCandidateGuids);
-        if (deleteEntitiesResult.right.size() > 0) {
-            onEntitiesDeleted(deleteEntitiesResult.right);
-        }
-        return deleteEntitiesResult.left;
+    private AtlasClient.EntityResult deleteGuids(List<String> deleteCandidateGuids) throws AtlasException {
+        AtlasClient.EntityResult entityResult = repository.deleteEntities(deleteCandidateGuids);
+        onEntitiesAddedUpdated(entityResult);
+        return entityResult;
     }
 
     private void onEntitiesDeleted(List<ITypedReferenceableInstance> entities) throws AtlasException {

repository/src/test/java/org/apache/atlas/repository/audit/AuditRepositoryTestBase.java

@@ -19,6 +19,7 @@
 package org.apache.atlas.repository.audit;
 
 import org.apache.atlas.EntityAuditEvent;
+import org.apache.atlas.typesystem.Referenceable;
 import org.apache.commons.lang.RandomStringUtils;
 import org.testng.annotations.Test;
 
@@ -38,7 +39,7 @@ public class AuditRepositoryTestBase {
     @Test
     public void testAddEvents() throws Exception {
         EntityAuditEvent event = new EntityAuditEvent(rand(), System.currentTimeMillis(), "u1",
-                EntityAuditEvent.EntityAuditAction.ENTITY_CREATE, "d1");
+                EntityAuditEvent.EntityAuditAction.ENTITY_CREATE, "d1", new Referenceable(rand()));
 
         eventRepository.putEvents(event);
 
@@ -54,17 +55,18 @@ public class AuditRepositoryTestBase {
         String id2 = "id2" + rand();
         String id3 = "id3" + rand();
         long ts = System.currentTimeMillis();
+        Referenceable entity = new Referenceable(rand());
         List<EntityAuditEvent> expectedEvents = new ArrayList<>(3);
         for (int i = 0; i < 3; i++) {
             //Add events for both ids
             EntityAuditEvent event = new EntityAuditEvent(id2, ts - i, "user" + i,
-                            EntityAuditEvent.EntityAuditAction.ENTITY_UPDATE, "details" + i);
+                            EntityAuditEvent.EntityAuditAction.ENTITY_UPDATE, "details" + i, entity);
             eventRepository.putEvents(event);
             expectedEvents.add(event);
             eventRepository.putEvents(new EntityAuditEvent(id1, ts - i, "user" + i,
-                    EntityAuditEvent.EntityAuditAction.TAG_DELETE, "details" + i));
+                    EntityAuditEvent.EntityAuditAction.TAG_DELETE, "details" + i, entity));
             eventRepository.putEvents(new EntityAuditEvent(id3, ts - i, "user" + i,
-                    EntityAuditEvent.EntityAuditAction.TAG_ADD, "details" + i));
+                    EntityAuditEvent.EntityAuditAction.TAG_ADD, "details" + i, entity));
         }
 
         //Use ts for which there is no event - ts + 2

repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryDeleteTestBase.java

@@ -23,7 +23,6 @@ import com.google.common.collect.ImmutableSet;
 import com.thinkaurelius.titan.core.TitanGraph;
 import com.thinkaurelius.titan.core.util.TitanCleanup;
 import com.tinkerpop.blueprints.Vertex;
-
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.RepositoryMetadataModule;
@@ -50,7 +49,6 @@ import org.apache.atlas.typesystem.types.Multiplicity;
 import org.apache.atlas.typesystem.types.StructTypeDefinition;
 import org.apache.atlas.typesystem.types.TraitType;
 import org.apache.atlas.typesystem.types.TypeSystem;
-import org.apache.atlas.typesystem.types.TypeUtils;
 import org.apache.atlas.typesystem.types.utils.TypesUtil;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
@@ -60,7 +58,6 @@ import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
 import javax.inject.Inject;
-
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -71,6 +68,7 @@ import java.util.Map;
 import static org.apache.atlas.TestUtils.COLUMNS_ATTR_NAME;
 import static org.apache.atlas.TestUtils.COLUMN_TYPE;
 import static org.apache.atlas.TestUtils.NAME;
+import static org.apache.atlas.TestUtils.PII;
 import static org.apache.atlas.TestUtils.PROCESS_TYPE;
 import static org.apache.atlas.TestUtils.TABLE_TYPE;
 import static org.apache.atlas.TestUtils.createColumnEntity;
@@ -78,8 +76,6 @@ import static org.apache.atlas.TestUtils.createDBEntity;
 import static org.apache.atlas.TestUtils.createTableEntity;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotEquals;
-import static org.testng.Assert.assertNotNull;
-import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
 
@@ -145,7 +141,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         assertEquals(instance.getId()._getId(), id);
 
         //delete entity should mark it as deleted
-        List<String> results = deleteEntities(id);
+        List<String> results = deleteEntities(id).getDeletedEntities();
         assertEquals(results.get(0), id);
         assertEntityDeleted(id);
 
@@ -167,6 +163,26 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
     }
 
     @Test
+    public void testDeleteEntityWithTraits() throws Exception {
+        Referenceable entity = createDBEntity();
+        String id = createInstance(entity);
+
+        TraitType dataType = typeSystem.getDataType(TraitType.class, PII);
+        ITypedStruct trait = dataType.convert(new Struct(TestUtils.PII), Multiplicity.REQUIRED);
+        repositoryService.addTrait(id, trait);
+
+        ITypedReferenceableInstance instance = repositoryService.getEntityDefinition(id);
+        assertTrue(instance.getTraits().contains(PII));
+
+        deleteEntities(id);
+        assertEntityDeleted(id);
+        assertTestDeleteEntityWithTraits(id);
+    }
+
+    protected abstract void assertTestDeleteEntityWithTraits(String guid)
+            throws EntityNotFoundException, RepositoryException, Exception;
+
+    @Test
     public void testDeleteReference() throws Exception {
         //Deleting column should update table
         Referenceable db = createDBEntity();
@@ -179,13 +195,16 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         table.set(COLUMNS_ATTR_NAME, Arrays.asList(new Id(colId, 0, COLUMN_TYPE)));
         String tableId = createInstance(table);
 
-        deleteEntities(colId);
+        AtlasClient.EntityResult entityResult = deleteEntities(colId);
+        assertEquals(entityResult.getDeletedEntities().size(), 1);
+        assertEquals(entityResult.getDeletedEntities().get(0), colId);
+        assertEquals(entityResult.getUpdateEntities().size(), 1);
+        assertEquals(entityResult.getUpdateEntities().get(0), tableId);
+
         assertEntityDeleted(colId);
 
         ITypedReferenceableInstance tableInstance = repositoryService.getEntityDefinition(tableId);
-        List<ITypedReferenceableInstance> columns =
-                (List<ITypedReferenceableInstance>) tableInstance.get(COLUMNS_ATTR_NAME);
-        assertNull(columns);
+        assertColumnForTestDeleteReference(tableInstance);
 
         //Deleting table should update process
         Referenceable process = new Referenceable(PROCESS_TYPE);
@@ -195,18 +214,23 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
 
         deleteEntities(tableId);
         assertEntityDeleted(tableId);
-        assertTestDeleteReference(processInstance);
+
+        assertTableForTestDeleteReference(tableId);
+        assertProcessForTestDeleteReference(processInstance);
     }
 
-    protected abstract void assertTestDeleteReference(ITypedReferenceableInstance processInstance) throws Exception;
+    protected abstract void assertTableForTestDeleteReference(String tableId) throws Exception;
+
+    protected abstract void assertColumnForTestDeleteReference(ITypedReferenceableInstance tableInstance)
+            throws AtlasException;
+
+    protected abstract void assertProcessForTestDeleteReference(ITypedReferenceableInstance processInstance) throws Exception;
 
     protected abstract void assertEntityDeleted(String id) throws Exception;
 
-    private List<String> deleteEntities(String... id) throws Exception {
+    private AtlasClient.EntityResult deleteEntities(String... id) throws Exception {
         RequestContext.createContext();
-        List<String> response = repositoryService.deleteEntities(Arrays.asList(id)).left;
-        assertNotNull(response);
-        return response;
+        return repositoryService.deleteEntities(Arrays.asList(id));
     }
 
     private String createInstance(Referenceable entity) throws Exception {
@@ -228,21 +252,41 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         table1Entity.set(COLUMNS_ATTR_NAME, ImmutableList.of(col1, col2, col3));
         createInstance(table1Entity);
 
-        // Retrieve the table entities from the auditRepository,
-        // to get their guids and the composite column guids.
+        // Retrieve the table entities from the Repository, to get their guids and the composite column guids.
         ITypedReferenceableInstance tableInstance = repositoryService.getEntityDefinition(TestUtils.TABLE_TYPE,
                 NAME, table1Entity.get(NAME));
-        List<IReferenceableInstance> table1Columns = (List<IReferenceableInstance>) tableInstance.get(COLUMNS_ATTR_NAME);
+        List<IReferenceableInstance> columns = (List<IReferenceableInstance>) tableInstance.get(COLUMNS_ATTR_NAME);
+
+        //Delete column
+        String colId = columns.get(0).getId()._getId();
+        String tableId = tableInstance.getId()._getId();
+
+        AtlasClient.EntityResult entityResult = deleteEntities(colId);
+        assertEquals(entityResult.getDeletedEntities().size(), 1);
+        assertEquals(entityResult.getDeletedEntities().get(0), colId);
+        assertEquals(entityResult.getUpdateEntities().size(), 1);
+        assertEquals(entityResult.getUpdateEntities().get(0), tableId);
+        assertEntityDeleted(colId);
+
+        tableInstance = repositoryService.getEntityDefinition(TestUtils.TABLE_TYPE, NAME, table1Entity.get(NAME));
+        assertDeletedColumn(tableInstance);
 
-        // Delete the table entities.  The deletion should cascade
-        // to their composite columns.
-        List<String> deletedGuids = deleteEntities(tableInstance.getId()._getId());
+        //update by removing a column
+        tableInstance.set(COLUMNS_ATTR_NAME, ImmutableList.of(col3));
+        entityResult = updatePartial(tableInstance);
+        colId = columns.get(1).getId()._getId();
+        assertEquals(entityResult.getDeletedEntities().size(), 1);
+        assertEquals(entityResult.getDeletedEntities().get(0), colId);
+        assertEntityDeleted(colId);
+
+        // Delete the table entities.  The deletion should cascade to their composite columns.
+        tableInstance = repositoryService.getEntityDefinition(TestUtils.TABLE_TYPE, NAME, table1Entity.get(NAME));
+        List<String> deletedGuids = deleteEntities(tableInstance.getId()._getId()).getDeletedEntities();
+        assertEquals(deletedGuids.size(), 2);
 
         // Verify that deleteEntities() response has guids for tables and their composite columns.
         Assert.assertTrue(deletedGuids.contains(tableInstance.getId()._getId()));
-        for (IReferenceableInstance column : table1Columns) {
-            Assert.assertTrue(deletedGuids.contains(column.getId()._getId()));
-        }
+        Assert.assertTrue(deletedGuids.contains(columns.get(2).getId()._getId()));
 
         // Verify that tables and their composite columns have been deleted from the graph Repository.
         for (String guid : deletedGuids) {
@@ -251,6 +295,8 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         assertTestDeleteEntities(tableInstance);
     }
 
+    protected abstract void assertDeletedColumn(ITypedReferenceableInstance tableInstance) throws AtlasException;
+
     protected abstract void assertTestDeleteEntities(ITypedReferenceableInstance tableInstance) throws Exception;
 
     /**
@@ -276,12 +322,13 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         vertexCount = getVertices(Constants.ENTITY_TYPE_PROPERTY_KEY, "SecurityClearance").size();
         Assert.assertEquals(vertexCount, 1);
 
-        List<String> deletedEntities = deleteEntities(hrDeptGuid);
+        List<String> deletedEntities = deleteEntities(hrDeptGuid).getDeletedEntities();
         assertTrue(deletedEntities.contains(hrDeptGuid));
+        assertEntityDeleted(hrDeptGuid);
 
         // Verify Department entity and its contained Person entities were deleted.
-        assertEntityDeleted(hrDeptGuid);
         for (String employeeGuid : employeeGuids) {
+            assertTrue(deletedEntities.contains(employeeGuid));
             assertEntityDeleted(employeeGuid);
         }
 
@@ -341,15 +388,16 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         object = mapOwnerVertex.getProperty(atlasEdgeLabel.getQualifiedMapKey());
         Assert.assertNotNull(object);
 
-        List<String> deletedEntities = deleteEntities(mapOwnerGuid);
+        List<String> deletedEntities = deleteEntities(mapOwnerGuid).getDeletedEntities();
         Assert.assertEquals(deletedEntities.size(), 2);
-        Assert.assertTrue(deletedEntities.containsAll(guids));
+        Assert.assertTrue(deletedEntities.contains(mapOwnerGuid));
+        Assert.assertTrue(deletedEntities.contains(mapValueGuid));
 
         assertEntityDeleted(mapOwnerGuid);
         assertEntityDeleted(mapValueGuid);
     }
 
-    private TypeUtils.Pair<List<String>, List<String>> updatePartial(ITypedReferenceableInstance entity) throws RepositoryException {
+    private AtlasClient.EntityResult updatePartial(ITypedReferenceableInstance entity) throws RepositoryException {
         RequestContext.createContext();
         return repositoryService.updatePartial(entity);
     }
@@ -379,7 +427,9 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         ClassType personType = typeSystem.getDataType(ClassType.class, "Person");
         ITypedReferenceableInstance maxEntity = personType.createInstance(max.getId());
         maxEntity.set("mentor", johnGuid);
-        updatePartial(maxEntity);
+        AtlasClient.EntityResult entityResult = updatePartial(maxEntity);
+        assertEquals(entityResult.getUpdateEntities().size(), 1);
+        assertTrue(entityResult.getUpdateEntities().contains(maxGuid));
 
         // Verify the update was applied correctly - john should now be max's mentor.
         max = repositoryService.getEntityDefinition(maxGuid);
@@ -394,7 +444,9 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
 
         // Update max's mentor reference to jane.
         maxEntity.set("mentor", janeGuid);
-        updatePartial(maxEntity);
+        entityResult = updatePartial(maxEntity);
+        assertEquals(entityResult.getUpdateEntities().size(), 1);
+        assertTrue(entityResult.getUpdateEntities().contains(maxGuid));
 
         // Verify the update was applied correctly - jane should now be max's mentor.
         max = repositoryService.getEntityDefinition(maxGuid);
@@ -411,7 +463,11 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         Id juliusGuid = julius.getId();
         maxEntity = personType.createInstance(max.getId());
         maxEntity.set("manager", juliusGuid);
-        updatePartial(maxEntity);
+        entityResult = updatePartial(maxEntity);
+        //TODO ATLAS-499 should have updated julius' subordinates
+        assertEquals(entityResult.getUpdateEntities().size(), 2);
+        assertTrue(entityResult.getUpdateEntities().contains(maxGuid));
+        assertTrue(entityResult.getUpdateEntities().contains(janeGuid._getId()));
 
         // Verify the update was applied correctly - julius should now be max's manager.
         max = repositoryService.getEntityDefinition(maxGuid);
@@ -456,41 +512,38 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         }
         Assert.assertTrue(subordinateIds.contains(maxGuid));
 
-        List<String> deletedEntities = deleteEntities(maxGuid);
-        Assert.assertTrue(deletedEntities.contains(maxGuid));
-        assertEntityDeleted(maxGuid);
 
-        // Verify that the Department.employees reference to the deleted employee
-        // was disconnected.
-        hrDept = repositoryService.getEntityDefinition(hrDeptGuid);
-        refValue = hrDept.get("employees");
-        Assert.assertTrue(refValue instanceof List);
-        List<Object> employees = (List<Object>)refValue;
-        Assert.assertEquals(employees.size(), 3);
-        for (Object listValue : employees) {
-            Assert.assertTrue(listValue instanceof ITypedReferenceableInstance);
-            ITypedReferenceableInstance employee = (ITypedReferenceableInstance) listValue;
-            Assert.assertNotEquals(employee.getId()._getId(), maxGuid);
-        }
+        AtlasClient.EntityResult entityResult = deleteEntities(maxGuid);
+        ITypedReferenceableInstance john = repositoryService.getEntityDefinition("Manager", "name", "John");
 
-        // Verify that max's Person.mentor unidirectional reference to john was disconnected.
-        ITypedReferenceableInstance john = repositoryService.getEntityDefinition(johnGuid);
-        refValue = john.get("mentor");
-        Assert.assertNull(refValue);
+        assertEquals(entityResult.getDeletedEntities().size(), 1);
+        assertTrue(entityResult.getDeletedEntities().contains(maxGuid));
+        assertEquals(entityResult.getUpdateEntities().size(), 3);
+        assertTrue(entityResult.getUpdateEntities().containsAll(Arrays.asList(jane.getId()._getId(), hrDeptGuid,
+                john.getId()._getId())));
+        assertEntityDeleted(maxGuid);
 
-        assertTestDisconnectBidirectionalReferences(janeGuid);
+        assertMaxForTestDisconnectBidirectionalReferences(nameGuidMap);
 
         // Now delete jane - this should disconnect the manager reference from her
         // subordinate.
-        deletedEntities = deleteEntities(janeGuid);
-        Assert.assertTrue(deletedEntities.contains(janeGuid));
+        entityResult = deleteEntities(janeGuid);
+        assertEquals(entityResult.getDeletedEntities().size(), 1);
+        assertTrue(entityResult.getDeletedEntities().contains(janeGuid));
+        assertEquals(entityResult.getUpdateEntities().size(), 2);
+        assertTrue(entityResult.getUpdateEntities().containsAll(Arrays.asList(hrDeptGuid, john.getId()._getId())));
+
         assertEntityDeleted(janeGuid);
 
-        john = repositoryService.getEntityDefinition(johnGuid);
-        Assert.assertNull(john.get("manager"));
+        john = repositoryService.getEntityDefinition("Person", "name", "John");
+        assertJohnForTestDisconnectBidirectionalReferences(john, janeGuid);
     }
 
-    protected abstract void assertTestDisconnectBidirectionalReferences(String janeGuid) throws Exception;
+    protected abstract void assertJohnForTestDisconnectBidirectionalReferences(ITypedReferenceableInstance john,
+                                                                               String janeGuid) throws Exception;
+
+    protected abstract void assertMaxForTestDisconnectBidirectionalReferences(Map<String, String> nameGuidMap)
+            throws Exception;
 
     /**
      * Verify deleting entity that is the target of a unidirectional class array reference
@@ -503,30 +556,27 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         // Get the guid for one of the table's columns.
         ITypedReferenceableInstance table = repositoryService.getEntityDefinition(TestUtils.TABLE_TYPE, "name", TestUtils.TABLE_NAME);
         String tableGuid = table.getId()._getId();
-        Object refValues = table.get("columns");
-        Assert.assertTrue(refValues instanceof List);
-        List<Object> refList = (List<Object>) refValues;
-        Assert.assertEquals(refList.size(), 5);
-        Assert.assertTrue(refList.get(0) instanceof ITypedReferenceableInstance);
-        ITypedReferenceableInstance column = (ITypedReferenceableInstance) refList.get(0);
-        String columnGuid = column.getId()._getId();
+        List<ITypedReferenceableInstance> columns = (List<ITypedReferenceableInstance>) table.get("columns");
+        Assert.assertEquals(columns.size(), 5);
+        String columnGuid = columns.get(0).getId()._getId();
 
         // Delete the column.
-        List<String> deletedEntities = deleteEntities(columnGuid);
-        Assert.assertTrue(deletedEntities.contains(columnGuid));
+        AtlasClient.EntityResult entityResult = deleteEntities(columnGuid);
+        assertEquals(entityResult.getDeletedEntities().size(), 1);
+        Assert.assertTrue(entityResult.getDeletedEntities().contains(columnGuid));
+        assertEquals(entityResult.getUpdateEntities().size(), 1);
+        Assert.assertTrue(entityResult.getUpdateEntities().contains(tableGuid));
         assertEntityDeleted(columnGuid);
 
         // Verify table.columns reference to the deleted column has been disconnected.
         table = repositoryService.getEntityDefinition(tableGuid);
-        refList = (List<Object>) table.get("columns");
-        Assert.assertEquals(refList.size(), 4);
-        for (Object refValue : refList) {
-            Assert.assertTrue(refValue instanceof ITypedReferenceableInstance);
-            column = (ITypedReferenceableInstance)refValue;
-            Assert.assertFalse(column.getId()._getId().equals(columnGuid));
-        }
+        assertTestDisconnectUnidirectionalArrayReferenceFromClassType(
+                (List<ITypedReferenceableInstance>) table.get("columns"), columnGuid);
     }
 
+    protected abstract void assertTestDisconnectUnidirectionalArrayReferenceFromClassType(
+            List<ITypedReferenceableInstance> columns, String columnGuid);
+
     /**
      * Verify deleting entities that are the target of a unidirectional class array reference
      * from a struct or trait instance.
@@ -559,8 +609,8 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         Referenceable structTargetEntity = new Referenceable("StructTarget");
         Referenceable traitTargetEntity = new Referenceable("TraitTarget");
         Referenceable structContainerEntity = new Referenceable("StructContainer");
-        Referenceable structInstance = new Referenceable("TestStruct");
-        Referenceable nestedStructInstance = new Referenceable("NestedStruct");
+        Struct structInstance = new Struct("TestStruct");
+        Struct nestedStructInstance = new Struct("NestedStruct");
         Referenceable traitInstance = new Referenceable("TestTrait");
         structContainerEntity.set("struct", structInstance);
         structInstance.set("target", ImmutableList.of(structTargetEntity));
@@ -623,19 +673,19 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         Assert.assertEquals(refList.get(0).getId()._getId(), traitTargetGuid);
 
         // Delete the entities that are targets of the struct and trait instances.
-        List<String> deletedEntities = deleteEntities(structTargetGuid, traitTargetGuid);
+        AtlasClient.EntityResult entityResult = deleteEntities(structTargetGuid, traitTargetGuid);
+        Assert.assertEquals(entityResult.getDeletedEntities().size(), 2);
+        Assert.assertTrue(entityResult.getDeletedEntities().containsAll(Arrays.asList(structTargetGuid, traitTargetGuid)));
         assertEntityDeleted(structTargetGuid);
         assertEntityDeleted(traitTargetGuid);
-        Assert.assertEquals(deletedEntities.size(), 2);
-        Assert.assertTrue(deletedEntities.containsAll(Arrays.asList(structTargetGuid, traitTargetGuid)));
 
         assertTestDisconnectUnidirectionalArrayReferenceFromStructAndTraitTypes(structContainerGuid);
 
         // Delete the entity which contains nested structs and has the TestTrait trait.
-        deletedEntities = deleteEntities(structContainerGuid);
+        entityResult = deleteEntities(structContainerGuid);
+        Assert.assertEquals(entityResult.getDeletedEntities().size(), 1);
+        Assert.assertTrue(entityResult.getDeletedEntities().contains(structContainerGuid));
         assertEntityDeleted(structContainerGuid);
-        Assert.assertEquals(deletedEntities.size(), 1);
-        Assert.assertTrue(deletedEntities.contains(structContainerGuid));
 
         // Verify all TestStruct struct vertices were removed.
         assertVerticesDeleted(getVertices(Constants.ENTITY_TYPE_PROPERTY_KEY, "TestStruct"));
@@ -890,13 +940,14 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         return list;
     }
 
-    private Map<String, String> getEmployeeNameGuidMap(ITypedReferenceableInstance hrDept) throws AtlasException {
-
+    private Map<String, String> getEmployeeNameGuidMap(final ITypedReferenceableInstance hrDept) throws AtlasException {
         Object refValue = hrDept.get("employees");
         Assert.assertTrue(refValue instanceof List);
         List<Object> employees = (List<Object>)refValue;
         Assert.assertEquals(employees.size(), 4);
-        Map<String, String> nameGuidMap = new HashMap<String, String>();
+        Map<String, String> nameGuidMap = new HashMap<String, String>() {{
+            put("hr", hrDept.getId()._getId());
+        }};
 
         for (Object listValue : employees) {
             Assert.assertTrue(listValue instanceof ITypedReferenceableInstance);

repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedRepositoryHardDeleteTest.java

@@ -21,8 +21,10 @@ package org.apache.atlas.repository.graph;
 import com.tinkerpop.blueprints.Vertex;
 
 import org.apache.atlas.AtlasClient;
+import org.apache.atlas.AtlasException;
 import org.apache.atlas.TestUtils;
 import org.apache.atlas.repository.Constants;
+import org.apache.atlas.typesystem.IReferenceableInstance;
 import org.apache.atlas.typesystem.IStruct;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.ITypedStruct;
@@ -32,8 +34,11 @@ import org.apache.atlas.typesystem.types.TypeSystem;
 import org.testng.Assert;
 
 import java.util.List;
+import java.util.Map;
 
+import static org.apache.atlas.TestUtils.COLUMNS_ATTR_NAME;
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertNull;
 import static org.testng.Assert.fail;
 import static org.testng.AssertJUnit.assertNotNull;
@@ -45,7 +50,24 @@ public class GraphBackedRepositoryHardDeleteTest extends GraphBackedMetadataRepo
     }
 
     @Override
-    protected void assertTestDeleteReference(ITypedReferenceableInstance processInstance) throws Exception {
+    protected void assertTestDeleteEntityWithTraits(String guid) {
+        //entity is deleted. So, no assertions
+    }
+
+    @Override
+    protected void assertTableForTestDeleteReference(String tableId) {
+        //entity is deleted. So, no assertions
+    }
+
+    @Override
+    protected void assertColumnForTestDeleteReference(ITypedReferenceableInstance tableInstance) throws AtlasException {
+        List<ITypedReferenceableInstance> columns =
+                (List<ITypedReferenceableInstance>) tableInstance.get(COLUMNS_ATTR_NAME);
+        assertNull(columns);
+    }
+
+    @Override
+    protected void assertProcessForTestDeleteReference(ITypedReferenceableInstance processInstance) throws Exception {
         //assert that outputs is empty
         ITypedReferenceableInstance newProcess =
                 repositoryService.getEntityDefinition(processInstance.getId()._getId());
@@ -63,6 +85,11 @@ public class GraphBackedRepositoryHardDeleteTest extends GraphBackedMetadataRepo
     }
 
     @Override
+    protected void assertDeletedColumn(ITypedReferenceableInstance tableInstance) throws AtlasException {
+        assertEquals(((List<IReferenceableInstance>) tableInstance.get(COLUMNS_ATTR_NAME)).size(), 2);
+    }
+
+    @Override
     protected void assertTestDeleteEntities(ITypedReferenceableInstance tableInstance) {
         int vertexCount = getVertices(Constants.ENTITY_TYPE_PROPERTY_KEY, TestUtils.TABLE_TYPE).size();
         assertEquals(vertexCount, 0);
@@ -85,12 +112,42 @@ public class GraphBackedRepositoryHardDeleteTest extends GraphBackedMetadataRepo
     }
 
     @Override
-    protected void assertTestDisconnectBidirectionalReferences(String janeGuid) throws Exception {
+    protected void assertJohnForTestDisconnectBidirectionalReferences(ITypedReferenceableInstance john,
+                                                                      String janeGuid) throws Exception {
+        assertNull(john.get("manager"));
+    }
+
+    @Override
+    protected void assertMaxForTestDisconnectBidirectionalReferences(Map<String, String> nameGuidMap)
+            throws Exception {
+        // Verify that the Department.employees reference to the deleted employee
+        // was disconnected.
+        ITypedReferenceableInstance hrDept = repositoryService.getEntityDefinition(nameGuidMap.get("hr"));
+        List<ITypedReferenceableInstance> employees = (List<ITypedReferenceableInstance>) hrDept.get("employees");
+        Assert.assertEquals(employees.size(), 3);
+        String maxGuid = nameGuidMap.get("Max");
+        for (ITypedReferenceableInstance employee : employees) {
+            Assert.assertNotEquals(employee.getId()._getId(), maxGuid);
+        }
+
         // Verify that the Manager.subordinates reference to the deleted employee
         // Max was disconnected.
-        ITypedReferenceableInstance jane = repositoryService.getEntityDefinition(janeGuid);
+        ITypedReferenceableInstance jane = repositoryService.getEntityDefinition(nameGuidMap.get("Jane"));
         List<ITypedReferenceableInstance> subordinates = (List<ITypedReferenceableInstance>) jane.get("subordinates");
         assertEquals(subordinates.size(), 1);
+
+        // Verify that max's Person.mentor unidirectional reference to john was disconnected.
+        ITypedReferenceableInstance john = repositoryService.getEntityDefinition(nameGuidMap.get("John"));
+        assertNull(john.get("mentor"));
+    }
+
+    @Override
+    protected void assertTestDisconnectUnidirectionalArrayReferenceFromClassType(
+            List<ITypedReferenceableInstance> columns, String columnGuid) {
+        assertEquals(columns.size(), 4);
+        for (ITypedReferenceableInstance column : columns) {
+            assertFalse(column.getId()._getId().equals(columnGuid));
+        }
     }
 
     @Override

repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedRepositorySoftDeleteTest.java

@@ -19,10 +19,11 @@
 package org.apache.atlas.repository.graph;
 
 import com.tinkerpop.blueprints.Vertex;
-
 import org.apache.atlas.AtlasClient;
+import org.apache.atlas.AtlasException;
 import org.apache.atlas.TestUtils;
 import org.apache.atlas.repository.Constants;
+import org.apache.atlas.typesystem.IReferenceableInstance;
 import org.apache.atlas.typesystem.IStruct;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.ITypedStruct;
@@ -33,8 +34,12 @@ import org.testng.Assert;
 import java.util.List;
 import java.util.Map;
 
+import static org.apache.atlas.TestUtils.COLUMNS_ATTR_NAME;
+import static org.apache.atlas.TestUtils.NAME;
+import static org.apache.atlas.TestUtils.PII;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
 
 public class GraphBackedRepositorySoftDeleteTest extends GraphBackedMetadataRepositoryDeleteTestBase {
     @Override
@@ -43,7 +48,38 @@ public class GraphBackedRepositorySoftDeleteTest extends GraphBackedMetadataRepo
     }
 
     @Override
-    protected void assertTestDeleteReference(ITypedReferenceableInstance expected) throws Exception {
+    protected void assertTestDeleteEntityWithTraits(String guid) throws Exception {
+        ITypedReferenceableInstance instance = repositoryService.getEntityDefinition(guid);
+        assertTrue(instance.getTraits().contains(PII));
+    }
+
+    @Override
+    protected void assertTableForTestDeleteReference(String tableId) throws Exception  {
+        ITypedReferenceableInstance table = repositoryService.getEntityDefinition(tableId);
+        assertNotNull(table.get(NAME));
+        assertNotNull(table.get("description"));
+        assertNotNull(table.get("type"));
+        assertNotNull(table.get("tableType"));
+        assertNotNull(table.get("created"));
+
+        Id dbId = (Id) table.get("database");
+        assertNotNull(dbId);
+
+        ITypedReferenceableInstance db = repositoryService.getEntityDefinition(dbId.getId()._getId());
+        assertNotNull(db);
+        assertEquals(db.getId().getState(), Id.EntityState.ACTIVE);
+    }
+
+    @Override
+    protected void assertColumnForTestDeleteReference(ITypedReferenceableInstance tableInstance) throws AtlasException {
+        List<ITypedReferenceableInstance> columns =
+                (List<ITypedReferenceableInstance>) tableInstance.get(COLUMNS_ATTR_NAME);
+        assertEquals(columns.size(), 1);
+        assertEquals(columns.get(0).getId().getState(), Id.EntityState.DELETED);
+    }
+
+    @Override
+    protected void assertProcessForTestDeleteReference(ITypedReferenceableInstance expected) throws Exception {
         ITypedReferenceableInstance process = repositoryService.getEntityDefinition(expected.getId()._getId());
         List<ITypedReferenceableInstance> outputs =
                 (List<ITypedReferenceableInstance>) process.get(AtlasClient.PROCESS_ATTRIBUTE_OUTPUTS);
@@ -59,6 +95,13 @@ public class GraphBackedRepositorySoftDeleteTest extends GraphBackedMetadataRepo
     }
 
     @Override
+    protected void assertDeletedColumn(ITypedReferenceableInstance tableInstance) throws AtlasException {
+        List<IReferenceableInstance> columns = (List<IReferenceableInstance>) tableInstance.get(COLUMNS_ATTR_NAME);
+        assertEquals(columns.size(), 3);
+        assertEquals(columns.get(0).getId().getState(), Id.EntityState.DELETED);
+    }
+
+    @Override
     protected void assertTestDeleteEntities(ITypedReferenceableInstance expected) throws Exception {
         //Assert that the deleted table can be fully constructed back
         ITypedReferenceableInstance table = repositoryService.getEntityDefinition(expected.getId()._getId());
@@ -67,6 +110,7 @@ public class GraphBackedRepositorySoftDeleteTest extends GraphBackedMetadataRepo
         List<ITypedReferenceableInstance> expectedColumns =
                 (List<ITypedReferenceableInstance>) table.get(TestUtils.COLUMNS_ATTR_NAME);
         assertEquals(columns.size(), expectedColumns.size());
+        assertNotNull(table.get("database"));
     }
 
     @Override
@@ -85,11 +129,57 @@ public class GraphBackedRepositorySoftDeleteTest extends GraphBackedMetadataRepo
     }
 
     @Override
-    protected void assertTestDisconnectBidirectionalReferences(String janeGuid) throws Exception {
+    protected void assertJohnForTestDisconnectBidirectionalReferences(ITypedReferenceableInstance john, String janeGuid)
+            throws Exception {
+        Id mgr = (Id) john.get("manager");
+        assertNotNull(mgr);
+        assertEquals(mgr._getId(), janeGuid);
+        assertEquals(mgr.getState(), Id.EntityState.DELETED);
+    }
+
+    @Override
+    protected void assertMaxForTestDisconnectBidirectionalReferences(Map<String, String> nameGuidMap) throws Exception {
+        // Verify that the Department.employees reference to the deleted employee
+        // was disconnected.
+        ITypedReferenceableInstance hrDept = repositoryService.getEntityDefinition(nameGuidMap.get("hr"));
+        List<ITypedReferenceableInstance> employees = (List<ITypedReferenceableInstance>) hrDept.get("employees");
+        Assert.assertEquals(employees.size(), 4);
+        String maxGuid = nameGuidMap.get("Max");
+        for (ITypedReferenceableInstance employee : employees) {
+            if (employee.getId()._getId().equals(maxGuid)) {
+                assertEquals(employee.getId().getState(), Id.EntityState.DELETED);
+            }
+        }
+
         // Verify that the Manager.subordinates still references deleted employee
-        ITypedReferenceableInstance jane = repositoryService.getEntityDefinition(janeGuid);
+        ITypedReferenceableInstance jane = repositoryService.getEntityDefinition(nameGuidMap.get("Jane"));
         List<ITypedReferenceableInstance> subordinates = (List<ITypedReferenceableInstance>) jane.get("subordinates");
         assertEquals(subordinates.size(), 2);
+        for (ITypedReferenceableInstance subordinate : subordinates) {
+            if (subordinate.getId()._getId().equals(maxGuid)) {
+                assertEquals(subordinate.getId().getState(), Id.EntityState.DELETED);
+            }
+        }
+
+        // Verify that max's Person.mentor unidirectional reference to john was disconnected.
+        ITypedReferenceableInstance john = repositoryService.getEntityDefinition(nameGuidMap.get("John"));
+        Id mentor = (Id) john.get("mentor");
+        assertEquals(mentor._getId(), maxGuid);
+        assertEquals(mentor.getState(), Id.EntityState.DELETED);
+    }
+
+    @Override
+    protected void assertTestDisconnectUnidirectionalArrayReferenceFromClassType(
+            List<ITypedReferenceableInstance> columns, String columnGuid) {
+        Assert.assertEquals(columns.size(), 5);
+        for (ITypedReferenceableInstance column : columns) {
+            if (column.getId()._getId().equals(columnGuid)) {
+                assertEquals(column.getId().getState(), Id.EntityState.DELETED);
+            } else {
+                assertEquals(column.getId().getState(), Id.EntityState.ACTIVE);
+            }
+        }
+
     }
 
     @Override
@@ -122,7 +212,6 @@ public class GraphBackedRepositorySoftDeleteTest extends GraphBackedMetadataRepo
 
     @Override
     protected void assertTestDeleteTargetOfMultiplicityRequiredReference() throws Exception {
-
         // No-op - it's ok that no exception was thrown if soft deletes are enabled.
     }
 }

repository/src/test/java/org/apache/atlas/service/DefaultMetadataServiceTest.java

@@ -73,6 +73,7 @@ import java.util.Map;
 
 import static org.apache.atlas.TestUtils.COLUMNS_ATTR_NAME;
 import static org.apache.atlas.TestUtils.COLUMN_TYPE;
+import static org.apache.atlas.TestUtils.PII;
 import static org.apache.atlas.TestUtils.TABLE_TYPE;
 import static org.apache.atlas.TestUtils.createColumnEntity;
 import static org.apache.atlas.TestUtils.createDBEntity;
@@ -80,6 +81,7 @@ import static org.apache.atlas.TestUtils.createTableEntity;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
 
 @Guice(modules = RepositoryMetadataModule.class)
@@ -154,23 +156,21 @@ public class DefaultMetadataServiceTest {
         String entityjson = InstanceSerialization.toJson(entity, true);
         JSONArray entitiesJson = new JSONArray();
         entitiesJson.put(entityjson);
-        String response = metadataService.createEntities(entitiesJson.toString());
-        JSONArray guids = new JSONArray(response);
-        if (guids != null && guids.length() > 0) {
-            return guids.getString(guids.length() - 1);
+        List<String> guids = metadataService.createEntities(entitiesJson.toString());
+        if (guids != null && guids.size() > 0) {
+            return guids.get(guids.size() - 1);
         }
         return null;
     }
 
-    private String updateInstance(Referenceable entity) throws Exception {
+    private AtlasClient.EntityResult updateInstance(Referenceable entity) throws Exception {
         RequestContext.createContext();
         ParamChecker.notNull(entity, "Entity");
         ParamChecker.notNull(entity.getId(), "Entity");
         String entityjson = InstanceSerialization.toJson(entity, true);
         JSONArray entitiesJson = new JSONArray();
         entitiesJson.put(entityjson);
-        String response = metadataService.updateEntities(entitiesJson.toString());
-        return new JSONArray(response).getString(0);
+        return metadataService.updateEntities(entitiesJson.toString());
     }
 
     @Test(expectedExceptions = TypeNotFoundException.class)
@@ -201,6 +201,32 @@ public class DefaultMetadataServiceTest {
     }
 
     @Test
+    public void testAddDeleteTrait() throws Exception {
+        Referenceable entity = createDBEntity();
+        String id = createInstance(entity);
+
+        //add trait
+        Struct tag = new Struct(TestUtils.PII);
+        metadataService.addTrait(id, InstanceSerialization.toJson(tag, true));
+
+        List<String> traits = metadataService.getTraitNames(id);
+        assertEquals(traits.size(), 1);
+        assertEquals(traits.get(0), PII);
+
+        //delete trait
+        metadataService.deleteTrait(id, PII);
+        traits = metadataService.getTraitNames(id);
+        assertEquals(traits.size(), 0);
+
+        //add trait again
+        metadataService.addTrait(id, InstanceSerialization.toJson(tag, true));
+
+        traits = metadataService.getTraitNames(id);
+        assertEquals(traits.size(), 1);
+        assertEquals(traits.get(0), PII);
+    }
+
+    @Test
     public void testEntityAudit() throws Exception {
         //create entity
         Referenceable entity = createDBEntity();
@@ -221,7 +247,7 @@ public class DefaultMetadataServiceTest {
         assertAuditEvents(id, EntityAuditEvent.EntityAuditAction.ENTITY_DELETE);
     }
 
-    private List<String> deleteEntities(String... guids) throws AtlasException {
+    private AtlasClient.EntityResult deleteEntities(String... guids) throws AtlasException {
         RequestContext.createContext();
         return metadataService.deleteEntities(Arrays.asList(guids));
     }
@@ -350,7 +376,7 @@ public class DefaultMetadataServiceTest {
         Assert.assertTrue(partsMap.get("part2").equalsContents(((Map<String, Struct>)tableDefinition.get("partitionsMap")).get("part2")));
 
         //Test map pointing to a class
-        final Map<String, Struct> columnsMap = new HashMap<>();
+        final Map<String, Referenceable> columnsMap = new HashMap<>();
         Referenceable col0Type = new Referenceable(TestUtils.COLUMN_TYPE,
             new HashMap<String, Object>() {{
                 put(NAME, "test1");
@@ -393,17 +419,33 @@ public class DefaultMetadataServiceTest {
         verifyMapUpdates(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME), null, TestUtils.COLUMNS_MAP);
     }
 
-    private void verifyMapUpdates(String typeName, String uniqAttrName, String uniqAttrValue, Map<String, Struct> expectedMap, String mapAttrName) throws AtlasException {
+    private void verifyMapUpdates(String typeName, String uniqAttrName, String uniqAttrValue,
+                                  Map<String, Referenceable> expectedMap, String mapAttrName) throws AtlasException {
         String json =
             metadataService.getEntityDefinition(typeName, uniqAttrName, uniqAttrValue);
         Referenceable tableDefinition = InstanceSerialization.fromJsonReferenceable(json, true);
+        Map<String, Referenceable> actualMap = (Map<String, Referenceable>) tableDefinition.get(mapAttrName);
 
-        if(expectedMap == null) {
-            Assert.assertNull(tableDefinition.get(TestUtils.COLUMNS_MAP));
+        if (expectedMap == null && actualMap != null) {
+            //all are marked as deleted in case of soft delete
+            for (String key : actualMap.keySet()) {
+                assertEquals(actualMap.get(key).getId().state, Id.EntityState.DELETED);
+            }
+        } else if(expectedMap == null) {
+            //hard delete case
+            assertNull(actualMap);
         } else {
-            Assert.assertEquals(((Map<String, Referenceable>)tableDefinition.get(mapAttrName)).size(), expectedMap.size());
+            assertTrue(actualMap.size() >= expectedMap.size());
+
             for (String key : expectedMap.keySet()) {
-                Assert.assertTrue(((Map<String, Referenceable>) tableDefinition.get(mapAttrName)).get(key).equalsContents(expectedMap.get(key)));
+                assertTrue(actualMap.get(key).equalsContents(expectedMap.get(key)));
+            }
+
+            //rest of the keys are marked as deleted
+            List<String> extraKeys = new ArrayList<>(actualMap.keySet());
+            extraKeys.removeAll(expectedMap.keySet());
+            for (String key : extraKeys) {
+                assertEquals(actualMap.get(key).getId().getState(), Id.EntityState.DELETED);
             }
         }
     }
@@ -438,14 +480,13 @@ public class DefaultMetadataServiceTest {
         Assert.assertEquals(actualColumns, updatedColNameList);
     }
 
-    private void updateEntityPartial(String guid, Referenceable entity) throws AtlasException {
+    private AtlasClient.EntityResult updateEntityPartial(String guid, Referenceable entity) throws AtlasException {
         RequestContext.createContext();
-        metadataService.updateEntityPartialByGuid(guid, entity);
+        return metadataService.updateEntityPartialByGuid(guid, entity);
     }
 
     @Test
     public void testUpdateEntityArrayOfClass() throws Exception {
-        
         //test array of class with id
         final List<Referenceable> columns = new ArrayList<>();
         Map<String, Object> values = new HashMap<>();
@@ -456,63 +497,67 @@ public class DefaultMetadataServiceTest {
         Referenceable tableUpdated = new Referenceable(TestUtils.TABLE_TYPE, new HashMap<String, Object>() {{
             put(COLUMNS_ATTR_NAME, columns);
         }});
-        updateEntityPartial(tableId._getId(), tableUpdated);
 
+        AtlasClient.EntityResult entityResult = updateEntityPartial(tableId._getId(), tableUpdated);
+        assertEquals(entityResult.getCreatedEntities().size(), 1);  //col1 created
+        assertEquals(entityResult.getUpdateEntities().size(), 1);  //table updated
+        assertEquals(entityResult.getUpdateEntities().get(0), tableId._getId());
         verifyArrayUpdates(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME), columns, COLUMNS_ATTR_NAME);
 
-        //Partial update. Add col5 But also update col1
+        //Partial update. Add col2 But also update col1
         Map<String, Object> valuesCol5 = new HashMap<>();
-        valuesCol5.put(NAME, "col5");
+        valuesCol5.put(NAME, "col2");
         valuesCol5.put("type", "type");
         Referenceable col2 = new Referenceable(TestUtils.COLUMN_TYPE, valuesCol5);
         //update col1
         col1.set("type", "type1");
-
-        //add col5
-        final List<Referenceable> updateColumns = Arrays.asList(col1, col2);
+        columns.add(col2);
 
         tableUpdated = new Referenceable(TestUtils.TABLE_TYPE, new HashMap<String, Object>() {{
-            put(COLUMNS_ATTR_NAME, updateColumns);
+            put(COLUMNS_ATTR_NAME, columns);
         }});
-        updateEntityPartial(tableId._getId(), tableUpdated);
+        entityResult = updateEntityPartial(tableId._getId(), tableUpdated);
+        assertEquals(entityResult.getCreatedEntities().size(), 1);  //col2 created
+        assertEquals(entityResult.getUpdateEntities().size(), 2);  //table, col1 updated
 
-        verifyArrayUpdates(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME), updateColumns, COLUMNS_ATTR_NAME);
+        verifyArrayUpdates(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME), columns, COLUMNS_ATTR_NAME);
 
-        //Complete update. Add  array elements - col3,4
+        //Complete update. Add  array elements - col3,col4
         Map<String, Object> values1 = new HashMap<>();
         values1.put(NAME, "col3");
         values1.put("type", "type");
-        Referenceable ref1 = new Referenceable(TestUtils.COLUMN_TYPE, values1);
-        columns.add(ref1);
+        Referenceable col3 = new Referenceable(TestUtils.COLUMN_TYPE, values1);
+        columns.add(col3);
 
         Map<String, Object> values2 = new HashMap<>();
         values2.put(NAME, "col4");
         values2.put("type", "type");
-        Referenceable ref2 = new Referenceable(TestUtils.COLUMN_TYPE, values2);
-        columns.add(ref2);
+        Referenceable col4 = new Referenceable(TestUtils.COLUMN_TYPE, values2);
+        columns.add(col4);
 
         table.set(COLUMNS_ATTR_NAME, columns);
-        updateInstance(table);
+        entityResult = updateInstance(table);
+        assertEquals(entityResult.getCreatedEntities().size(), 2);  //col3, col4 created
 
         verifyArrayUpdates(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME), columns, COLUMNS_ATTR_NAME);
 
         //Swap elements
         columns.clear();
-        columns.add(ref2);
-        columns.add(ref1);
+        columns.add(col4);
+        columns.add(col3);
 
         table.set(COLUMNS_ATTR_NAME, columns);
-        updateInstance(table);
-
+        entityResult = updateInstance(table);
+        assertEquals(entityResult.getDeletedEntities().size(), 2);  //col1, col2 are deleted
         verifyArrayUpdates(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME), columns, COLUMNS_ATTR_NAME);
 
         //drop a single column
         columns.clear();
-        columns.add(ref1);
+        columns.add(col3);
 
         table.set(COLUMNS_ATTR_NAME, columns);
-        updateInstance(table);
-
+        entityResult = updateInstance(table);
+        assertEquals(entityResult.getDeletedEntities().size(), 1);  //col4 deleted
         verifyArrayUpdates(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME), columns, COLUMNS_ATTR_NAME);
 
         //Remove a class reference/Id and insert another reference
@@ -520,34 +565,46 @@ public class DefaultMetadataServiceTest {
         values.clear();
         columns.clear();
 
-        values.put(NAME, "col2");
+        values.put(NAME, "col5");
         values.put("type", "type");
-        col1 = new Referenceable(TestUtils.COLUMN_TYPE, values);
-        columns.add(col1);
+        Referenceable col5 = new Referenceable(TestUtils.COLUMN_TYPE, values);
+        columns.add(col5);
         table.set(COLUMNS_ATTR_NAME, columns);
-        updateInstance(table);
+        entityResult = updateInstance(table);
+        assertEquals(entityResult.getCreatedEntities().size(), 1);  //col5 created
+        assertEquals(entityResult.getDeletedEntities().size(), 1);  //col3 deleted
 
         verifyArrayUpdates(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME), columns, COLUMNS_ATTR_NAME);
 
         //Update array column to null
         table.setNull(COLUMNS_ATTR_NAME);
-        String newtableId = updateInstance(table);
-        Assert.assertEquals(newtableId, tableId._getId());
-
+        entityResult = updateInstance(table);
+        assertEquals(entityResult.getDeletedEntities().size(), 1);
         verifyArrayUpdates(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME), null, COLUMNS_ATTR_NAME);
     }
 
-    private void verifyArrayUpdates(String typeName, String uniqAttrName, String uniqAttrValue, List<? extends Struct> expectedArray, String arrAttrName) throws AtlasException {
-        String json =
-            metadataService.getEntityDefinition(typeName, uniqAttrName, uniqAttrValue);
+    private void verifyArrayUpdates(String typeName, String uniqAttrName, String uniqAttrValue,
+                                    List<Referenceable> expectedArray, String arrAttrName) throws AtlasException {
+        String json = metadataService.getEntityDefinition(typeName, uniqAttrName, uniqAttrValue);
         Referenceable entityDefinition = InstanceSerialization.fromJsonReferenceable(json, true);
-
-        if (expectedArray == null) {
-            Assert.assertNull(entityDefinition.get(arrAttrName));
+        List<Referenceable> actualArray = (List<Referenceable>) entityDefinition.get(arrAttrName);
+        if (expectedArray == null && actualArray != null) {
+            //all are marked as deleted in case of soft delete
+            for (int index = 0; index < actualArray.size(); index++) {
+                assertEquals(actualArray.get(index).getId().state, Id.EntityState.DELETED);
+            }
+        } else if(expectedArray == null) {
+            //hard delete case
+            assertNull(actualArray);
         } else {
-            Assert.assertEquals(((List<Referenceable>)entityDefinition.get(arrAttrName)).size(), expectedArray.size());
-            for (int index = 0; index < expectedArray.size(); index++) {
-                Assert.assertTrue(((List<Referenceable>) entityDefinition.get(arrAttrName)).get(index).equalsContents(expectedArray.get(index)));
+            int index;
+            for (index = 0; index < expectedArray.size(); index++) {
+                Assert.assertTrue(actualArray.get(index).equalsContents(expectedArray.get(index)));
+            }
+
+            //Rest of the entities in the list are marked as deleted
+            for (; index < actualArray.size(); index++) {
+                assertEquals(actualArray.get(index).getId().state, Id.EntityState.DELETED);
             }
         }
     }
@@ -560,7 +617,7 @@ public class DefaultMetadataServiceTest {
         serdeInstance.set("description", "testDesc");
         table.set("serde1", serdeInstance);
 
-        String newtableId = updateInstance(table);
+        String newtableId = updateInstance(table).getUpdateEntities().get(0);
         Assert.assertEquals(newtableId, tableId._getId());
 
         String tableDefinitionJson =
@@ -664,7 +721,7 @@ public class DefaultMetadataServiceTest {
         List<Struct> partitions = new ArrayList<Struct>(){{ add(partition1); add(partition2); }};
         table.set("partitions", partitions);
 
-        String newtableId = updateInstance(table);
+        String newtableId = updateInstance(table).getUpdateEntities().get(0);
         Assert.assertEquals(newtableId, tableId._getId());
 
         String tableDefinitionJson =
@@ -673,15 +730,14 @@ public class DefaultMetadataServiceTest {
 
         Assert.assertNotNull(tableDefinition.get("partitions"));
         List<Struct> partitionsActual = (List<Struct>) tableDefinition.get("partitions");
-        Assert.assertEquals(partitionsActual.size(), 2);
-        Assert.assertTrue(partitions.get(0).equalsContents(partitionsActual.get(0)));
+        assertPartitions(partitionsActual, partitions);
 
         //add a new element to array of struct
         final Struct partition3 = new Struct(TestUtils.PARTITION_STRUCT_TYPE);
         partition3.set(NAME, "part3");
         partitions.add(partition3);
         table.set("partitions", partitions);
-        newtableId = updateInstance(table);
+        newtableId = updateInstance(table).getUpdateEntities().get(0);
         Assert.assertEquals(newtableId, tableId._getId());
 
         tableDefinitionJson =
@@ -690,13 +746,12 @@ public class DefaultMetadataServiceTest {
 
         Assert.assertNotNull(tableDefinition.get("partitions"));
         partitionsActual = (List<Struct>) tableDefinition.get("partitions");
-        Assert.assertEquals(partitionsActual.size(), 3);
-        Assert.assertTrue(partitions.get(2).equalsContents(partitionsActual.get(2)));
+        assertPartitions(partitionsActual, partitions);
 
         //remove one of the struct values
         partitions.remove(1);
         table.set("partitions", partitions);
-        newtableId = updateInstance(table);
+        newtableId = updateInstance(table).getUpdateEntities().get(0);
         Assert.assertEquals(newtableId, tableId._getId());
 
         tableDefinitionJson =
@@ -705,13 +760,11 @@ public class DefaultMetadataServiceTest {
 
         Assert.assertNotNull(tableDefinition.get("partitions"));
         partitionsActual = (List<Struct>) tableDefinition.get("partitions");
-        Assert.assertEquals(partitionsActual.size(), 2);
-        Assert.assertTrue(partitions.get(0).equalsContents(partitionsActual.get(0)));
-        Assert.assertTrue(partitions.get(1).equalsContents(partitionsActual.get(1)));
+        assertPartitions(partitionsActual, partitions);
 
         //Update struct value within array of struct
-        partition1.set(NAME, "part4");
-        newtableId = updateInstance(table);
+        partitions.get(0).set(NAME, "part4");
+        newtableId = updateInstance(table).getUpdateEntities().get(0);
         Assert.assertEquals(newtableId, tableId._getId());
 
         tableDefinitionJson =
@@ -720,15 +773,14 @@ public class DefaultMetadataServiceTest {
 
         Assert.assertNotNull(tableDefinition.get("partitions"));
         partitionsActual = (List<Struct>) tableDefinition.get("partitions");
-        Assert.assertEquals(partitionsActual.size(), 2);
-        Assert.assertTrue(partitions.get(0).equalsContents(partitionsActual.get(0)));
+        assertPartitions(partitionsActual, partitions);
 
         //add a repeated element to array of struct
         final Struct partition4 = new Struct(TestUtils.PARTITION_STRUCT_TYPE);
         partition4.set(NAME, "part4");
         partitions.add(partition4);
         table.set("partitions", partitions);
-        newtableId = updateInstance(table);
+        newtableId = updateInstance(table).getUpdateEntities().get(0);
         Assert.assertEquals(newtableId, tableId._getId());
 
         tableDefinitionJson =
@@ -737,15 +789,12 @@ public class DefaultMetadataServiceTest {
 
         Assert.assertNotNull(tableDefinition.get("partitions"));
         partitionsActual = (List<Struct>) tableDefinition.get("partitions");
-        Assert.assertEquals(partitionsActual.size(), 3);
-        Assert.assertEquals(partitionsActual.get(2).get(NAME), "part4");
-        Assert.assertEquals(partitionsActual.get(0).get(NAME), "part4");
-        Assert.assertTrue(partitions.get(2).equalsContents(partitionsActual.get(2)));
+        assertPartitions(partitionsActual, partitions);
 
 
         // Remove all elements. Should set array attribute to null
         partitions.clear();
-        newtableId = updateInstance(table);
+        newtableId = updateInstance(table).getUpdateEntities().get(0);
         Assert.assertEquals(newtableId, tableId._getId());
 
         tableDefinitionJson =
@@ -755,6 +804,13 @@ public class DefaultMetadataServiceTest {
         Assert.assertNull(tableDefinition.get("partitions"));
     }
 
+    private void assertPartitions(List<Struct> partitionsActual, List<Struct> partitions) {
+        assertEquals(partitionsActual.size(), partitions.size());
+        for (int index = 0; index < partitions.size(); index++) {
+            assertTrue(partitionsActual.get(index).equalsContents(partitions.get(index)));
+        }
+    }
+
 
     @Test(expectedExceptions = ValueConversionException.class)
     public void testUpdateRequiredAttrToNull() throws Exception {
@@ -772,7 +828,6 @@ public class DefaultMetadataServiceTest {
 
     @Test
     public void testUpdateOptionalAttrToNull() throws Exception {
-
         String tableDefinitionJson =
             metadataService.getEntityDefinition(TestUtils.TABLE_TYPE, NAME, (String) table.get(NAME));
         Referenceable tableDefinition = InstanceSerialization.fromJsonReferenceable(tableDefinitionJson, true);
@@ -782,7 +837,7 @@ public class DefaultMetadataServiceTest {
         //Update optional attribute
             table.setNull("created");
 
-        String newtableId = updateInstance(table);
+        String newtableId = updateInstance(table).getUpdateEntities().get(0);
         Assert.assertEquals(newtableId, tableId._getId());
 
         tableDefinitionJson =
@@ -853,35 +908,39 @@ public class DefaultMetadataServiceTest {
 
         // Register an EntityChangeListener to verify the notification mechanism
         // is working for deleteEntities().
-        DeleteEntitiesChangeListener listener = new DeleteEntitiesChangeListener();
+        EntitiesChangeListener listener = new EntitiesChangeListener();
         metadataService.registerListener(listener);
 
+        //Delete one column
+        String columnId = table1Columns.get(0).getId()._getId();
+        AtlasClient.EntityResult entityResult = deleteEntities(columnId);
+        //column is deleted and table is updated
+        assertEquals(entityResult.getDeletedEntities().get(0), columnId);
+        assertEquals(entityResult.getUpdateEntities().get(0), table1Entity.getId()._getId());
+
+        //verify listener was called for updates and deletes
+        assertEquals(entityResult.getDeletedEntities(), listener.getDeletedEntities());
+        assertEquals(entityResult.getUpdateEntities(), listener.getUpdatedEntities());
+
         // Delete the table entities.  The deletion should cascade
         // to their composite columns.
-        List<String> deletedGuids = deleteEntities(table1Entity.getId()._getId());
+        entityResult = deleteEntities(table1Entity.getId()._getId());
 
         // Verify that deleteEntities() response has guids for tables and their composite columns.
-        Assert.assertTrue(deletedGuids.contains(table1Entity.getId()._getId()));
-        for (IReferenceableInstance column : table1Columns) {
-            Assert.assertTrue(deletedGuids.contains(column.getId()._getId()));
-        }
+        Assert.assertTrue(entityResult.getDeletedEntities().contains(table1Entity.getId()._getId()));
+        Assert.assertTrue(entityResult.getDeletedEntities().contains(table1Columns.get(1).getId()._getId()));
+        Assert.assertTrue(entityResult.getDeletedEntities().contains(table1Columns.get(2).getId()._getId()));
 
         // Verify that tables and their composite columns have been deleted from the repository.
         assertEntityDeleted(TABLE_TYPE, NAME, table1Entity.get(NAME));
-        assertEntityDeleted(COLUMN_TYPE, NAME, col1.get(NAME));
         assertEntityDeleted(COLUMN_TYPE, NAME, col2.get(NAME));
         assertEntityDeleted(COLUMN_TYPE, NAME, col3.get(NAME));
 
         // Verify that the listener was notified about the deleted entities.
-        Collection<ITypedReferenceableInstance> deletedEntitiesFromListener = listener.getDeletedEntities();
+        List<String> deletedEntitiesFromListener = listener.getDeletedEntities();
         Assert.assertNotNull(deletedEntitiesFromListener);
-        Assert.assertEquals(deletedEntitiesFromListener.size(), deletedGuids.size());
-        List<String> deletedGuidsFromListener = new ArrayList<>(deletedGuids.size());
-        for (ITypedReferenceableInstance deletedEntity : deletedEntitiesFromListener) {
-            deletedGuidsFromListener.add(deletedEntity.getId()._getId());
-        }
-        Assert.assertEquals(deletedGuidsFromListener.size(), deletedGuids.size());
-        Assert.assertTrue(deletedGuidsFromListener.containsAll(deletedGuids));
+        Assert.assertEquals(deletedEntitiesFromListener.size(), entityResult.getDeletedEntities().size());
+        Assert.assertTrue(deletedEntitiesFromListener.containsAll(entityResult.getDeletedEntities()));
     }
 
     private void assertEntityDeleted(String typeName, String attributeName, Object attributeValue)
@@ -915,12 +974,13 @@ public class DefaultMetadataServiceTest {
 
         // Register an EntityChangeListener to verify the notification mechanism
         // is working for deleteEntityByUniqueAttribute().
-        DeleteEntitiesChangeListener listener = new DeleteEntitiesChangeListener();
+        EntitiesChangeListener listener = new EntitiesChangeListener();
         metadataService.registerListener(listener);
 
         // Delete the table entities.  The deletion should cascade
         // to their composite columns.
-        List<String> deletedGuids = metadataService.deleteEntityByUniqueAttribute(TestUtils.TABLE_TYPE, NAME, (String) table1Entity.get(NAME));
+        List<String> deletedGuids = metadataService.deleteEntityByUniqueAttribute(TestUtils.TABLE_TYPE, NAME,
+                (String) table1Entity.get(NAME)).getDeletedEntities();
 
         // Verify that deleteEntities() response has guids for tables and their composite columns.
         Assert.assertTrue(deletedGuids.contains(table1Entity.getId()._getId()));
@@ -936,15 +996,10 @@ public class DefaultMetadataServiceTest {
         assertEntityDeleted(COLUMN_TYPE, NAME, col3.get(NAME));
 
         // Verify that the listener was notified about the deleted entities.
-        Collection<ITypedReferenceableInstance> deletedEntitiesFromListener = listener.getDeletedEntities();
+        List<String> deletedEntitiesFromListener = listener.getDeletedEntities();
         Assert.assertNotNull(deletedEntitiesFromListener);
         Assert.assertEquals(deletedEntitiesFromListener.size(), deletedGuids.size());
-        List<String> deletedGuidsFromListener = new ArrayList<>(deletedGuids.size());
-        for (ITypedReferenceableInstance deletedEntity : deletedEntitiesFromListener) {
-            deletedGuidsFromListener.add(deletedEntity.getId()._getId());
-        }
-        Assert.assertEquals(deletedGuidsFromListener.size(), deletedGuids.size());
-        Assert.assertTrue(deletedGuidsFromListener.containsAll(deletedGuids));
+        Assert.assertTrue(deletedEntitiesFromListener.containsAll(deletedGuids));
     }
 
     @Test
@@ -1024,10 +1079,10 @@ public class DefaultMetadataServiceTest {
         }
     }
 
-    private static class DeleteEntitiesChangeListener implements EntityChangeListener {
-        
-        private Collection<ITypedReferenceableInstance> deletedEntities_;
-        
+    private static class EntitiesChangeListener implements EntityChangeListener {
+        private List<String> deletedEntities = new ArrayList<>();
+        private List<String> updatedEntities = new ArrayList<>();
+
         @Override
         public void onEntitiesAdded(Collection<ITypedReferenceableInstance> entities)
             throws AtlasException {
@@ -1036,6 +1091,10 @@ public class DefaultMetadataServiceTest {
         @Override
         public void onEntitiesUpdated(Collection<ITypedReferenceableInstance> entities)
             throws AtlasException {
+            updatedEntities.clear();
+            for (ITypedReferenceableInstance entity : entities) {
+                updatedEntities.add(entity.getId()._getId());
+            }
         }
 
         @Override
@@ -1051,11 +1110,18 @@ public class DefaultMetadataServiceTest {
         @Override
         public void onEntitiesDeleted(Collection<ITypedReferenceableInstance> entities)
             throws AtlasException {
-            deletedEntities_ = entities;
+            deletedEntities.clear();
+            for (ITypedReferenceableInstance entity : entities) {
+                deletedEntities.add(entity.getId()._getId());
+            }
         }
         
-        public Collection<ITypedReferenceableInstance> getDeletedEntities() {
-            return deletedEntities_;
+        public List<String> getDeletedEntities() {
+            return deletedEntities;
+        }
+
+        public List<String> getUpdatedEntities() {
+            return updatedEntities;
         }
     }
 }

server-api/src/main/java/org/apache/atlas/RequestContext.java

@@ -49,7 +49,16 @@ public class RequestContext {
     private RequestContext() {
     }
 
+    //To handle gets from background threads where createContext() is not called
+    //createContext called for every request in the filter
     public static RequestContext get() {
+        if (CURRENT_CONTEXT.get() == null) {
+            synchronized (RequestContext.class) {
+                if (CURRENT_CONTEXT.get() == null) {
+                    createContext();
+                }
+            }
+        }
         return CURRENT_CONTEXT.get();
     }
 
@@ -72,15 +81,19 @@ public class RequestContext {
         this.user = user;
     }
 
-    public void recordCreatedEntities(Collection<String> createdEntityIds) {
+    public void recordEntityCreate(Collection<String> createdEntityIds) {
         this.createdEntityIds.addAll(createdEntityIds);
     }
 
-    public void recordUpdatedEntities(Collection<String> updatedEntityIds) {
+    public void recordEntityUpdate(Collection<String> updatedEntityIds) {
         this.updatedEntityIds.addAll(updatedEntityIds);
     }
 
-    public void recordDeletedEntity(String entityId, String typeName) throws AtlasException {
+    public void recordEntityUpdate(String entityId) {
+        this.updatedEntityIds.add(entityId);
+    }
+
+    public void recordEntityDelete(String entityId, String typeName) throws AtlasException {
         ClassType type = typeSystem.getDataType(ClassType.class, typeName);
         ITypedReferenceableInstance entity = type.createInstance(new Id(entityId, 0, typeName));
         if (deletedEntityIds.add(entityId)) {

server-api/src/main/java/org/apache/atlas/services/MetadataService.java

@@ -18,6 +18,7 @@
 
 package org.apache.atlas.services;
 
+import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.EntityAuditEvent;
 import org.apache.atlas.listener.EntityChangeListener;
@@ -80,7 +81,7 @@ public interface MetadataService {
      * @param entityDefinition definition
      * @return json array of guids of entities created
      */
-    String createEntities(String entityDefinition) throws AtlasException;
+    List<String> createEntities(String entityDefinition) throws AtlasException;
 
     /**
      * Get a typed entity instance.
@@ -136,7 +137,7 @@ public interface MetadataService {
      * @param value    property value
      * @return json array of guids of entities created/updated
      */
-    String updateEntityAttributeByGuid(String guid, String attribute, String value) throws AtlasException;
+    AtlasClient.EntityResult updateEntityAttributeByGuid(String guid, String attribute, String value) throws AtlasException;
 
     /**
      * Supports Partial updates of an entity. Users can update a subset of attributes for an entity identified by its guid
@@ -146,7 +147,7 @@ public interface MetadataService {
      * @return json array of guids of entities created/updated
      * @throws AtlasException
      */
-    String updateEntityPartialByGuid(String guid, Referenceable entity) throws AtlasException;
+    AtlasClient.EntityResult updateEntityPartialByGuid(String guid, Referenceable entity) throws AtlasException;
 
     /**
      * Batch API - Adds/Updates the given entity id(guid).
@@ -154,7 +155,7 @@ public interface MetadataService {
      * @param entityJson entity json
      * @return json array of guids of entities created/updated
      */
-    String updateEntities(String entityJson) throws AtlasException;
+    AtlasClient.EntityResult updateEntities(String entityJson) throws AtlasException;
 
     // Trait management functions
 
@@ -168,8 +169,9 @@ public interface MetadataService {
      * @return Guid of updated entity
      * @throws AtlasException
      */
-    String updateEntityByUniqueAttribute(String typeName, String uniqueAttributeName, String attrValue,
-                                         Referenceable updatedEntity) throws AtlasException;
+    AtlasClient.EntityResult updateEntityByUniqueAttribute(String typeName, String uniqueAttributeName,
+                                                           String attrValue,
+                                                           Referenceable updatedEntity) throws AtlasException;
 
     /**
      * Gets the list of trait names for a given entity represented by a guid.
@@ -207,10 +209,10 @@ public interface MetadataService {
      * Delete the specified entities from the repository
      * 
      * @param guids entity guids to be deleted
-     * @return List of guids for deleted entities 
+     * @return List of guids for deleted entities
      * @throws AtlasException
      */
-    List<String> deleteEntities(List<String> guids) throws AtlasException;
+    AtlasClient.EntityResult deleteEntities(List<String> guids) throws AtlasException;
     
     /**
      * Register a listener for entity change.
@@ -235,7 +237,8 @@ public interface MetadataService {
      * @return List of guids for deleted entities (including their composite references)
      * @throws AtlasException
      */
-    List<String> deleteEntityByUniqueAttribute(String typeName, String uniqueAttributeName, String attrValue) throws AtlasException;
+    AtlasClient.EntityResult deleteEntityByUniqueAttribute(String typeName, String uniqueAttributeName,
+                                                           String attrValue) throws AtlasException;
 
     /**
      * Returns entity audit events for entity id in the decreasing order of timestamp

typesystem/src/main/java/org/apache/atlas/typesystem/Referenceable.java

@@ -100,7 +100,7 @@ public class Referenceable extends Struct implements IReferenceableInstance {
      * @throws AtlasException if the referenceable can not be created
      */
     public Referenceable(IReferenceableInstance instance) throws AtlasException {
-        this(instance.getId()._getId(), instance.getTypeName(), instance.getValuesMap(), instance.getTraits(),
+        this(instance.getId(), instance.getTypeName(), instance.getValuesMap(), instance.getTraits(),
             getTraits(instance));
     }
 

webapp/src/main/java/org/apache/atlas/LocalAtlasClient.java

@@ -83,13 +83,13 @@ public class LocalAtlasClient extends AtlasClient {
             }
         };
         JSONObject response = entityOperation.run();
-        List<String> results = extractResults(response, GUID, new ExtractOperation<String, String>());
+        EntityResult results = extractEntityResult(response);
         LOG.debug("Create entities returned results: {}", results);
-        return results;
+        return results.getCreatedEntities();
     }
 
     @Override
-    protected List<String> updateEntities(final JSONArray entities) throws AtlasServiceException {
+    protected EntityResult updateEntities(final JSONArray entities) throws AtlasServiceException {
         LOG.debug("Updating entities: {}", entities);
         EntityOperation entityOperation = new EntityOperation(API.UPDATE_ENTITY) {
             @Override
@@ -98,7 +98,7 @@ public class LocalAtlasClient extends AtlasClient {
             }
         };
         JSONObject response = entityOperation.run();
-        List<String> results = extractResults(response, GUID, new ExtractOperation<String, String>());
+        EntityResult results = extractEntityResult(response);
         LOG.debug("Update entities returned results: {}", results);
         return results;
     }
@@ -130,7 +130,7 @@ public class LocalAtlasClient extends AtlasClient {
     }
 
     @Override
-    public String updateEntity(final String entityType, final String uniqueAttributeName,
+    public EntityResult updateEntity(final String entityType, final String uniqueAttributeName,
                                final String uniqueAttributeValue, Referenceable entity) throws AtlasServiceException {
         final String entityJson = InstanceSerialization.toJson(entity, true);
         LOG.debug("Updating entity type: {}, attributeName: {}, attributeValue: {}, entity: {}", entityType,
@@ -143,13 +143,13 @@ public class LocalAtlasClient extends AtlasClient {
             }
         };
         JSONObject response = entityOperation.run();
-        String result = getString(response, GUID);
+        EntityResult result = extractEntityResult(response);
         LOG.debug("Update entity returned result: {}", result);
         return result;
     }
 
     @Override
-    public List<String> deleteEntity(final String entityType, final String uniqueAttributeName,
+    public EntityResult deleteEntity(final String entityType, final String uniqueAttributeName,
                                      final String uniqueAttributeValue) throws AtlasServiceException {
         LOG.debug("Deleting entity type: {}, attributeName: {}, attributeValue: {}", entityType, uniqueAttributeName,
                 uniqueAttributeValue);
@@ -160,7 +160,7 @@ public class LocalAtlasClient extends AtlasClient {
             }
         };
         JSONObject response = entityOperation.run();
-        List<String> results = extractResults(response, GUID, new ExtractOperation<String, String>());
+        EntityResult results = extractEntityResult(response);
         LOG.debug("Delete entities returned results: {}", results);
         return results;
     }
@@ -191,18 +191,18 @@ public class LocalAtlasClient extends AtlasClient {
     }
 
     @Override
-    public void updateEntityAttribute(final String guid, final String attribute, String value) throws AtlasServiceException {
+    public EntityResult updateEntityAttribute(final String guid, final String attribute, String value) throws AtlasServiceException {
         throw new IllegalStateException("Not supported in LocalAtlasClient");
     }
 
     @Override
-    public void updateEntity(String guid, Referenceable entity) throws AtlasServiceException {
+    public EntityResult updateEntity(String guid, Referenceable entity) throws AtlasServiceException {
         throw new IllegalStateException("Not supported in LocalAtlasClient");
     }
 
 
     @Override
-    public List<String> deleteEntities(final String ... guids) throws AtlasServiceException {
+    public EntityResult deleteEntities(final String ... guids) throws AtlasServiceException {
         throw new IllegalStateException("Not supported in LocalAtlasClient");
     }
 

notification/src/main/java/org/apache/atlas/notification/entity/NotificationEntityChangeListener.java → webapp/src/main/java/org/apache/atlas/notification/NotificationEntityChangeListener.java

@@ -15,22 +15,30 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.atlas.notification.entity;
+package org.apache.atlas.notification;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.inject.Inject;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.listener.EntityChangeListener;
-import org.apache.atlas.notification.NotificationInterface;
+import org.apache.atlas.notification.entity.EntityNotification;
+import org.apache.atlas.notification.entity.EntityNotificationImpl;
 import org.apache.atlas.typesystem.IReferenceableInstance;
 import org.apache.atlas.typesystem.IStruct;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.Referenceable;
+import org.apache.atlas.typesystem.Struct;
+import org.apache.atlas.typesystem.types.FieldMapping;
+import org.apache.atlas.typesystem.types.TraitType;
 import org.apache.atlas.typesystem.types.TypeSystem;
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * Listen to the repository for entity changes and produce entity change notifications.
@@ -87,6 +95,53 @@ public class NotificationEntityChangeListener implements EntityChangeListener {
     // ----- helper methods -------------------------------------------------
 
 
+    // ----- helper methods ----------------------------------------------------
+    @VisibleForTesting
+    public static List<IStruct> getAllTraits(IReferenceableInstance entityDefinition,
+                                              TypeSystem typeSystem) throws AtlasException {
+        List<IStruct> traitInfo = new LinkedList<>();
+        for (String traitName : entityDefinition.getTraits()) {
+            IStruct trait = entityDefinition.getTrait(traitName);
+            String typeName = trait.getTypeName();
+            Map<String, Object> valuesMap = trait.getValuesMap();
+            traitInfo.add(new Struct(typeName, valuesMap));
+            traitInfo.addAll(getSuperTraits(typeName, valuesMap, typeSystem));
+        }
+        return traitInfo;
+    }
+
+    private static List<IStruct> getSuperTraits(
+            String typeName, Map<String, Object> values, TypeSystem typeSystem) throws AtlasException {
+
+        List<IStruct> superTypes = new LinkedList<>();
+
+        TraitType traitDef = typeSystem.getDataType(TraitType.class, typeName);
+        Set<String> superTypeNames = traitDef.getAllSuperTypeNames();
+
+        for (String superTypeName : superTypeNames) {
+            TraitType superTraitDef = typeSystem.getDataType(TraitType.class, superTypeName);
+
+            Map<String, Object> superTypeValues = new HashMap<>();
+
+            FieldMapping fieldMapping = superTraitDef.fieldMapping();
+
+            if (fieldMapping != null) {
+                Set<String> superTypeAttributeNames = fieldMapping.fields.keySet();
+
+                for (String superTypeAttributeName : superTypeAttributeNames) {
+                    if (values.containsKey(superTypeAttributeName)) {
+                        superTypeValues.put(superTypeAttributeName, values.get(superTypeAttributeName));
+                    }
+                }
+            }
+            IStruct superTrait = new Struct(superTypeName, superTypeValues);
+            superTypes.add(superTrait);
+            superTypes.addAll(getSuperTraits(superTypeName, values, typeSystem));
+        }
+
+        return superTypes;
+    }
+
     // send notification of entity change
     private void notifyOfEntityEvent(Collection<ITypedReferenceableInstance> entityDefinitions,
                                      EntityNotification.OperationType operationType) throws AtlasException {
@@ -96,7 +151,7 @@ public class NotificationEntityChangeListener implements EntityChangeListener {
             Referenceable entity = new Referenceable(entityDefinition);
 
             EntityNotificationImpl notification =
-                    new EntityNotificationImpl(entity, operationType, typeSystem);
+                    new EntityNotificationImpl(entity, operationType, getAllTraits(entity, typeSystem));
 
             messages.add(notification);
         }

webapp/src/main/java/org/apache/atlas/web/resources/EntityResource.java

@@ -19,7 +19,6 @@
 package org.apache.atlas.web.resources;
 
 import com.google.common.base.Preconditions;
-
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.EntityAuditEvent;
@@ -59,9 +58,7 @@ import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.core.UriBuilder;
 import javax.ws.rs.core.UriInfo;
-
 import java.net.URI;
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 
@@ -119,19 +116,11 @@ public class EntityResource {
 
             LOG.debug("submitting entities {} ", AtlasClient.toString(new JSONArray(entities)));
 
-            final String guids = metadataService.createEntities(entities);
+            final List<String> guids = metadataService.createEntities(entities);
+            JSONObject response = getResponse(new AtlasClient.EntityResult(guids, null, null));
 
             UriBuilder ub = uriInfo.getAbsolutePathBuilder();
-            URI locationURI = ub.path(guids).build();
-
-            JSONObject response = new JSONObject();
-            response.put(AtlasClient.REQUEST_ID, Servlets.getRequestId());
-            JSONArray guidArray = new JSONArray(guids);
-            response.put(AtlasClient.GUID, guidArray);
-            if (guidArray.length() > 0) {
-                response.put(AtlasClient.DEFINITION,
-                        new JSONObject(metadataService.getEntityDefinition(new JSONArray(guids).getString(0))));
-            }
+            URI locationURI = guids.isEmpty() ? null : ub.path(guids.get(0)).build();
 
             return Response.created(locationURI).entity(response).build();
 
@@ -150,6 +139,18 @@ public class EntityResource {
         }
     }
 
+    private JSONObject getResponse(AtlasClient.EntityResult entityResult) throws AtlasException, JSONException {
+        JSONObject response = new JSONObject();
+        response.put(AtlasClient.REQUEST_ID, Servlets.getRequestId());
+        response.put(AtlasClient.ENTITIES, new JSONObject(entityResult.toString()).get(AtlasClient.ENTITIES));
+        String sampleEntityId = getSample(entityResult);
+        if (sampleEntityId != null) {
+            String entityDefinition = metadataService.getEntityDefinition(sampleEntityId);
+            response.put(AtlasClient.DEFINITION, new JSONObject(entityDefinition));
+        }
+        return response;
+    }
+
     /**
      * Complete update of a set of entities - the values not specified will be replaced with null/removed
      * Adds/Updates given entities identified by its GUID or unique attribute
@@ -163,14 +164,8 @@ public class EntityResource {
             final String entities = Servlets.getRequestPayload(request);
             LOG.debug("updating entities {} ", AtlasClient.toString(new JSONArray(entities)));
 
-            final String guids = metadataService.updateEntities(entities);
-
-            JSONObject response = new JSONObject();
-            response.put(AtlasClient.REQUEST_ID, Servlets.getRequestId());
-            JSONArray guidsArray = new JSONArray(guids);
-            response.put(AtlasClient.GUID, guidsArray);
-            String entityDefinition = metadataService.getEntityDefinition(guidsArray.getString(0));
-            response.put(AtlasClient.DEFINITION, new JSONObject(entityDefinition));
+            AtlasClient.EntityResult entityResult = metadataService.updateEntities(entities);
+            JSONObject response = getResponse(entityResult);
             return Response.ok(response).build();
         } catch(EntityExistsException e) {
             LOG.error("Unique constraint violation", e);
@@ -187,6 +182,25 @@ public class EntityResource {
         }
     }
 
+    private String getSample(AtlasClient.EntityResult entityResult) {
+        String sample = getSample(entityResult.getCreatedEntities());
+        if (sample == null) {
+            sample = getSample(entityResult.getUpdateEntities());
+        }
+        if (sample == null) {
+            sample = getSample(entityResult.getDeletedEntities());
+        }
+        return sample;
+    }
+
+
+    private String getSample(List<String> list) {
+        if (list != null && list.size() > 0) {
+            return list.get(0);
+        }
+        return null;
+    }
+
     /**
      * Adds/Updates given entity identified by its unique attribute( entityType, attributeName and value)
      * Updates support only partial update of an entity - Adds/updates any new values specified
@@ -214,11 +228,10 @@ public class EntityResource {
 
             Referenceable updatedEntity =
                 InstanceSerialization.fromJsonReferenceable(entities, true);
-            final String guid = metadataService.updateEntityByUniqueAttribute(entityType, attribute, value, updatedEntity);
+            AtlasClient.EntityResult entityResult =
+                    metadataService.updateEntityByUniqueAttribute(entityType, attribute, value, updatedEntity);
 
-            JSONObject response = new JSONObject();
-            response.put(AtlasClient.REQUEST_ID, Thread.currentThread().getName());
-            response.put(AtlasClient.GUID, guid);
+            JSONObject response = getResponse(entityResult);
             return Response.ok(response).build();
         } catch (ValueConversionException ve) {
             LOG.error("Unable to persist entity instance due to a desrialization error ", ve);
@@ -268,10 +281,8 @@ public class EntityResource {
 
             Referenceable updatedEntity =
                     InstanceSerialization.fromJsonReferenceable(entityJson, true);
-            metadataService.updateEntityPartialByGuid(guid, updatedEntity);
-
-            JSONObject response = new JSONObject();
-            response.put(AtlasClient.REQUEST_ID, Thread.currentThread().getName());
+            AtlasClient.EntityResult entityResult = metadataService.updateEntityPartialByGuid(guid, updatedEntity);
+            JSONObject response = getResponse(entityResult);
             return Response.ok(response).build();
         } catch (EntityNotFoundException e) {
             LOG.error("An entity with GUID={} does not exist", guid, e);
@@ -301,12 +312,8 @@ public class EntityResource {
             String value = Servlets.getRequestPayload(request);
             Preconditions.checkNotNull(value, "Entity value cannot be null");
 
-            metadataService.updateEntityAttributeByGuid(guid, property, value);
-
-            JSONObject response = new JSONObject();
-            response.put(AtlasClient.REQUEST_ID, Thread.currentThread().getName());
-            response.put(AtlasClient.GUID, guid);
-
+            AtlasClient.EntityResult entityResult = metadataService.updateEntityAttributeByGuid(guid, property, value);
+            JSONObject response = getResponse(entityResult);
             return Response.ok(response).build();
         } catch (EntityNotFoundException e) {
             LOG.error("An entity with GUID={} does not exist", guid, e);
@@ -340,19 +347,13 @@ public class EntityResource {
         @QueryParam("value") String value) {
         
         try {
-            List<String> deletedGuids = new ArrayList<>();
+            AtlasClient.EntityResult entityResult;
             if (guids != null && !guids.isEmpty()) {
-                deletedGuids = metadataService.deleteEntities(guids);
+                entityResult = metadataService.deleteEntities(guids);
             } else {
-                deletedGuids = metadataService.deleteEntityByUniqueAttribute(entityType, attribute, value);
-            }
-            JSONObject response = new JSONObject();
-            response.put(AtlasClient.REQUEST_ID, Servlets.getRequestId());
-            JSONArray guidArray = new JSONArray(deletedGuids.size());
-            for (String guid : deletedGuids) {
-                guidArray.put(guid);
+                entityResult = metadataService.deleteEntityByUniqueAttribute(entityType, attribute, value);
             }
-            response.put(AtlasClient.GUID, guidArray);
+            JSONObject response = getResponse(entityResult);
             return Response.ok(response).build();
         } catch (EntityNotFoundException e) {
             if(guids != null || !guids.isEmpty()) {
@@ -386,7 +387,6 @@ public class EntityResource {
 
             JSONObject response = new JSONObject();
             response.put(AtlasClient.REQUEST_ID, Servlets.getRequestId());
-            response.put(AtlasClient.GUID, guid);
 
             Response.Status status = Response.Status.NOT_FOUND;
             if (entityDefinition != null) {
@@ -518,7 +518,6 @@ public class EntityResource {
 
             JSONObject response = new JSONObject();
             response.put(AtlasClient.REQUEST_ID, Servlets.getRequestId());
-            response.put(AtlasClient.GUID, guid);
             response.put(AtlasClient.RESULTS, new JSONArray(traitNames));
             response.put(AtlasClient.COUNT, traitNames.size());
 
@@ -555,7 +554,6 @@ public class EntityResource {
 
             JSONObject response = new JSONObject();
             response.put(AtlasClient.REQUEST_ID, Servlets.getRequestId());
-            response.put(AtlasClient.GUID, guid);
 
             return Response.created(locationURI).entity(response).build();
         } catch (EntityNotFoundException | TypeNotFoundException e) {
@@ -588,7 +586,6 @@ public class EntityResource {
 
             JSONObject response = new JSONObject();
             response.put(AtlasClient.REQUEST_ID, Servlets.getRequestId());
-            response.put(AtlasClient.GUID, guid);
             response.put(TRAIT_NAME, traitName);
 
             return Response.ok(response).build();

webapp/src/main/java/org/apache/atlas/web/service/ServiceModule.java

@@ -23,7 +23,7 @@ import com.google.inject.multibindings.Multibinder;
 import org.apache.atlas.kafka.KafkaNotification;
 import org.apache.atlas.listener.EntityChangeListener;
 import org.apache.atlas.notification.NotificationHookConsumer;
-import org.apache.atlas.notification.entity.NotificationEntityChangeListener;
+import org.apache.atlas.notification.NotificationEntityChangeListener;
 import org.apache.atlas.service.Service;
 
 public class ServiceModule extends AbstractModule {

webapp/src/test/java/org/apache/atlas/LocalAtlasClientTest.java

@@ -23,7 +23,6 @@ import org.apache.atlas.typesystem.Referenceable;
 import org.apache.atlas.web.resources.EntityResource;
 import org.apache.atlas.web.service.ServiceState;
 import org.apache.commons.lang.RandomStringUtils;
-import org.codehaus.jettison.json.JSONArray;
 import org.codehaus.jettison.json.JSONObject;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
@@ -36,6 +35,7 @@ import javax.ws.rs.core.Response;
 import java.util.Arrays;
 import java.util.List;
 
+import static org.apache.atlas.AtlasClient.ENTITIES;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyListOf;
 import static org.mockito.Matchers.anyString;
@@ -64,7 +64,8 @@ public class LocalAtlasClientTest {
         when(entityResource.submit(any(HttpServletRequest.class))).thenReturn(response);
         final String guid = random();
         when(response.getEntity()).thenReturn(new JSONObject() {{
-            put(AtlasClient.GUID, new JSONArray(Arrays.asList(guid)));
+            put(ENTITIES, new JSONObject(
+                    new AtlasClient.EntityResult(Arrays.asList(guid), null, null).toString()).get(ENTITIES));
         }});
 
         LocalAtlasClient atlasClient = new LocalAtlasClient(serviceState, entityResource);
@@ -119,12 +120,14 @@ public class LocalAtlasClientTest {
         when(entityResource.updateByUniqueAttribute(anyString(), anyString(), anyString(),
                 any(HttpServletRequest.class))).thenReturn(response);
         when(response.getEntity()).thenReturn(new JSONObject() {{
-            put(AtlasClient.GUID, guid);
+            put(ENTITIES, new JSONObject(
+                    new AtlasClient.EntityResult(null, Arrays.asList(guid), null).toString()).get(ENTITIES));
         }});
 
         LocalAtlasClient atlasClient = new LocalAtlasClient(serviceState, entityResource);
-        String actualId = atlasClient.updateEntity(random(), random(), random(), new Referenceable(random()));
-        assertEquals(actualId, guid);
+        AtlasClient.EntityResult
+                entityResult = atlasClient.updateEntity(random(), random(), random(), new Referenceable(random()));
+        assertEquals(entityResult.getUpdateEntities(), Arrays.asList(guid));
     }
 
     @Test
@@ -132,14 +135,14 @@ public class LocalAtlasClientTest {
         final String guid = random();
         Response response = mock(Response.class);
         when(response.getEntity()).thenReturn(new JSONObject() {{
-            put(AtlasClient.GUID, new JSONArray(Arrays.asList(guid)));
+            put(ENTITIES, new JSONObject(
+                    new AtlasClient.EntityResult(null, null, Arrays.asList(guid)).toString()).get(ENTITIES));
         }});
 
         when(entityResource.deleteEntities(anyListOf(String.class), anyString(), anyString(), anyString())).thenReturn(response);
         LocalAtlasClient atlasClient = new LocalAtlasClient(serviceState, entityResource);
-        List<String> results = atlasClient.deleteEntity(random(), random(), random());
-        assertEquals(results.size(), 1);
-        assertEquals(results.get(0), guid);
+        AtlasClient.EntityResult entityResult = atlasClient.deleteEntity(random(), random(), random());
+        assertEquals(entityResult.getDeletedEntities(), Arrays.asList(guid));
     }
 
     private String random() {

webapp/src/test/java/org/apache/atlas/notification/EntityNotificationIT.java

@@ -22,7 +22,6 @@ import com.google.common.collect.ImmutableSet;
 import com.google.inject.Inject;
 import com.sun.jersey.api.client.ClientResponse;
 import com.sun.jersey.api.client.WebResource;
-
 import org.apache.atlas.notification.entity.EntityNotification;
 import org.apache.atlas.typesystem.IReferenceableInstance;
 import org.apache.atlas.typesystem.IStruct;
@@ -43,7 +42,6 @@ import org.testng.annotations.Test;
 
 import javax.ws.rs.HttpMethod;
 import javax.ws.rs.core.Response;
-
 import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;

webapp/src/test/java/org/apache/atlas/notification/NotificationEntityChangeListenerTest.java

+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.atlas.notification;
+
+import org.apache.atlas.typesystem.IStruct;
+import org.apache.atlas.typesystem.Referenceable;
+import org.apache.atlas.typesystem.Struct;
+import org.apache.atlas.typesystem.types.TraitType;
+import org.apache.atlas.typesystem.types.TypeSystem;
+import org.testng.annotations.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+public class NotificationEntityChangeListenerTest {
+    @Test
+    public void testGetAllTraitsSuperTraits() throws Exception {
+
+        TypeSystem typeSystem = mock(TypeSystem.class);
+
+        String traitName = "MyTrait";
+        IStruct myTrait = new Struct(traitName);
+
+        String superTraitName = "MySuperTrait";
+
+        TraitType traitDef = mock(TraitType.class);
+        Set<String> superTypeNames = Collections.singleton(superTraitName);
+
+        TraitType superTraitDef = mock(TraitType.class);
+        Set<String> superSuperTypeNames = Collections.emptySet();
+
+        Referenceable entity = getEntity("id", myTrait);
+
+        when(typeSystem.getDataType(TraitType.class, traitName)).thenReturn(traitDef);
+        when(typeSystem.getDataType(TraitType.class, superTraitName)).thenReturn(superTraitDef);
+
+        when(traitDef.getAllSuperTypeNames()).thenReturn(superTypeNames);
+        when(superTraitDef.getAllSuperTypeNames()).thenReturn(superSuperTypeNames);
+
+        List<IStruct> allTraits = NotificationEntityChangeListener.getAllTraits(entity, typeSystem);
+
+        assertEquals(2, allTraits.size());
+
+        for (IStruct trait : allTraits) {
+            String typeName = trait.getTypeName();
+            assertTrue(typeName.equals(traitName) || typeName.equals(superTraitName));
+        }
+    }
+
+    private Referenceable getEntity(String id, IStruct... traits) {
+        String typeName = "typeName";
+        Map<String, Object> values = new HashMap<>();
+
+        List<String> traitNames = new LinkedList<>();
+        Map<String, IStruct> traitMap = new HashMap<>();
+
+        for (IStruct trait : traits) {
+            String traitName = trait.getTypeName();
+
+            traitNames.add(traitName);
+            traitMap.put(traitName, trait);
+        }
+        return new Referenceable(id, typeName, values, traitNames, traitMap);
+    }
+}

webapp/src/test/java/org/apache/atlas/web/resources/EntityJerseyResourceIT.java

@@ -71,6 +71,7 @@ import java.util.Map;
 import java.util.UUID;
 
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.fail;
 
 
@@ -154,7 +155,10 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
 
         JSONObject response = new JSONObject(responseAsString);
         Assert.assertNotNull(response.get(AtlasClient.REQUEST_ID));
-        Assert.assertNotNull(response.get(AtlasClient.GUID));
+
+        AtlasClient.EntityResult entityResult = AtlasClient.EntityResult.fromString(response.toString());
+        assertEquals(entityResult.getCreatedEntities().size(), 1);
+        assertNotNull(entityResult.getCreatedEntities().get(0));
     }
 
     @Test
@@ -376,7 +380,9 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
     }
 
     private void addProperty(String guid, String property, String value) throws AtlasServiceException {
-        serviceClient.updateEntityAttribute(guid, property, value);
+        AtlasClient.EntityResult entityResult = serviceClient.updateEntityAttribute(guid, property, value);
+        assertEquals(entityResult.getUpdateEntities().size(), 1);
+        assertEquals(entityResult.getUpdateEntities().get(0), guid);
     }
 
     private ClientResponse getEntityDefinition(String guid) {
@@ -482,7 +488,6 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
 
         JSONObject response = new JSONObject(responseAsString);
         Assert.assertNotNull(response.get(AtlasClient.REQUEST_ID));
-        Assert.assertNotNull(response.get("GUID"));
 
         final JSONArray list = response.getJSONArray(AtlasClient.RESULTS);
         Assert.assertEquals(list.length(), 7);
@@ -513,7 +518,6 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
 
         JSONObject response = new JSONObject(responseAsString);
         Assert.assertNotNull(response.get(AtlasClient.REQUEST_ID));
-        Assert.assertNotNull(response.get(AtlasClient.GUID));
 
         assertEntityAudit(guid, EntityAuditEvent.EntityAuditAction.TAG_ADD);
     }
@@ -561,7 +565,6 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
 
         JSONObject response = new JSONObject(responseAsString);
         Assert.assertNotNull(response.get(AtlasClient.REQUEST_ID));
-        Assert.assertNotNull(response.get(AtlasClient.GUID));
 
         // verify the response
         clientResponse = getEntityDefinition(guid);
@@ -612,7 +615,6 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
 
         JSONObject response = new JSONObject(responseAsString);
         Assert.assertNotNull(response.get(AtlasClient.REQUEST_ID));
-        Assert.assertNotNull(response.get("GUID"));
         Assert.assertNotNull(response.get("traitName"));
         assertEntityAudit(guid, EntityAuditEvent.EntityAuditAction.TAG_DELETE);
     }
@@ -635,7 +637,8 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
                 "trait=" + traitName + " should be defined in type system before it can be deleted");
         Assert.assertNotNull(response.get(AtlasClient.STACKTRACE));
     }
-@Test(dependsOnMethods = "testSubmitEntity()")
+
+    @Test(dependsOnMethods = "testSubmitEntity()")
     public void testDeleteExistentTraitNonExistentForEntity() throws Exception {
     
         final String guid = tableId._getId();
@@ -704,7 +707,9 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
         }});
 
         LOG.debug("Updating entity= " + tableUpdated);
-        serviceClient.updateEntity(tableId._getId(), tableUpdated);
+        AtlasClient.EntityResult entityResult = serviceClient.updateEntity(tableId._getId(), tableUpdated);
+        assertEquals(entityResult.getUpdateEntities().size(), 1);
+        assertEquals(entityResult.getUpdateEntities().get(0), tableId._getId());
 
         ClientResponse response = getEntityDefinition(tableId._getId());
         String definition = getEntityDefinition(response);
@@ -722,8 +727,10 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
         }});
 
         LOG.debug("Updating entity= " + tableUpdated);
-        serviceClient.updateEntity(BaseResourceIT.HIVE_TABLE_TYPE, "name", (String) tableInstance.get("name"),
-                tableUpdated);
+        entityResult = serviceClient.updateEntity(BaseResourceIT.HIVE_TABLE_TYPE, "name",
+                (String) tableInstance.get("name"), tableUpdated);
+        assertEquals(entityResult.getUpdateEntities().size(), 1);
+        assertEquals(entityResult.getUpdateEntities().get(0), tableId._getId());
 
         response = getEntityDefinition(tableId._getId());
         definition = getEntityDefinition(response);
@@ -732,7 +739,6 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
 
         Assert.assertTrue(refs.get(0).equalsContents(columns.get(0)));
         Assert.assertEquals(refs.get(0).get("dataType"), "int");
-
     }
 
     @Test(dependsOnMethods = "testSubmitEntity")
@@ -765,9 +771,8 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
         // ATLAS-586: verify response entity can be parsed by GSON.
         String entity = clientResponse.getEntity(String.class);
         Gson gson = new Gson();
-        UpdateEntitiesResponse updateEntitiesResponse = null;
         try {
-            updateEntitiesResponse = gson.fromJson(entity, UpdateEntitiesResponse.class);
+            UpdateEntitiesResponse updateEntitiesResponse = gson.fromJson(entity, UpdateEntitiesResponse.class);
         }
         catch (JsonSyntaxException e) {
             Assert.fail("Response entity from " + service.path(ENTITIES).getURI() + " not parseable by GSON", e);
@@ -785,7 +790,7 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
     
     private static class UpdateEntitiesResponse {
         String requestId;
-        String[] GUID;
+        AtlasClient.EntityResult entities;
         AtlasEntity definition;
     }
     
@@ -811,15 +816,9 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
             queryParam(AtlasClient.GUID.toLowerCase(), db1Id._getId()).
             queryParam(AtlasClient.GUID.toLowerCase(), db2Id._getId()).
             accept(Servlets.JSON_MEDIA_TYPE).type(Servlets.JSON_MEDIA_TYPE).method(HttpMethod.DELETE, ClientResponse.class);
+
         JSONObject response = getEntity(clientResponse);
-        final String deletedGuidsJson = response.getString(AtlasClient.GUID);
-        Assert.assertNotNull(deletedGuidsJson);
-        JSONArray guidsArray = new JSONArray(deletedGuidsJson);
-        Assert.assertEquals(guidsArray.length(), 2);
-        List<String> deletedGuidsList = new ArrayList<>(2);
-        for (int index = 0; index < guidsArray.length(); index++) {
-            deletedGuidsList.add(guidsArray.getString(index));
-        }
+        List<String> deletedGuidsList = AtlasClient.EntityResult.fromString(response.toString()).getDeletedEntities();
         Assert.assertTrue(deletedGuidsList.contains(db1Id._getId()));
         Assert.assertTrue(deletedGuidsList.contains(db2Id._getId()));
 
@@ -843,7 +842,8 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
         Id db2Id = createInstance(db2);
         
         // Delete the database entities
-        List<String> deletedGuidsList = serviceClient.deleteEntities(db1Id._getId(), db2Id._getId());
+        List<String> deletedGuidsList =
+                serviceClient.deleteEntities(db1Id._getId(), db2Id._getId()).getDeletedEntities();
         
         // Verify that deleteEntities() response has database entity guids 
         Assert.assertEquals(deletedGuidsList.size(), 2);
@@ -867,7 +867,7 @@ public class EntityJerseyResourceIT extends BaseResourceIT {
         Id db1Id = createInstance(db1);
 
         // Delete the database entity
-        List<String> deletedGuidsList = serviceClient.deleteEntity(DATABASE_TYPE, "name", dbName);
+        List<String> deletedGuidsList = serviceClient.deleteEntity(DATABASE_TYPE, "name", dbName).getDeletedEntities();
 
         // Verify that deleteEntities() response has database entity guids
         Assert.assertEquals(deletedGuidsList.size(), 1);

webapp/src/test/java/org/apache/atlas/web/service/CuratorFactoryTest.java

@@ -23,21 +23,16 @@ import org.apache.atlas.ha.HAConfiguration;
 import org.apache.commons.configuration.Configuration;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.framework.api.ACLProvider;
-import org.apache.zookeeper.ZooDefs;
 import org.apache.zookeeper.data.ACL;
 import org.mockito.ArgumentMatcher;
 import org.mockito.Mock;
-import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
-import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
-import static junit.framework.TestCase.assertEquals;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.argThat;
 import static org.mockito.Matchers.eq;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;