Examples with NiFiAtlasClient org.apache.nifi.atlas.NiFiAtlasClient used on opensource projects

Example 1 with NiFiAtlasClient

use of org.apache.nifi.atlas.NiFiAtlasClient in project nifi by apache.

the class ReportLineageToAtlas method createNiFiAtlasClient.

/**
 * In order to avoid authentication expiration issues (i.e. Kerberos ticket and DelegationToken expiration),
 * create Atlas client instance at every onTrigger execution.
 */
private NiFiAtlasClient createNiFiAtlasClient(ReportingContext context) {
    List<String> urls = new ArrayList<>();
    parseAtlasUrls(context.getProperty(ATLAS_URLS), urls::add);
    try {
        return new NiFiAtlasClient(atlasAuthN.createClient(urls.toArray(new String[] {})));
    } catch (final NullPointerException e) {
        throw new ProcessException(String.format("Failed to initialize Atlas client due to %s." + " Make sure 'atlas-application.properties' is in the directory specified with %s" + " or under root classpath if not specified.", e, ATLAS_CONF_DIR.getDisplayName()), e);
    }
}

Example 2 with NiFiAtlasClient

use of org.apache.nifi.atlas.NiFiAtlasClient in project nifi by apache.

the class ReportLineageToAtlas method onTrigger.

@Override
public void onTrigger(ReportingContext context) {
    final String clusterNodeId = context.getClusterNodeIdentifier();
    final boolean isClustered = context.isClustered();
    if (isClustered && isEmpty(clusterNodeId)) {
        // Clustered, but this node's ID is unknown. Not ready for processing yet.
        return;
    }
    // If standalone or being primary node in a NiFi cluster, this node is responsible for doing primary tasks.
    final boolean isResponsibleForPrimaryTasks = !isClustered || getNodeTypeProvider().isPrimary();
    final NiFiAtlasClient atlasClient = createNiFiAtlasClient(context);
    // Create Entity defs in Atlas if there's none yet.
    if (!isTypeDefCreated) {
        try {
            if (isResponsibleForPrimaryTasks) {
                // Create NiFi type definitions in Atlas type system.
                atlasClient.registerNiFiTypeDefs(false);
            } else {
                // Otherwise, just check existence of NiFi type definitions.
                if (!atlasClient.isNiFiTypeDefsRegistered()) {
                    getLogger().debug("NiFi type definitions are not ready in Atlas type system yet.");
                    return;
                }
            }
            isTypeDefCreated = true;
        } catch (AtlasServiceException e) {
            throw new RuntimeException("Failed to check and create NiFi flow type definitions in Atlas due to " + e, e);
        }
    }
    // Regardless of whether being a primary task node, each node has to analyse NiFiFlow.
    // Assuming each node has the same flow definition, that is guaranteed by NiFi cluster management mechanism.
    final NiFiFlow nifiFlow = createNiFiFlow(context, atlasClient);
    if (isResponsibleForPrimaryTasks) {
        try {
            atlasClient.registerNiFiFlow(nifiFlow);
        } catch (AtlasServiceException e) {
            throw new RuntimeException("Failed to register NiFI flow. " + e, e);
        }
    }
    // NOTE: There is a race condition between the primary node and other nodes.
    // If a node notifies an event related to a NiFi component which is not yet created by NiFi primary node,
    // then the notification message will fail due to having a reference to a non-existing entity.
    nifiAtlasHook.setAtlasClient(atlasClient);
    consumeNiFiProvenanceEvents(context, nifiFlow);
}