org.apache.atlas.type.AtlasType Java Examples

private void cleanupIndexForAttribute(AtlasGraphManagement management, String typeName, AtlasAttributeDef attributeDef) {
    final String propertyName = GraphHelper.encodePropertyKey(typeName + "." + attributeDef.getName());
    String attribTypeName = attributeDef.getTypeName();
    boolean isBuiltInType = AtlasTypeUtil.isBuiltInType(attribTypeName);
    boolean isArrayType = AtlasTypeUtil.isArrayType(attribTypeName);
    boolean isMapType = AtlasTypeUtil.isMapType(attribTypeName);

    try {
        AtlasType atlasType = typeRegistry.getType(attribTypeName);

        if (isMapType || isArrayType || isClassificationType(atlasType) || isEntityType(atlasType)) {
            LOG.warn("Ignoring non-indexable attribute {}", attribTypeName);
        } else if (isBuiltInType || isEnumType(atlasType)) {
            cleanupIndex(management, propertyName);
        } else if (isStructType(atlasType)) {
            AtlasStructDef structDef = typeRegistry.getStructDefByName(attribTypeName);
            cleanupIndices(management, structDef);
        }
    } catch (AtlasBaseException e) {
        LOG.error("No type exists for {}", attribTypeName, e);
    }
}
 

protected void validateEntity(AtlasEntityExtInfo entityExtInfo, AtlasStruct actual, AtlasStruct expected) throws AtlasBaseException, AtlasException {
    if (expected == null) {
        Assert.assertNull(actual, "expected null instance. Found " + actual);

        return;
    }

    Assert.assertNotNull(actual, "found null instance");

    AtlasStructType entityType = (AtlasStructType) typeRegistry.getType(actual.getTypeName());
    for (String attrName : expected.getAttributes().keySet()) {
        Object expectedVal = expected.getAttribute(attrName);
        Object actualVal   = actual.getAttribute(attrName);

        AtlasType attrType = entityType.getAttributeType(attrName);
        validateAttribute(entityExtInfo, actualVal, expectedVal, attrType, attrName);
    }
}
 

public static TypesDef toTypesDef(AtlasType type, AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
    final TypesDef ret;

    if (type instanceof AtlasEnumType) {
        ret = TypeConverterUtil.enumToTypesDef((AtlasEnumType) type);
    } else if (type instanceof AtlasEntityType) {
        ret = TypeConverterUtil.entityToTypesDef((AtlasEntityType) type, typeRegistry);
    } else if (type instanceof AtlasClassificationType) {
        ret = TypeConverterUtil.classificationToTypesDef((AtlasClassificationType) type, typeRegistry);
    } else if (type instanceof AtlasStructType) {
        ret = TypeConverterUtil.structToTypesDef((AtlasStructType) type, typeRegistry);
    } else {
        ret = new TypesDef();
    }

    return ret;
}
 

@Test
public void testUTF8() throws Exception {
    String              attrName            = randomUTF8();
    String              attrValue           = randomUTF8();
    String              classType           = randomString(); //Type names cannot be arbitrary UTF8 characters. See org.apache.atlas.type.AtlasTypeUtil#validateType()
    ClassTypeDefinition classTypeDefinition = TypesUtil.createClassTypeDef(classType, null, Collections.<String>emptySet(), TypesUtil.createUniqueRequiredAttrDef(attrName, AtlasBaseTypeDef.ATLAS_TYPE_STRING));
    TypesDef            typesDef            = new TypesDef(Collections.<EnumTypeDefinition>emptyList(), Collections.<StructTypeDefinition>emptyList(), Collections.<TraitTypeDefinition>emptyList(), Collections.singletonList(classTypeDefinition));

    createType(typesDef);

    Referenceable entityToCreate  = new Referenceable(classType, Collections.singletonMap(attrName, attrValue));
    Id            guid            = createInstance(entityToCreate);
    ObjectNode    response        = atlasClientV1.callAPIWithBodyAndParams(AtlasClient.API_V1.GET_ENTITY, null, guid._getId());
    Object        objResponse     = response.get(AtlasClient.DEFINITION);
    String        jsonResponse    = AtlasType.toJson(objResponse);
    Referenceable createdEntity   = AtlasType.fromV1Json(jsonResponse, Referenceable.class);
    Object        entityAttrValue = createdEntity.get(attrName);

    Assert.assertEquals(entityAttrValue, attrValue,
                        "attrName=" + attrName + "; attrValue=" + attrValue + "; entityToCreate=" + entityToCreate + "; entityId=" + guid + "; getEntityResponse_Obj=" + objResponse + "; getEntityResponse_Json=" + jsonResponse + "; getEntityResponse_Entity=" + createdEntity);
}
 

@Test
public void serializeAtlasImportRequstFromJsonWithEmptyTransforms() {
    String jsonData = "{ \"options\": { \"transforms\": \"{ }\" } }";

    AtlasImportRequest request = AtlasType.fromJson(jsonData, AtlasImportRequest.class);

    assertNotNull(request);
    assertNotNull(request.getOptions());
    assertNotNull(request.getOptions().get(AtlasImportRequest.TRANSFORMS_KEY));

    ImportTransforms tr = ImportTransforms.fromJson(request.getOptions().get(AtlasImportRequest.TRANSFORMS_KEY));

    assertNotNull(tr);
    assertNotNull(tr.getTransforms());
    assertEquals(tr.getTransforms().size(), 0);
}
 

protected void deleteEdge(AtlasEdge edge, boolean updateInverseAttribute, boolean force) throws AtlasBaseException {
    //update inverse attribute
    if (updateInverseAttribute) {
        AtlasEdgeLabel atlasEdgeLabel = new AtlasEdgeLabel(edge.getLabel());

        AtlasType parentType = typeRegistry.getType(atlasEdgeLabel.getTypeName());

        if (parentType instanceof AtlasEntityType) {
            AtlasEntityType parentEntityType = (AtlasEntityType) parentType;

            AtlasStructType.AtlasAttribute attribute = parentEntityType.getAttribute(atlasEdgeLabel.getAttributeName());
            if (attribute.getInverseRefAttribute() != null) {
                deleteEdgeBetweenVertices(edge.getInVertex(), edge.getOutVertex(), attribute.getInverseRefAttribute());
            }
        }
    }

    deleteEdge(edge, force);
}
 

private void mapAttributes(AtlasEdge edge, AtlasRelationshipWithExtInfo relationshipWithExtInfo) throws AtlasBaseException {
    AtlasRelationship relationship = relationshipWithExtInfo.getRelationship();
    AtlasType         objType      = typeRegistry.getType(relationship.getTypeName());

    if (!(objType instanceof AtlasRelationshipType)) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, relationship.getTypeName());
    }

    AtlasRelationshipType relationshipType = (AtlasRelationshipType) objType;

    for (AtlasAttribute attribute : relationshipType.getAllAttributes().values()) {
        // mapping only primitive attributes
        Object attrValue = mapVertexToPrimitive(edge, attribute.getVertexPropertyName(), attribute.getAttributeDef());

        relationship.setAttribute(attribute.getName(), attrValue);
    }
}
 

@Test
//API should accept single entity (or jsonarray of entities)
public void testSubmitSingleEntity() throws Exception {
    Referenceable databaseInstance = new Referenceable(DATABASE_TYPE_BUILTIN);
    String dbName = randomString();
    databaseInstance.set("name", dbName);
    databaseInstance.set(QUALIFIED_NAME, dbName);
    databaseInstance.set("clusterName", randomString());
    databaseInstance.set("description", randomString());
    databaseInstance.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, dbName);
    databaseInstance.set("owner", "user1");
    databaseInstance.set("clusterName", "cl1");
    databaseInstance.set("parameters", Collections.EMPTY_MAP);
    databaseInstance.set("location", "/tmp");

    ObjectNode response = atlasClientV1.callAPIWithBody(AtlasClient.API_V1.CREATE_ENTITY, AtlasType.toV1Json(databaseInstance));
    assertNotNull(response);
    Assert.assertNotNull(response.get(AtlasClient.REQUEST_ID));

    EntityResult entityResult = EntityResult.fromString(response.toString());
    assertEquals(entityResult.getCreatedEntities().size(), 1);
    assertNotNull(entityResult.getCreatedEntities().get(0));
}
 

@Test
public void serializeAtlasImportRequstFromJsonWith1Transform() {
    String jsonData = "{ \"options\": { \"transforms\": \"{ \\\"hive_db\\\": { \\\"qualifiedName\\\": [ \\\"replace:@cl1:@cl2\\\" ] } }\" } }";

    AtlasImportRequest request = AtlasType.fromJson(jsonData, AtlasImportRequest.class);

    assertNotNull(request);
    assertNotNull(request.getOptions());
    assertNotNull(request.getOptions().get(AtlasImportRequest.TRANSFORMS_KEY));

    ImportTransforms tr = ImportTransforms.fromJson(request.getOptions().get(AtlasImportRequest.TRANSFORMS_KEY));

    assertNotNull(tr);
    assertNotNull(tr.getTransforms());
    assertEquals(tr.getTransforms().size(), 1);
    assertTrue(tr.getTransforms().containsKey("hive_db"));
    assertEquals(tr.getTransforms("hive_db").entrySet().size(), 1);
    assertTrue(tr.getTransforms("hive_db").containsKey("qualifiedName"));
    assertEquals(tr.getTransforms("hive_db").get("qualifiedName").size(), 1);
}
 

@Test
public void serializeOptions_VerifyAccessors() {
    String guid = "\"abcd\"";
    String pos = "\"1\"";
    String trueVal = "\"true\"";

    String jsonData = "{ \"options\": " +
            "               {" +
            "\"startGuid\":" + guid + "," +
            "\"startPosition\":" + pos + "," +
            "\"updateTypeDefinition\":" + trueVal +
            "}" +
            "}";

    AtlasImportRequest request = AtlasType.fromJson(jsonData, AtlasImportRequest.class);

    assertNotNull(request);
    assertNotNull(request.getStartGuid());
    assertNotNull(request.getStartPosition());
    assertNotNull(request.getUpdateTypeDefs());

    assertEquals(request.getStartGuid(), guid.replace("\"", ""));
    assertEquals(request.getStartPosition(), pos.replace("\"", ""));
    assertEquals(request.getUpdateTypeDefs(), trueVal.replace("\"", ""));
}
 

@Test
public void serializeAtlasImportRequstFromJsonWithEmptyTransforms() {
    String jsonData = "{ \"options\": { \"transforms\": \"{ }\" } }";

    AtlasImportRequest request = AtlasType.fromJson(jsonData, AtlasImportRequest.class);

    assertNotNull(request);
    assertNotNull(request.getOptions());
    assertNotNull(request.getOptions().get(AtlasImportRequest.TRANSFORMS_KEY));

    ImportTransforms tr = ImportTransforms.fromJson(request.getOptions().get(AtlasImportRequest.TRANSFORMS_KEY));

    assertNotNull(tr);
    assertNotNull(tr.getTransforms());
    assertEquals(tr.getTransforms().size(), 0);
}
 

@Test
public void serializeOptions_VerifyAccessors() {
    String guid = "\"abcd\"";
    String pos = "\"1\"";
    String trueVal = "\"true\"";

    String jsonData = "{ \"options\": " +
            "               {" +
            "\"startGuid\":" + guid + "," +
            "\"startPosition\":" + pos + "," +
            "\"updateTypeDefinition\":" + trueVal +
            "}" +
            "}";

    AtlasImportRequest request = AtlasType.fromJson(jsonData, AtlasImportRequest.class);

    assertNotNull(request);
    assertNotNull(request.getStartGuid());
    assertNotNull(request.getStartPosition());
    assertNotNull(request.getUpdateTypeDefs());

    assertEquals(request.getStartGuid(), guid.replace("\"", ""));
    assertEquals(request.getStartPosition(), pos.replace("\"", ""));
    assertEquals(request.getUpdateTypeDefs(), trueVal.replace("\"", ""));
}
 

public static TypesDef toTypesDef(AtlasType type, AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
    final TypesDef ret;

    if (type instanceof AtlasEnumType) {
        ret = TypeConverterUtil.enumToTypesDef((AtlasEnumType) type);
    } else if (type instanceof AtlasEntityType) {
        ret = TypeConverterUtil.entityToTypesDef((AtlasEntityType) type, typeRegistry);
    } else if (type instanceof AtlasClassificationType) {
        ret = TypeConverterUtil.classificationToTypesDef((AtlasClassificationType) type, typeRegistry);
    } else if (type instanceof AtlasStructType) {
        ret = TypeConverterUtil.structToTypesDef((AtlasStructType) type, typeRegistry);
    } else {
        ret = new TypesDef();
    }

    return ret;
}
 

@Override
public void onClassificationsAdded(AtlasEntity entity, List<AtlasClassification> classifications) throws AtlasBaseException {
    if (CollectionUtils.isNotEmpty(classifications)) {
        MetricRecorder metric = RequestContext.get().startMetricRecord("entityAudit");

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

        for (AtlasClassification classification : classifications) {
            if (entity.getGuid().equals(classification.getEntityGuid())) {
                events.add(createEvent(entity, CLASSIFICATION_ADD, "Added classification: " + AtlasType.toJson(classification)));
            } else {
                events.add(createEvent(entity, PROPAGATED_CLASSIFICATION_ADD, "Added propagated classification: " + AtlasType.toJson(classification)));
            }
        }

        auditRepository.putEventsV2(events);

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

void visitCollectionReferences(AtlasType elemType, Object val) throws AtlasBaseException {
    if (elemType == null || val == null || isPrimitive(elemType.getTypeCategory())) {
        return;
    }

    Iterator it = null;

    if (val instanceof Collection) {
        it = ((Collection) val).iterator();
    } else if (val instanceof Iterable) {
        it = ((Iterable) val).iterator();
    } else if (val instanceof Iterator) {
        it = (Iterator) val;
    }

    if (it != null) {
        while (it.hasNext()) {
            Object elem = it.next();
            visitAttribute(elemType, elem);
        }
    }
}
 

private void addType(AtlasType type, ExportContext context) {
    if (type.getTypeCategory() == TypeCategory.PRIMITIVE) {
        return;
    }

    if (type instanceof AtlasArrayType) {
        AtlasArrayType arrayType = (AtlasArrayType)type;

        addType(arrayType.getElementType(), context);
    } else if (type instanceof AtlasMapType) {
        AtlasMapType mapType = (AtlasMapType)type;

        addType(mapType.getKeyType(), context);
        addType(mapType.getValueType(), context);
    } else if (type instanceof AtlasEntityType) {
        addEntityType((AtlasEntityType)type, context);
    } else if (type instanceof AtlasClassificationType) {
        addClassificationType((AtlasClassificationType)type, context);
    } else if (type instanceof AtlasStructType) {
        addStructType((AtlasStructType)type, context);
    } else if (type instanceof AtlasEnumType) {
        addEnumType((AtlasEnumType)type, context);
    }
}
 

@Test
public void serializeAtlasImportRequstFromJsonWith1Transform() {
    String jsonData = "{ \"options\": { \"transforms\": \"{ \\\"hive_db\\\": { \\\"qualifiedName\\\": [ \\\"replace:@cl1:@cl2\\\" ] } }\" } }";

    AtlasImportRequest request = AtlasType.fromJson(jsonData, AtlasImportRequest.class);

    assertNotNull(request);
    assertNotNull(request.getOptions());
    assertNotNull(request.getOptions().get(AtlasImportRequest.TRANSFORMS_KEY));

    ImportTransforms tr = ImportTransforms.fromJson(request.getOptions().get(AtlasImportRequest.TRANSFORMS_KEY));

    assertNotNull(tr);
    assertNotNull(tr.getTransforms());
    assertEquals(tr.getTransforms().size(), 1);
    assertTrue(tr.getTransforms().containsKey("hive_db"));
    assertEquals(tr.getTransforms("hive_db").entrySet().size(), 1);
    assertTrue(tr.getTransforms("hive_db").containsKey("qualifiedName"));
    assertEquals(tr.getTransforms("hive_db").get("qualifiedName").size(), 1);
}
 

public ITypedStruct getTrait(AtlasClassification classification) throws AtlasBaseException {
    AtlasFormatConverter converter          = instanceFormatters.getConverter(TypeCategory.CLASSIFICATION);
    AtlasType            classificationType = typeRegistry.getType(classification.getTypeName());
    Struct               trait               = (Struct)converter.fromV2ToV1(classification, classificationType, new ConverterContext());

    try {
        return metadataService.createTraitInstance(trait);
    } catch (AtlasException e) {
        LOG.error("Exception while getting a typed reference for the entity ", e);
        throw toAtlasBaseException(e);
    }
}
 

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

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

    return ret;
}
 

@BeforeTest
public void setup() throws AtlasException, InterruptedException, AtlasBaseException {
    MockitoAnnotations.initMocks(this);
    AtlasType mockType = mock(AtlasType.class);
    when(typeRegistry.getType(anyString())).thenReturn(mockType);
    AtlasEntity.AtlasEntitiesWithExtInfo mockEntity = mock(AtlasEntity.AtlasEntitiesWithExtInfo.class);
    when(instanceConverter.toAtlasEntities(anyList())).thenReturn(mockEntity);
    kafkaNotification = startKafkaServer();
}
 

AtlasEntity serDeserEntity(AtlasEntity entity) throws IOException {
    //Convert from json to object and back to trigger the case where it gets translated to a map for attributes instead of AtlasEntity
    String      jsonString = AtlasType.toJson(entity);
    AtlasEntity newEntity  = AtlasType.fromJson(jsonString, AtlasEntity.class);

    return newEntity;
}
 

private void addType(String typeName, ExportContext context) {
    AtlasType type = null;

    try {
        type = typeRegistry.getType(typeName);

        addType(type, context);
    } catch (AtlasBaseException excp) {
        LOG.error("unknown type {}", typeName);
    }
}
 

private <T> T convertFromJson(Class<T> clazz, String jsonData) throws AtlasBaseException {
    T t;
    try {
        t = AtlasType.fromJson(jsonData, clazz);
        if (t == null) {
            LOG.error("Error converting file to JSON.");
            return null;
        }
    } catch (Exception e) {
        throw new AtlasBaseException("Error converting file to JSON.", e);
    }

    return t;
}
 

@Test
public void testDeserializeLegacyMessage() throws Exception {
    Referenceable       entity              = generateEntityWithTrait();
    EntityUpdateRequest message             = new EntityUpdateRequest("user1", entity);
    String              jsonMsg             = AtlasType.toV1Json(message);
    HookNotification    deserializedMessage = deserialize(Collections.singletonList(jsonMsg));

    assertEqualMessage(deserializedMessage, message);
}
 

@Test(dependsOnMethods = "verifySavedSearchesForUser")
public void verifyQueryConversionFromJSON() throws AtlasBaseException {
    List<AtlasUserSavedSearch> list = userProfileService.getSavedSearches("first-0");

    for (int i = 0; i < NUM_SEARCHES; i++) {
        SearchParameters sp   = list.get(i).getSearchParameters();
        String           json = AtlasType.toJson(sp);

        assertEquals(AtlasType.toJson(getActualSearchParameters()).replace("\n", "").replace(" ", ""), json);
    }
}
 

@Override
public void onTraitsAdded(Referenceable entity, Collection<? extends Struct> traits) throws AtlasException {
    if (traits != null) {
        MetricRecorder metric = RequestContext.get().startMetricRecord("entityAudit");

        for (Struct trait : traits) {
            EntityAuditEvent event = createEvent(entity, EntityAuditAction.TAG_ADD,
                                                 "Added trait: " + AtlasType.toV1Json(trait));

            auditRepository.putEventsV1(event);
        }

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

void visitEntity(AtlasEntityType entityType, AtlasEntity entity) throws AtlasBaseException {
    visitStruct(entityType, entity);

    for (AtlasAttribute attribute : entityType.getRelationshipAttributes().values()) {
        AtlasType attrType = attribute.getAttributeType();
        Object    attrVal  = entity.getRelationshipAttribute(attribute.getName());

        visitAttribute(attrType, attrVal);
    }
}
 

@Override
public Object fromV2ToV1(Object v2Obj, AtlasType type, ConverterContext converterContext) throws AtlasBaseException {
    Id ret = null;

    if (v2Obj != null) {

        if (v2Obj instanceof Map) {
            Map    v2Map    = (Map) v2Obj;
            String idStr    = (String)v2Map.get(AtlasObjectId.KEY_GUID);
            String typeName = type.getTypeName();

            if (StringUtils.isEmpty(idStr)) {
                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND);
            }

            ret = new Id(idStr, 0, typeName);
        } else if (v2Obj instanceof AtlasObjectId) { // transient-id
            ret = new Id(((AtlasObjectId) v2Obj).getGuid(), 0, type.getTypeName());
        } else if (v2Obj instanceof AtlasEntity) {
            AtlasEntity entity = (AtlasEntity) v2Obj;
            ret = new Id(((AtlasObjectId) v2Obj).getGuid(), 0, type.getTypeName());
        } else {
            throw new AtlasBaseException(AtlasErrorCode.TYPE_CATEGORY_INVALID, type.getTypeCategory().name());
        }
    }
    return ret;
}
 

private AtlasEntityType getEntityType(String typeName) throws AtlasBaseException {
    AtlasType objType = typeRegistry.getType(typeName);

    if (!(objType instanceof AtlasEntityType)) {
        throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, typeName);
    }

    return (AtlasEntityType)objType;
}
 

void visitStruct(AtlasStructType structType, AtlasStruct struct) throws AtlasBaseException {
    for (AtlasAttribute attribute : structType.getAllAttributes().values()) {
        AtlasType attrType = attribute.getAttributeType();
        Object    attrVal  = struct.getAttribute(attribute.getName());

        visitAttribute(attrType, attrVal);
    }
}