Examples with ProcessContext org.apache.nifi.processor.ProcessContext used on opensource projects

Example 66 with ProcessContext

use of org.apache.nifi.processor.ProcessContext in project nifi by apache.

the class QueryElasticsearchHttp method getPage.

private int getPage(final Response getResponse, final URL url, final ProcessContext context, final ProcessSession session, FlowFile flowFile, final ComponentLog logger, final long startNanos, boolean targetIsContent) throws IOException {
    List<FlowFile> page = new ArrayList<>();
    final int statusCode = getResponse.code();
    if (isSuccess(statusCode)) {
        ResponseBody body = getResponse.body();
        final byte[] bodyBytes = body.bytes();
        JsonNode responseJson = parseJsonResponse(new ByteArrayInputStream(bodyBytes));
        JsonNode hits = responseJson.get("hits").get("hits");
        for (int i = 0; i < hits.size(); i++) {
            JsonNode hit = hits.get(i);
            String retrievedId = hit.get("_id").asText();
            String retrievedIndex = hit.get("_index").asText();
            String retrievedType = hit.get("_type").asText();
            FlowFile documentFlowFile = null;
            if (flowFile != null) {
                documentFlowFile = targetIsContent ? session.create(flowFile) : session.clone(flowFile);
            } else {
                documentFlowFile = session.create();
            }
            JsonNode source = hit.get("_source");
            documentFlowFile = session.putAttribute(documentFlowFile, "es.id", retrievedId);
            documentFlowFile = session.putAttribute(documentFlowFile, "es.index", retrievedIndex);
            documentFlowFile = session.putAttribute(documentFlowFile, "es.type", retrievedType);
            if (targetIsContent) {
                documentFlowFile = session.putAttribute(documentFlowFile, "filename", retrievedId);
                documentFlowFile = session.putAttribute(documentFlowFile, "mime.type", "application/json");
                documentFlowFile = session.write(documentFlowFile, out -> {
                    out.write(source.toString().getBytes());
                });
            } else {
                Map<String, String> attributes = new HashMap<>();
                for (Iterator<Entry<String, JsonNode>> it = source.fields(); it.hasNext(); ) {
                    Entry<String, JsonNode> entry = it.next();
                    attributes.put(ATTRIBUTE_PREFIX + entry.getKey(), entry.getValue().asText());
                }
                documentFlowFile = session.putAllAttributes(documentFlowFile, attributes);
            }
            page.add(documentFlowFile);
        }
        logger.debug("Elasticsearch retrieved " + responseJson.size() + " documents, routing to success");
        session.transfer(page, REL_SUCCESS);
    } else {
        try {
            // 5xx -> RETRY, but a server error might last a while, so yield
            if (statusCode / 100 == 5) {
                throw new RetryableException(String.format("Elasticsearch returned code %s with message %s, transferring flow file to retry. This is likely a server problem, yielding...", statusCode, getResponse.message()));
            } else if (context.hasIncomingConnection()) {
                // 1xx, 3xx, 4xx -> NO RETRY
                throw new UnretryableException(String.format("Elasticsearch returned code %s with message %s, transferring flow file to failure", statusCode, getResponse.message()));
            } else {
                logger.warn("Elasticsearch returned code {} with message {}", new Object[] { statusCode, getResponse.message() });
            }
        } finally {
            if (!page.isEmpty()) {
                session.remove(page);
                page.clear();
            }
        }
    }
    // emit provenance event
    final long millis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
    if (!page.isEmpty()) {
        if (context.hasNonLoopConnection()) {
            page.forEach(f -> session.getProvenanceReporter().fetch(f, url.toExternalForm(), millis));
        } else {
            page.forEach(f -> session.getProvenanceReporter().receive(f, url.toExternalForm(), millis));
        }
    }
    return page.size();
}

Example 67 with ProcessContext

use of org.apache.nifi.processor.ProcessContext in project nifi by apache.

the class ScrollElasticsearchHttp method getPage.

private void getPage(final Response getResponse, final URL url, final ProcessContext context, final ProcessSession session, FlowFile flowFile, final ComponentLog logger, final long startNanos) throws IOException {
    final int statusCode = getResponse.code();
    if (isSuccess(statusCode)) {
        ResponseBody body = getResponse.body();
        final byte[] bodyBytes = body.bytes();
        JsonNode responseJson = parseJsonResponse(new ByteArrayInputStream(bodyBytes));
        String scrollId = responseJson.get("_scroll_id").asText();
        StringBuilder builder = new StringBuilder();
        builder.append("{ \"hits\" : [");
        JsonNode hits = responseJson.get("hits").get("hits");
        if (hits.size() == 0) {
            finishQuery(context.getStateManager());
            session.remove(flowFile);
            return;
        }
        for (int i = 0; i < hits.size(); i++) {
            JsonNode hit = hits.get(i);
            String retrievedIndex = hit.get("_index").asText();
            String retrievedType = hit.get("_type").asText();
            JsonNode source = hit.get("_source");
            flowFile = session.putAttribute(flowFile, "es.index", retrievedIndex);
            flowFile = session.putAttribute(flowFile, "es.type", retrievedType);
            flowFile = session.putAttribute(flowFile, "mime.type", "application/json");
            builder.append(source.toString());
            if (i < hits.size() - 1) {
                builder.append(", ");
            }
        }
        builder.append("] }");
        logger.debug("Elasticsearch retrieved " + responseJson.size() + " documents, routing to success");
        flowFile = session.write(flowFile, out -> {
            out.write(builder.toString().getBytes());
        });
        session.transfer(flowFile, REL_SUCCESS);
        saveScrollId(context.getStateManager(), scrollId);
        // emit provenance event
        final long millis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
        session.getProvenanceReporter().receive(flowFile, url.toExternalForm(), millis);
    } else {
        // 5xx -> RETRY, but a server error might last a while, so yield
        if (statusCode / 100 == 5) {
            logger.warn("Elasticsearch returned code {} with message {}, removing the flow file. This is likely a server problem, yielding...", new Object[] { statusCode, getResponse.message() });
            session.remove(flowFile);
            context.yield();
        } else {
            logger.warn("Elasticsearch returned code {} with message {}", new Object[] { statusCode, getResponse.message() });
            session.remove(flowFile);
        }
    }
}

Example 68 with ProcessContext

use of org.apache.nifi.processor.ProcessContext in project nifi by apache.

the class AbstractEmailProcessor method transfer.

/**
 * Disposes the message by converting it to a {@link FlowFile} transferring
 * it to the REL_SUCCESS relationship.
 */
private void transfer(Message emailMessage, ProcessContext context, ProcessSession processSession) {
    long start = System.nanoTime();
    FlowFile flowFile = processSession.create();
    flowFile = processSession.append(flowFile, out -> {
        try {
            emailMessage.writeTo(out);
        } catch (MessagingException e) {
            throw new IOException(e);
        }
    });
    long executionDuration = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start);
    String fromAddressesString = "";
    try {
        Address[] fromAddresses = emailMessage.getFrom();
        if (fromAddresses != null) {
            fromAddressesString = Arrays.asList(fromAddresses).toString();
        }
    } catch (MessagingException e) {
        this.logger.warn("Failed to retrieve 'From' attribute from Message.");
    }
    processSession.getProvenanceReporter().receive(flowFile, this.displayUrl, "Received message from " + fromAddressesString, executionDuration);
    this.getLogger().info("Successfully received {} from {} in {} millis", new Object[] { flowFile, fromAddressesString, executionDuration });
    processSession.transfer(flowFile, REL_SUCCESS);
}

Example 69 with ProcessContext

use of org.apache.nifi.processor.ProcessContext in project nifi by apache.

the class ExecuteFlumeSinkTest method testValidators.

@Test
public void testValidators() {
    TestRunner runner = TestRunners.newTestRunner(ExecuteFlumeSink.class);
    Collection<ValidationResult> results;
    ProcessContext pc;
    results = new HashSet<>();
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    Assert.assertEquals(1, results.size());
    for (ValidationResult vr : results) {
        logger.debug(vr.toString());
        Assert.assertTrue(vr.toString().contains("is invalid because Sink Type is required"));
    }
    // non-existent class
    results = new HashSet<>();
    runner.setProperty(ExecuteFlumeSink.SINK_TYPE, "invalid.class.name");
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    Assert.assertEquals(1, results.size());
    for (ValidationResult vr : results) {
        logger.debug(vr.toString());
        Assert.assertTrue(vr.toString().contains("is invalid because unable to load sink"));
    }
    // class doesn't implement Sink
    results = new HashSet<>();
    runner.setProperty(ExecuteFlumeSink.SINK_TYPE, AvroSource.class.getName());
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    Assert.assertEquals(1, results.size());
    for (ValidationResult vr : results) {
        logger.debug(vr.toString());
        Assert.assertTrue(vr.toString().contains("is invalid because unable to create sink"));
    }
    results = new HashSet<>();
    runner.setProperty(ExecuteFlumeSink.SINK_TYPE, NullSink.class.getName());
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    Assert.assertEquals(0, results.size());
}

Example 70 with ProcessContext

use of org.apache.nifi.processor.ProcessContext in project nifi by apache.

the class ExecuteFlumeSourceTest method testValidators.

@Test
public void testValidators() {
    TestRunner runner = TestRunners.newTestRunner(ExecuteFlumeSource.class);
    Collection<ValidationResult> results;
    ProcessContext pc;
    results = new HashSet<>();
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    Assert.assertEquals(1, results.size());
    for (ValidationResult vr : results) {
        logger.debug(vr.toString());
        Assert.assertTrue(vr.toString().contains("is invalid because Source Type is required"));
    }
    // non-existent class
    results = new HashSet<>();
    runner.setProperty(ExecuteFlumeSource.SOURCE_TYPE, "invalid.class.name");
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    Assert.assertEquals(1, results.size());
    for (ValidationResult vr : results) {
        logger.debug(vr.toString());
        Assert.assertTrue(vr.toString().contains("is invalid because unable to load source"));
    }
    // class doesn't implement Source
    results = new HashSet<>();
    runner.setProperty(ExecuteFlumeSource.SOURCE_TYPE, NullSink.class.getName());
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    Assert.assertEquals(1, results.size());
    for (ValidationResult vr : results) {
        logger.debug(vr.toString());
        Assert.assertTrue(vr.toString().contains("is invalid because unable to create source"));
    }
    results = new HashSet<>();
    runner.setProperty(ExecuteFlumeSource.SOURCE_TYPE, AvroSource.class.getName());
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    Assert.assertEquals(0, results.size());
}