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

Example 31 with ProcessContext

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

the class ExecuteSparkInteractive method onTrigger.

@Override
public void onTrigger(ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final ComponentLog log = getLogger();
    final LivySessionService livySessionService = context.getProperty(LIVY_CONTROLLER_SERVICE).asControllerService(LivySessionService.class);
    final Map<String, String> livyController = livySessionService.getSession();
    if (livyController == null || livyController.isEmpty()) {
        log.debug("No Spark session available (yet), routing flowfile to wait");
        session.transfer(flowFile, REL_WAIT);
        return;
    }
    final long statusCheckInterval = context.getProperty(STATUS_CHECK_INTERVAL).evaluateAttributeExpressions(flowFile).asTimePeriod(TimeUnit.MILLISECONDS);
    Charset charset;
    try {
        charset = Charset.forName(context.getProperty(CHARSET).evaluateAttributeExpressions(flowFile).getValue());
    } catch (Exception e) {
        log.warn("Illegal character set name specified, defaulting to UTF-8");
        charset = StandardCharsets.UTF_8;
    }
    String sessionId = livyController.get("sessionId");
    String livyUrl = livyController.get("livyUrl");
    String code = context.getProperty(CODE).evaluateAttributeExpressions(flowFile).getValue();
    if (StringUtils.isEmpty(code)) {
        try (InputStream inputStream = session.read(flowFile)) {
            // If no code was provided, assume it is in the content of the incoming flow file
            code = IOUtils.toString(inputStream, charset);
        } catch (IOException ioe) {
            log.error("Error reading input flowfile, penalizing and routing to failure", new Object[] { flowFile, ioe.getMessage() }, ioe);
            flowFile = session.penalize(flowFile);
            session.transfer(flowFile, REL_FAILURE);
            return;
        }
    }
    code = StringEscapeUtils.escapeJavaScript(code);
    String payload = "{\"code\":\"" + code + "\"}";
    try {
        final JSONObject result = submitAndHandleJob(livyUrl, livySessionService, sessionId, payload, statusCheckInterval);
        log.debug("ExecuteSparkInteractive Result of Job Submit: " + result);
        if (result == null) {
            session.transfer(flowFile, REL_FAILURE);
        } else {
            try {
                final JSONObject output = result.getJSONObject("data");
                flowFile = session.write(flowFile, out -> out.write(output.toString().getBytes()));
                flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), LivySessionService.APPLICATION_JSON);
                session.transfer(flowFile, REL_SUCCESS);
            } catch (JSONException je) {
                // The result doesn't contain the data, just send the output object as the flow file content to failure (after penalizing)
                log.error("Spark Session returned an error, sending the output JSON object as the flow file content to failure (after penalizing)");
                flowFile = session.write(flowFile, out -> out.write(result.toString().getBytes()));
                flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), LivySessionService.APPLICATION_JSON);
                flowFile = session.penalize(flowFile);
                session.transfer(flowFile, REL_FAILURE);
            }
        }
    } catch (IOException ioe) {
        log.error("Failure processing flowfile {} due to {}, penalizing and routing to failure", new Object[] { flowFile, ioe.getMessage() }, ioe);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
    }
}

Example 32 with ProcessContext

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

the class TestExecuteProcess method testRedirectErrorStreamWithExpressions.

@Test
public void testRedirectErrorStreamWithExpressions() {
    final TestRunner runner = TestRunners.newTestRunner(ExecuteProcess.class);
    runner.setProperty(ExecuteProcess.COMMAND, "ls");
    runner.setProperty(ExecuteProcess.COMMAND_ARGUMENTS, "${literal('does-not-exist'):toUpper()}");
    runner.setProperty(ExecuteProcess.REDIRECT_ERROR_STREAM, "true");
    ProcessContext processContext = runner.getProcessContext();
    ExecuteProcess processor = (ExecuteProcess) runner.getProcessor();
    processor.updateScheduledTrue();
    processor.setupExecutor(processContext);
    processor.onTrigger(processContext, runner.getProcessSessionFactory());
    if (isCommandFailed(runner))
        return;
    final List<LogMessage> warnMessages = runner.getLogger().getWarnMessages();
    assertEquals("If redirect error stream is true " + "the output should be sent as a content of flow-file.", 0, warnMessages.size());
    final List<MockFlowFile> succeeded = runner.getFlowFilesForRelationship(ExecuteProcess.REL_SUCCESS);
    assertEquals(1, succeeded.size());
    assertTrue(new String(succeeded.get(0).toByteArray()).contains("DOES-NOT-EXIST"));
    assertEquals(succeeded.get(0).getAttribute(ExecuteProcess.ATTRIBUTE_COMMAND), "ls");
    assertEquals(succeeded.get(0).getAttribute(ExecuteProcess.ATTRIBUTE_COMMAND_ARGS), "DOES-NOT-EXIST");
}

Example 33 with ProcessContext

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

the class TestFTP method testValidators.

@Test
public void testValidators() {
    TestRunner runner = TestRunners.newTestRunner(PutFTP.class);
    Collection<ValidationResult> results;
    ProcessContext pc;
    /* Set the basic required values */
    results = new HashSet<>();
    runner.setProperty(FTPTransfer.USERNAME, "${el-username}");
    runner.setProperty(FTPTransfer.HOSTNAME, "static-hostname");
    runner.setProperty(FTPTransfer.PORT, "${el-portNumber}");
    results = new HashSet<>();
    runner.setProperty(FTPTransfer.REMOTE_PATH, "static-remote-target");
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    assertEquals(0, results.size());
    results = new HashSet<>();
    runner.setProperty(FTPTransfer.REMOTE_PATH, "${el-remote-target}");
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    assertEquals(0, results.size());
    results = new HashSet<>();
    runner.setProperty(FTPTransfer.USERNAME, "static-username");
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    assertEquals(0, results.size());
    /* Try an invalid expression */
    results = new HashSet<>();
    runner.setProperty(FTPTransfer.USERNAME, "");
    runner.enqueue(new byte[0]);
    pc = runner.getProcessContext();
    if (pc instanceof MockProcessContext) {
        results = ((MockProcessContext) pc).validate();
    }
    assertEquals(1, results.size());
}

Example 34 with ProcessContext

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

the class ConsumeWindowsEventLog method subscribe.

/**
 * Register subscriber via native call
 *
 * @param context the process context
 */
private String subscribe(ProcessContext context) throws URISyntaxException {
    String channel = context.getProperty(CHANNEL).getValue();
    String query = context.getProperty(QUERY).getValue();
    renderedXMLs = new LinkedBlockingQueue<>(context.getProperty(MAX_EVENT_QUEUE_SIZE).asInteger());
    provenanceUri = new URI("winlog", name, "/" + channel, query, null).toASCIIString();
    evtSubscribeCallback = new EventSubscribeXmlRenderingCallback(getLogger(), s -> {
        try {
            renderedXMLs.put(s);
        } catch (InterruptedException e) {
            throw new IllegalStateException("Got interrupted while waiting to add to queue.", e);
        }
    }, context.getProperty(MAX_BUFFER_SIZE).asInteger(), wEvtApi, kernel32, errorLookup);
    subscriptionHandle = wEvtApi.EvtSubscribe(null, null, channel, query, null, null, evtSubscribeCallback, WEvtApi.EvtSubscribeFlags.SUBSCRIBE_TO_FUTURE | WEvtApi.EvtSubscribeFlags.EVT_SUBSCRIBE_STRICT);
    if (!isSubscribed()) {
        return UNABLE_TO_SUBSCRIBE + errorLookup.getLastError();
    }
    return null;
}

Example 35 with ProcessContext

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

the class TestJmsConsumer method testMap2FlowFileBytesMessage.

/**
 * Test BytesMessage to FlowFile conversion
 *
 * @throws java.lang.Exception ex
 */
@Test
public void testMap2FlowFileBytesMessage() throws Exception {
    TestRunner runner = TestRunners.newTestRunner(GetJMSQueue.class);
    BytesMessage bytesMessage = new ActiveMQBytesMessage();
    String sourceString = "Apache NiFi is an easy to use, powerful, and reliable system to process and distribute data.!";
    byte[] payload = sourceString.getBytes("UTF-8");
    bytesMessage.writeBytes(payload);
    bytesMessage.reset();
    ProcessContext context = runner.getProcessContext();
    ProcessSession session = runner.getProcessSessionFactory().createSession();
    ProcessorInitializationContext pic = new MockProcessorInitializationContext(runner.getProcessor(), (MockProcessContext) runner.getProcessContext());
    JmsProcessingSummary summary = JmsConsumer.map2FlowFile(context, session, bytesMessage, true, pic.getLogger());
    assertEquals("BytesMessage content length should equal to FlowFile content size", payload.length, summary.getLastFlowFile().getSize());
    final byte[] buffer = new byte[payload.length];
    runner.clearTransferState();
    session.read(summary.getLastFlowFile(), new InputStreamCallback() {

        @Override
        public void process(InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer, false);
        }
    });
    String contentString = new String(buffer, "UTF-8");
    assertEquals("", sourceString, contentString);
}