Java Code Examples for org.apache.atlas.model.instance.AtlasEntity#setGuid()

@SuppressWarnings("unchecked")
private void createEntityByNotification(AtlasEntity entity) {
    final String key = toTypedQname(entity);
    final AtlasEntity exEntity = atlasEntitiesByTypedQname.get(key);

    if (exEntity != null) {
        convertReferenceableToObjectId(entity.getAttributes()).forEach((k, v) -> {
            Object r = v;
            final Object exAttr = exEntity.getAttribute(k);
            if (exAttr != null && exAttr instanceof Collection) {
                ((Collection) exAttr).addAll((Collection) v);
                r = exAttr;
            }
            exEntity.setAttribute(k, r);
        });
    } else {
        String guid = String.valueOf(guidSeq.getAndIncrement());
        entity.setGuid(guid);
        atlasEntitiesByTypedQname.put(key, entity);
        atlasEntitiesByGuid.put(guid, entity);
    }
}
 

@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);
}
 

public void prepareForPostUpdate() {
    if (postUpdateEntities != null) {
        ListIterator<AtlasEntity> iter = postUpdateEntities.listIterator();

        while (iter.hasNext()) {
            AtlasEntity entity       = iter.next();
            String      assignedGuid = getAssignedGuid(entity.getGuid());

            // no need to perform partial-update for entities that are created/deleted while processing this message
            if (createdEntities.contains(assignedGuid) || deletedEntities.contains(assignedGuid)) {
                iter.remove();
            } else {
                entity.setGuid(assignedGuid);

                if (entity.getAttributes() != null) {
                    setAssignedGuids(entity.getAttributes().values());
                }

                if (entity.getRelationshipAttributes() != null) {
                    setAssignedGuids(entity.getRelationshipAttributes().values());
                }
            }
        }
    }
}
 

@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);
}
 

private AtlasEntity mapSystemAttributes(AtlasVertex entityVertex, AtlasEntity entity) {
    if (LOG.isDebugEnabled()) {
        LOG.debug("Mapping system attributes for type {}", entity.getTypeName());
    }

    entity.setGuid(GraphHelper.getGuid(entityVertex));
    entity.setTypeName(GraphHelper.getTypeName(entityVertex));
    entity.setStatus(GraphHelper.getStatus(entityVertex));
    entity.setVersion(GraphHelper.getVersion(entityVertex).longValue());

    entity.setCreatedBy(GraphHelper.getCreatedByAsString(entityVertex));
    entity.setUpdatedBy(GraphHelper.getModifiedByAsString(entityVertex));

    entity.setCreateTime(new Date(GraphHelper.getCreatedTime(entityVertex)));
    entity.setUpdateTime(new Date(GraphHelper.getModifiedTime(entityVertex)));

    return entity;
}
 

private AtlasEntity registerNiFiFlowEntity(final NiFiFlow nifiFlow) throws AtlasServiceException {
    final List<AtlasEntity> entities = new ArrayList<>();
    final AtlasEntity.AtlasEntitiesWithExtInfo atlasEntities = new AtlasEntity.AtlasEntitiesWithExtInfo(entities);

    if (!nifiFlow.isMetadataUpdated()) {
        // Nothing has been changed, return existing entity.
        return nifiFlow.getExEntity();
    }

    // Create parent flow entity using existing NiFiFlow entity if available, so that common properties are taken over.
    final AtlasEntity flowEntity = nifiFlow.getExEntity() != null ? new AtlasEntity(nifiFlow.getExEntity()) : new AtlasEntity();
    flowEntity.setTypeName(TYPE_NIFI_FLOW);
    flowEntity.setVersion(1L);
    flowEntity.setAttribute(ATTR_NAME, nifiFlow.getFlowName());
    flowEntity.setAttribute(ATTR_QUALIFIED_NAME, nifiFlow.toQualifiedName(nifiFlow.getRootProcessGroupId()));
    flowEntity.setAttribute(ATTR_URL, nifiFlow.getUrl());
    flowEntity.setAttribute(ATTR_DESCRIPTION, nifiFlow.getDescription());

    // If flowEntity is not persisted yet, then store nifi_flow entity to make nifiFlowId available for other entities.
    if (flowEntity.getGuid().startsWith("-")) {
        entities.add(flowEntity);
        final EntityMutationResponse mutationResponse = atlasClient.createEntities(atlasEntities);
        logger.debug("Registered a new nifi_flow entity, mutation response={}", mutationResponse);
        final String assignedNiFiFlowGuid = mutationResponse.getGuidAssignments().get(flowEntity.getGuid());
        flowEntity.setGuid(assignedNiFiFlowGuid);
        nifiFlow.setAtlasGuid(assignedNiFiFlowGuid);
    }

    return flowEntity;
}
 

@Override
@GraphTransaction
public EntityMutationResponse updateEntity(AtlasObjectId objectId, AtlasEntityWithExtInfo updatedEntityInfo, boolean isPartialUpdate) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> updateEntity({}, {}, {})", objectId, updatedEntityInfo, isPartialUpdate);
    }

    if (objectId == null || updatedEntityInfo == null || updatedEntityInfo.getEntity() == null) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "null entity-id/entity");
    }

    final String guid;

    if (AtlasTypeUtil.isAssignedGuid(objectId.getGuid())) {
        guid = objectId.getGuid();
    } else {
        AtlasEntityType entityType = typeRegistry.getEntityTypeByName(objectId.getTypeName());

        if (entityType == null) {
            throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_TYPENAME, objectId.getTypeName());
        }

        guid = AtlasGraphUtilsV2.getGuidByUniqueAttributes(graph, typeRegistry.getEntityTypeByName(objectId.getTypeName()), objectId.getUniqueAttributes());
    }

    AtlasEntity entity = updatedEntityInfo.getEntity();

    entity.setGuid(guid);

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

@Test
public void recordsEntityEntries() throws AtlasBaseException {
    ByteArrayOutputStream byteOutputStream = new ByteArrayOutputStream();
    ZipSink zs = new ZipSink(byteOutputStream);

    AtlasEntity entity = new AtlasEntity();
    entity.setGuid(String.format(knownEntityGuidFormat, 0));

    zs.add(entity);
    assertTrue(zs.hasEntity(String.format(knownEntityGuidFormat, 0)));

    zs.close();
}
 

@Test
public void recordsEntityEntries() throws AtlasBaseException {
    ByteArrayOutputStream byteOutputStream = new ByteArrayOutputStream();
    ZipSink zs = new ZipSink(byteOutputStream);

    AtlasEntity entity = new AtlasEntity();
    entity.setGuid(String.format(knownEntityGuidFormat, 0));

    zs.add(entity);
    assertTrue(zs.hasEntity(String.format(knownEntityGuidFormat, 0)));

    zs.close();
}
 

@Test
public void testEntityInvalidValue() throws Exception {
    AtlasEntity databaseInstance = new AtlasEntity(DATABASE_TYPE_V2);
    String dbName = randomString();
    String nullString = null;
    String emptyString = "";
    databaseInstance.setAttribute("name", dbName);
    databaseInstance.setAttribute("description", nullString);
    AtlasEntityHeader created = createEntity(databaseInstance);

    // null valid value for required attr - description
    assertNull(created);

    databaseInstance.setAttribute("description", emptyString);
    created = createEntity(databaseInstance);

    // empty string valid value for required attr
    assertNotNull(created);

    databaseInstance.setGuid(created.getGuid());
    databaseInstance.setAttribute("owner", nullString);
    databaseInstance.setAttribute("locationUri", emptyString);

    created = updateEntity(databaseInstance);

    // null/empty string valid value for optional attr
    assertNotNull(created);
}
 

@Test
public void testSubmitEntityWithBadDateFormat() throws Exception {
    AtlasEntity       hiveDBEntity = createHiveDBInstanceV2("db" + randomString());
    AtlasEntityHeader hiveDBHeader = createEntity(hiveDBEntity);
    hiveDBEntity.setGuid(hiveDBHeader.getGuid());

    AtlasEntity tableInstance = createHiveTableInstanceV2(hiveDBEntity, "table" + randomString());
    //Dates with an invalid format are simply nulled out.  This does not produce
    //an error.  See AtlasBuiltInTypes.AtlasDateType.getNormalizedValue().
    tableInstance.setAttribute("lastAccessTime", 1107201407);
    AtlasEntityHeader tableEntityHeader = createEntity(tableInstance);
    assertNotNull(tableEntityHeader);
}
 

private AtlasEntity createHiveDB(String dbName) {
    AtlasEntity hiveDBInstanceV2 = createHiveDBInstanceV2(dbName);
    AtlasEntityHeader entityHeader = createEntity(hiveDBInstanceV2);
    assertNotNull(entityHeader);
    assertNotNull(entityHeader.getGuid());
    hiveDBInstanceV2.setGuid(entityHeader.getGuid());
    return hiveDBInstanceV2;
}
 

private AtlasEntity createHiveTable(AtlasEntity dbInstanceV2, String tableName) throws Exception {
    AtlasEntity hiveTableInstanceV2 = createHiveTableInstanceV2(dbInstanceV2, tableName);
    AtlasEntityHeader createdHeader = createEntity(hiveTableInstanceV2);
    assertNotNull(createdHeader);
    assertNotNull(createdHeader.getGuid());
    hiveTableInstanceV2.setGuid(createdHeader.getGuid());
    tableEntity = hiveTableInstanceV2;
    return hiveTableInstanceV2;
}
 

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;
}
 

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;
}
 

private String getAuditEventDetail(AtlasEntity entity, EntityAuditActionV2 action) {
    Map<String, Object> prunedAttributes = pruneEntityAttributesForAudit(entity);

    String auditPrefix  = getV2AuditPrefix(action);
    String auditString  = auditPrefix + AtlasType.toJson(entity);
    byte[] auditBytes   = auditString.getBytes(StandardCharsets.UTF_8);
    long   auditSize    = auditBytes != null ? auditBytes.length : 0;
    long   auditMaxSize = auditRepository.repositoryMaxSize();

    if (auditMaxSize >= 0 && auditSize > auditMaxSize) { // don't store attributes in audit
        LOG.warn("audit record too long: entityType={}, guid={}, size={}; maxSize={}. entity attribute values not stored in audit",
                entity.getTypeName(), entity.getGuid(), auditSize, auditMaxSize);

        Map<String, Object> attrValues    = entity.getAttributes();
        Map<String, Object> relAttrValues = entity.getRelationshipAttributes();

        entity.setAttributes(null);
        entity.setRelationshipAttributes(null);

        auditString = auditPrefix + AtlasType.toJson(entity);
        auditBytes  = auditString.getBytes(StandardCharsets.UTF_8); // recheck auditString size
        auditSize   = auditBytes != null ? auditBytes.length : 0;

        if (auditMaxSize >= 0 && auditSize > auditMaxSize) { // don't store classifications and meanings as well
            LOG.warn("audit record still too long: entityType={}, guid={}, size={}; maxSize={}. audit will have only summary details",
                    entity.getTypeName(), entity.getGuid(), auditSize, auditMaxSize);

            AtlasEntity shallowEntity = new AtlasEntity();

            shallowEntity.setGuid(entity.getGuid());
            shallowEntity.setTypeName(entity.getTypeName());
            shallowEntity.setCreateTime(entity.getCreateTime());
            shallowEntity.setUpdateTime(entity.getUpdateTime());
            shallowEntity.setCreatedBy(entity.getCreatedBy());
            shallowEntity.setUpdatedBy(entity.getUpdatedBy());
            shallowEntity.setStatus(entity.getStatus());
            shallowEntity.setVersion(entity.getVersion());

            auditString = auditPrefix + AtlasType.toJson(shallowEntity);
        }

        entity.setAttributes(attrValues);
        entity.setRelationshipAttributes(relAttrValues);
    }

    restoreEntityAttributes(entity, prunedAttributes);

    return auditString;
}
 

private EntityMutationContext preCreateOrUpdate(EntityStream entityStream, EntityGraphMapper entityGraphMapper, boolean isPartialUpdate) throws AtlasBaseException {
    MetricRecorder metric = RequestContext.get().startMetricRecord("preCreateOrUpdate");

    EntityGraphDiscovery        graphDiscoverer  = new AtlasEntityGraphDiscoveryV2(graph, typeRegistry, entityStream, entityGraphMapper);
    EntityGraphDiscoveryContext discoveryContext = graphDiscoverer.discoverEntities();
    EntityMutationContext       context          = new EntityMutationContext(discoveryContext);
    RequestContext              requestContext   = RequestContext.get();

    for (String guid : discoveryContext.getReferencedGuids()) {
        AtlasEntity entity = entityStream.getByGuid(guid);

        if (entity != null) { // entity would be null if guid is not in the stream but referenced by an entity in the stream
            AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());

            if (entityType == null) {
                throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
            }

            compactAttributes(entity, entityType);

            AtlasVertex vertex = getResolvedEntityVertex(discoveryContext, entity);

            if (vertex != null) {
                if (!isPartialUpdate) {
                    graphDiscoverer.validateAndNormalize(entity);

                    // change entity 'isInComplete' to 'false' during full update
                    if (isEntityIncomplete(vertex)) {
                        vertex.removeProperty(IS_INCOMPLETE_PROPERTY_KEY);

                        entity.setIsIncomplete(FALSE);
                    }
                } else {
                    graphDiscoverer.validateAndNormalizeForUpdate(entity);
                }

                String guidVertex = AtlasGraphUtilsV2.getIdFromVertex(vertex);

                if (!StringUtils.equals(guidVertex, guid)) { // if entity was found by unique attribute
                    entity.setGuid(guidVertex);

                    requestContext.recordEntityGuidUpdate(entity, guid);
                }

                context.addUpdated(guid, entity, entityType, vertex);
            } else {
                graphDiscoverer.validateAndNormalize(entity);

                //Create vertices which do not exist in the repository
                if (RequestContext.get().isImportInProgress() && AtlasTypeUtil.isAssignedGuid(entity.getGuid())) {
                    vertex = entityGraphMapper.createVertexWithGuid(entity, entity.getGuid());
                } else {
                     vertex = entityGraphMapper.createVertex(entity);
                }

                discoveryContext.addResolvedGuid(guid, vertex);

                discoveryContext.addResolvedIdByUniqAttribs(getAtlasObjectId(entity), vertex);

                String generatedGuid = AtlasGraphUtilsV2.getIdFromVertex(vertex);

                entity.setGuid(generatedGuid);

                requestContext.recordEntityGuidUpdate(entity, guid);

                context.addCreated(guid, entity, entityType, vertex);
            }

            // during import, update the system attributes
            if (RequestContext.get().isImportInProgress()) {
                Status newStatus = entity.getStatus();

                if (newStatus != null) {
                    Status currStatus = AtlasGraphUtilsV2.getState(vertex);

                    if (currStatus == Status.ACTIVE && newStatus == Status.DELETED) {
                        if (LOG.isDebugEnabled()) {
                            LOG.debug("entity-delete via import - guid={}", guid);
                        }

                        context.addEntityToDelete(vertex);
                    } else if (currStatus == Status.DELETED && newStatus == Status.ACTIVE) {
                        LOG.warn("Import is attempting to activate deleted entity (guid={}).", guid);
                        entityGraphMapper.importActivateEntity(vertex, entity);
                        context.addCreated(guid, entity, entityType, vertex);
                    }
                }

                entityGraphMapper.updateSystemAttributes(vertex, entity);
            }
        }
    }

    RequestContext.get().endMetricRecord(metric);

    return context;
}
 

@Test
public void testUpdateFlowPath() throws AtlasServiceException {
    final NotificationSender sender = new NotificationSender();
    final Referenceable fileC = createRef("fs_path", "/tmp/file-c.txt@test");
    final Referenceable fileD = createRef("fs_path", "/tmp/file-d.txt@test");

    // New in/out fileC and fileD are found for path1.
    final Referenceable newPath1Lineage = createRef(TYPE_NIFI_FLOW_PATH, "path1@test");
    newPath1Lineage.set(ATTR_INPUTS, singleton(fileC));
    newPath1Lineage.set(ATTR_OUTPUTS, singleton(fileD));

    final List<HookNotification> messages = asList(
            new HookNotificationV1.EntityCreateRequest(NIFI_USER, fileC),
            new HookNotificationV1.EntityCreateRequest(NIFI_USER, fileD),
            new HookNotificationV1.EntityPartialUpdateRequest(NIFI_USER, TYPE_NIFI_FLOW_PATH, ATTR_QUALIFIED_NAME, "path1@test", newPath1Lineage)
    );

    final NiFiAtlasClient atlasClient = mock(NiFiAtlasClient.class);
    sender.setAtlasClient(atlasClient);

    // Existing nifi_flow_path
    final AtlasEntity path1Entity = new AtlasEntity(TYPE_NIFI_FLOW_PATH, ATTR_QUALIFIED_NAME, "path1@test");
    path1Entity.setGuid("path1-guid");
    path1Entity.setAttribute(ATTR_INPUTS, singleton(createGuidReference("fs_path", "fileA-guid")));
    path1Entity.setAttribute(ATTR_OUTPUTS, singleton(createGuidReference("fs_path", "fileB-guid")));

    final AtlasEntity fileAEntity = new AtlasEntity("fs_path", ATTR_QUALIFIED_NAME, "file-a.txt@test");
    fileAEntity.setGuid("fileA-guid");

    final AtlasEntity fileBEntity = new AtlasEntity("fs_path", ATTR_QUALIFIED_NAME, "file-b.txt@test");
    fileBEntity.setGuid("fileA-guid");

    final AtlasEntity.AtlasEntityWithExtInfo path1Ext = new AtlasEntity.AtlasEntityWithExtInfo(path1Entity);
    final AtlasEntity.AtlasEntityWithExtInfo fileAExt = new AtlasEntity.AtlasEntityWithExtInfo(fileAEntity);
    final AtlasEntity.AtlasEntityWithExtInfo fileBExt = new AtlasEntity.AtlasEntityWithExtInfo(fileBEntity);
    when(atlasClient.searchEntityDef(eq(new AtlasObjectId(TYPE_NIFI_FLOW_PATH, ATTR_QUALIFIED_NAME,"path1@test")))).thenReturn(path1Ext);
    when(atlasClient.searchEntityDef(eq(new AtlasObjectId("fileA-guid")))).thenReturn(fileAExt);
    when(atlasClient.searchEntityDef(eq(new AtlasObjectId("fileB-guid")))).thenReturn(fileBExt);

    final Notifier notifier = new Notifier();
    sender.send(messages, notifier);

    assertCreateMessage(notifier, 0, fileC, fileD);
    final Referenceable updatedPath1 = createRef(TYPE_NIFI_FLOW_PATH, "path1@test");
    updatedPath1.set(ATTR_INPUTS, asList(new Referenceable("fileA-guid", "fs_path", Collections.emptyMap()), fileC));
    updatedPath1.set(ATTR_OUTPUTS, asList(new Referenceable("fileB-guid", "fs_path", Collections.emptyMap()), fileD));
    assertUpdateFlowPathMessage(notifier, 1, updatedPath1);
}
 

@Test
public void testPartialUpdateArrayAttr() throws Exception {
    // Create a table entity, with 3 reference column entities
    init();
    final AtlasEntity      dbEntity           = TestUtilsV2.createDBEntity();
    EntityMutationResponse dbCreationResponse = entityStore.createOrUpdate(new AtlasEntityStream(dbEntity), false);

    final AtlasEntity        tableEntity      = TestUtilsV2.createTableEntity(dbEntity);
    AtlasEntitiesWithExtInfo entitiesInfo     = new AtlasEntitiesWithExtInfo(tableEntity);

    final AtlasEntity columnEntity1 = TestUtilsV2.createColumnEntity(tableEntity);
    columnEntity1.setAttribute("description", "desc for col1");
    entitiesInfo.addReferredEntity(columnEntity1);

    final AtlasEntity columnEntity2 = TestUtilsV2.createColumnEntity(tableEntity);
    columnEntity2.setAttribute("description", "desc for col2");
    entitiesInfo.addReferredEntity(columnEntity2);

    final AtlasEntity columnEntity3 = TestUtilsV2.createColumnEntity(tableEntity);
    columnEntity3.setAttribute("description", "desc for col3");
    entitiesInfo.addReferredEntity(columnEntity3);

    tableEntity.setAttribute(COLUMNS_ATTR_NAME, Arrays.asList(AtlasTypeUtil.getAtlasObjectId(columnEntity1),
                                                              AtlasTypeUtil.getAtlasObjectId(columnEntity2),
                                                              AtlasTypeUtil.getAtlasObjectId(columnEntity3)));

    init();

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

    final AtlasEntityHeader column1Created = tblCreationResponse.getCreatedEntityByTypeNameAndAttribute(COLUMN_TYPE, NAME, (String) columnEntity1.getAttribute(NAME));
    final AtlasEntityHeader column2Created = tblCreationResponse.getCreatedEntityByTypeNameAndAttribute(COLUMN_TYPE, NAME, (String) columnEntity2.getAttribute(NAME));
    final AtlasEntityHeader column3Created = tblCreationResponse.getCreatedEntityByTypeNameAndAttribute(COLUMN_TYPE, NAME, (String) columnEntity3.getAttribute(NAME));

    // update only description attribute of all 3 columns
    AtlasEntity col1 = new AtlasEntity(COLUMN_TYPE);
    col1.setGuid(column1Created.getGuid());
    col1.setAttribute("description", "desc for col1:updated");

    AtlasEntity col2 = new AtlasEntity(COLUMN_TYPE);
    col2.setGuid(column2Created.getGuid());
    col2.setAttribute("description", "desc for col2:updated");

    AtlasEntity col3 = new AtlasEntity(COLUMN_TYPE);
    col3.setGuid(column3Created.getGuid());
    col3.setAttribute("description", "desc for col3:updated");

    final AtlasEntity tableEntity1 = new AtlasEntity(TABLE_TYPE);
    tableEntity1.setGuid(createdTblHeader.getGuid());
    tableEntity1.setAttribute(COLUMNS_ATTR_NAME, Arrays.asList(AtlasTypeUtil.getAtlasObjectId(col1),
            AtlasTypeUtil.getAtlasObjectId(col2),
            AtlasTypeUtil.getAtlasObjectId(col3)));
    AtlasEntitiesWithExtInfo tableInfo = new AtlasEntitiesWithExtInfo(tableEntity1);
    tableInfo.addReferredEntity(col1.getGuid(), col1);
    tableInfo.addReferredEntity(col2.getGuid(), col2);
    tableInfo.addReferredEntity(col3.getGuid(), col3);

    init();

    final EntityMutationResponse tblUpdateResponse = entityStore.createOrUpdate(new AtlasEntityStream(tableInfo), true);
    final AtlasEntityHeader      updatedTblHeader  = tblUpdateResponse.getFirstEntityPartialUpdated();
    final AtlasEntity            updatedTblEntity2 = getEntityFromStore(updatedTblHeader);
    List<AtlasEntityHeader>      updatedColHeaders = tblUpdateResponse.getPartialUpdatedEntitiesByTypeName(COLUMN_TYPE);

    final AtlasEntity updatedCol1Entity = getEntityFromStore(updatedColHeaders.get(0));
    final AtlasEntity updatedCol2Entity = getEntityFromStore(updatedColHeaders.get(1));
    final AtlasEntity updatedCol3Entity = getEntityFromStore(updatedColHeaders.get(2));

    assertEquals(col1.getAttribute("description"), updatedCol1Entity.getAttribute("description"));
    assertEquals(col2.getAttribute("description"), updatedCol2Entity.getAttribute("description"));
    assertEquals(col3.getAttribute("description"), updatedCol3Entity.getAttribute("description"));
}
 

private AtlasEntity getTopicEntityWithGuid(String guid) {
    AtlasEntity ret = new AtlasEntity(KafkaDataTypes.KAFKA_TOPIC.getName());
    ret.setGuid(guid);
    return ret;
}