org.apache.atlas.model.instance.EntityMutationResponse Java Examples

@Override
@GraphTransaction
public EntityMutationResponse updateByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes,
                                                       AtlasEntityWithExtInfo updatedEntityInfo) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> updateByUniqueAttributes({}, {})", entityType.getTypeName(), uniqAttributes);
    }

    if (updatedEntityInfo == null || updatedEntityInfo.getEntity() == null) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "no entity to update.");
    }

    String      guid   = AtlasGraphUtilsV2.getGuidByUniqueAttributes(graph, entityType, uniqAttributes);
    AtlasEntity entity = updatedEntityInfo.getEntity();

    entity.setGuid(guid);

    AtlasAuthorizationUtils.verifyAccess(new AtlasEntityAccessRequest(typeRegistry, AtlasPrivilege.ENTITY_UPDATE, new AtlasEntityHeader(entity)), "update entity ByUniqueAttributes");

    return createOrUpdate(new AtlasEntityStream(updatedEntityInfo), true, false, false);
}
 

/*******
 * Entity Partial Update - Add/Update entity attribute identified by its GUID.
 * Supports only uprimitive attribute type and entity references.
 * does not support updation of complex types like arrays, maps
 * Null updates are not possible
 *******/
@PUT
@Path("/guid/{guid}")
public EntityMutationResponse partialUpdateEntityAttrByGuid(@PathParam("guid") String guid,
                                                            @QueryParam("name") String attrName,
                                                            Object attrValue) throws Exception {
    Servlets.validateQueryParamLength("guid", guid);
    Servlets.validateQueryParamLength("name", attrName);

    AtlasPerfTracer perf = null;

    try {
        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.partialUpdateEntityAttrByGuid(" + guid + "," + attrName + ")");
        }

        return entitiesStore.updateEntityAttributeByGuid(guid, attrName, attrValue);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

/**
 * Delete an entity identified by its type and unique attributes.
 *
 * In addition to the typeName path parameter, attribute key-value pair(s) can be provided in the following format
 *
 * attr:<attrName>=<attrValue>
 *
 * NOTE: The attrName and attrValue should be unique across entities, eg. qualifiedName
 *
 * The REST request would look something like this
 *
 * DELETE /v2/entity/uniqueAttribute/type/aType?attr:aTypeAttribute=someValue
 *
 * @param  typeName - entity type to be deleted
 * @param  servletRequest - request containing unique attributes/values
 * @return EntityMutationResponse
 */
@DELETE
@Path("/uniqueAttribute/type/{typeName}")
public EntityMutationResponse deleteByUniqueAttribute(@PathParam("typeName") String typeName,
                                                      @Context HttpServletRequest servletRequest) throws AtlasBaseException {
    Servlets.validateQueryParamLength("typeName", typeName);

    AtlasPerfTracer perf = null;

    try {
        Map<String, Object> attributes = getAttributes(servletRequest);

        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.deleteByUniqueAttribute(" + typeName + "," + attributes + ")");
        }

        AtlasEntityType entityType = ensureEntityType(typeName);

        return entitiesStore.deleteByUniqueAttributes(entityType, attributes);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

/**
 * Bulk API to delete list of entities identified by its GUIDs
 */
@DELETE
@Path("/bulk")
public EntityMutationResponse deleteByGuids(@QueryParam("guid") final List<String> guids) throws AtlasBaseException {
    if (CollectionUtils.isNotEmpty(guids)) {
        for (String guid : guids) {
            Servlets.validateQueryParamLength("guid", guid);
        }
    }

    AtlasPerfTracer perf = null;

    try {
        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.deleteByGuids(" + guids  + ")");
        }

        return entitiesStore.deleteByIds(guids);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

@Test
public void testDeleteEntitiesDoesNotLookupDeletedEntity() throws Exception {
    List<String> guids = Collections.singletonList(DELETED_GUID);
    List<AtlasEntityHeader> deletedEntities = Collections.singletonList(new AtlasEntityHeader(null, DELETED_GUID, null));

    // Create EntityResult with a deleted guid and no other guids.
    EntityMutationResponse  resp    = new EntityMutationResponse();
    List<AtlasEntityHeader> headers = toAtlasEntityHeaders(guids);

    if (CollectionUtils.isNotEmpty(headers)) {
        for (AtlasEntityHeader entity : headers) {
            resp.addEntity(EntityMutations.EntityOperation.DELETE, entity);
        }
    }

    when(entitiesStore.deleteByIds(guids)).thenReturn(resp);

    EntityMutationResponse response = entitiesStore.deleteByIds(guids);

    List<AtlasEntityHeader> responseDeletedEntities = response.getDeletedEntities();

    Assert.assertEquals(responseDeletedEntities, deletedEntities);
}
 

@Test
public void testDeleteEntitiesDoesNotLookupDeletedEntity() throws Exception {
    List<String> guids = Collections.singletonList(DELETED_GUID);
    List<AtlasEntityHeader> deletedEntities = Collections.singletonList(new AtlasEntityHeader(null, DELETED_GUID, null));

    // Create EntityResult with a deleted guid and no other guids.
    EntityMutationResponse  resp    = new EntityMutationResponse();
    List<AtlasEntityHeader> headers = toAtlasEntityHeaders(guids);

    if (CollectionUtils.isNotEmpty(headers)) {
        for (AtlasEntityHeader entity : headers) {
            resp.addEntity(EntityMutations.EntityOperation.DELETE, entity);
        }
    }

    when(entitiesStore.deleteByIds(guids)).thenReturn(resp);

    EntityMutationResponse response = entitiesStore.deleteByIds(guids);

    List<AtlasEntityHeader> responseDeletedEntities = response.getDeletedEntities();

    Assert.assertEquals(responseDeletedEntities, deletedEntities);
}
 

@BeforeClass
public void setUp() throws Exception {
    new GraphBackedSearchIndexer(typeRegistry);

    // create employee relationship types
    AtlasTypesDef employeeTypes = getDepartmentEmployeeTypes();
    typeDefStore.createTypesDef(employeeTypes);

    AtlasEntitiesWithExtInfo employeeInstances = getDepartmentEmployeeInstances();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(employeeInstances), false);

    for (AtlasEntityHeader entityHeader : response.getCreatedEntities()) {
        employeeNameIdMap.put((String) entityHeader.getAttribute(NAME), getAtlasObjectId(entityHeader));
    }

    init();
    AtlasTypesDef testTypes = getInverseReferenceTestTypes();
    typeDefStore.createTypesDef(testTypes);
}
 

/*******
 * Entity Partial Update - Allows a subset of attributes to be updated on
 * an entity which is identified by its type and unique attribute  eg: Referenceable.qualifiedName.
 * Null updates are not possible
 *
 * In addition to the typeName path parameter, attribute key-value pair(s) can be provided in the following format
 *
 * attr:<attrName>=<attrValue>
 *
 * NOTE: The attrName and attrValue should be unique across entities, eg. qualifiedName
 *
 * The REST request would look something like this
 *
 * PUT /v2/entity/uniqueAttribute/type/aType?attr:aTypeAttribute=someValue
 *

 *******/
@PUT
@Path("/uniqueAttribute/type/{typeName}")
public EntityMutationResponse partialUpdateEntityByUniqueAttrs(@PathParam("typeName") String typeName,
                                                               @Context HttpServletRequest servletRequest,
                                                               AtlasEntityWithExtInfo entityInfo) throws Exception {
    Servlets.validateQueryParamLength("typeName", typeName);

    AtlasPerfTracer perf = null;

    try {
        Map<String, Object> uniqueAttributes = getAttributes(servletRequest);

        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "EntityREST.partialUpdateEntityByUniqueAttrs(" + typeName + "," + uniqueAttributes + ")");
        }

        AtlasEntityType entityType = ensureEntityType(typeName);

        validateUniqueAttribute(entityType, uniqueAttributes);

        return entitiesStore.updateByUniqueAttributes(entityType, uniqueAttributes, entityInfo);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

private void createEntities() throws Exception {
    dbEntities.clear();

    for (int i = 1; i <= 2; i++) {
        AtlasEntity dbEntity = TestUtilsV2.createDBEntity();

        final EntityMutationResponse response = entityREST.createOrUpdate(new AtlasEntity.AtlasEntitiesWithExtInfo(dbEntity));

        assertNotNull(response);
        List<AtlasEntityHeader> entitiesMutated = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);

        assertNotNull(entitiesMutated);
        Assert.assertEquals(entitiesMutated.size(), 1);
        assertNotNull(entitiesMutated.get(0));
        dbEntity.setGuid(entitiesMutated.get(0).getGuid());

        dbEntities.add(dbEntity);
    }
}
 

@Test
public void deleteByGuidsSoft() throws Exception {
    RequestContext.get().setDeleteType(DeleteType.SOFT);

    createEntities();
    List<String> guids = new ArrayList<>();
    guids.add(dbEntities.get(0).getGuid());
    guids.add(dbEntities.get(1).getGuid());

    EntityMutationResponse response = entityREST.deleteByGuids(guids);

    assertNotNull(response);
    assertNotNull(response.getDeletedEntities());

    for (String guid : guids) {
        assertSoftDelete(guid);
    }
}
 

@Override
@GraphTransaction
public EntityMutationResponse updateByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes,
                                                       AtlasEntityWithExtInfo updatedEntityInfo) throws AtlasBaseException {

    if (LOG.isDebugEnabled()) {
        LOG.debug("==> updateByUniqueAttributes({}, {})", entityType.getTypeName(), uniqAttributes);
    }

    if (updatedEntityInfo == null || updatedEntityInfo.getEntity() == null) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "no entity to update.");
    }

    String guid = AtlasGraphUtilsV1.getGuidByUniqueAttributes(entityType, uniqAttributes);

    AtlasEntity entity = updatedEntityInfo.getEntity();

    entity.setGuid(guid);

    return createOrUpdate(new AtlasEntityStream(updatedEntityInfo), true);
}
 

@Override
protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
    final AtlasEntity.AtlasEntitiesWithExtInfo withExtInfo = readInputJSON(req, AtlasEntity.AtlasEntitiesWithExtInfo.class);
    final Map<String, String> guidAssignments = new HashMap<>();
    withExtInfo.getEntities().forEach(entity -> {
        atlasEntitiesByTypedQname.put(toTypedQname(entity), entity);
        String guid = entity.getGuid();
        if (!AtlasUtils.isGuidAssigned(guid)) {
            final String _guid = String.valueOf(guidSeq.getAndIncrement());
            guidAssignments.put(guid, _guid);
            entity.setGuid(_guid);
            guid = _guid;
        }
        atlasEntitiesByGuid.put(guid, entity);
    });
    final EntityMutationResponse mutationResponse = new EntityMutationResponse();
    mutationResponse.setGuidAssignments(guidAssignments);
    respondWithJson(resp, mutationResponse);
}
 

@BeforeClass
public void setUp() throws Exception {
    metadataService = TestUtils.addSessionCleanupWrapper(metadataService);

    AtlasTypesDef[] testTypesDefs = new AtlasTypesDef[] { TestUtilsV2.defineDeptEmployeeTypes(),
                                                          TestUtilsV2.defineInverseReferenceTestTypes()
                                                        };

    for (AtlasTypesDef typesDef : testTypesDefs) {
        AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(typesDef, typeRegistry);

        if (!typesToCreate.isEmpty()) {
            typeDefStore.createTypesDef(typesToCreate);
        }
    }

    deptEntity = TestUtilsV2.createDeptEg2();
    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(deptEntity), false);
    for (AtlasEntityHeader entityHeader : response.getCreatedEntities()) {
        nameIdMap.put((String)entityHeader.getAttribute(NAME), AtlasTypeUtil.getAtlasObjectId(entityHeader));
    }
}
 

private void createDimensionTaggedEntityAndDelete() throws AtlasBaseException {
    AtlasEntity entityToDelete = new AtlasEntity(HIVE_TABLE_TYPE);
    entityToDelete.setAttribute("name", "entity to be deleted");
    entityToDelete.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, "entity.tobedeleted");

    List<AtlasClassification> cls = new ArrayList<>();
    cls.add(new AtlasClassification(DIMENSION_CLASSIFICATION));
    entityToDelete.setClassifications(cls);

    //create entity
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(new AtlasEntity.AtlasEntitiesWithExtInfo(entityToDelete)), false);
    AtlasEntityHeader entityHeader = response.getCreatedEntities().get(0);
    dimensionTagDeleteGuid = entityHeader.getGuid();

    //delete entity
    entityStore.deleteById(dimensionTagDeleteGuid);
}
 

private void createDimensionalTaggedEntityWithAttr() throws AtlasBaseException {
    AtlasEntity entityToDelete = new AtlasEntity(HIVE_TABLE_TYPE);
    entityToDelete.setAttribute("name", "Entity1");
    entityToDelete.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, "entity.one");

    List<AtlasClassification> cls = new ArrayList<>();
    cls.add(new AtlasClassification(DIMENSIONAL_CLASSIFICATION, new HashMap<String, Object>() {{
        put("attr1", "Test");
    }}));
    entityToDelete.setClassifications(cls);

    //create entity
    final EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(new AtlasEntity.AtlasEntitiesWithExtInfo(entityToDelete)), false);
    AtlasEntityHeader entityHeader = response.getCreatedEntities().get(0);
    dimensionalTagGuid = entityHeader.getGuid();

}
 

private EntityMutationResponse createOrUpdate(EntityStream entityStream, boolean isPartialUpdate, boolean replaceClassifications) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> createOrUpdate()");
    }

    if (entityStream == null || !entityStream.hasNext()) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "no entities to create/update.");
    }

    // Create/Update entities
    EntityMutationContext context = preCreateOrUpdate(entityStream, entityGraphMapper, isPartialUpdate);

    EntityMutationResponse ret = entityGraphMapper.mapAttributesAndClassifications(context, isPartialUpdate, replaceClassifications);

    ret.setGuidAssignments(context.getGuidAssignments());

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== createOrUpdate()");
    }

    // Notify the change listeners
    entityChangeNotifier.onEntitiesMutated(ret, entityStream instanceof EntityImportStream);

    return ret;
}
 

public void onEntitiesMutated(EntityMutationResponse entityMutationResponse, boolean isImport) throws AtlasBaseException {
    if (CollectionUtils.isEmpty(entityChangeListeners) || instanceConverter == null) {
        return;
    }

    List<AtlasEntityHeader> createdEntities          = entityMutationResponse.getCreatedEntities();
    List<AtlasEntityHeader> updatedEntities          = entityMutationResponse.getUpdatedEntities();
    List<AtlasEntityHeader> partiallyUpdatedEntities = entityMutationResponse.getPartialUpdatedEntities();
    List<AtlasEntityHeader> deletedEntities          = entityMutationResponse.getDeletedEntities();

    // complete full text mapping before calling toITypedReferenceable(), from notifyListners(), to
    // include all vertex updates in the current graph-transaction
    doFullTextMapping(createdEntities);
    doFullTextMapping(updatedEntities);
    doFullTextMapping(partiallyUpdatedEntities);

    notifyListeners(createdEntities, EntityOperation.CREATE, isImport);
    notifyListeners(updatedEntities, EntityOperation.UPDATE, isImport);
    notifyListeners(partiallyUpdatedEntities, EntityOperation.PARTIAL_UPDATE, isImport);
    notifyListeners(deletedEntities, EntityOperation.DELETE, isImport);
}
 

@Test(dependsOnMethods = "typeCreationFromFile")
public void entityCreationUsingSoftRef() throws IOException, AtlasBaseException {
    final int EXPECTED_ENTITY_COUNT = 6;
    AtlasEntity.AtlasEntityWithExtInfo dbEntity = AtlasType.fromJson(
            TestResourceFileUtils.getJson(RDBMS_DB_FILE), AtlasEntity.AtlasEntityWithExtInfo.class);

    EntityMutationResponse  response = entityStore.createOrUpdate(new AtlasEntityStream(dbEntity), false);

    assertNotNull(response);
    assertTrue(response.getCreatedEntities().size() == EXPECTED_ENTITY_COUNT);
    assertGraphStructure(response.getCreatedEntities().get(0).getGuid(),
            response.getCreatedEntities().get(1).getGuid(), RDBMS_SD_PROPERTY);

    dbGuid = response.getCreatedEntities().get(0).getGuid();
    storageGuid = response.getCreatedEntities().get(1).getGuid();
}
 

@Test(dependsOnMethods = "entityCreationUsingSoftRef")
public void deletetingCollections() throws AtlasBaseException {
    AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = entityStore.getById(dbGuid);

    assertNotNull(entityWithExtInfo);
    List list = (List)entityWithExtInfo.getEntity().getAttribute(RDBMS_DB_TABLES_PROPERTY);
    list.remove(1);

    Map map = (Map) entityWithExtInfo.getEntity().getAttribute(RDBMS_DB_REGIONS_PROPERTY);
    map.remove("east");

    EntityMutationResponse  response = entityStore.createOrUpdate(new AtlasEntityStream(entityWithExtInfo), true);
    assertNotNull(response);
    assertTrue(response.getPartialUpdatedEntities().size() > 0);
    assertAttribute(dbGuid, storageGuid, 1, 1);
}
 

@Test(dependsOnMethods = "testCreate")
public void testMapOfPrimitivesUpdate() throws Exception {
    AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
    AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);

    entitiesInfo.addReferredEntity(tableEntity);

    //Add a new entry
    Map<String, String> paramsMap = (Map<String, String>) tableEntity.getAttribute("parametersMap");
    paramsMap.put("newParam", "value");
    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 1);
    AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //Remove an entry
    paramsMap.remove("key1");
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 1);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
}
 

protected AtlasEntityHeader modifyEntity(AtlasEntity atlasEntity, boolean update) {
    EntityMutationResponse entity = null;
    try {
        if (!update) {
            entity = atlasClientV2.createEntity(new AtlasEntityWithExtInfo(atlasEntity));
            assertNotNull(entity);
            assertNotNull(entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE));
            assertTrue(entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).size() > 0);
            return entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).get(0);
        } else {
            entity = atlasClientV2.updateEntity(new AtlasEntityWithExtInfo(atlasEntity));
            assertNotNull(entity);
            assertNotNull(entity.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
            assertTrue(entity.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).size() > 0);
            return entity.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).get(0);
        }

    } catch (AtlasServiceException e) {
        LOG.error("Entity {} failed", update ? "update" : "creation", entity);
    }
    return null;
}
 

private void createTestEntity() throws Exception {
    AtlasEntity dbEntity = TestUtilsV2.createDBEntity();

    final EntityMutationResponse response = entityREST.createOrUpdate(new AtlasEntitiesWithExtInfo(dbEntity));

    Assert.assertNotNull(response);
    List<AtlasEntityHeader> entitiesMutated = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);

    Assert.assertNotNull(entitiesMutated);
    Assert.assertEquals(entitiesMutated.size(), 1);
    Assert.assertNotNull(entitiesMutated.get(0));
    dbEntity.setGuid(entitiesMutated.get(0).getGuid());

    this.dbEntity = dbEntity;
}
 

public void testReactivation(String tableEntityGuid, int columnCount) throws AtlasBaseException, IOException {
    importSeedData();

    AtlasEntity.AtlasEntityWithExtInfo entity = entityStore.getById(tableEntityGuid);
    EntityMutationResponse response = createColumn(entity.getEntity());
    String columnGuid = response.getCreatedEntities().get(0).getGuid();
    assertNotNull(columnGuid);
    columnCount++;

    entityStore.deleteById(tableEntityGuid);
    entity = entityStore.getById(tableEntityGuid);
    assertEquals(entity.getEntity().getStatus(), AtlasEntity.Status.DELETED);

    importSeedData();

    AtlasEntity atlasEntity = entityStore.getById(tableEntityGuid).getEntity();

    assertEquals(atlasEntity.getStatus(), AtlasEntity.Status.ACTIVE);
    List<AtlasRelatedObjectId> columns = (List<AtlasRelatedObjectId>) atlasEntity.getRelationshipAttribute("columns");
    assertEquals(columns.size(), columnCount);

    int activeColumnCount = 0;
    int deletedColumnCount = 0;
    for (AtlasRelatedObjectId column : columns) {
        if (column.getGuid().equals(columnGuid)){
            assertEquals(column.getEntityStatus(), AtlasEntity.Status.DELETED);
            assertEquals(column.getRelationshipStatus(), AtlasRelationship.Status.DELETED);
            deletedColumnCount++;
        }else{
            assertEquals(column.getEntityStatus(), AtlasEntity.Status.ACTIVE);
            assertEquals(column.getRelationshipStatus(), AtlasRelationship.Status.ACTIVE);
            activeColumnCount++;
        }
    }
    assertEquals(activeColumnCount, --columnCount);
    assertEquals(deletedColumnCount, 1);
}
 

@Test(dependsOnMethods = "deletetingCollections")
public void addingCollections() throws AtlasBaseException {
    AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = entityStore.getById(dbGuid);

    assertNotNull(entityWithExtInfo);
    addNewTables(entityWithExtInfo);
    addNewRegions(entityWithExtInfo);

    EntityMutationResponse  response = entityStore.createOrUpdate(new AtlasEntityStream(entityWithExtInfo), true);
    assertNotNull(response);
    assertTrue(response.getPartialUpdatedEntities().size() > 0);
    assertAttribute(dbGuid, storageGuid, 3, 3);
}
 

@Test
public void testSetObjectIdAttrToNull() throws Exception {
    final AtlasEntity dbEntity  = TestUtilsV2.createDBEntity();
    final AtlasEntity db2Entity = TestUtilsV2.createDBEntity();

    entityStore.createOrUpdate(new AtlasEntityStream(dbEntity), false);
    entityStore.createOrUpdate(new AtlasEntityStream(db2Entity), false);

    final AtlasEntity tableEntity = TestUtilsV2.createTableEntity(dbEntity);

    tableEntity.setAttribute("databaseComposite", AtlasTypeUtil.getAtlasObjectId(db2Entity));

    final EntityMutationResponse tblCreationResponse = entityStore.createOrUpdate(new AtlasEntityStream(tableEntity), false);
    final AtlasEntityHeader      createdTblHeader    = tblCreationResponse.getCreatedEntityByTypeNameAndAttribute(TABLE_TYPE, NAME, (String) tableEntity.getAttribute(NAME));
    final AtlasEntity            createdTblEntity    = getEntityFromStore(createdTblHeader);

    init();

    createdTblEntity.setAttribute("databaseComposite", null);

    final EntityMutationResponse tblUpdateResponse = entityStore.createOrUpdate(new AtlasEntityStream(createdTblEntity), true);
    final AtlasEntityHeader      updatedTblHeader  = tblUpdateResponse.getFirstEntityPartialUpdated();
    final AtlasEntity            updatedTblEntity  = getEntityFromStore(updatedTblHeader);
    final AtlasEntity            deletedDb2Entity  = getEntityFromStore(db2Entity.getGuid());

    assertEquals(deletedDb2Entity.getStatus(), AtlasEntity.Status.DELETED);
}
 

private void addProperty(String guid, String property, Object value) throws AtlasServiceException {

        AtlasEntity entityByGuid = getEntityByGuid(guid);
        entityByGuid.setAttribute(property, value);
        EntityMutationResponse response = atlasClientV2.updateEntity(new AtlasEntityWithExtInfo(entityByGuid));
        assertNotNull(response);
        assertNotNull(response.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
    }
 

@Override
@GraphTransaction
public EntityMutationResponse purgeByIds(Set<String> guids) throws AtlasBaseException {
    if (CollectionUtils.isEmpty(guids)) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "Guid(s) not specified");
    }

    AtlasAuthorizationUtils.verifyAccess(new AtlasAdminAccessRequest(AtlasPrivilege.ADMIN_PURGE), "purge entity: guids=", guids);
    Collection<AtlasVertex> purgeCandidates = new ArrayList<>();

    for (String guid : guids) {
        AtlasVertex vertex = AtlasGraphUtilsV2.findDeletedByGuid(graph, guid);

        if (vertex == null) {
            // Entity does not exist - treat as non-error, since the caller
            // wanted to delete the entity and it's already gone.
            LOG.warn("Purge request ignored for non-existent/active entity with guid " + guid);

            continue;
        }

        purgeCandidates.add(vertex);
    }

    if (purgeCandidates.isEmpty()) {
        LOG.info("No purge candidate entities were found for guids: " + guids + " which is already deleted");
    }

    EntityMutationResponse ret = purgeVertices(purgeCandidates);

    // Notify the change listeners
    entityChangeNotifier.onEntitiesMutated(ret, false);

    return ret;
}
 

@Test
public void testDeleteAndCreate() throws Exception {
    init();
    final AtlasEntity dbEntity = TestUtilsV2.createDBEntity();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(dbEntity), false);

    init();
    //delete entity should mark it as deleted
    EntityMutationResponse deleteResponse = entityStore.deleteById(response.getFirstEntityCreated().getGuid());
    AtlasEntityHeader dbEntityCreated = response.getFirstEntityCreated();
    assertEquals(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).get(0).getGuid(), dbEntityCreated.getGuid());

    //get entity by unique attribute should throw EntityNotFoundException
    try {
        metadataService.getEntityDefinition(TestUtils.DATABASE_TYPE, "name", (String) response.getFirstEntityCreated().getAttribute("name"));
        fail("Expected EntityNotFoundException");
    } catch(EntityNotFoundException e) {
        //expected
    }

    init();
    //Create the same entity again, should create new entity
    AtlasEntity newDBEntity = TestUtilsV2.createDBEntity((String) dbEntity.getAttribute(NAME));
    EntityMutationResponse newCreationResponse = entityStore.createOrUpdate(new AtlasEntityStream(newDBEntity), false);
    assertNotEquals(newCreationResponse.getFirstEntityCreated().getGuid(), response.getFirstEntityCreated().getGuid());

    //get by unique attribute should return the new entity
    ITypedReferenceableInstance instance = metadataService.getEntityDefinitionReference(TestUtils.DATABASE_TYPE, "name", (String) dbEntity.getAttribute("name"));
    assertEquals(instance.getId()._getId(), newCreationResponse.getFirstEntityCreated().getGuid());
}
 

@Test(dependsOnMethods = "testSubmitEntity")
public void testCompleteUpdate() throws Exception {
    final List<AtlasEntity> columns = new ArrayList<>();
    Map<String, Object> values1 = new HashMap<>();
    values1.put("name", "col3");
    values1.put(NAME, "qualifiedName.col3");
    values1.put("type", "string");
    values1.put("comment", "col3 comment");

    Map<String, Object> values2 = new HashMap<>();
    values2.put("name", "col4");
    values2.put(NAME, "qualifiedName.col4");
    values2.put("type", "string");
    values2.put("comment", "col4 comment");

    AtlasEntity colEntity1 = new AtlasEntity(BaseResourceIT.COLUMN_TYPE_V2, values1);
    AtlasEntity colEntity2 = new AtlasEntity(BaseResourceIT.COLUMN_TYPE_V2, values2);
    columns.add(colEntity1);
    columns.add(colEntity2);
    AtlasEntity hiveTable = createHiveTable();
    hiveTable.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));

    AtlasEntityWithExtInfo entityInfo = new AtlasEntityWithExtInfo(hiveTable);
    entityInfo.addReferredEntity(colEntity1);
    entityInfo.addReferredEntity(colEntity2);

    EntityMutationResponse updateEntityResult = atlasClientV2.updateEntity(entityInfo);
    assertNotNull(updateEntityResult);
    assertNotNull(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
    assertNotNull(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE));
    //2 columns are being created, and 1 hiveTable is being updated
    assertEquals(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).size(), 1);
    assertEquals(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).size(), 2);

    AtlasEntity entityByGuid = getEntityByGuid(hiveTable.getGuid());
    List<AtlasObjectId> refs = (List<AtlasObjectId>) entityByGuid.getAttribute("columns");
    assertEquals(refs.size(), 2);
}
 

@Test
public void  testUpdateGetDeleteEntityByUniqueAttribute() throws Exception {
    AtlasEntity            dbEntity = TestUtilsV2.createDBEntity();
    EntityMutationResponse response = entityREST.createOrUpdate(new AtlasEntitiesWithExtInfo(dbEntity));
    String                 dbGuid   = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).get(0).getGuid();

    Assert.assertTrue(AtlasTypeUtil.isAssignedGuid(dbGuid));

    final String prevDBName    = (String) dbEntity.getAttribute(TestUtilsV2.NAME);
    final String updatedDBName = prevDBName + ":updated";

    dbEntity.setAttribute(TestUtilsV2.NAME, updatedDBName);

    response = entityREST.partialUpdateEntityByUniqueAttrs(TestUtilsV2.DATABASE_TYPE, toHttpServletRequest(TestUtilsV2.NAME, prevDBName), new AtlasEntityWithExtInfo(dbEntity));

    Assert.assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.PARTIAL_UPDATE).get(0).getGuid(), dbGuid);

    //Get By unique attribute
    AtlasEntityWithExtInfo entity = entityREST.getByUniqueAttributes(TestUtilsV2.DATABASE_TYPE, toHttpServletRequest(TestUtilsV2.NAME, updatedDBName));
    Assert.assertNotNull(entity);
    Assert.assertNotNull(entity.getEntity().getGuid());
    Assert.assertEquals(entity.getEntity().getGuid(), dbGuid);
    TestEntitiesREST.verifyAttributes(entity.getEntity().getAttributes(), dbEntity.getAttributes());

    final EntityMutationResponse deleteResponse = entityREST.deleteByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, toHttpServletRequest(TestUtilsV2.NAME, (String) dbEntity.getAttribute(TestUtilsV2.NAME)));

    Assert.assertNotNull(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE));
    Assert.assertEquals(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 1);
    Assert.assertEquals(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).get(0).getGuid(), dbGuid);
}