Java Code Examples for org.apache.atlas.model.instance.AtlasEntity#AtlasEntityWithExtInfo

@Test(dataProvider = "stocks-legacy")
public void importLegacy(InputStream inputStream) throws IOException, AtlasBaseException {
    loadBaseModel();
    loadFsModel();
    loadHiveModel();

    runImportWithNoParameters(importService, inputStream);
    List<AtlasEntityHeader> result = getImportedEntities("hive_db", "886c5e9c-3ac6-40be-8201-fb0cebb64783");
    assertEquals(result.size(), 1);

    AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = getEntity(result.get(0));
    Map<String, Object> relationshipAttributes = entityWithExtInfo.getEntity().getRelationshipAttributes();
    assertNotNull(relationshipAttributes);
    assertNotNull(relationshipAttributes.get("tables"));

    List<AtlasRelatedObjectId> relatedList = (List<AtlasRelatedObjectId>) relationshipAttributes.get("tables");
    AtlasRelatedObjectId relatedObjectId = relatedList.get(0);
    assertNotNull(relatedObjectId.getRelationshipGuid());
}
 

@Override
protected void processItem(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) {
    int delta = MapUtils.isEmpty(entityWithExtInfo.getReferredEntities())
            ? 1
            : entityWithExtInfo.getReferredEntities().size() + 1;

    long currentCount = counter.addAndGet(delta);
    currentBatch.addAndGet(delta);

    try {
        processEntity(entityWithExtInfo, currentCount);
        attemptCommit();
    } catch (Exception e) {
        LOG.info("Invalid entities. Possible data loss: Please correct and re-submit!", e);
    }
}
 

@Test
public void recordsEntityWithExtInfoEntries() throws AtlasBaseException {
    final int max_entries = 3;
    ByteArrayOutputStream byteOutputStream = new ByteArrayOutputStream();
    ZipSink zs = new ZipSink(byteOutputStream);

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

    AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = new AtlasEntity.AtlasEntityWithExtInfo(entity);
    addReferredEntities(entityWithExtInfo, max_entries);

    zs.add(entityWithExtInfo);
    for (int i = 0; i <= max_entries; i++) {
        String g = String.format(knownEntityGuidFormat, i);
        assertTrue(zs.hasEntity(g));
    }

    zs.close();
}
 

@Test
public void loadFile() throws IOException, AtlasBaseException {
    final int EXPECTED_ENTITY_COUNT = 3;

    InputStream inputStream = ZipFileResourceTestUtils.getFileInputStream("zip-direct-2.zip");
    ZipSourceDirect zipSourceDirect = new ZipSourceDirect(inputStream, EXPECTED_ENTITY_COUNT);

    assertNotNull(zipSourceDirect);
    assertNotNull(zipSourceDirect.getTypesDef());
    assertTrue(zipSourceDirect.getTypesDef().getEntityDefs().size() > 0);
    assertNotNull(zipSourceDirect.getExportResult());

    int count = 0;
    AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo;
    while((entityWithExtInfo = zipSourceDirect.getNextEntityWithExtInfo()) != null) {
        assertNotNull(entityWithExtInfo);
        count++;
    }

    assertEquals(count, EXPECTED_ENTITY_COUNT);
}
 

@Test(dataProvider = "zipFileStocksFloat")
public void attemptToSerializeLongFloats(ZipSource zipSource) throws IOException, AtlasBaseException {
    Assert.assertTrue(zipSource.hasNext());
    assertTrue(zipSource.hasNext());
    assertTrue(zipSource.hasNext());

    AtlasEntity.AtlasEntityWithExtInfo e = zipSource.getNextEntityWithExtInfo();
    assertNotNull(e);
    assertTrue(e.getEntity().getClassifications().size() > 0);
    assertNotNull(e.getEntity().getClassifications().get(0).getAttribute("fv"));
    assertEquals(e.getEntity().getClassifications().get(0).getAttribute("fv").toString(), "3.4028235E+38");

    assertTrue(zipSource.hasNext());
}
 

private void assertEntities(List<String> entityGuids, String tagName) throws AtlasBaseException {
    for (String guid : entityGuids) {
        AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = this.entityStore.getById(guid);

        assertNotNull(entityWithExtInfo);
        assertTag(entityWithExtInfo, tagName);
    }
}
 

public static void createAtlasEntity(AtlasEntityStoreV2 entityStoreV1, AtlasEntity.AtlasEntityWithExtInfo atlasEntity) {
    try {
        EntityMutationResponse response = entityStoreV1.createOrUpdateForImport(new AtlasEntityStreamForImport(atlasEntity, null));
        assertNotNull(response);
        assertTrue((response.getCreatedEntities() != null && response.getCreatedEntities().size() > 0) ||
                (response.getMutatedEntities() != null && response.getMutatedEntities().size() > 0));
    } catch (AtlasBaseException e) {
        throw new SkipException(String.format("createAtlasEntity: could not loaded '%s'.", atlasEntity.getEntity().getTypeName()));
    }
}
 

private void applyTransformers(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) {
    if (entityWithExtInfo == null) {
        return;
    }

    transform(entityWithExtInfo.getEntity());

    if (MapUtils.isNotEmpty(entityWithExtInfo.getReferredEntities())) {
        for (AtlasEntity e : entityWithExtInfo.getReferredEntities().values()) {
            transform(e);
        }
    }
}
 

private AtlasEntity.AtlasEntityWithExtInfo getHDFSPathEntity(AtlasEntityType hdfsPathEntityType, String path, String clusterName) {
    try {
        return entityStore.getByUniqueAttributes(hdfsPathEntityType, getUniqueAttributes(path, clusterName));
    } catch (AtlasBaseException e) {
        return null;
    }
}
 

@Override
public AtlasEntity.AtlasEntityWithExtInfo toEntityWithExtInfo(final AtlasGlossaryTerm obj) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasGlossaryTermDTO.toEntityWithExtInfo()", obj);
    }

    Objects.requireNonNull(obj, "atlasGlossaryTerm");
    AtlasEntity entity = toEntity(obj);
    AtlasEntity.AtlasEntityWithExtInfo ret = new AtlasEntity.AtlasEntityWithExtInfo(entity);

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

@Override
public AtlasEntity.AtlasEntityWithExtInfo getNextEntityWithExtInfo() {
    try {
        return getEntityWithExtInfo(moveNext());
    } catch (AtlasBaseException e) {
        LOG.error("getNextEntityWithExtInfo", e);
        return null;
    }
}
 

public AtlasEntity.AtlasEntityWithExtInfo getNextEntityWithExtInfo() {
    if (navigateResidualList == false) {
        return stream.getNextEntityWithExtInfo();
    } else {
        stream.setPositionUsingEntityGuid(residualList.get(currentResidualListIndex++));
        return stream.getNextEntityWithExtInfo();
    }
}
 

@Override
protected void assertTestDeleteEntities(final AtlasEntity.AtlasEntityWithExtInfo tableInstance) throws Exception {
    //Assert that the deleted table can be fully constructed back
    List<IReferenceableInstance> columns = (List<IReferenceableInstance>) tableInstance.getEntity().getAttribute(COLUMNS_ATTR_NAME);
    assertEquals(columns.size(), 3);
    assertNotNull(tableInstance.getEntity().getAttribute("database"));
}
 

@Override
public AtlasEntity.AtlasEntityWithExtInfo toEntityWithExtInfo(ExportImportAuditEntry obj) throws AtlasBaseException {
    return new AtlasEntity.AtlasEntityWithExtInfo(toEntity(obj));
}
 

private AtlasEntity getEntity(String guid) throws AtlasBaseException {
    AtlasEntity.AtlasEntityWithExtInfo extInfo = getEntityWithExtInfo(guid);
    return (extInfo != null) ? extInfo.getEntity() : null;
}
 

/**
 * Verify deleting entities that are the target of class map references.
 */
@Test
public void testDisconnectMapReferenceFromClassType() throws Exception {
    // Define type for map value.
    AtlasStructDef.AtlasAttributeDef[] mapValueAttributes = new AtlasStructDef.AtlasAttributeDef[]{
        new AtlasStructDef.AtlasAttributeDef("biMapOwner", "MapOwner",
            true,
            AtlasStructDef.AtlasAttributeDef.Cardinality.SINGLE, 0, 1,
            false, false,
            new ArrayList<AtlasStructDef.AtlasConstraintDef>() {{
                add(new AtlasStructDef.AtlasConstraintDef(
                    AtlasStructDef.AtlasConstraintDef.CONSTRAINT_TYPE_INVERSE_REF, new HashMap<String, Object>() {{
                    put(AtlasStructDef.AtlasConstraintDef.CONSTRAINT_PARAM_ATTRIBUTE, "biMap");
                }}));
            }})};

    AtlasEntityDef mapValueContainerDef =
        new AtlasEntityDef("MapValue", "MapValue_desc", "1.0",
            Arrays.asList(mapValueAttributes), Collections.<String>emptySet());

    // Define type with unidirectional and bidirectional map references,
    // where the map value is a class reference to MapValue.

    AtlasStructDef.AtlasAttributeDef[] mapOwnerAttributes = new AtlasStructDef.AtlasAttributeDef[]{
        new AtlasStructDef.AtlasAttributeDef("map", "map<string,MapValue>",
            true,
            AtlasStructDef.AtlasAttributeDef.Cardinality.SINGLE, 0, 1,
            false, false,
            Collections.<AtlasStructDef.AtlasConstraintDef>emptyList()),
        new AtlasStructDef.AtlasAttributeDef("biMap", "map<string,MapValue>",
            true,
            AtlasStructDef.AtlasAttributeDef.Cardinality.SINGLE, 0, 1,
            false, false,
            new ArrayList<AtlasStructDef.AtlasConstraintDef>() {{
                add(new AtlasStructDef.AtlasConstraintDef(
                    AtlasStructDef.AtlasConstraintDef.CONSTRAINT_TYPE_INVERSE_REF, new HashMap<String, Object>() {{
                    put(AtlasStructDef.AtlasConstraintDef.CONSTRAINT_PARAM_ATTRIBUTE, "biMapOwner");
                }}));
            }})};

    AtlasEntityDef mapOwnerContainerDef =
        new AtlasEntityDef("MapOwner", "MapOwner_desc", "1.0",
            Arrays.asList(mapOwnerAttributes), Collections.<String>emptySet());

    AtlasTypesDef typesDef = AtlasTypeUtil.getTypesDef(ImmutableList.<AtlasEnumDef>of(),
        ImmutableList.<AtlasStructDef>of(),
        ImmutableList.<AtlasClassificationDef>of(),
        ImmutableList.<AtlasEntityDef>of(mapValueContainerDef, mapOwnerContainerDef));

    typeDefStore.createTypesDef(typesDef);

    // Create instances of MapOwner and MapValue.
    // Set MapOwner.map and MapOwner.biMap with one entry that references MapValue instance.
    AtlasEntity mapOwnerInstance = new AtlasEntity("MapOwner");
    AtlasEntity mapValueInstance = new AtlasEntity("MapValue");

    mapOwnerInstance.setAttribute("map", Collections.singletonMap("value1", AtlasTypeUtil.getAtlasObjectId(mapValueInstance)));
    mapOwnerInstance.setAttribute("biMap", Collections.singletonMap("value1", AtlasTypeUtil.getAtlasObjectId(mapValueInstance)));
    // Set biMapOwner reverse reference on MapValue.
    mapValueInstance.setAttribute("biMapOwner", AtlasTypeUtil.getAtlasObjectId(mapOwnerInstance));

    AtlasEntity.AtlasEntitiesWithExtInfo entities = new AtlasEntity.AtlasEntitiesWithExtInfo();
    entities.addReferredEntity(mapValueInstance);
    entities.addEntity(mapOwnerInstance);

    final EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entities), false);
    Assert.assertEquals(response.getCreatedEntities().size(), 2);
    final List<AtlasEntityHeader> mapOwnerCreated = response.getCreatedEntitiesByTypeName("MapOwner");
    AtlasEntity.AtlasEntityWithExtInfo mapOwnerEntity = entityStore.getById(mapOwnerCreated.get(0).getGuid());

    String edgeLabel = AtlasGraphUtilsV1.getAttributeEdgeLabel(typeRegistry.getEntityTypeByName("MapOwner"), "map");
    String mapEntryLabel = edgeLabel + "." + "value1";
    AtlasEdgeLabel atlasEdgeLabel = new AtlasEdgeLabel(mapEntryLabel);

    // Verify MapOwner.map attribute has expected value.
    String mapValueGuid = null;
    AtlasVertex mapOwnerVertex = null;
    for (String mapAttrName : Arrays.asList("map", "biMap")) {
        Object object = mapOwnerEntity.getEntity().getAttribute(mapAttrName);
        Assert.assertNotNull(object);
        Assert.assertTrue(object instanceof Map);
        Map<String, AtlasObjectId> map = (Map<String, AtlasObjectId>)object;
        Assert.assertEquals(map.size(), 1);
        AtlasObjectId value1Id = map.get("value1");
        Assert.assertNotNull(value1Id);
        mapValueGuid = value1Id.getGuid();
        mapOwnerVertex = GraphHelper.getInstance().getVertexForGUID(mapOwnerEntity.getEntity().getGuid());
        object = mapOwnerVertex.getProperty(atlasEdgeLabel.getQualifiedMapKey(), Object.class);
        Assert.assertNotNull(object);
    }

    // Delete the map value instance.
    // This should disconnect the references from the map owner instance.
    entityStore.deleteById(mapValueGuid);
    assertEntityDeleted(mapValueGuid);
    assertTestDisconnectMapReferenceFromClassType(mapOwnerEntity.getEntity().getGuid());
}
 

@Override
protected void assertEntityDeleted(final String id) throws Exception {
    final AtlasEntity.AtlasEntityWithExtInfo byId = entityStore.getById(id);
    assertEquals(byId.getEntity().getStatus(), AtlasEntity.Status.DELETED);
}
 

@Override
public ExportImportAuditEntry from(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) {
    return from(entityWithExtInfo.getEntity());
}
 

@Override
public AtlasEntity next() {
    AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = getNextEntityWithExtInfo();

    return entityWithExtInfo != null ? entityWithExtInfo.getEntity() : null;
}
 

/**
 * Verify deleting an entity which is contained by another
 * entity through a bi-directional composite reference.
 *
 * @throws Exception
 */
@Test
public void testDisconnectBidirectionalReferences() throws Exception {
    AtlasEntity.AtlasEntitiesWithExtInfo hrDept = TestUtilsV2.createDeptEg2();
    init();
    final EntityMutationResponse hrDeptCreationResponse = entityStore.createOrUpdate(new AtlasEntityStream(hrDept), false);

    final AtlasEntityHeader deptCreated = hrDeptCreationResponse.getFirstCreatedEntityByTypeName(DEPARTMENT_TYPE);
    final AtlasEntityHeader maxEmployee = hrDeptCreationResponse.getCreatedEntityByTypeNameAndAttribute(TestUtilsV2.EMPLOYEE_TYPE, NAME, "Max");
    final AtlasEntityHeader johnEmployee = hrDeptCreationResponse.getCreatedEntityByTypeNameAndAttribute(TestUtilsV2.EMPLOYEE_TYPE, NAME, "John");
    final AtlasEntityHeader janeEmployee = hrDeptCreationResponse.getCreatedEntityByTypeNameAndAttribute(TestUtilsV2.MANAGER_TYPE, NAME, "Jane");
    final AtlasEntityHeader juliusEmployee = hrDeptCreationResponse.getCreatedEntityByTypeNameAndAttribute(TestUtilsV2.MANAGER_TYPE, NAME, "Julius");

    ITypedReferenceableInstance hrDeptInstance = metadataService.getEntityDefinition(deptCreated.getGuid());
    Map<String, String> nameGuidMap = getEmployeeNameGuidMap(hrDeptInstance);

    // Verify that Max is one of Jane's subordinates.
    ITypedReferenceableInstance jane = metadataService.getEntityDefinition(janeEmployee.getGuid());
    Object refValue = jane.get("subordinates");
    Assert.assertTrue(refValue instanceof List);
    List<Object> subordinates = (List<Object>)refValue;
    Assert.assertEquals(subordinates.size(), 2);
    List<String> subordinateIds = new ArrayList<>(2);
    for (Object listValue : subordinates) {
        Assert.assertTrue(listValue instanceof ITypedReferenceableInstance);
        ITypedReferenceableInstance employee = (ITypedReferenceableInstance) listValue;
        subordinateIds.add(employee.getId()._getId());
    }
    Assert.assertTrue(subordinateIds.contains(maxEmployee.getGuid()));

    init();
    EntityMutationResponse entityResult = entityStore.deleteById(maxEmployee.getGuid());
    ITypedReferenceableInstance john = metadataService.getEntityDefinitionReference(TestUtilsV2.EMPLOYEE_TYPE, NAME, "John");

    assertEquals(entityResult.getDeletedEntities().size(), 1);
    assertEquals(entityResult.getDeletedEntities().get(0).getGuid(), maxEmployee.getGuid());
    assertEquals(entityResult.getUpdatedEntities().size(), 3);

    assertEquals(extractGuids(entityResult.getUpdatedEntities()), Arrays.asList(janeEmployee.getGuid(), deptCreated.getGuid(), johnEmployee.getGuid()));
    assertEntityDeleted(maxEmployee.getGuid());

    assertMaxForTestDisconnectBidirectionalReferences(nameGuidMap);

    // Now delete jane - this should disconnect the manager reference from her
    // subordinate.
    init();
    entityResult = entityStore.deleteById(janeEmployee.getGuid());
    assertEquals(entityResult.getDeletedEntities().size(), 1);
    assertEquals(entityResult.getDeletedEntities().get(0).getGuid(), janeEmployee.getGuid());
    assertEquals(entityResult.getUpdatedEntities().size(), 2);
    assertEquals(extractGuids(entityResult.getUpdatedEntities()), Arrays.asList(deptCreated.getGuid(), johnEmployee.getGuid()));

    assertEntityDeleted(janeEmployee.getGuid());

    final AtlasEntity.AtlasEntityWithExtInfo johnUpdated = entityStore.getById(johnEmployee.getGuid());
    assertJohnForTestDisconnectBidirectionalReferences(johnUpdated, janeEmployee.getGuid());
}