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

private void clearRefAttributesAndMove(AtlasEntity entity, PreprocessorContext context) {
    switch (entity.getTypeName()) {
        case TYPE_RDBMS_INSTANCE:
            context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_DATABASES, RELATIONSHIP_TYPE_RDBMS_INSTANCE_DATABASES, ATTRIBUTE_INSTANCE);
        break;

        case TYPE_RDBMS_DB:
            context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_TABLES, RELATIONSHIP_TYPE_RDBMS_DB_TABLES, ATTRIBUTE_DB);
        break;

        case TYPE_RDBMS_TABLE:
            context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_COLUMNS, RELATIONSHIP_TYPE_RDBMS_TABLE_COLUMNS, ATTRIBUTE_TABLE);
            context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_INDEXES, RELATIONSHIP_TYPE_RDBMS_TABLE_INDEXES, ATTRIBUTE_TABLE);
            context.removeRefAttributeAndRegisterToMove(entity, ATTRIBUTE_FOREIGN_KEYS, RELATIONSHIP_TYPE_RDBMS_TABLE_FOREIGN_KEYS, ATTRIBUTE_TABLE);
        break;
    }
}
 

private void verifyOzoneEntities(String scheme, String path, Map<String, AtlasEntity> knownEntities) {
    for (AtlasEntity knownEntity : knownEntities.values()) {
        switch (knownEntity.getTypeName()){
            case OZONE_KEY:
                verifyOzoneKeyEntity(path, knownEntity);
                break;

            case OZONE_VOLUME:
                assertEquals(knownEntity.getAttribute(ATTRIBUTE_QUALIFIED_NAME), scheme + "volume1" + QNAME_METADATA_NAMESPACE);
                assertEquals(knownEntity.getAttribute(ATTRIBUTE_NAME), "volume1");
                break;

            case OZONE_BUCKET:
                assertEquals(knownEntity.getAttribute(ATTRIBUTE_QUALIFIED_NAME), scheme + "volume1.bucket1" + QNAME_METADATA_NAMESPACE);
                assertEquals(knownEntity.getAttribute(ATTRIBUTE_NAME), "bucket1");
                break;
        }
    }
}
 

@Override
public void validateAndNormalizeForUpdate(AtlasEntity entity) throws AtlasBaseException {
    List<String> messages = new ArrayList<>();

    if (! AtlasTypeUtil.isValidGuid(entity.getGuid())) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, "invalid guid " + entity.getGuid());
    }

    AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());

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

    type.validateValueForUpdate(entity, entity.getTypeName(), messages);

    if (!messages.isEmpty()) {
        throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, messages);
    }

    type.getNormalizedValueForUpdate(entity);
}
 

@Test
public void testHiveTableClearAttributeHandler() {
    // clear replicatedTo attribute for hive_table entities
    AttributeTransform p1 = new AttributeTransform(Collections.singletonMap("hive_table.replicatedTo", "HAS_VALUE:"),
                                                   Collections.singletonMap("hive_table.replicatedTo", "CLEAR:"));

    List<BaseEntityHandler> handlers = initializeHandlers(Collections.singletonList(p1));

    List<AtlasEntity> entities = getAllEntities();

    for (AtlasEntity entity : entities) {
        String  replicatedTo = (String) entity.getAttribute("replicatedTo");

        if (entity.getTypeName() == HIVE_TABLE) {
            assertTrue(StringUtils.isNotEmpty(replicatedTo));
        }

        applyTransforms(entity, handlers);

        String transformedValue = (String) entity.getAttribute("replicatedTo");

        if (entity.getTypeName() == HIVE_TABLE) {
            assertTrue(StringUtils.isEmpty(transformedValue));
        }
    }
}
 

public static AtlasObjectId getAtlasObjectId(AtlasEntity entity, AtlasTypeRegistry typeRegistry) {
    String              typeName       = entity.getTypeName();
    AtlasEntityType     entityType     = typeRegistry.getEntityTypeByName(typeName);
    Map<String, Object> uniqAttributes = null;

    if (entityType != null && MapUtils.isNotEmpty(entityType.getUniqAttributes())) {
        for (AtlasAttribute attribute : entityType.getUniqAttributes().values()) {
            Object attrValue = entity.getAttribute(attribute.getName());

            if (attrValue != null) {
                if (uniqAttributes == null) {
                    uniqAttributes = new HashMap<>();
                }

                uniqAttributes.put(attribute.getName(), attrValue);
            }
        }
    }

    return new AtlasObjectId(entity.getGuid(), typeName, uniqAttributes);
}
 

@Override
public void validateAndNormalizeForUpdate(AtlasEntity entity) throws AtlasBaseException {
    List<String> messages = new ArrayList<>();

    if (!AtlasTypeUtil.isValidGuid(entity.getGuid())) {
        throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, "invalid guid " + entity.getGuid());
    }

    AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());

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

    validateCustomAttributes(entity);

    validateLabels(entity.getLabels());

    type.validateValueForUpdate(entity, entity.getTypeName(), messages);

    if (!messages.isEmpty()) {
        throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, messages);
    }

    type.getNormalizedValueForUpdate(entity);
}
 

void walkEntityGraph(AtlasEntity entity) throws AtlasBaseException {
    if (entity == null) {
        return;
    }

    AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());

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

    recordObjectReference(entity.getGuid());

    visitEntity(type, entity);
}
 

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

private void renameColumns(List<AtlasObjectId> columns, AtlasEntityExtInfo oldEntityExtInfo, String newTableQualifiedName, List<HookNotification> notifications) {
    if (CollectionUtils.isNotEmpty(columns)) {
        for (AtlasObjectId columnId : columns) {
            AtlasEntity   oldColumn   = oldEntityExtInfo.getEntity(columnId.getGuid());
            AtlasObjectId oldColumnId = new AtlasObjectId(oldColumn.getTypeName(), ATTRIBUTE_QUALIFIED_NAME, oldColumn.getAttribute(ATTRIBUTE_QUALIFIED_NAME));
            AtlasEntity   newColumn   = new AtlasEntity(oldColumn.getTypeName(), ATTRIBUTE_QUALIFIED_NAME, getColumnQualifiedName(newTableQualifiedName, (String) oldColumn.getAttribute(ATTRIBUTE_NAME)));

            notifications.add(new EntityPartialUpdateRequestV2(getUserName(), oldColumnId, new AtlasEntityWithExtInfo(newColumn)));
        }
    }
}
 

private void mapRelationshipAttributes(AtlasVertex entityVertex, AtlasEntity entity, AtlasEntityExtInfo entityExtInfo, boolean isMinExtInfo) throws AtlasBaseException {
    AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());

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

    for (String attributeName : entityType.getRelationshipAttributes().keySet()) {
        mapVertexToRelationshipAttribute(entityVertex, entityType, attributeName, entity, entityExtInfo, isMinExtInfo);
    }
}
 

private AtlasEntityHeader constructHeader(AtlasEntity entity, final AtlasEntityType type, AtlasVertex vertex) {
    AtlasEntityHeader header = new AtlasEntityHeader(entity.getTypeName());

    header.setGuid(getIdFromVertex(vertex));

    for (AtlasAttribute attribute : type.getUniqAttributes().values()) {
        header.setAttribute(attribute.getName(), entity.getAttribute(attribute.getName()));
    }

    return header;
}
 

public EntityGraphDiscoveryContext resolveEntityReferences(EntityGraphDiscoveryContext context) throws AtlasBaseException {
    if (context == null) {
        throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "IDBasedEntityResolver.resolveEntityReferences(): context is null");
    }

    EntityStream entityStream = context.getEntityStream();

    for (String guid : context.getReferencedGuids()) {
        boolean isAssignedGuid = AtlasTypeUtil.isAssignedGuid(guid);
        AtlasVertex vertex = isAssignedGuid ? AtlasGraphUtilsV1.findByGuid(guid) : null;

        if (vertex == null && !(entityStream instanceof EntityImportStream)) { // if not found in the store, look if the entity is present in the stream
            AtlasEntity entity = entityStream.getByGuid(guid);

            if (entity != null) { // look for the entity in the store using unique-attributes
                AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());

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

                vertex = AtlasGraphUtilsV1.findByUniqueAttributes(entityType, entity.getAttributes());
            } else if (!isAssignedGuid) { // for local-guids, entity must be in the stream
                throw new AtlasBaseException(AtlasErrorCode.REFERENCED_ENTITY_NOT_FOUND, guid);
            }
        }

        if (vertex != null) {
            context.addResolvedGuid(guid, vertex);
        } else {
            if (isAssignedGuid && !(entityStream instanceof EntityImportStream)) {
                throw new AtlasBaseException(AtlasErrorCode.REFERENCED_ENTITY_NOT_FOUND, guid);
            } else {
                context.addLocalGuidReference(guid);
            }
        }
    }

    return context;
}
 

public static AtlasObjectId getAtlasObjectId(AtlasEntity entity) {
    return new AtlasObjectId(entity.getGuid(), entity.getTypeName());
}
 

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

private AtlasObjectId getAtlasObjectId(AtlasEntity e) {
    return new AtlasObjectId(e.getGuid(), e.getTypeName());
}
 

public static AtlasObjectId getAtlasObjectId(AtlasEntity entity) {
    return new AtlasObjectId(entity.getGuid(), entity.getTypeName());
}
 

public static AtlasObjectId getObjectId(AtlasEntity entity) {
    String qualifiedName = (String) entity.getAttribute(ATTRIBUTE_QUALIFIED_NAME);
    AtlasObjectId ret = new AtlasObjectId(entity.getGuid(), entity.getTypeName(), Collections.singletonMap(ATTRIBUTE_QUALIFIED_NAME, qualifiedName));

    return ret;
}
 

private boolean isProcessEntity(AtlasEntity entity) {
    String typeName = entity.getTypeName();
    AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);

    return entityType.isSubTypeOf(AtlasBaseTypeDef.ATLAS_TYPE_PROCESS);
}
 

public static Map<String, AtlasEntity> createDeptEg1() {
    Map<String, AtlasEntity> deptEmpEntities = new HashMap<>();

    AtlasEntity hrDept = new AtlasEntity(DEPARTMENT_TYPE);
    AtlasEntity john = new AtlasEntity(EMPLOYEE_TYPE);

    AtlasEntity jane = new AtlasEntity("Manager");
    AtlasEntity johnAddr = new AtlasEntity("Address");
    AtlasEntity janeAddr = new AtlasEntity("Address");
    AtlasEntity julius = new AtlasEntity("Manager");
    AtlasEntity juliusAddr = new AtlasEntity("Address");
    AtlasEntity max = new AtlasEntity(EMPLOYEE_TYPE);
    AtlasEntity maxAddr = new AtlasEntity("Address");

    AtlasObjectId deptId = new AtlasObjectId(hrDept.getGuid(), hrDept.getTypeName());
    hrDept.setAttribute("name", "hr");
    john.setAttribute("name", "John");
    john.setAttribute("department", deptId);
    johnAddr.setAttribute("street", "Stewart Drive");
    johnAddr.setAttribute("city", "Sunnyvale");
    john.setAttribute("address", johnAddr);

    john.setAttribute("birthday",new Date(1950, 5, 15));
    john.setAttribute("hasPets", true);
    john.setAttribute("numberOfCars", 1);
    john.setAttribute("houseNumber", 153);
    john.setAttribute("carMileage", 13364);
    john.setAttribute("shares", 15000);
    john.setAttribute("salary", 123345.678);
    john.setAttribute("age", 50);
    john.setAttribute("numberOfStarsEstimate", new BigInteger("1000000000000000000000"));
    john.setAttribute("approximationOfPi", new BigDecimal("3.141592653589793238462643383279502884197169399375105820974944592307816406286"));

    jane.setAttribute("name", "Jane");
    jane.setAttribute("department", deptId);
    janeAddr.setAttribute("street", "Great America Parkway");
    janeAddr.setAttribute("city", "Santa Clara");
    jane.setAttribute("address", janeAddr);
    janeAddr.setAttribute("street", "Great America Parkway");

    julius.setAttribute("name", "Julius");
    julius.setAttribute("department", deptId);
    juliusAddr.setAttribute("street", "Madison Ave");
    juliusAddr.setAttribute("city", "Newtonville");
    julius.setAttribute("address", juliusAddr);
    julius.setAttribute("subordinates", ImmutableList.of());

    AtlasObjectId janeId = AtlasTypeUtil.getAtlasObjectId(jane);
    AtlasObjectId johnId = AtlasTypeUtil.getAtlasObjectId(john);

    //TODO - Change to MANAGER_TYPE for JULIUS
    AtlasObjectId maxId = new AtlasObjectId(max.getGuid(), EMPLOYEE_TYPE);
    AtlasObjectId juliusId = new AtlasObjectId(julius.getGuid(), EMPLOYEE_TYPE);

    max.setAttribute("name", "Max");
    max.setAttribute("department", deptId);
    maxAddr.setAttribute("street", "Ripley St");
    maxAddr.setAttribute("city", "Newton");
    max.setAttribute("address", maxAddr);
    max.setAttribute("manager", janeId);
    max.setAttribute("mentor", juliusId);
    max.setAttribute("birthday",new Date(1979, 3, 15));
    max.setAttribute("hasPets", true);
    max.setAttribute("age", 36);
    max.setAttribute("numberOfCars", 2);
    max.setAttribute("houseNumber", 17);
    max.setAttribute("carMileage", 13);
    max.setAttribute("shares", Long.MAX_VALUE);
    max.setAttribute("salary", Double.MAX_VALUE);
    max.setAttribute("numberOfStarsEstimate", new BigInteger("1000000000000000000000000000000"));
    max.setAttribute("approximationOfPi", new BigDecimal("3.1415926535897932"));

    john.setAttribute("manager", janeId);
    john.setAttribute("mentor", maxId);
    hrDept.setAttribute("employees", ImmutableList.of(johnId, janeId, juliusId, maxId));

    jane.setAttribute("subordinates", ImmutableList.of(johnId, maxId));

    deptEmpEntities.put(jane.getGuid(), jane);
    deptEmpEntities.put(john.getGuid(), john);
    deptEmpEntities.put(julius.getGuid(), julius);
    deptEmpEntities.put(max.getGuid(), max);
    deptEmpEntities.put(deptId.getGuid(), hrDept);
    return deptEmpEntities;
}
 

private PreprocessorContext preProcessNotificationMessage(AtlasKafkaMessage<HookNotification> kafkaMsg) {
    PreprocessorContext context = null;

    if (preprocessEnabled) {
        context = new PreprocessorContext(kafkaMsg, typeRegistry, hiveTablesToIgnore, hiveTablesToPrune, hiveTablesCache, hiveDummyDatabasesToIgnore, hiveDummyTablesToIgnore, hiveTablePrefixesToIgnore, hiveTypesRemoveOwnedRefAttrs, rdbmsTypesRemoveOwnedRefAttrs, updateHiveProcessNameWithQualifiedName);

        if (context.isHivePreprocessEnabled()) {
            preprocessHiveTypes(context);
        }

        if (skipHiveColumnLineageHive20633) {
            skipHiveColumnLineage(context);
        }

        if (rdbmsTypesRemoveOwnedRefAttrs) {
            rdbmsTypeRemoveOwnedRefAttrs(context);
        }

        context.moveRegisteredReferredEntities();

        if (context.isHivePreprocessEnabled() && CollectionUtils.isNotEmpty(context.getEntities()) && context.getEntities().size() > 1) {
            // move hive_process and hive_column_lineage entities to end of the list
            List<AtlasEntity> entities = context.getEntities();
            int               count    = entities.size();

            for (int i = 0; i < count; i++) {
                AtlasEntity entity = entities.get(i);

                switch (entity.getTypeName()) {
                    case TYPE_HIVE_PROCESS:
                    case TYPE_HIVE_COLUMN_LINEAGE:
                        entities.remove(i--);
                        entities.add(entity);
                        count--;
                        break;
                }
            }

            if (entities.size() - count > 0) {
                LOG.info("preprocess: moved {} hive_process/hive_column_lineage entities to end of list (listSize={}). topic-offset={}, partition={}", entities.size() - count, entities.size(), kafkaMsg.getOffset(), kafkaMsg.getPartition());
            }
        }
    }

    return context;
}