Java Code Examples for org.apache.atlas.type.AtlasType#toV1Json()

@Test
public void testCreateEntity() throws Exception {
    setupRetryParams();
    AtlasClient atlasClient = new AtlasClient(service, configuration);

    WebResource.Builder builder = setupBuilder(AtlasClient.API_V1.CREATE_ENTITY, service);
    ClientResponse response = mock(ClientResponse.class);
    when(response.getStatus()).thenReturn(Response.Status.CREATED.getStatusCode());

    String jsonResponse = AtlasType.toV1Json(new EntityResult(Arrays.asList("id"), null, null));
    when(response.getEntity(String.class)).thenReturn(jsonResponse.toString());
    when(response.getLength()).thenReturn(jsonResponse.length());

    String entityJson = AtlasType.toV1Json(new Referenceable("type"));
    when(builder.method(anyString(), Matchers.<Class>any(), anyString())).thenReturn(response);

    List<String> ids = atlasClient.createEntity(entityJson);
    assertEquals(ids.size(), 1);
    assertEquals(ids.get(0), "id");
}
 

/**
 * Supports Partial updates
 * Updates properties set in the definition for the entity corresponding to guid
 * @param entityType Type of the entity being updated
 * @param uniqueAttributeName Attribute Name that uniquely identifies the entity
 * @param uniqueAttributeValue Attribute Value that uniquely identifies the entity
 * @param entity entity definition
 */
public EntityResult updateEntity(final String entityType, final String uniqueAttributeName,
                                 final String uniqueAttributeValue,
                                 Referenceable entity) throws AtlasServiceException {
    final API api        = API_V1.UPDATE_ENTITY_PARTIAL;
    String    entityJson = AtlasType.toV1Json(entity);
    LOG.debug("Updating entity type: {}, attributeName: {}, attributeValue: {}, entity: {}", entityType,
              uniqueAttributeName, uniqueAttributeValue, entityJson);
    ObjectNode response = callAPIWithRetries(api, entityJson, new ResourceCreator() {
        @Override
        public WebResource createResource() {
            WebResource resource = getResource(api, QUALIFIED_NAME);
            resource = resource.queryParam(TYPE, entityType);
            resource = resource.queryParam(ATTRIBUTE_NAME, uniqueAttributeName);
            resource = resource.queryParam(ATTRIBUTE_VALUE, uniqueAttributeValue);
            return resource;
        }
    });
    EntityResult result = extractEntityResult(response);
    LOG.debug("Update entity returned result: {}", result);
    return result;
}
 

private List<AtlasKafkaMessage<HookNotification>> testReceiveHelper(EntityUpdateRequest message, String topic) throws Exception {

        String                               json    = AtlasType.toV1Json(new AtlasNotificationMessage<>(new MessageVersion("1.0.0"), message));
        TopicPartition                       tp      = new TopicPartition(topic, 0);
        List<ConsumerRecord<String, String>> klist   = Collections.singletonList(new ConsumerRecord<>(topic, 0, 0L, "mykey", json));
        Map                                  mp      = Collections.singletonMap(tp, klist);
        ConsumerRecords                      records = new ConsumerRecords(mp);

        when(kafkaConsumer.poll(100)).thenReturn(records);

        kafkaConsumer.assign(Collections.singletonList(tp));

        AtlasKafkaConsumer                        consumer    = new AtlasKafkaConsumer(NotificationType.HOOK, kafkaConsumer, false, 100L);
        List<AtlasKafkaMessage<HookNotification>> messageList = consumer.receive();
        return messageList;
    }
 

private String getAuditEventDetail(Referenceable entity, EntityAuditAction action) throws AtlasException {
    Map<String, Object> prunedAttributes = pruneEntityAttributesForAudit(entity);

    String auditPrefix  = getV1AuditPrefix(action);
    String auditString  = auditPrefix + AtlasType.toV1Json(entity);
    byte[] auditBytes   = auditString.getBytes(StandardCharsets.UTF_8);
    long   auditSize    = auditBytes != null ? auditBytes.length : 0;
    long   auditMaxSize = auditRepository.repositoryMaxSize();

    if (auditMaxSize >= 0 && auditSize > auditMaxSize) { // don't store attributes in audit
        LOG.warn("audit record too long: entityType={}, guid={}, size={}; maxSize={}. entity attribute values not stored in audit",
                entity.getTypeName(), entity.getId()._getId(), auditSize, auditMaxSize);

        Map<String, Object> attrValues = entity.getValuesMap();

        entity.setValues(null);

        auditString = auditPrefix + AtlasType.toV1Json(entity);
        auditBytes  = auditString.getBytes(StandardCharsets.UTF_8); // recheck auditString size
        auditSize   = auditBytes != null ? auditBytes.length : 0;

        if (auditMaxSize >= 0 && auditSize > auditMaxSize) { // don't store classifications as well
            LOG.warn("audit record still too long: entityType={}, guid={}, size={}; maxSize={}. audit will have only summary details",
                    entity.getTypeName(), entity.getId()._getId(), auditSize, auditMaxSize);

            Referenceable shallowEntity = new Referenceable(entity.getId(), entity.getTypeName(), null, entity.getSystemAttributes(), null, null);

            auditString = auditPrefix + AtlasType.toJson(shallowEntity);
        }

        entity.setValues(attrValues);
    }

    restoreEntityAttributes(entity, prunedAttributes);

    return auditString;
}
 

@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
public void testNewMessageSerDe() throws Exception {
    Referenceable entity1 = new Referenceable("sometype");
    Referenceable entity2 = new Referenceable("newtype");

    entity1.set("attr", "value");
    entity1.set("complex", new Referenceable("othertype"));

    String user = "user";

    EntityCreateRequest request           = new EntityCreateRequest(user, entity1, entity2);
    String              notificationJson  = AtlasType.toV1Json(request);
    HookNotification    actualNotification = deserializer.deserialize(notificationJson);

    assertEquals(actualNotification.getType(), HookNotificationType.ENTITY_CREATE);
    assertEquals(actualNotification.getUser(), user);
    assertTrue(actualNotification instanceof EntityCreateRequest);

    EntityCreateRequest createRequest = (EntityCreateRequest) actualNotification;

    assertEquals(createRequest.getEntities().size(), 2);

    Referenceable actualEntity1 = createRequest.getEntities().get(0);

    assertEquals(actualEntity1.getTypeName(), "sometype");
    assertEquals(((Referenceable)actualEntity1.get("complex")).getTypeName(), "othertype");
    assertEquals(createRequest.getEntities().get(1).getTypeName(), "newtype");
}
 

@JsonIgnore
public String getEntityDefinitionString() {
    if (entityDefinition != null) {
        return AtlasType.toV1Json(entityDefinition);
    }
    return null;
}
 

@Test
public void testBackwardCompatibility() throws Exception {
    //Code to generate the json, use it for hard-coded json used later in this test
    Referenceable entity = new Referenceable("sometype");
    entity.set("attr", "value");

    EntityCreateRequest request                  = new EntityCreateRequest(null, entity);
    String              notificationJsonFromCode = AtlasType.toV1Json(request);

    System.out.println(notificationJsonFromCode);

    //Json without user and assert that the string can be deserialised
    String notificationJson = "{\n"
            + "  \"entities\": [\n"
            + "    {\n"
            + "      \"jsonClass\": \"org.apache.atlas.typesystem.json.InstanceSerialization$_Reference\",\n"
            + "      \"id\": {\n"
            + "        \"jsonClass\": \"org.apache.atlas.typesystem.json.InstanceSerialization$_Id\",\n"
            + "        \"id\": \"-1459493350903186000\",\n"
            + "        \"version\": 0,\n"
            + "        \"typeName\": \"sometype\",\n"
            + "        \"state\": \"ACTIVE\"\n"
            + "      },\n"
            + "      \"typeName\": \"sometype\",\n"
            + "      \"values\": {\n"
            + "        \"attr\": \"value\"\n"
            + "      },\n"
            + "      \"traitNames\": [],\n"
            + "      \"traits\": {}\n"
            + "    }\n"
            + "  ],\n"
            + "  \"type\": \"ENTITY_CREATE\"\n"
            + "}";


    HookNotification actualNotification = deserializer.deserialize(notificationJson);

    assertEquals(actualNotification.getType(), HookNotificationType.ENTITY_CREATE);
    assertEquals(actualNotification.getUser(), HookNotification.UNKNOW_USER);
}
 

void createTypes() throws Exception {
    TypesDef typesDef = createTypeDefinitions();

    String typesAsJSON = AtlasType.toV1Json(typesDef);
    System.out.println("typesAsJSON = " + typesAsJSON);
    metadataServiceClient.createType(typesAsJSON);

    // verify types created
    verifyTypesCreated();
}
 

private void createTrait(String traitName, String ... superTraitNames) throws Exception {
    TraitTypeDefinition traitDef = TypesUtil.createTraitTypeDef(traitName, null, new HashSet<>(Arrays.asList(superTraitNames)));
    TypesDef            typesDef = new TypesDef(Collections.<EnumTypeDefinition>emptyList(),
                                                Collections.<StructTypeDefinition>emptyList(),
                                                Collections.singletonList(traitDef),
                                                Collections.<ClassTypeDefinition>emptyList());
    String traitDefinitionJSON = AtlasType.toV1Json(typesDef);

    LOG.debug("Trait definition = {}", traitDefinitionJSON);

    createType(traitDefinitionJSON);
}
 

private List<AtlasEntityDef> convertV1toV2(List<ClassTypeDefinition> types) throws AtlasBaseException {
    List<ClassTypeDefinition> classTypeList       = new ArrayList(types);
    TypesDef                  toConvert           = new TypesDef(Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), classTypeList);
    String                    json                = AtlasType.toV1Json(toConvert);
    AtlasTypeRegistry         emptyRegistry       = new AtlasTypeRegistry();
    AtlasTypesDef             converted           = TypeConverterUtil.toAtlasTypesDef(json, emptyRegistry);
    List<AtlasEntityDef>      convertedEntityDefs = converted.getEntityDefs();

    return convertedEntityDefs;
}
 

@Test
public void testAddTraitWithAttribute() throws Exception {
    String dbName = "db" + randomString();
    String tableName = "table" + randomString();
    Referenceable hiveDBInstance = createHiveDBInstanceBuiltIn(dbName);
    Id dbId = createInstance(hiveDBInstance);
    Referenceable hiveTableInstance = createHiveTableInstanceBuiltIn(dbName, tableName, dbId);
    Id id = createInstance(hiveTableInstance);

    final String guid = id._getId();
    try {
        Assert.assertNotNull(UUID.fromString(guid));
    } catch (IllegalArgumentException e) {
        Assert.fail("Response is not a guid, " + guid);
    }

    final String traitName = "PII_Trait" + randomString();
    TraitTypeDefinition piiTrait = TypesUtil
            .createTraitTypeDef(traitName, null, Collections.<String>emptySet(),
                    TypesUtil.createRequiredAttrDef("type", AtlasBaseTypeDef.ATLAS_TYPE_STRING));
    String traitDefinitionAsJSON = AtlasType.toV1Json(piiTrait);
    LOG.debug("traitDefinitionAsJSON = {}", traitDefinitionAsJSON);

    TypesDef typesDef = new TypesDef(Collections.emptyList(), Collections.emptyList(), Collections.singletonList(piiTrait), Collections.emptyList());

    createType(typesDef);

    Struct traitInstance = new Struct(traitName);
    traitInstance.set("type", "SSN");
    atlasClientV1.addTrait(guid, traitInstance);

    // verify the response
    Referenceable entity = atlasClientV1.getEntity(guid);
    Assert.assertNotNull(entity);
    Assert.assertEquals(entity.getId()._getId(), guid);

    assertNotNull(entity.getTrait(traitName));
    assertEquals(entity.getTrait(traitName).get("type"), traitInstance.get("type"));
}
 

@Test
  public void testNextVersionMismatch() throws Exception {
      Referenceable                        entity  = getEntity(TRAIT_NAME);
      EntityUpdateRequest                  message = new EntityUpdateRequest("user1", entity);
      String                               json    = AtlasType.toV1Json(new AtlasNotificationMessage<>(new MessageVersion("2.0.0"), message));
      TopicPartition                       tp      = new TopicPartition(ATLAS_HOOK_TOPIC,0);
      List<ConsumerRecord<String, String>> klist   = Collections.singletonList(new ConsumerRecord<>(ATLAS_HOOK_TOPIC, 0, 0L, "mykey", json));
      Map                                  mp      = Collections.singletonMap(tp,klist);
      ConsumerRecords                      records = new ConsumerRecords(mp);

      kafkaConsumer.assign(Collections.singletonList(tp));

      when(kafkaConsumer.poll(100L)).thenReturn(records);

      AtlasKafkaConsumer consumer =new AtlasKafkaConsumer(NotificationType.HOOK, kafkaConsumer ,false, 100L);

      try {
          List<AtlasKafkaMessage<HookNotification>> messageList = consumer.receive();

          assertTrue(messageList.size() > 0);

          HookNotification consumedMessage  = messageList.get(0).getMessage();

          fail("Expected VersionMismatchException!");
      } catch (IncompatibleVersionException e) {
          e.printStackTrace();
      }
}
 

@Test
public void testAddTrait() throws Exception {
    String dbName = "db" + randomString();
    String tableName = "table" + randomString();
    Referenceable hiveDBInstance = createHiveDBInstanceBuiltIn(dbName);
    Id dbId = createInstance(hiveDBInstance);
    Referenceable hiveTableInstance = createHiveTableInstanceBuiltIn(dbName, tableName, dbId);
    Id id = createInstance(hiveTableInstance);

    final String guid = id._getId();
    try {
        Assert.assertNotNull(UUID.fromString(guid));
    } catch (IllegalArgumentException e) {
        Assert.fail("Response is not a guid, " + guid);
    }

    String traitName = "PII_Trait" + randomString();
    TraitTypeDefinition piiTrait =
            TypesUtil.createTraitTypeDef(traitName, null, Collections.<String>emptySet());
    String traitDefinitionAsJSON = AtlasType.toV1Json(piiTrait);
    LOG.debug("traitDefinitionAsJSON = {}", traitDefinitionAsJSON);

    TypesDef typesDef = new TypesDef(Collections.emptyList(), Collections.emptyList(), Collections.singletonList(piiTrait), Collections.emptyList());

    createType(typesDef);

    Struct traitInstance = new Struct(traitName);

    atlasClientV1.addTrait(guid, traitInstance);
    assertEntityAudit(guid, EntityAuditEvent.EntityAuditAction.TAG_ADD);
}
 

private String[] addTraits() throws Exception {
    String[] traitNames = {"class_trait", "secure_trait", "pii_trait", "ssn_trait", "salary_trait", "sox_trait",};

    for (String traitName : traitNames) {
        TraitTypeDefinition traitTypeDef =
                TypesUtil.createTraitTypeDef(traitName, null, Collections.<String>emptySet());
        TypesDef typesDef = new TypesDef(Collections.emptyList(), Collections.emptyList(), Collections.singletonList(traitTypeDef), Collections.emptyList());

        String json = AtlasType.toV1Json(typesDef);
        createType(json);
    }

    return traitNames;
}
 

@Test
public void testSubmit() throws Exception {
    for (HierarchicalTypeDefinition typeDefinition : typeDefinitions) {
        try{
            atlasClientV1.getType(typeDefinition.getTypeName());
        } catch (AtlasServiceException ase){
            TypesDef typesDef = null;

            if (typeDefinition instanceof ClassTypeDefinition) {
                typesDef = new TypesDef(Collections.emptyList(), Collections.emptyList(),
                        Collections.emptyList(), Collections.singletonList((ClassTypeDefinition) typeDefinition));
            } else if (typeDefinition instanceof TraitTypeDefinition) {
                typesDef = new TypesDef(Collections.emptyList(), Collections.emptyList(),
                        Collections.singletonList((TraitTypeDefinition) typeDefinition), Collections.emptyList());
            }

            String typesAsJSON = AtlasType.toV1Json(typesDef);

            System.out.println("typesAsJSON = " + typesAsJSON);

            ObjectNode response = atlasClientV1.callAPIWithBody(AtlasClient.API_V1.CREATE_TYPE, typesAsJSON);
            Assert.assertNotNull(response);


            ArrayNode typesAdded = (ArrayNode) response.get(AtlasClient.TYPES);
            assertEquals(typesAdded.size(), 1);
            assertEquals(typesAdded.get(0).get(NAME).asText(), typeDefinition.getTypeName());
            Assert.assertNotNull(response.get(AtlasClient.REQUEST_ID));}
    }
}
 

private Id createInstance(Referenceable referenceable) throws Exception {
    String typeName = referenceable.getTypeName();

    String entityJSON = AtlasType.toV1Json(referenceable);
    System.out.println("Submitting new entity= " + entityJSON);
    List<String> guids = metadataServiceClient.createEntity(entityJSON);
    System.out.println("created instance for type " + typeName + ", guid: " + guids);

    // return the Id for created instance with guid
    if (guids.size() > 0) {
        return new Id(guids.get(guids.size() - 1), referenceable.getId().getVersion(), referenceable.getTypeName());
    }

    return null;
}
 

@Override
public String toString() { return AtlasType.toV1Json(this); }
 

@Test
public void testDeserializeCompressedMessage() throws Exception {
    Referenceable       entity     = generateLargeEntityWithTrait();
    EntityUpdateRequest message    = new EntityUpdateRequest("user1", entity);
    List<String>       jsonMsgList = new ArrayList<>();

    AbstractNotification.createNotificationMessages(message, jsonMsgList);

    assertTrue(jsonMsgList.size() == 1);

    String compressedMsg   = jsonMsgList.get(0);
    String uncompressedMsg = AtlasType.toV1Json(message);

    assertTrue(compressedMsg.length() < uncompressedMsg.length(), "Compressed message (" + compressedMsg.length() + ") should be shorter than uncompressed message (" + uncompressedMsg.length() + ")");

    HookNotification deserializedMessage = deserialize(jsonMsgList);

    assertEqualMessage(deserializedMessage, message);
}
 

@Override
public String toString() {
    return AtlasType.toV1Json(this);
}