Java Code Examples for org.apache.atlas.v1.model.instance.Referenceable#set()

@Test
public void testMessageHandleFailureConsumerContinues() throws Exception {
    //send invalid message - update with invalid type
    sendHookMessage(new EntityPartialUpdateRequest(TEST_USER, randomString(), null, null, new Referenceable(randomString())));

    //send valid message
    final Referenceable entity = new Referenceable(DATABASE_TYPE_BUILTIN);
    final String        dbName = "db" + randomString();

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

    sendHookMessage(new EntityCreateRequest(TEST_USER, entity));

    waitFor(MAX_WAIT_TIME, new Predicate() {
        @Override
        public boolean evaluate() throws Exception {
            ArrayNode results = searchByDSL(String.format("%s where name='%s'", DATABASE_TYPE_BUILTIN, entity.get(NAME)));

            return results.size() == 1;
        }
    });
}
 

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

Id loadProcess(String name, String user, List<Id> inputTables, List<Id> outputTables, String queryText,
        String queryPlan, String queryId, String queryGraph, String... traitNames) throws Exception {
    Referenceable referenceable = new Referenceable(HIVE_PROCESS_TYPE, traitNames);
    referenceable.set("name", name);
    referenceable.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
    referenceable.set("userName", user);
    referenceable.set("startTime", System.currentTimeMillis());
    referenceable.set("endTime", System.currentTimeMillis() + 10000);

    referenceable.set("inputs", inputTables);
    referenceable.set("outputs", outputTables);

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

    return createInstance(referenceable);
}
 

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

private Referenceable toReferenceable(NiFiFlow nifiFlow) {
    final Referenceable flowRef = new Referenceable(TYPE_NIFI_FLOW);
    flowRef.set(ATTR_NAME, nifiFlow.getFlowName());
    flowRef.set(ATTR_QUALIFIED_NAME, nifiFlow.getQualifiedName());
    flowRef.set(ATTR_URL, nifiFlow.getUrl());
    return flowRef;
}
 

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

@Test
public void testDeleteEntitiesViaClientApi() throws Exception {
    // Create 2 database entities
    Referenceable db1 = new Referenceable(DATABASE_TYPE_BUILTIN);
    String dbName = randomString();
    db1.set("name", dbName);
    db1.set("description", randomString());
    db1.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, dbName);
    db1.set("owner", "user1");
    db1.set(CLUSTER_NAME, "cl1");
    db1.set("parameters", Collections.EMPTY_MAP);
    db1.set("location", "/tmp");
    Id db1Id = createInstance(db1);
    Referenceable db2 = new Referenceable(DATABASE_TYPE_BUILTIN);
    String dbName2 = randomString();
    db2.set("name", dbName2);
    db2.set(QUALIFIED_NAME, dbName2);
    db2.set(CLUSTER_NAME, randomString());
    db2.set("description", randomString());
    db2.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, dbName2);
    db2.set("owner", "user2");
    db2.set("clusterName", "cl1");
    db2.set("parameters", Collections.EMPTY_MAP);
    db2.set("location", "/tmp");
    Id db2Id = createInstance(db2);

    // Delete the database entities
    List<String> deletedGuidsList = atlasClientV1.deleteEntities(db1Id._getId(), db2Id._getId()).getDeletedEntities();

    // Verify that deleteEntities() response has database entity guids
    Assert.assertEquals(deletedGuidsList.size(), 2);
    Assert.assertTrue(deletedGuidsList.contains(db1Id._getId()));
    Assert.assertTrue(deletedGuidsList.contains(db2Id._getId()));

    // Verify entities were deleted from the repository.
    for (String guid : deletedGuidsList) {
        Referenceable entity = atlasClientV1.getEntity(guid);
        assertEquals(entity.getId().getState(), Id.EntityState.DELETED);
    }
}
 

@Override
public DataSetRefs analyze(AnalysisContext context, ProvenanceEventRecord event) {
    final Referenceable ref = new Referenceable(TYPE);
    final URI uri = parseUri(event.getTransitUri());
    final String namespace = context.getNamespaceResolver().fromHostNames(uri.getHost());
    final String path = uri.getPath();
    ref.set(ATTR_NAME, path);
    ref.set(ATTR_PATH, path);
    // The attribute 'clusterName' is in the 'hdfs_path' 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, path));

    return singleDataSetRef(event.getComponentId(), event.getEventType(), ref);
}
 

private Referenceable createDirectoryRef(URI uri, String namespace) {
    final Referenceable ref = new Referenceable(TYPE_DIRECTORY);

    ref.set(ATTR_QUALIFIED_NAME, toQualifiedName(namespace, uri.toString().toLowerCase()));
    ref.set(ATTR_NAME, uri.getPath().toLowerCase());
    ref.set(ATTR_OBJECT_PREFIX, uri.getPath().toLowerCase());
    ref.set(ATTR_BUCKET, createBucketRef(uri, namespace));

    return ref;
}
 

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

private Referenceable toReferenceable(NiFiFlowPath flowPath, Referenceable flowRef, String namespace, String nifiUrl) {
    final Referenceable flowPathRef = new Referenceable(TYPE_NIFI_FLOW_PATH);
    flowPathRef.set(ATTR_NAME, flowPath.getName());
    flowPathRef.set(ATTR_QUALIFIED_NAME, flowPath.getId() + "@" + namespace);
    flowPathRef.set(ATTR_NIFI_FLOW, flowRef);
    flowPathRef.set(ATTR_URL, flowPath.createDeepLinkURL(nifiUrl));
    // Referenceable has to have GUID assigned, otherwise it will not be stored due to lack of required attribute.
    // If a Referencible has GUID, Atlas does not validate all required attributes.
    flowPathRef.set(ATTR_INPUTS, flowPath.getInputs().stream().map(this::toReferenceable).collect(Collectors.toList()));
    flowPathRef.set(ATTR_OUTPUTS,  flowPath.getOutputs().stream().map(this::toReferenceable).collect(Collectors.toList()));
    return flowPathRef;
}
 

Referenceable createEntity() {
    final Referenceable entity = new Referenceable(AtlasClient.DATA_SET_SUPER_TYPE);

    entity.set(NAME, "db" + randomString());
    entity.set(DESCRIPTION, randomString());
    entity.set(QUALIFIED_NAME, randomString());

    return entity;
}
 

private Referenceable generateLargeEntityWithTrait() {
    Referenceable ret = EntityNotificationTest.getEntity("id", new Struct("MyTrait", Collections.<String, Object>emptyMap()));

    // add 100 attributes, each with value of size 10k
    // Json Size=1,027,984; GZipped Size=16,387 ==> will compress, but not split
    String attrValue = RandomStringUtils.randomAlphanumeric(10 * 1024); // use the same value for all attributes - to aid better compression
    for (int i = 0; i < 100; i++) {
        ret.set("attr_" + i, attrValue);
    }

    return ret;
}
 

@Override
public DataSetRefs analyze(AnalysisContext context, ProvenanceEventRecord event) {
    final Referenceable ref = new Referenceable(TYPE);

    final String transitUri = event.getTransitUri();
    if (transitUri == null) {
        return null;
    }

    final Matcher uriMatcher = URI_PATTERN.matcher(transitUri);
    if (!uriMatcher.matches()) {
        logger.warn("Unexpected transit URI: {}", new Object[]{transitUri});
        return null;
    }

    final String[] hostNames = splitHostNames(uriMatcher.group(1));
    final String namespace = context.getNamespaceResolver().fromHostNames(hostNames);

    final String topicName = uriMatcher.group(2);

    ref.set(ATTR_NAME, topicName);
    ref.set(ATTR_TOPIC, topicName);
    ref.set(ATTR_QUALIFIED_NAME, toQualifiedName(namespace, topicName));
    ref.set(ATTR_URI, transitUri);

    return singleDataSetRef(event.getComponentId(), event.getEventType(), ref);
}
 

protected Referenceable createHiveDBInstanceBuiltIn(String dbName) {
    Referenceable databaseInstance = new Referenceable(DATABASE_TYPE_BUILTIN);

    databaseInstance.set(NAME, dbName);
    databaseInstance.set(QUALIFIED_NAME, dbName);
    databaseInstance.set(CLUSTER_NAME, randomString());
    databaseInstance.set(DESCRIPTION, "foo database");

    return databaseInstance;
}
 

Id view(String name, Id dbId, List<Id> inputTables, String... traitNames) throws AtlasBaseException {
    try {
        Referenceable referenceable = new Referenceable(VIEW_TYPE, traitNames);
        referenceable.set("name", name);
        referenceable.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
        referenceable.set("db", dbId);

        referenceable.set(INPUT_TABLES_ATTRIBUTE, inputTables);

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

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

    atlasClientV1.createEntity(entity);

    final Referenceable newEntity = new Referenceable(DATABASE_TYPE_BUILTIN);

    newEntity.set(NAME, dbName);
    newEntity.set(DESCRIPTION, randomString());
    newEntity.set("owner", randomString());
    newEntity.set(QUALIFIED_NAME, qualifiedName);
    newEntity.set(CLUSTER_NAME, clusterName);

    //updating unique attribute
    sendHookMessage(new EntityUpdateRequest(TEST_USER, newEntity));

    waitFor(MAX_WAIT_TIME, new Predicate() {
        @Override
        public boolean evaluate() throws Exception {
            ArrayNode results = searchByDSL(String.format("%s where qualifiedName='%s'", DATABASE_TYPE_BUILTIN, newEntity.get(QUALIFIED_NAME)));

            return results.size() == 1;
        }
    });

    Referenceable actualEntity = atlasClientV1.getEntity(DATABASE_TYPE_BUILTIN, QUALIFIED_NAME, qualifiedName);

    assertEquals(actualEntity.get(DESCRIPTION), newEntity.get(DESCRIPTION));
    assertEquals(actualEntity.get("owner"), newEntity.get("owner"));
}
 

@Override
public DataSetRefs analyze(AnalysisContext context, ProvenanceEventRecord event) {

    if (!ProvenanceEventType.SEND.equals(event.getEventType())
            && !ProvenanceEventType.RECEIVE.equals(event.getEventType())) {
        return null;
    }

    final boolean isRemoteInputPort = event.getComponentType().equals("Remote Input Port");
    final String type = isRemoteInputPort ? TYPE_NIFI_INPUT_PORT : TYPE_NIFI_OUTPUT_PORT;

    final S2SPort s2SPort = analyzeS2SPort(event, context.getNamespaceResolver());

    // Find connections that connects to/from the remote port.
    final String componentId = event.getComponentId();
    final List<ConnectionStatus> connections = isRemoteInputPort
            ? context.findConnectionTo(componentId)
            : context.findConnectionFrom(componentId);
    if (connections == null || connections.isEmpty()) {
        logger.warn("Connection was not found: {}", new Object[]{event});
        return null;
    }

    // The name of remote port can be retrieved from any connection, use the first one.
    final ConnectionStatus connection = connections.get(0);
    final Referenceable ref = new Referenceable(type);
    ref.set(ATTR_NAME, isRemoteInputPort ? connection.getDestinationName() : connection.getSourceName());
    ref.set(ATTR_QUALIFIED_NAME, toQualifiedName(s2SPort.namespace, s2SPort.targetPortId));

    return singleDataSetRef(event.getComponentId(), event.getEventType(), ref);
}
 

public static List<Referenceable> createFeedCreationEntity(Feed feed, ConfigurationStore falconStore) throws FalconException, URISyntaxException {
    LOG.info("Creating feed : {}", feed.getName());

    List<Referenceable> entities = new ArrayList<>();

    if (feed.getClusters() != null) {
        List<Referenceable> replicationInputs = new ArrayList<>();
        List<Referenceable> replicationOutputs = new ArrayList<>();

        for (org.apache.falcon.entity.v0.feed.Cluster feedCluster : feed.getClusters().getClusters()) {
            org.apache.falcon.entity.v0.cluster.Cluster cluster = falconStore.get(EntityType.CLUSTER,
                    feedCluster.getName());

            // set cluster
            Referenceable clusterReferenceable = getClusterEntityReference(cluster.getName(), cluster.getColo());
            entities.add(clusterReferenceable);

            // input as hive_table or hdfs_path, output as falcon_feed dataset
            List<Referenceable> inputs = new ArrayList<>();
            List<Referenceable> inputReferenceables = getInputEntities(cluster, feed);
            if (inputReferenceables != null) {
                entities.addAll(inputReferenceables);
                inputs.add(inputReferenceables.get(inputReferenceables.size() - 1));
            }

            List<Referenceable> outputs = new ArrayList<>();
            Referenceable feedEntity = createFeedEntity(feed, clusterReferenceable);
            if (feedEntity != null) {
                entities.add(feedEntity);
                outputs.add(feedEntity);
            }

            if (!inputs.isEmpty() || !outputs.isEmpty()) {
                Referenceable feedCreateEntity = new Referenceable(FalconDataTypes.FALCON_FEED_CREATION.getName());
                String feedQualifiedName = getFeedQualifiedName(feed.getName(), cluster.getName());

                feedCreateEntity.set(AtlasClient.NAME, feed.getName());
                feedCreateEntity.set(AtlasClient.DESCRIPTION, "Feed creation - " + feed.getName());
                feedCreateEntity.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, feedQualifiedName);

                if (!inputs.isEmpty()) {
                    feedCreateEntity.set(AtlasClient.PROCESS_ATTRIBUTE_INPUTS, inputs);
                }
                if (!outputs.isEmpty()) {
                    feedCreateEntity.set(AtlasClient.PROCESS_ATTRIBUTE_OUTPUTS, outputs);
                }

                feedCreateEntity.set(FalconBridge.STOREDIN, clusterReferenceable);
                entities.add(feedCreateEntity);
            }

            if (ClusterType.SOURCE == feedCluster.getType()) {
                replicationInputs.add(feedEntity);
            } else if (ClusterType.TARGET == feedCluster.getType()) {
                replicationOutputs.add(feedEntity);
            }
        }

        if (!replicationInputs.isEmpty() && !replicationInputs.isEmpty()) {
            Referenceable feedReplicationEntity = new Referenceable(FalconDataTypes
                    .FALCON_FEED_REPLICATION.getName());

            feedReplicationEntity.set(AtlasClient.NAME, feed.getName());
            feedReplicationEntity.set(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, feed.getName());

            feedReplicationEntity.set(AtlasClient.PROCESS_ATTRIBUTE_INPUTS, replicationInputs);
            feedReplicationEntity.set(AtlasClient.PROCESS_ATTRIBUTE_OUTPUTS, replicationOutputs);
            entities.add(feedReplicationEntity);
        }

    }
    return entities;
}
 

@SuppressWarnings("unchecked")
private void mergeRefs(Referenceable r1, Referenceable r2) {
    r1.set(ATTR_INPUTS, mergeRefs((Collection<Referenceable>) r1.get(ATTR_INPUTS), (Collection<Referenceable>) r2.get(ATTR_INPUTS)));
    r1.set(ATTR_OUTPUTS, mergeRefs((Collection<Referenceable>) r1.get(ATTR_OUTPUTS), (Collection<Referenceable>) r2.get(ATTR_OUTPUTS)));
}