org.apache.atlas.v1.model.instance.Referenceable Java Examples

Id loadProcess(String name, String description, String user, List<Id> inputTables, List<Id> outputTables,
               String queryText, String queryPlan, String queryId, String queryGraph, String... traitNames)
        throws AtlasBaseException {
    try {
        Referenceable referenceable = new Referenceable(LOAD_PROCESS_TYPE, traitNames);
        // super type attributes
        referenceable.set(AtlasClient.NAME, name);
        referenceable.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
        referenceable.set("description", description);
        referenceable.set(INPUTS_ATTRIBUTE, inputTables);
        referenceable.set(OUTPUTS_ATTRIBUTE, outputTables);

        referenceable.set("user", user);
        referenceable.set("startTime", System.currentTimeMillis());
        referenceable.set("endTime", System.currentTimeMillis() + 10000);

        referenceable.set("queryText", queryText);
        referenceable.set("queryPlan", queryPlan);
        referenceable.set("queryId", queryId);
        referenceable.set("queryGraph", queryGraph);

        return createInstance(referenceable);
    } catch (Exception e) {
        throw new AtlasBaseException(AtlasErrorCode.QUICK_START, e, String.format("%s process entity creation failed", name));
    }
}
 

private void verifyFeedLineage(String feedName, String clusterName, String feedId, String dbName, String tableName)
        throws Exception{
    //verify that lineage from hive table to falcon feed is created
    String processId = assertEntityIsRegistered(FalconDataTypes.FALCON_FEED_CREATION.getName(),
            AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
            FalconBridge.getFeedQualifiedName(feedName, clusterName));
    Referenceable processEntity = atlasClient.getEntity(processId);
    assertEquals(((List<Id>)processEntity.get("outputs")).get(0).getId(), feedId);

    String inputId = ((List<Id>) processEntity.get("inputs")).get(0).getId();
    Referenceable tableEntity = atlasClient.getEntity(inputId);
    assertEquals(tableEntity.getTypeName(), HiveDataTypes.HIVE_TABLE.getName());
    assertEquals(tableEntity.get(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME),
            HiveMetaStoreBridge.getTableQualifiedName(clusterName, dbName, tableName));

}
 

/**
 * Get an entity given the entity id
 * @param entityType entity type name
 * @param attribute qualified name of the entity
 * @param value
 * @return result object
 * @throws AtlasServiceException
 */
public Referenceable getEntity(final String entityType, final String attribute, final String value)
        throws AtlasServiceException {
    final API api = API_V1.GET_ENTITY;
    ObjectNode jsonResponse = callAPIWithRetries(api, null, new ResourceCreator() {
        @Override
        public WebResource createResource() {
            WebResource resource = getResource(api);
            resource = resource.queryParam(TYPE, entityType);
            resource = resource.queryParam(ATTRIBUTE_NAME, attribute);
            resource = resource.queryParam(ATTRIBUTE_VALUE, value);
            return resource;
        }
    });
    String entityInstanceDefinition =  AtlasType.toJson(jsonResponse.get(AtlasClient.DEFINITION));
    return AtlasType.fromV1Json(entityInstanceDefinition, Referenceable.class);
}
 

@Test
public void testDeleteByQualifiedName() throws Exception {
    final Referenceable entity        = new Referenceable(DATABASE_TYPE_BUILTIN);
    final String        dbName        = "db" + randomString();
    final String        clusterName   = randomString();
    final String        qualifiedName = dbName + "@" + clusterName;

    entity.set(NAME, dbName);
    entity.set(DESCRIPTION, randomString());
    entity.set(QUALIFIED_NAME, qualifiedName);
    entity.set(CLUSTER_NAME, clusterName);

    final String dbId = atlasClientV1.createEntity(entity).get(0);

    sendHookMessage(new EntityDeleteRequest(TEST_USER, DATABASE_TYPE_BUILTIN, QUALIFIED_NAME, qualifiedName));

    waitFor(MAX_WAIT_TIME, new Predicate() {
        @Override
        public boolean evaluate() throws Exception {
            Referenceable getEntity = atlasClientV1.getEntity(dbId);

            return getEntity.getId().getState() == Id.EntityState.DELETED;
        }
    });
}
 

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

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

@Override
public DataSetRefs analyze(AnalysisContext context, ProvenanceEventRecord event) {
    if (event.getTransitUri() == null) {
        return null;
    }

    final URI uri = parseUri(event.getTransitUri());
    final String namespace = context.getNamespaceResolver().fromHostNames(uri.getHost());
    final Set<Tuple<String, String>> outputTables = parseTableNames(null, event.getAttribute(ATTR_OUTPUT_TABLES));
    if (outputTables.isEmpty()) {
        return null;
    }

    final DataSetRefs refs = new DataSetRefs(event.getComponentId());
    outputTables.forEach(tableName -> {
        final Referenceable ref = createTableRef(namespace, tableName);
        refs.addOutput(ref);
    });
    return refs;
}
 

/**
 * For the given type, finds an unique attribute and checks if there is an existing instance with the same
 * unique value
 *
 * @param classType
 * @param instance
 * @return
 * @throws AtlasException
 */
public AtlasVertex getVertexForInstanceByUniqueAttribute(AtlasEntityType classType, Referenceable instance)
    throws AtlasException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("Checking if there is an instance with the same unique attributes for instance {}", instance.toShortString());
    }

    AtlasVertex result = null;
    for (AtlasAttribute attributeInfo : classType.getUniqAttributes().values()) {
        String propertyKey = attributeInfo.getQualifiedName();
        try {
            result = findVertex(propertyKey, instance.get(attributeInfo.getName()),
                    ENTITY_TYPE_PROPERTY_KEY, classType.getTypeName(),
                    STATE_PROPERTY_KEY, Id.EntityState.ACTIVE.name());
            if (LOG.isDebugEnabled()) {
                LOG.debug("Found vertex by unique attribute : {}={}", propertyKey, instance.get(attributeInfo.getName()));
            }
        } catch (EntityNotFoundException e) {
            //Its ok if there is no entity with the same unique value
        }
    }

    return result;
}
 

private static Referenceable createFeedEntity(Feed feed, Referenceable clusterReferenceable) {
    LOG.info("Creating feed dataset: {}", feed.getName());

    Referenceable feedEntity = new Referenceable(FalconDataTypes.FALCON_FEED.getName());
    feedEntity.set(AtlasClient.NAME, feed.getName());
    feedEntity.set(AtlasClient.DESCRIPTION, feed.getDescription());
    String feedQualifiedName =
            getFeedQualifiedName(feed.getName(), (String) clusterReferenceable.get(AtlasClient.NAME));
    feedEntity.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, feedQualifiedName);
    feedEntity.set(FalconBridge.FREQUENCY, feed.getFrequency().toString());
    feedEntity.set(FalconBridge.STOREDIN, clusterReferenceable);
    if (feed.getACL() != null) {
        feedEntity.set(AtlasClient.OWNER, feed.getACL().getOwner());
    }

    if (StringUtils.isNotEmpty(feed.getTags())) {
        feedEntity.set(FalconBridge.TAGS,
                EventUtil.convertKeyValueStringToMap(feed.getTags()));
    }

    if (feed.getGroups() != null) {
        feedEntity.set(FalconBridge.GROUPS, feed.getGroups());
    }

    return feedEntity;
}
 

@SuppressWarnings("unchecked")
private void assertUpdateFlowPathMessage(Notifier notifier, int notificationIndex, Referenceable ... expects) {
    assertTrue(notifier.notifications.size() > notificationIndex);
    final List<HookNotification> messages = notifier.notifications.get(notificationIndex);
    assertEquals(expects.length, messages.size());
    for (int i = 0; i < expects.length; i++) {
        final Referenceable expect = expects[i];
        final HookNotificationV1.EntityPartialUpdateRequest actual = (HookNotificationV1.EntityPartialUpdateRequest) messages.get(i);
        assertEquals(expect.getTypeName(), actual.getTypeName());
        assertEquals(ATTR_QUALIFIED_NAME, actual.getAttribute());
        assertEquals(expect.get(ATTR_QUALIFIED_NAME), actual.getAttributeValue());

        final Collection expIn = (Collection) expect.get(ATTR_INPUTS);
        final Collection expOut = (Collection) expect.get(ATTR_OUTPUTS);
        assertTrue(expIn.containsAll((Collection) actual.getEntity().get(ATTR_INPUTS)));
        assertTrue(expOut.containsAll((Collection) actual.getEntity().get(ATTR_OUTPUTS)));
    }
}
 

@Test
public void testGetTraitNames() 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);
    }

    List<String> traits = atlasClientV1.listTraits(guid);
    assertNotNull(traits);
    Assert.assertEquals(traits.size(), 7);
}
 

@Test
public void testPublishKafka() {
    final String processorName = "PublishKafka";
    final String transitUri = "PLAINTEXT://0.example.com:6667/topicA";
    final ProvenanceEventRecord record = Mockito.mock(ProvenanceEventRecord.class);
    when(record.getComponentType()).thenReturn(processorName);
    when(record.getTransitUri()).thenReturn(transitUri);
    when(record.getEventType()).thenReturn(ProvenanceEventType.SEND);

    final NamespaceResolvers namespaceResolvers = Mockito.mock(NamespaceResolvers.class);
    when(namespaceResolvers.fromHostNames(matches(".+\\.example\\.com"))).thenReturn("namespace1");

    final AnalysisContext context = Mockito.mock(AnalysisContext.class);
    when(context.getNamespaceResolver()).thenReturn(namespaceResolvers);

    final NiFiProvenanceEventAnalyzer analyzer = NiFiProvenanceEventAnalyzerFactory.getAnalyzer(processorName, transitUri, record.getEventType());
    assertNotNull(analyzer);

    final DataSetRefs refs = analyzer.analyze(context, record);
    assertEquals(0, refs.getInputs().size());
    assertEquals(1, refs.getOutputs().size());
    Referenceable ref = refs.getOutputs().iterator().next();
    assertEquals("topicA", ref.get(ATTR_NAME));
    assertEquals("topicA", ref.get("topic"));
    assertEquals("topicA@namespace1", ref.get(ATTR_QUALIFIED_NAME));
}
 

private boolean hasAnyAssignedAttribute(org.apache.atlas.v1.model.instance.Referenceable rInstance) {
    boolean ret = false;

    Map<String, Object> attributes = rInstance.getValues();

    if (MapUtils.isNotEmpty(attributes)) {
        for (Map.Entry<String, Object> attribute : attributes.entrySet()) {
            if (attribute.getValue() != null) {
                ret = true;
                break;
            }
        }
    }

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

    Referenceable entity = atlasClientV1.getEntity(guid);
    Assert.assertNotNull(entity);
}
 

private Predicate<Referenceable> distinctReferenceable() {
    final Set<String> keys = new HashSet<>();
    return r -> {
        final String key = AtlasUtils.toTypedQualifiedName(r.getTypeName(), (String) r.get(ATTR_QUALIFIED_NAME));
        return keys.add(key);
    };
}
 

private Referenceable createBucketRef(URI uri, String namespace) {
    final Referenceable ref = new Referenceable(TYPE_BUCKET);

    ref.set(ATTR_QUALIFIED_NAME, toQualifiedName(namespace, String.format("%s://%s", uri.getScheme(), uri.getAuthority())));
    ref.set(ATTR_NAME, uri.getAuthority());

    return ref;
}
 

private static Referenceable getClusterEntityReference(final String clusterName,
                                                       final String colo) {
    LOG.info("Getting reference for entity {}", clusterName);
    Referenceable clusterRef = new Referenceable(FalconDataTypes.FALCON_CLUSTER.getName());
    clusterRef.set(AtlasClient.NAME, String.format("%s", clusterName));
    clusterRef.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, clusterName);
    clusterRef.set(FalconBridge.COLO, colo);
    return clusterRef;
}
 

@Test
public void testRemoteInputPortHTTP() {
    final String componentType = "Remote Input Port";
    final String transitUri = "http://0.example.com:8080/nifi-api/data-transfer/input-ports/port-guid/transactions/tx-guid/flow-files";
    final ProvenanceEventRecord sendEvent = Mockito.mock(ProvenanceEventRecord.class);
    when(sendEvent.getEventId()).thenReturn(123L);
    when(sendEvent.getComponentId()).thenReturn("port-guid");
    when(sendEvent.getComponentType()).thenReturn(componentType);
    when(sendEvent.getTransitUri()).thenReturn(transitUri);
    when(sendEvent.getEventType()).thenReturn(ProvenanceEventType.SEND);

    final NamespaceResolvers namespaceResolvers = Mockito.mock(NamespaceResolvers.class);
    when(namespaceResolvers.fromHostNames(matches(".+\\.example\\.com"))).thenReturn("namespace1");

    final List<ConnectionStatus> connections = new ArrayList<>();
    final ConnectionStatus connection = new ConnectionStatus();
    connection.setDestinationId("port-guid");
    connection.setDestinationName("inputPortA");
    connections.add(connection);

    final AnalysisContext context = Mockito.mock(AnalysisContext.class);
    when(context.getNamespaceResolver()).thenReturn(namespaceResolvers);
    when(context.findConnectionTo(matches("port-guid"))).thenReturn(connections);

    final NiFiProvenanceEventAnalyzer analyzer = NiFiProvenanceEventAnalyzerFactory.getAnalyzer(componentType, transitUri, sendEvent.getEventType());
    assertNotNull(analyzer);

    final DataSetRefs refs = analyzer.analyze(context, sendEvent);
    assertEquals(0, refs.getInputs().size());
    assertEquals(1, refs.getOutputs().size());
    assertEquals(1, refs.getComponentIds().size());
    // Should report connected componentId.
    assertTrue(refs.getComponentIds().contains("port-guid"));

    Referenceable ref = refs.getOutputs().iterator().next();
    assertEquals(TYPE_NIFI_INPUT_PORT, ref.getTypeName());
    assertEquals("inputPortA", ref.get(ATTR_NAME));
    assertEquals("port-guid@namespace1", ref.get(ATTR_QUALIFIED_NAME));
}
 

private void restoreAttributes(Map<String, Object> prunedAttributes, Referenceable attributeEntity) throws AtlasException {
    Object                      obj          = prunedAttributes.get(attributeEntity.getId()._getId());

    if (obj instanceof Map) {
        restoreEntityAttributes(attributeEntity, (Map) obj);
    }
}
 

private Referenceable getReferenceable(List<Referenceable> refs, String name) {
    Referenceable ret = null;

    for (Referenceable ref : refs) {
        Map<String, Object> values     = ref.getValuesMap();
        String              entityName = (String) values.get("name");

        if (StringUtils.equalsIgnoreCase(name, entityName)) {
            ret = ref;
            break;
        }
    }

    return ret;
}
 

protected Referenceable createDatabaseRef(String namespace, String databaseName) {
    final Referenceable ref = new Referenceable(TYPE_DATABASE);
    ref.set(ATTR_NAME, databaseName);
    // The attribute 'clusterName' is in the 'hive_db' Atlas entity so it cannot be changed.
    //  Using 'namespace' as value for lack of better solution.
    ref.set(ATTR_CLUSTER_NAME, namespace);
    ref.set(ATTR_QUALIFIED_NAME, toQualifiedName(namespace, databaseName));
    return ref;
}
 

@Test
public void testViewIsAdded() throws AtlasServiceException {

    Referenceable view = atlasClientV1.getEntity(QuickStart.VIEW_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, QuickStart.PRODUCT_DIM_VIEW);

    assertEquals(QuickStart.PRODUCT_DIM_VIEW, view.get(AtlasClient.NAME));

    Id productDimId = getTable(QuickStart.PRODUCT_DIM_TABLE).getId();
    Id inputTableId = ((List<Id>) view.get(QuickStart.INPUT_TABLES_ATTRIBUTE)).get(0);
    assertEquals(productDimId, inputTableId);
}
 

@SuppressWarnings("unchecked")
private void assertCreateMessage(Notifier notifier, int notificationIndex, Referenceable ... expects) {
    assertTrue(notifier.notifications.size() > notificationIndex);
    final List<HookNotification> messages = notifier.notifications.get(notificationIndex);
    assertEquals(1, messages.size());
    final HookNotificationV1.EntityCreateRequest message = (HookNotificationV1.EntityCreateRequest) messages.get(0);
    assertEquals(expects.length, message.getEntities().size());

    // The use of 'flatMap' at NotificationSender does not preserve actual entities order.
    // Use typed qname map to assert regardless of ordering.
    final Map<String, Referenceable> entities = message.getEntities().stream().collect(Collectors.toMap(
            ref -> AtlasUtils.toTypedQualifiedName(ref.getTypeName(), (String) ref.get(ATTR_QUALIFIED_NAME)), ref -> ref));

    boolean hasFlowPathSeen = false;
    for (int i = 0; i < expects.length; i++) {
        final Referenceable expect = expects[i];
        final String typeName = expect.getTypeName();
        final Referenceable actual = entities.get(AtlasUtils.toTypedQualifiedName(typeName, (String) expect.get(ATTR_QUALIFIED_NAME)));
        assertNotNull(actual);
        assertEquals(typeName, actual.getTypeName());
        assertEquals(expect.get(ATTR_QUALIFIED_NAME), actual.get(ATTR_QUALIFIED_NAME));

        if (TYPE_NIFI_FLOW_PATH.equals(typeName)) {
            assertIOReferences(expect, actual, ATTR_INPUTS);
            assertIOReferences(expect, actual, ATTR_OUTPUTS);
            hasFlowPathSeen = true;
        } else {
            assertFalse("Types other than nifi_flow_path should be created before any nifi_flow_path entity.", hasFlowPathSeen);
        }
    }
}
 

Referenceable column(String name, String type, String comment, String... traitNames) throws Exception {
    Referenceable referenceable = new Referenceable(COLUMN_TYPE, traitNames);
    referenceable.set(NAME, name);
    referenceable.set(QUALIFIED_NAME, name);
    referenceable.set("type", type);
    referenceable.set("comment", comment);

    return referenceable;
}
 

static void populateColumnReferenceableMap(Map<String, Referenceable> m,
                                           Referenceable r) {
    if (r.getTypeName().equals(HiveDataTypes.HIVE_TABLE.getName())) {
        String qName = (String) r.get(ATTRIBUTE_QUALIFIED_NAME);
        String[] qNameComps = extractComponents(qName);
        for (Referenceable col : (List<Referenceable>) r.get(ATTRIBUTE_COLUMNS)) {
            String cName = (String) col.get(ATTRIBUTE_QUALIFIED_NAME);
            String[] colQNameComps = extractComponents(cName);
            String colQName = colQNameComps[0] + "." + colQNameComps[1] + "." + colQNameComps[2];
            m.put(colQName, col);
        }
        String tableQName = qNameComps[0] + "." + qNameComps[1];
        m.put(tableQName, r);
    }
}
 

@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 testGetTraitDefinitionForEntity() 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());

    TypesDef typesDef = new TypesDef(Collections.emptyList(), Collections.emptyList(), Collections.singletonList(piiTrait), Collections.emptyList());
    String traitDefinitionAsJSON = AtlasType.toV1Json(typesDef);
    LOG.debug("traitDefinitionAsJSON = {}", traitDefinitionAsJSON);
    createType(AtlasType.toV1Json(typesDef));

    Struct traitInstance = new Struct(traitName);
    atlasClientV1.addTrait(guid, traitInstance);
    Struct traitDef = atlasClientV1.getTraitDefinition(guid, traitName);
    Assert.assertEquals(traitDef.getTypeName(), traitName);

    List<Struct> allTraitDefs = atlasClientV1.listTraitDefinitions(guid);
    System.out.println(allTraitDefs.toString());
    Assert.assertEquals(allTraitDefs.size(), 8);
}
 

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

@Test
public void testGenerateFlowFile() {
    final String processorName = "GenerateFlowFile";
    final String processorId = "processor-1234";
    final ProvenanceEventRecord record = Mockito.mock(ProvenanceEventRecord.class);
    when(record.getComponentType()).thenReturn(processorName);
    when(record.getComponentId()).thenReturn(processorId);
    when(record.getEventType()).thenReturn(ProvenanceEventType.CREATE);

    final NamespaceResolvers namespaceResolvers = Mockito.mock(NamespaceResolvers.class);
    when(namespaceResolvers.fromHostNames(matches(".+\\.example\\.com"))).thenReturn("namespace1");

    final List<ConnectionStatus> connections = new ArrayList<>();

    final AnalysisContext context = Mockito.mock(AnalysisContext.class);
    when(context.getNamespaceResolver()).thenReturn(namespaceResolvers);
    when(context.findConnectionTo(processorId)).thenReturn(connections);
    when(context.getNiFiNamespace()).thenReturn("test_namespace");

    final NiFiProvenanceEventAnalyzer analyzer = NiFiProvenanceEventAnalyzerFactory.getAnalyzer(processorName, null, record.getEventType());
    assertNotNull(analyzer);

    final DataSetRefs refs = analyzer.analyze(context, record);
    assertEquals(1, refs.getInputs().size());
    assertEquals(0, refs.getOutputs().size());
    Referenceable ref = refs.getInputs().iterator().next();
    assertEquals("nifi_data", ref.getTypeName());
    assertEquals("GenerateFlowFile", ref.get(ATTR_NAME));
    assertEquals("processor-1234@test_namespace", ref.get(ATTR_QUALIFIED_NAME));
}
 

public void put(Object value, Referenceable instance, int index) {
    IndexedInstance wrapper = new IndexedInstance(instance, index);
    Collection<IndexedInstance> existingValues = valueMap_.get(value);
    if(existingValues == null) {
        //only expect 1 value
        existingValues = new HashSet<>(1);
        valueMap_.put(value, existingValues);
    }
    existingValues.add(wrapper);
}