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

private AtlasVertex getResolvedEntityVertex(EntityGraphDiscoveryContext context, AtlasEntity entity) throws AtlasBaseException {
    AtlasObjectId objectId = getAtlasObjectId(entity);
    AtlasVertex   ret      = context.getResolvedEntityVertex(entity.getGuid());

    if (ret != null) {
        context.addResolvedIdByUniqAttribs(objectId, ret);
    } else {
        ret = context.getResolvedEntityVertex(objectId);

        if (ret != null) {
            context.addResolvedGuid(entity.getGuid(), ret);
        }
    }

    return ret;
}
 

AtlasEntity createProcessExecution(AtlasEntity hiveProcess, String name, String description, String user,
                          String queryText, String queryPlan, String queryId, String queryGraph, String... classificationNames) throws Exception {

    AtlasEntity entity = new AtlasEntity(LOAD_PROCESS_EXECUTION_TYPE);
    Long startTime = System.currentTimeMillis();
    Long endTime = System.currentTimeMillis() + 10000;
    // set attributes
    entity.setAttribute("name", name);
    entity.setAttribute(REFERENCEABLE_ATTRIBUTE_NAME, name + CLUSTER_SUFFIX + startTime.toString() + endTime.toString());
    entity.setAttribute("description", description);
    entity.setAttribute("user", user);
    entity.setAttribute("startTime", startTime);
    entity.setAttribute("endTime", endTime);
    entity.setAttribute("queryText", queryText);
    entity.setAttribute("queryPlan", queryPlan);
    entity.setAttribute("queryId", queryId);
    entity.setAttribute("queryGraph", queryGraph);
    entity.setRelationshipAttribute("process", AtlasTypeUtil.toAtlasRelatedObjectId(hiveProcess));

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

    return createInstance(entity);
}
 

@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
public void testGetPathEntityHdfsPath() {
    Map<String, AtlasEntity> knownEntities = new HashMap<>();
    AtlasEntityWithExtInfo extInfo = new AtlasEntityWithExtInfo();

    PathExtractorContext extractorContext = new PathExtractorContext(METADATA_NAMESPACE);

    Path path = new Path(HDFS_PATH);
    AtlasEntityWithExtInfo entityWithExtInfo = AtlasPathExtractorUtil.getPathEntity(path, extractorContext);
    AtlasEntity entity = entityWithExtInfo.getEntity();

    assertNotNull(entity);
    assertEquals(entity.getTypeName(), HDFS_PATH_TYPE);
    verifyHDFSEntity(entity, false);

    assertNull(extInfo.getReferredEntities());
    assertEquals(extractorContext.getKnownEntities().size(), 1);
    extractorContext.getKnownEntities().values().forEach(x -> verifyHDFSEntity(x, false));
}
 

@Override
public AtlasEntity.AtlasEntityWithExtInfo getEntityWithExtInfo(String json) throws AtlasBaseException {
    if (StringUtils.isEmpty(json)) {
        return null;
    }

    AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = convertFromJson(AtlasEntity.AtlasEntityWithExtInfo.class, json);

    if (importTransform != null) {
        entityWithExtInfo = importTransform.apply(entityWithExtInfo);
    }

    if (entityHandlers != null) {
        applyTransformers(entityWithExtInfo);
    }

    return entityWithExtInfo;
}
 

public void testCreateTableHiveProcessNameAttribute() throws Exception {
	//test if \n is trimmed from name attribute of the process entity
    String tableName = tableName();
    String processNameQuery = String.format("create table %s (id string)", tableName);
    //add \n at the beginning of the query
    String query = String.format("%n%n%s", processNameQuery);

    runCommand(query);

    String dbId = assertDatabaseIsRegistered(DEFAULT_DB);
    String tableId = assertTableIsRegistered(DEFAULT_DB, tableName);

    //verify lineage is created and the name attribute is the query without \n
    String      processId      = assertEntityIsRegistered(HiveDataTypes.HIVE_PROCESS.getName(), AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, getTableProcessQualifiedName(DEFAULT_DB, tableName), null);
    AtlasEntity processsEntity = atlasClientV2.getEntityByGuid(processId).getEntity();

    assertEquals(processsEntity.getAttribute("name"), processNameQuery);
}
 

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

    AtlasExportRequest request = getRequest();
    InputStream inputStream = runExportWithParameters(exportService, request);

    ZipSource source = new ZipSource(inputStream);
    AtlasEntity.AtlasEntityWithExtInfo entities = ZipFileResourceTestUtils.getEntities(source, expectedEntityCount);

    int count = 0;
    for (Map.Entry<String, AtlasEntity> entry : entities.getReferredEntities().entrySet()) {
        assertNotNull(entry.getValue());
        if(entry.getValue().getTypeName().equals("hive_process")) {
            fail("Process entities should not be part of export!");
        }
        count++;
    }

    assertEquals(count, expectedEntityCount);
}
 

private void createOrUpdateNamespaceInstance(HBaseOperationContext hbaseOperationContext) {
    AtlasEntity nameSpace = buildNameSpace(hbaseOperationContext);

    switch (hbaseOperationContext.getOperation()) {
        case CREATE_NAMESPACE:
            LOG.info("Create NameSpace {}", nameSpace.getAttribute(REFERENCEABLE_ATTRIBUTE_NAME));

            hbaseOperationContext.addMessage(new EntityCreateRequestV2(hbaseOperationContext.getUser(), new AtlasEntitiesWithExtInfo(nameSpace)));
            break;

        case ALTER_NAMESPACE:
            LOG.info("Modify NameSpace {}", nameSpace.getAttribute(REFERENCEABLE_ATTRIBUTE_NAME));

            hbaseOperationContext.addMessage(new EntityUpdateRequestV2(hbaseOperationContext.getUser(), new AtlasEntitiesWithExtInfo(nameSpace)));
            break;
    }
}
 

private List<AtlasEntity> updateClassificationText(AtlasClassification classification, Collection<AtlasVertex> propagatedVertices) throws AtlasBaseException {
    List<AtlasEntity> propagatedEntities = new ArrayList<>();

    if(CollectionUtils.isNotEmpty(propagatedVertices)) {
        for(AtlasVertex vertex : propagatedVertices) {
            AtlasEntity entity = instanceConverter.getAndCacheEntity(graphHelper.getGuid(vertex), ENTITY_CHANGE_NOTIFY_IGNORE_RELATIONSHIP_ATTRIBUTES);

            if (isActive(entity)) {
                vertex.setProperty(CLASSIFICATION_TEXT_KEY, fullTextMapperV2.getClassificationTextForEntity(entity));
                propagatedEntities.add(entity);
            }
        }
    }

    return propagatedEntities;
}
 

@Test (dependsOnMethods = "addInvalidKeysToEntityCustomAttributes")
public void addInvalidValuesToEntityCustomAttributes() throws AtlasBaseException {
    AtlasEntity tblEntity = getEntityFromStore(tblEntityGuid);

    // value length is greater than 500
    Map<String, String> invalidCustomAttributes = new HashMap<>();
    invalidCustomAttributes.put("key1", randomAlphanumeric(500));
    invalidCustomAttributes.put("key2", randomAlphanumeric(501));

    tblEntity.setCustomAttributes(invalidCustomAttributes);

    try {
        entityStore.createOrUpdate(new AtlasEntityStream(tblEntity), false);
    } catch (AtlasBaseException ex) {
        assertEquals(ex.getAtlasErrorCode(), INVALID_CUSTOM_ATTRIBUTE_VALUE);
    }
}
 

@Test
public void addScopedClassification() throws AtlasBaseException {
    AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo = getAtlasEntityWithExtInfo();
    AtlasEntity entity = entityWithExtInfo.getReferredEntities().get("2");

    int existingClassificationsCount =  entityWithExtInfo.getEntity().getClassifications() != null ? entity.getClassifications().size() : 0;
    ImportTransforms t = ImportTransforms.fromJson(jsonAddClasificationScoped);

    assertTrue(t.getTransforms().size() > 0);

    ImportTransformer.AddClassification classification = (ImportTransformer.AddClassification) t.getTransforms().get("hive_column").get("*").get(0);
    AtlasObjectId objectId = new AtlasObjectId("hive_column", ATTR_NAME_QUALIFIED_NAME, String.format(COLUMN_QUALIFIED_NAME_FORMAT, 2));
    classification.addFilter(objectId);
    t.apply(entityWithExtInfo);

    assertNotNull(t);

    assertNull(entityWithExtInfo.getEntity().getClassifications());
    assertNull(entityWithExtInfo.getReferredEntities().get("0").getClassifications());
    assertEquals(entityWithExtInfo.getReferredEntities().get("1").getClassifications().size(), existingClassificationsCount + 1);
    assertNull(entityWithExtInfo.getReferredEntities().get("2").getClassifications());
}
 

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

public List<AtlasEntity> getEntitiesWithModifiedTimestamp(AtlasEntityWithExtInfo entityWithExtInfo) {
    if(fetchType != ExportFetchType.INCREMENTAL) {
        return new ArrayList<>();
    }

    List<AtlasEntity> ret = new ArrayList<>();
    if(doesTimestampQualify(entityWithExtInfo.getEntity())) {
        ret.add(entityWithExtInfo.getEntity());
        return ret;
    }

    for (AtlasEntity entity : entityWithExtInfo.getReferredEntities().values()) {
        if((doesTimestampQualify(entity))) {
            ret.add(entity);
        }
    }

    return ret;
}
 

public AtlasEntity.AtlasEntitiesWithExtInfo toAtlasEntities(List<Referenceable> referenceables) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> toAtlasEntities");
    }

    AtlasFormatConverter.ConverterContext context = new AtlasFormatConverter.ConverterContext();
    for (Referenceable referenceable : referenceables) {
        AtlasEntity entity = fromV1toV2Entity(referenceable, context);

        context.addEntity(entity);
    }
    if (LOG.isDebugEnabled()) {
        LOG.debug("<== toAtlasEntities");
    }

    return context.getEntities();
}
 

@Override
public void onTermAdded(AtlasGlossaryTerm term, List<AtlasRelatedObjectId> entities) throws AtlasBaseException {
    if (term != null && CollectionUtils.isNotEmpty(entities)) {
        MetricRecorder metric = RequestContext.get().startMetricRecord("entityAudit");

        List<EntityAuditEventV2> events = new ArrayList<>();

        for (AtlasRelatedObjectId relatedObjectId : entities) {
            AtlasEntity entity = instanceConverter.getAndCacheEntity(relatedObjectId.getGuid());

            if (entity != null) {
                events.add(createEvent(entity, TERM_ADD, "Added term: " + term.toAuditString()));
            }
        }

        auditRepository.putEventsV2(events);

        RequestContext.get().endMetricRecord(metric);
    }
}
 

private List<AtlasObjectId> getAtlasObjectIds(List<AtlasEntity> columns) {
    List<AtlasObjectId> objIds = new ArrayList<>();
    for (AtlasEntity e : columns) {
        AtlasObjectId oid = getAtlasObjectId(e);
        objIds.add(oid);
    }
    return objIds;
}
 

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

@Test(dependsOnMethods = "testSubmitEntity")
    public void testAddProperty() throws Exception {
        //add property
        String description = "bar table - new desc";
        addProperty(createHiveTable().getGuid(), "description", description);

        AtlasEntity entityByGuid = getEntityByGuid(createHiveTable().getGuid());
        Assert.assertNotNull(entityByGuid);

        entityByGuid.setAttribute("description", description);

        // TODO: This behavior should've been consistent across APIs
//        //invalid property for the type
//        try {
//            addProperty(table.getGuid(), "invalid_property", "bar table");
//            Assert.fail("Expected AtlasServiceException");
//        } catch (AtlasServiceException e) {
//            assertNotNull(e.getStatus());
//            assertEquals(e.getStatus(), ClientResponse.Status.BAD_REQUEST);
//        }

        //non-string property, update
        Object currentTime = new DateTime();
        addProperty(createHiveTable().getGuid(), "createTime", currentTime);

        entityByGuid = getEntityByGuid(createHiveTable().getGuid());
        Assert.assertNotNull(entityByGuid);
    }
 

public AtlasEntity getCreatedOrUpdatedEntity(String parentGuid) {
    AtlasEntity e = getCreatedEntity(parentGuid);
    if(e == null) {
        return getUpdatedEntity(parentGuid);
    }

    return e;
}
 

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

private static List<AtlasObjectId> getAtlasObjectIds(AtlasEntity... entities) {
    List<AtlasObjectId> ret = new ArrayList<>();

    if (ArrayUtils.isNotEmpty(entities)) {
        for (AtlasEntity entity : entities) {
            ret.add(getAtlasObjectId(entity));
        }
    }

    return ret;
}
 

public static AtlasEntity newEntity(AtlasEntityDef entityDef, AtlasTypeRegistry typesRegistry) {
    AtlasEntity ret = null;

    AtlasEntityType entityType = typesRegistry.getEntityTypeByName(entityDef.getName());

    if (entityType != null) {
        ret = entityType.createDefaultValue();
    } else {
        LOG.error("failed to get entity-type {}", entityDef.getName());
    }

    return ret;
}
 

@Test
public void testGetEntityByAttribute() throws Exception {
    AtlasEntity hiveDB = createHiveDB();
    String qualifiedName = (String) hiveDB.getAttribute(NAME);
    //get entity by attribute

    AtlasEntity byAttribute = atlasClientV2.getEntityByAttribute(DATABASE_TYPE_V2, toMap(NAME, qualifiedName)).getEntity();
    assertEquals(byAttribute.getTypeName(), DATABASE_TYPE_V2);
    assertEquals(byAttribute.getAttribute(NAME), qualifiedName);
}
 

private void purgeDeletedEntities(List<AtlasEntityHeader> entities) {
    if (entities != null) {
        for (ListIterator<AtlasEntityHeader> iter = entities.listIterator(); iter.hasNext(); ) {
            AtlasEntityHeader entity = iter.next();

            if (entity.getStatus() == AtlasEntity.Status.DELETED) {
                if (LOG.isDebugEnabled()) {
                    LOG.debug("purgeDeletedEntities(guid={}, status={}): REMOVED", entity.getGuid(), entity.getStatus());
                }

                iter.remove();
            }
        }
    }
}
 

private void assertTag(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo, String tagName) {
    if(entityWithExtInfo.getReferredEntities() == null || entityWithExtInfo.getReferredEntities().size() == 0) {
        return;
    }

    for (AtlasEntity entity : entityWithExtInfo.getReferredEntities().values()) {
        assertTag(entity, tagName);
    }
}
 

@Override
public int compare(AtlasEntity entity1, AtlasEntity entity2) {
    String name1 = (String)entity1.getAttribute(ATTRIBUTE_QUALIFIED_NAME);
    String name2 = (String)entity2.getAttribute(ATTRIBUTE_QUALIFIED_NAME);

    if (name1 == null) {
        return -1;
    }

    if (name2 == null) {
        return 1;
    }

    return name1.toLowerCase().compareTo(name2.toLowerCase());
}
 

private AtlasEntity getHiveColumnEntity(String clusterName, String dbName, String tableName, String columnName) {
    String qualifiedName = dbName + "." + tableName + "." + columnName + "@" + clusterName;

    AtlasEntity entity = new AtlasEntity(TransformationConstants.HIVE_COLUMN);

    entity.setAttribute("owner", "hive");
    entity.setAttribute(ATTR_NAME_QUALIFIED_NAME, qualifiedName);
    entity.setAttribute("name", columnName);
    entity.setAttribute("position", 1);
    entity.setAttribute("type", "string");

    return entity;
}
 

public Referenceable getReferenceable(AtlasEntity entity, final ConverterContext ctx) throws AtlasBaseException {
    AtlasFormatConverter converter  = instanceFormatters.getConverter(TypeCategory.ENTITY);
    AtlasType            entityType = typeRegistry.getType(entity.getTypeName());
    Referenceable        ref        = (Referenceable) converter.fromV2ToV1(entity, entityType, ctx);

    return ref;
}
 

private String getCustomAttributesString(AtlasEntity entity) {
    String              ret              = null;
    Map<String, String> customAttributes = entity.getCustomAttributes();

    if (customAttributes != null) {
        ret = AtlasType.toJson(customAttributes);
    }

    return ret;
}
 

AtlasEntity createView(String name, AtlasEntity db, List<AtlasEntity> inputTables, String... traitNames) throws Exception {
    AtlasEntity entity = new AtlasEntity(VIEW_TYPE);

    entity.setClassifications(toAtlasClassifications(traitNames));
    entity.setAttribute("name", name);
    entity.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
    entity.setAttribute("db", db);
    entity.setAttribute("inputTables", inputTables);

    return createInstance(entity, traitNames);
}