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

AtlasEntity createDatabase(String name, String description, String owner, String locationUri, String... classificationNames) throws Exception {
    AtlasEntity entity = new AtlasEntity(DATABASE_TYPE);

    // set attributes
    entity.setAttribute("name", name);
    entity.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, name + CLUSTER_SUFFIX);
    entity.setAttribute("description", description);
    entity.setAttribute("owner", owner);
    entity.setAttribute("locationuri", locationUri);
    entity.setAttribute("createTime", System.currentTimeMillis());

    // set classifications
    entity.setClassifications(toAtlasClassifications(classificationNames));

    return createInstance(entity);
}
 

@Test(dependsOnMethods = "testCreate")
public void testClassUpdate() throws Exception {

    init();
    //Create new db instance
    final AtlasEntity databaseInstance = TestUtilsV2.createDBEntity();

    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(databaseInstance), false);
    final AtlasEntityHeader dbCreated = response.getFirstCreatedEntityByTypeName(TestUtilsV2.DATABASE_TYPE);

    init();
    Map<String, AtlasEntity> tableCloneMap = new HashMap<>();
    AtlasEntity tableClone = new AtlasEntity(tblEntity.getEntity());
    tableClone.setAttribute("database", new AtlasObjectId(dbCreated.getGuid(), TestUtilsV2.DATABASE_TYPE));

    tableCloneMap.put(dbCreated.getGuid(), databaseInstance);
    tableCloneMap.put(tableClone.getGuid(), tableClone);

    response = entityStore.createOrUpdate(new InMemoryMapEntityStream(tableCloneMap), false);
    final AtlasEntityHeader tableDefinition = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    AtlasEntity updatedTableDefinition = getEntityFromStore(tableDefinition);
    assertNotNull(updatedTableDefinition.getAttribute("database"));
    Assert.assertEquals(((AtlasObjectId) updatedTableDefinition.getAttribute("database")).getGuid(), dbCreated.getGuid());
}
 

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();

}
 

AtlasEntity createTable(String name, String description, AtlasEntity db, AtlasEntity sd, String owner, String tableType,
                        List<AtlasEntity> columns, String... traitNames) throws Exception {
    AtlasEntity entity = new AtlasEntity(TABLE_TYPE);

    entity.setClassifications(toAtlasClassifications(traitNames));
    entity.setAttribute("name", name);
    entity.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
    entity.setAttribute("description", description);
    entity.setAttribute("owner", owner);
    entity.setAttribute("tableType", tableType);
    entity.setAttribute("createTime", System.currentTimeMillis());
    entity.setAttribute("lastAccessTime", System.currentTimeMillis());
    entity.setAttribute("retention", System.currentTimeMillis());
    entity.setAttribute("db", AtlasTypeUtil.getAtlasObjectId(db));
    entity.setAttribute("sd", AtlasTypeUtil.getAtlasObjectId(sd));
    entity.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));

    return createInstance(entity, traitNames);
}
 

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

    AtlasStruct serdeInstance = new AtlasStruct(TestUtilsV2.SERDE_TYPE, NAME, "serde1Name");
    serdeInstance.setAttribute("serde", "test");
    serdeInstance.setAttribute("description", "testDesc");
    tableEntity.setAttribute("serde1", serdeInstance);
    tableEntity.setAttribute("database", new AtlasObjectId(databaseEntity.getTypeName(), NAME, databaseEntity.getAttribute(NAME)));

    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //update struct attribute
    serdeInstance.setAttribute("serde", "testUpdated");
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //set to null
    tableEntity.setAttribute("description", null);
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    Assert.assertNull(updatedTable.getAttribute("description"));
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
}
 

AtlasEntity createStorageDescriptor(String location, String inputFormat, String outputFormat, boolean compressed)
        throws Exception {
    AtlasEntity entity = new AtlasEntity(STORAGE_DESC_TYPE);

    entity.setAttribute("location", location);
    entity.setAttribute("inputFormat", inputFormat);
    entity.setAttribute("outputFormat", outputFormat);
    entity.setAttribute("compressed", compressed);

    return createInstance(entity, null);
}
 

public static AtlasEntity createDBEntity(String dbName) {
    AtlasEntity entity = new AtlasEntity(DATABASE_TYPE);
    entity.setAttribute(NAME, dbName);
    entity.setAttribute("description", "us db");

    return entity;
}
 

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

private void processDynamicAttributes(AtlasEntity entity) throws AtlasBaseException {
    AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());

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

    for (AtlasAttribute attribute : entityType.getDynEvalAttributes()) {
        String              attributeName   = attribute.getName();
        List<TemplateToken> tokens          = entityType.getParsedTemplates().get(attributeName);

        if (tokens == null) {
            continue;
        }

        StringBuilder dynAttributeValue = new StringBuilder();

        boolean set = true;

        for (TemplateToken token : tokens) {
            String evaluated = token.eval(entity);
            if (evaluated != null) {
                dynAttributeValue.append(evaluated);
            } else {
                set = false;
                LOG.warn("Attribute {} for {} unable to be generated because of dynamic attribute token {}", attributeName, entityType, token.getValue());
                break;
            }

        }

        if (set) {
            entity.setAttribute(attributeName,dynAttributeValue.toString());
        }
    }
}
 

public void createDummyEntity(String name, String type, String... traitNames) throws AtlasBaseException {
    AtlasEntity entity = new AtlasEntity(type);
    entity.setAttribute("name", name);
    entity.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
    entity.setClassifications(Stream.of(traitNames).map(AtlasClassification::new).collect(Collectors.toList()));
    entityStore.createOrUpdate(new AtlasEntityStream(new AtlasEntity.AtlasEntitiesWithExtInfo(entity)), false);
}
 

@Override
public AtlasEntity toEntity(final AtlasGlossaryTerm obj) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasGlossaryTermDTO.toEntity()", obj);
    }
    Objects.requireNonNull(obj, "atlasGlossaryTerm");
    Objects.requireNonNull(obj.getQualifiedName(), "atlasGlossaryTerm qualifiedName must be specified");
    Objects.requireNonNull(obj.getAnchor(), "atlasGlossaryTerm anchor must be specified");

    AtlasEntity ret = getDefaultAtlasEntity(obj);

    ret.setAttribute("qualifiedName", obj.getQualifiedName());
    ret.setAttribute("name", obj.getName());
    ret.setAttribute("shortDescription", obj.getShortDescription());
    ret.setAttribute("longDescription", obj.getLongDescription());
    ret.setAttribute("examples", obj.getExamples());
    ret.setAttribute("abbreviation", obj.getAbbreviation());
    ret.setAttribute("usage", obj.getUsage());
    ret.setAttribute("anchor", new AtlasObjectId(obj.getAnchor().getGlossaryGuid()));
    ret.setAttribute("additionalAttributes", obj.getAdditionalAttributes());

    if (CollectionUtils.isNotEmpty(obj.getClassifications())) {
        if (LOG.isDebugEnabled()) {
            LOG.debug("Processing term classifications");
        }
        ret.setClassifications(obj.getClassifications());
    }

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasGlossaryTermDTO.toEntity() : {}", ret);
    }
    return ret;
}
 

@Test
public void testUpdateWithSerializedEntities() throws  Exception {
    //Check with serialization and deserialization of entity attributes for the case
    // where attributes which are de-serialized into a map
    AtlasEntity dbEntity    = TestUtilsV2.createDBEntity();
    AtlasEntity tableEntity = TestUtilsV2.createTableEntity(dbEntity);

    final AtlasEntity colEntity = TestUtilsV2.createColumnEntity(tableEntity);
    List<AtlasEntity> columns = new ArrayList<AtlasEntity>() {{ add(colEntity); }};
    tableEntity.setAttribute("columns", getObjIdList(columns));

    AtlasEntity newDBEntity = serDeserEntity(dbEntity);
    AtlasEntity newTableEntity = serDeserEntity(tableEntity);

    AtlasEntitiesWithExtInfo newEntities = new AtlasEntitiesWithExtInfo();
    newEntities.addEntity(newDBEntity);
    newEntities.addEntity(newTableEntity);
    for (AtlasEntity column : columns) {
        newEntities.addReferredEntity(serDeserEntity(column));
    }

    EntityMutationResponse response2 = entityREST.createOrUpdate(newEntities);

    List<AtlasEntityHeader> newGuids = response2.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);
    Assert.assertNotNull(newGuids);
    Assert.assertEquals(newGuids.size(), 3);
}
 

private AtlasEntity toStorageDescEntity(StorageDescriptor storageDesc, String tableQualifiedName, String sdQualifiedName, AtlasObjectId tableId ) throws AtlasHookException {
    AtlasEntity ret = new AtlasEntity(HiveDataTypes.HIVE_STORAGEDESC.getName());

    ret.setRelationshipAttribute(ATTRIBUTE_TABLE, AtlasTypeUtil.getAtlasRelatedObjectId(tableId, RELATIONSHIP_HIVE_TABLE_STORAGE_DESC));
    ret.setAttribute(ATTRIBUTE_QUALIFIED_NAME, sdQualifiedName);
    ret.setAttribute(ATTRIBUTE_PARAMETERS, storageDesc.getParameters());
    ret.setAttribute(ATTRIBUTE_LOCATION, HdfsNameServiceResolver.getPathWithNameServiceID(storageDesc.getLocation()));
    ret.setAttribute(ATTRIBUTE_INPUT_FORMAT, storageDesc.getInputFormat());
    ret.setAttribute(ATTRIBUTE_OUTPUT_FORMAT, storageDesc.getOutputFormat());
    ret.setAttribute(ATTRIBUTE_COMPRESSED, storageDesc.isCompressed());
    ret.setAttribute(ATTRIBUTE_NUM_BUCKETS, storageDesc.getNumBuckets());
    ret.setAttribute(ATTRIBUTE_STORED_AS_SUB_DIRECTORIES, storageDesc.isStoredAsSubDirectories());

    if (storageDesc.getBucketCols().size() > 0) {
        ret.setAttribute(ATTRIBUTE_BUCKET_COLS, storageDesc.getBucketCols());
    }

    if (storageDesc.getSerdeInfo() != null) {
        SerDeInfo serdeInfo = storageDesc.getSerdeInfo();

        LOG.debug("serdeInfo = {}", serdeInfo);
        // SkewedInfo skewedInfo = storageDesc.getSkewedInfo();

        AtlasStruct serdeInfoStruct = new AtlasStruct(HiveDataTypes.HIVE_SERDE.getName());

        serdeInfoStruct.setAttribute(ATTRIBUTE_NAME, serdeInfo.getName());
        serdeInfoStruct.setAttribute(ATTRIBUTE_SERIALIZATION_LIB, serdeInfo.getSerializationLib());
        serdeInfoStruct.setAttribute(ATTRIBUTE_PARAMETERS, serdeInfo.getParameters());

        ret.setAttribute(ATTRIBUTE_SERDE_INFO, serdeInfoStruct);
    }

    if (CollectionUtils.isNotEmpty(storageDesc.getSortCols())) {
        List<AtlasStruct> sortColsStruct = new ArrayList<>();

        for (Order sortcol : storageDesc.getSortCols()) {
            String hiveOrderName = HiveDataTypes.HIVE_ORDER.getName();
            AtlasStruct colStruct = new AtlasStruct(hiveOrderName);
            colStruct.setAttribute("col", sortcol.getCol());
            colStruct.setAttribute("order", sortcol.getOrder());

            sortColsStruct.add(colStruct);
        }

        ret.setAttribute(ATTRIBUTE_SORT_COLS, sortColsStruct);
    }

    return ret;
}
 

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

@Test(enabled = false)
    //Titan doesn't allow some reserved chars in property keys. Verify that atlas encodes these
    //See GraphHelper.encodePropertyKey()
    //TODO : Failing in typedef creation
    public void testSpecialCharacters() throws Exception {
        //Verify that type can be created with reserved characters in typename, attribute name
        final String typeName = "test_type_"+ RandomStringUtils.randomAlphanumeric(10);
        String strAttrName = randomStrWithReservedChars();
        String arrayAttrName = randomStrWithReservedChars();
        String mapAttrName = randomStrWithReservedChars();

        AtlasEntityDef typeDefinition =
            AtlasTypeUtil.createClassTypeDef(typeName, "Special chars test type", ImmutableSet.<String>of(),
                AtlasTypeUtil.createOptionalAttrDef(strAttrName, "string"),
                AtlasTypeUtil.createOptionalAttrDef(arrayAttrName, "array<string>"),
                AtlasTypeUtil.createOptionalAttrDef(mapAttrName, "map<string,string>"));

        AtlasTypesDef atlasTypesDef = new AtlasTypesDef(null, null, null, Arrays.asList(typeDefinition));
        typeDefStore.createTypesDef(atlasTypesDef);

        //verify that entity can be created with reserved characters in string value, array value and map key and value
        AtlasEntity entity = new AtlasEntity();
        entity.setAttribute(strAttrName, randomStrWithReservedChars());
        entity.setAttribute(arrayAttrName, new String[]{randomStrWithReservedChars()});
        entity.setAttribute(mapAttrName, new HashMap<String, String>() {{
            put(randomStrWithReservedChars(), randomStrWithReservedChars());
        }});

        AtlasEntityWithExtInfo entityWithExtInfo = new AtlasEntityWithExtInfo(entity);

        final EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entityWithExtInfo), false);
        final AtlasEntityHeader firstEntityCreated = response.getFirstEntityCreated();
        validateEntity(entityWithExtInfo, getEntityFromStore(firstEntityCreated));


        //Verify that search with reserved characters works - for string attribute
//        String query =
//            String.format("`%s` where `%s` = '%s'", typeName, strAttrName, entity.getAttribute(strAttrName));
//        String responseJson = discoveryService.searchByDSL(query, new QueryParams(1, 0));
//        JSONObject response = new JSONObject(responseJson);
//        assertEquals(response.getJSONArray("rows").length(), 1);
    }
 

@Test
public void testInverseReferenceAutoUpdate_NonCompositeManyToOne() throws Exception {
    AtlasEntityType bType = typeRegistry.getEntityTypeByName("B");
    AtlasEntity a1 = new AtlasEntity("A");
    a1.setAttribute(NAME, TestUtils.randomString());
    AtlasEntity a2 = new AtlasEntity("A");
    a2.setAttribute(NAME, TestUtils.randomString());
    AtlasEntity a3 = new AtlasEntity("A");
    a3.setAttribute(NAME, TestUtils.randomString());
    AtlasEntity b = new AtlasEntity("B");

    b.setAttribute(NAME, TestUtils.randomString());
    AtlasEntitiesWithExtInfo atlasEntitiesWithExtInfo = new AtlasEntitiesWithExtInfo();
    atlasEntitiesWithExtInfo.addEntity(a1);
    atlasEntitiesWithExtInfo.addEntity(a2);
    atlasEntitiesWithExtInfo.addEntity(a3);
    atlasEntitiesWithExtInfo.addEntity(b);
    AtlasEntityStream entityStream = new AtlasEntityStream(atlasEntitiesWithExtInfo);
    EntityMutationResponse response = entityStore.createOrUpdate(entityStream , false);

    AtlasEntity bForPartialUpdate = new AtlasEntity("B");
    bForPartialUpdate.setAttribute("manyA", ImmutableList.of(AtlasTypeUtil.getAtlasObjectId(a1), AtlasTypeUtil.getAtlasObjectId(a2)));
    init();
    response = entityStore.updateByUniqueAttributes(bType, Collections.<String, Object>singletonMap(NAME, b.getAttribute(NAME)), new AtlasEntityWithExtInfo(bForPartialUpdate));
    List<AtlasEntityHeader> partialUpdatedEntities = response.getPartialUpdatedEntities();
    // Verify 3 entities were updated:
    // * set b.manyA reference to a1 and a2
    // * set inverse a1.oneB reference to b
    // * set inverse a2.oneB reference to b
    assertEquals(partialUpdatedEntities.size(), 3);
    AtlasEntitiesWithExtInfo storedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), b.getGuid()));
    AtlasEntity storedEntity = storedEntities.getEntity(a1.getGuid());
    verifyReferenceValue(storedEntity, "oneB", b.getGuid());

    storedEntity = storedEntities.getEntity(a2.getGuid());
    verifyReferenceValue(storedEntity, "oneB", b.getGuid());

    storedEntity = storedEntities.getEntity(b.getGuid());
    verifyReferenceList(storedEntity, "manyA", ImmutableList.of(AtlasTypeUtil.getAtlasObjectId(a1), AtlasTypeUtil.getAtlasObjectId(a2)));

    bForPartialUpdate.setAttribute("manyA", ImmutableList.of(AtlasTypeUtil.getAtlasObjectId(a3)));
    init();
    response = entityStore.updateByUniqueAttributes(bType, Collections.<String, Object>singletonMap(NAME, b.getAttribute(NAME)), new AtlasEntityWithExtInfo(bForPartialUpdate));
    partialUpdatedEntities = response.getPartialUpdatedEntities();
    // Verify 4 entities were updated:
    // * set b.manyA reference to a3
    // * set inverse a3.oneB reference to b
    // * disconnect inverse a1.oneB reference to b
    // * disconnect inverse a2.oneB reference to b
    assertEquals(partialUpdatedEntities.size(), 4);

    init();
    storedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), a3.getGuid(), b.getGuid()));
    verifyReferenceValue(storedEntities.getEntity(a3.getGuid()), "oneB", b.getGuid());

    verify_testInverseReferenceAutoUpdate_NonCompositeManyToOne(storedEntities.getEntity(a1.getGuid()), storedEntities.getEntity(a2.getGuid()),
        storedEntities.getEntity(a3.getGuid()), storedEntities.getEntity(b.getGuid()));
}
 

@Test
public void testInverseReferenceAutoUpdate_NonComposite_OneToOne() throws Exception {
    AtlasEntityType bType = typeRegistry.getEntityTypeByName("B");
    AtlasEntity a1 = new AtlasEntity("A");
    a1.setAttribute(NAME, TestUtils.randomString());
    AtlasEntity a2 = new AtlasEntity("A");
    a2.setAttribute(NAME, TestUtils.randomString());
    AtlasEntity b = new AtlasEntity("B");
    b.setAttribute(NAME, TestUtils.randomString());
    AtlasEntitiesWithExtInfo atlasEntitiesWithExtInfo = new AtlasEntitiesWithExtInfo();
    atlasEntitiesWithExtInfo.addEntity(a1);
    atlasEntitiesWithExtInfo.addEntity(a2);
    atlasEntitiesWithExtInfo.addEntity(b);
    AtlasEntityStream entityStream = new AtlasEntityStream(atlasEntitiesWithExtInfo);
    EntityMutationResponse response = entityStore.createOrUpdate(entityStream , false);

    AtlasEntity bForPartialUpdate = new AtlasEntity("B");
    bForPartialUpdate.setAttribute("a", AtlasTypeUtil.getAtlasObjectId(a1));
    init();
    response = entityStore.updateByUniqueAttributes(bType, Collections.<String, Object>singletonMap(NAME, b.getAttribute(NAME)), new AtlasEntityWithExtInfo(bForPartialUpdate));
    List<AtlasEntityHeader> partialUpdatedEntities = response.getPartialUpdatedEntities();
    // Verify 2 entities were updated:
    // * set b.a reference to a1
    // * set inverse a1.b reference to b
    assertEquals(partialUpdatedEntities.size(), 2);
    AtlasEntitiesWithExtInfo storedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), b.getGuid()));
    AtlasEntity storedEntity = storedEntities.getEntity(a1.getGuid());
    verifyReferenceValue(storedEntity, "b", b.getGuid());
    storedEntity = storedEntities.getEntity(b.getGuid());
    verifyReferenceValue(storedEntity, "a", a1.getGuid());

    // Update b.a to reference a2.
    bForPartialUpdate.setAttribute("a", AtlasTypeUtil.getAtlasObjectId(a2));
    init();
    response = entityStore.updateByUniqueAttributes(bType, Collections.<String, Object>singletonMap(NAME, b.getAttribute(NAME)), new AtlasEntityWithExtInfo(bForPartialUpdate));
    partialUpdatedEntities = response.getPartialUpdatedEntities();
    // Verify 3 entities were updated:
    // * set b.a reference to a2
    // * set a2.b reference to b
    // * disconnect a1.b reference
    assertEquals(partialUpdatedEntities.size(), 3);
    storedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), b.getGuid()));
    storedEntity = storedEntities.getEntity(a2.getGuid());
    verifyReferenceValue(storedEntity, "b", b.getGuid());
    storedEntity = storedEntities.getEntity(b.getGuid());
    verifyReferenceValue(storedEntity, "a", a2.getGuid());
    storedEntity = storedEntities.getEntity(a1.getGuid());
    Object refValue = storedEntity.getAttribute("b");
    verify_testInverseReferenceAutoUpdate_NonComposite_OneToOne(storedEntities.getEntity(a1.getGuid()), storedEntities.getEntity(b.getGuid()));
}
 

private AtlasEntityWithExtInfo toTableEntity(AtlasEntity database, final Table hiveTable, AtlasEntityWithExtInfo table) throws AtlasHookException {
    if (table == null) {
        table = new AtlasEntityWithExtInfo(new AtlasEntity(HiveDataTypes.HIVE_TABLE.getName()));
    }

    AtlasEntity tableEntity        = table.getEntity();
    String      tableQualifiedName = getTableQualifiedName(metadataNamespace, hiveTable);
    long        createTime         = BaseHiveEvent.getTableCreateTime(hiveTable);
    long        lastAccessTime     = hiveTable.getLastAccessTime() > 0 ? hiveTable.getLastAccessTime() : createTime;

    tableEntity.setRelationshipAttribute(ATTRIBUTE_DB, AtlasTypeUtil.getAtlasRelatedObjectId(database, RELATIONSHIP_HIVE_TABLE_DB));
    tableEntity.setAttribute(ATTRIBUTE_QUALIFIED_NAME, tableQualifiedName);
    tableEntity.setAttribute(ATTRIBUTE_NAME, hiveTable.getTableName().toLowerCase());
    tableEntity.setAttribute(ATTRIBUTE_OWNER, hiveTable.getOwner());

    tableEntity.setAttribute(ATTRIBUTE_CREATE_TIME, createTime);
    tableEntity.setAttribute(ATTRIBUTE_LAST_ACCESS_TIME, lastAccessTime);
    tableEntity.setAttribute(ATTRIBUTE_RETENTION, hiveTable.getRetention());
    tableEntity.setAttribute(ATTRIBUTE_PARAMETERS, hiveTable.getParameters());
    tableEntity.setAttribute(ATTRIBUTE_COMMENT, hiveTable.getParameters().get(ATTRIBUTE_COMMENT));
    tableEntity.setAttribute(ATTRIBUTE_TABLE_TYPE, hiveTable.getTableType().name());
    tableEntity.setAttribute(ATTRIBUTE_TEMPORARY, hiveTable.isTemporary());

    if (hiveTable.getViewOriginalText() != null) {
        tableEntity.setAttribute(ATTRIBUTE_VIEW_ORIGINAL_TEXT, hiveTable.getViewOriginalText());
    }

    if (hiveTable.getViewExpandedText() != null) {
        tableEntity.setAttribute(ATTRIBUTE_VIEW_EXPANDED_TEXT, hiveTable.getViewExpandedText());
    }

    AtlasEntity       sdEntity = toStorageDescEntity(hiveTable.getSd(), tableQualifiedName, getStorageDescQFName(tableQualifiedName), AtlasTypeUtil.getObjectId(tableEntity));
    List<AtlasEntity> partKeys = toColumns(hiveTable.getPartitionKeys(), tableEntity, RELATIONSHIP_HIVE_TABLE_PART_KEYS);
    List<AtlasEntity> columns  = toColumns(hiveTable.getCols(), tableEntity, RELATIONSHIP_HIVE_TABLE_COLUMNS);

    tableEntity.setRelationshipAttribute(ATTRIBUTE_STORAGEDESC, AtlasTypeUtil.getAtlasRelatedObjectId(sdEntity, RELATIONSHIP_HIVE_TABLE_STORAGE_DESC));
    tableEntity.setRelationshipAttribute(ATTRIBUTE_PARTITION_KEYS, AtlasTypeUtil.getAtlasRelatedObjectIds(partKeys, RELATIONSHIP_HIVE_TABLE_PART_KEYS));
    tableEntity.setRelationshipAttribute(ATTRIBUTE_COLUMNS, AtlasTypeUtil.getAtlasRelatedObjectIds(columns, RELATIONSHIP_HIVE_TABLE_COLUMNS));

    table.addReferredEntity(database);
    table.addReferredEntity(sdEntity);

    if (partKeys != null) {
        for (AtlasEntity partKey : partKeys) {
            table.addReferredEntity(partKey);
        }
    }

    if (columns != null) {
        for (AtlasEntity column : columns) {
            table.addReferredEntity(column);
        }
    }

    table.setEntity(tableEntity);

    return table;
}
 

@Test
public void testRelationshipAttributeUpdate_NonComposite_ManyToMany() throws Exception {
    AtlasEntity a1 = new AtlasEntity("A");
    a1.setAttribute(NAME, "a1_name");

    AtlasEntity a2 = new AtlasEntity("A");
    a2.setAttribute(NAME, "a2_name");

    AtlasEntity a3 = new AtlasEntity("A");
    a3.setAttribute(NAME, "a3_name");

    AtlasEntity b1 = new AtlasEntity("B");
    b1.setAttribute(NAME, "b1_name");

    AtlasEntity b2 = new AtlasEntity("B");
    b2.setAttribute(NAME, "b2_name");

    AtlasEntitiesWithExtInfo entitiesWithExtInfo = new AtlasEntitiesWithExtInfo();
    entitiesWithExtInfo.addEntity(a1);
    entitiesWithExtInfo.addEntity(a2);
    entitiesWithExtInfo.addEntity(a3);
    entitiesWithExtInfo.addEntity(b1);
    entitiesWithExtInfo.addEntity(b2);
    entityStore.createOrUpdate(new AtlasEntityStream(entitiesWithExtInfo) , false);

    AtlasEntity b1PartialUpdate = new AtlasEntity("B");
    b1PartialUpdate.setRelationshipAttribute("manyToManyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2)));

    init();
    EntityMutationResponse response = entityStore.updateByUniqueAttributes(typeRegistry.getEntityTypeByName("B"),
                                                                           Collections.singletonMap(NAME, b1.getAttribute(NAME)),
                                                                           new AtlasEntityWithExtInfo(b1PartialUpdate));

    List<AtlasEntityHeader> updatedEntityHeaders = response.getPartialUpdatedEntities();
    assertEquals(updatedEntityHeaders.size(), 3);

    AtlasEntitiesWithExtInfo updatedEntities = entityStore.getByIds(ImmutableList.of(a1.getGuid(), a2.getGuid(), b1.getGuid()));

    AtlasEntity b1Entity = updatedEntities.getEntity(b1.getGuid());
    verifyRelationshipAttributeList(b1Entity, "manyToManyA", ImmutableList.of(getAtlasObjectId(a1), getAtlasObjectId(a2)));

    AtlasEntity a1Entity = updatedEntities.getEntity(a1.getGuid());
    verifyRelationshipAttributeList(a1Entity, "manyB", ImmutableList.of(getAtlasObjectId(b1)));

    AtlasEntity a2Entity = updatedEntities.getEntity(a2.getGuid());
    verifyRelationshipAttributeList(a2Entity, "manyB", ImmutableList.of(getAtlasObjectId(b1)));
}
 

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

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

    final AtlasEntity columnEntity1 = TestUtilsV2.createColumnEntity(tableEntity);
    entitiesInfo.addReferredEntity(columnEntity1);
    final AtlasEntity columnEntity2 = TestUtilsV2.createColumnEntity(tableEntity);
    entitiesInfo.addReferredEntity(columnEntity2);
    final AtlasEntity columnEntity3 = TestUtilsV2.createColumnEntity(tableEntity);
    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 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));

    // Retrieve the table entities from the Repository, to get their guids and the composite column guids.
    ITypedReferenceableInstance tableInstance = metadataService.getEntityDefinitionReference(TestUtils.TABLE_TYPE, NAME, (String) tableEntity.getAttribute(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();

    init();

    EntityMutationResponse deletionResponse = entityStore.deleteById(colId);
    assertEquals(deletionResponse.getDeletedEntities().size(), 1);
    assertEquals(deletionResponse.getDeletedEntities().get(0).getGuid(), colId);
    assertEquals(deletionResponse.getUpdatedEntities().size(), 1);
    assertEquals(deletionResponse.getUpdatedEntities().get(0).getGuid(), tableId);
    assertEntityDeleted(colId);

    final AtlasEntity.AtlasEntityWithExtInfo tableEntityCreated = entityStore.getById(tableId);
    assertDeletedColumn(tableEntityCreated);

    assertTestDisconnectUnidirectionalArrayReferenceFromClassType(
        (List<AtlasObjectId>) tableEntityCreated.getEntity().getAttribute(COLUMNS_ATTR_NAME), colId);

    //update by removing a column - col1
    final AtlasEntity tableEntity1 = TestUtilsV2.createTableEntity(dbEntity, (String) tableEntity.getAttribute(NAME));

    AtlasEntity.AtlasEntitiesWithExtInfo entitiesInfo1 = new AtlasEntity.AtlasEntitiesWithExtInfo(tableEntity1);
    final AtlasEntity columnEntity3New = TestUtilsV2.createColumnEntity(tableEntity1, (String) column3Created.getAttribute(NAME));
    tableEntity1.setAttribute(COLUMNS_ATTR_NAME, Arrays.asList(AtlasTypeUtil.getAtlasObjectId(columnEntity3New)));
    entitiesInfo1.addReferredEntity(columnEntity3New);

    init();
    deletionResponse = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo1), false);

    assertEquals(deletionResponse.getDeletedEntities().size(), 1);
    assertEquals(deletionResponse.getDeletedEntities().get(0).getGuid(), column2Created.getGuid());
    assertEntityDeleted(colId);

    // Delete the table entities.  The deletion should cascade to their composite columns.
    tableInstance = metadataService.getEntityDefinitionReference(TestUtils.TABLE_TYPE, NAME, (String) tableEntity.getAttribute(NAME));

    init();
    EntityMutationResponse tblDeletionResponse = entityStore.deleteById(tableInstance.getId()._getId());
    assertEquals(tblDeletionResponse.getDeletedEntities().size(), 2);

    final AtlasEntityHeader tableDeleted = tblDeletionResponse.getFirstDeletedEntityByTypeName(TABLE_TYPE);
    final AtlasEntityHeader colDeleted = tblDeletionResponse.getFirstDeletedEntityByTypeName(COLUMN_TYPE);

    // Verify that deleteEntities() response has guids for tables and their composite columns.
    Assert.assertTrue(tableDeleted.getGuid().equals(tableInstance.getId()._getId()));
    Assert.assertTrue(colDeleted.getGuid().equals(column3Created.getGuid()));

    // Verify that tables and their composite columns have been deleted from the graph Repository.
    assertEntityDeleted(tableDeleted.getGuid());
    assertEntityDeleted(colDeleted.getGuid());

}