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

Example 61 with ProcessContext

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

the class TestPublishKafkaRecord_0_10 method setup.

@Before
public void setup() throws InitializationException, IOException {
    mockPool = mock(PublisherPool.class);
    mockLease = mock(PublisherLease.class);
    Mockito.doCallRealMethod().when(mockLease).publish(any(FlowFile.class), any(RecordSet.class), any(RecordSetWriterFactory.class), any(RecordSchema.class), any(String.class), any(String.class));
    when(mockPool.obtainPublisher()).thenReturn(mockLease);
    runner = TestRunners.newTestRunner(new PublishKafkaRecord_0_10() {

        @Override
        protected PublisherPool createPublisherPool(final ProcessContext context) {
            return mockPool;
        }
    });
    runner.setProperty(PublishKafkaRecord_0_10.TOPIC, TOPIC_NAME);
    final String readerId = "record-reader";
    final MockRecordParser readerService = new MockRecordParser();
    readerService.addSchemaField("name", RecordFieldType.STRING);
    readerService.addSchemaField("age", RecordFieldType.INT);
    runner.addControllerService(readerId, readerService);
    runner.enableControllerService(readerService);
    final String writerId = "record-writer";
    final RecordSetWriterFactory writerService = new MockRecordWriter("name, age");
    runner.addControllerService(writerId, writerService);
    runner.enableControllerService(writerService);
    runner.setProperty(PublishKafkaRecord_0_10.RECORD_READER, readerId);
    runner.setProperty(PublishKafkaRecord_0_10.RECORD_WRITER, writerId);
}

Example 62 with ProcessContext

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

the class ImportSqoopTest method testImportSqoopCommand.

@Test
public void testImportSqoopCommand() {
    String systemProperties = "-Dsystem.prop1=prop1 -Dsystem.prop2=prop2 ";
    String additionalProperties = "--arg1=test1 --arg=test2";
    runner.setProperty(ImportSqoop.SQOOP_SYSTEM_PROPERTIES, systemProperties);
    runner.setProperty(ImportSqoop.TARGET_COMPRESSION_ALGORITHM, "SNAPPY");
    runner.setProperty(ImportSqoop.SQOOP_CONNECTION_SERVICE, "sqoop-connection-service");
    runner.setProperty(ImportSqoop.SQOOP_ADDITIONAL_ARGUMENTS, additionalProperties);
    runner.setProperty(ImportSqoop.SOURCE_TABLE_WHERE_CLAUSE, "where clause");
    runner.setProperty(ImportSqoop.SOURCE_LOAD_STRATEGY, "FULL_LOAD");
    runner.enqueue(new byte[0]);
    ProcessContext ctx = runner.getProcessContext();
    Collection<ValidationResult> validationResults = validate(runner);
    runner.run();
    // look at the failure results since we dont have sqoop running
    List<MockFlowFile> failResults = runner.getFlowFilesForRelationship(ImportSqoop.REL_FAILURE);
    MockFlowFile result = failResults.get(0);
    String sqoopCommand = result.getAttribute("sqoop.command.text");
    Assert.assertTrue(StringUtils.startsWithIgnoreCase(sqoopCommand, "sqoop import " + systemProperties));
    Assert.assertTrue(StringUtils.endsWithIgnoreCase(sqoopCommand.trim(), additionalProperties));
}

Example 63 with ProcessContext

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

the class ExecuteSparkContextJob method onTrigger.

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLog();
    FlowFile flowFile = null;
    try {
        if (context.hasIncomingConnection()) {
            flowFile = session.get();
            if (flowFile == null) {
                return;
            }
        }
    } catch (NoSuchMethodError e) {
        logger.error("Failed to get incoming", e);
    }
    FlowFile outgoing = (flowFile == null ? session.create() : flowFile);
    final String appName = context.getProperty(APP_NAME).evaluateAttributeExpressions(flowFile).getValue().trim();
    final String classPath = context.getProperty(CLASS_PATH).evaluateAttributeExpressions(flowFile).getValue().trim();
    final String contextName = context.getProperty(CONTEXT_NAME).evaluateAttributeExpressions(flowFile).getValue().trim();
    final SparkContextType contextType = SparkContextType.valueOf(context.getProperty(CONTEXT_TYPE).getValue());
    final JobService jobService = context.getProperty(JOB_SERVICE).asControllerService(JobService.class);
    final String resultsOutputLocation = context.getProperty(RESULTS_OUTPUT_LOCATION).getValue();
    final String numExecutors = context.getProperty(NUM_EXECUTORS).evaluateAttributeExpressions(flowFile).getValue().trim();
    final String memPerNode = context.getProperty(MEM_PER_NODE).evaluateAttributeExpressions(flowFile).getValue().trim();
    final String numCPUCores = context.getProperty(NUM_CPU_CORES).evaluateAttributeExpressions(flowFile).getValue().trim();
    final int contextTimeout = context.getProperty(CONTEXT_TIMEOUT).evaluateAttributeExpressions(flowFile).asInteger();
    final boolean async = context.getProperty(ASYNC).asBoolean();
    String args = "";
    if (context.getProperty(ARGS).isSet()) {
        args = context.getProperty(ARGS).evaluateAttributeExpressions(flowFile).getValue().trim();
    }
    boolean createSuccess = jobService.createContext(contextName, numExecutors, memPerNode, numCPUCores, contextType, contextTimeout, false);
    if (createSuccess) {
        final SparkJobResult result = jobService.executeSparkContextJob(appName, classPath, contextName, args, async);
        if (result.success) {
            if (Objects.equals(resultsOutputLocation, FLOW_FILE_ATTRIBUTE_VALUE)) {
                outgoing = session.putAttribute(outgoing, appName + ".result", result.result);
            } else {
                outgoing = session.write(outgoing, outputStream -> IOUtils.write(result.result, outputStream, "UTF-8"));
            }
            session.transfer(outgoing, REL_SUCCESS);
        } else {
            session.transfer(outgoing, REL_FAILURE);
        }
    } else {
        session.transfer(outgoing, REL_FAILURE);
    }
}

Example 64 with ProcessContext

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

the class TestDeleteElasticsearch5 method testDeleteNonRetryableException.

@Test
public void testDeleteNonRetryableException() throws IOException {
    mockDeleteProcessor = new DeleteElasticsearch5() {

        @Override
        protected DeleteRequestBuilder prepareDeleteRequest(String index, String docId, String docType) {
            return null;
        }

        @Override
        protected DeleteResponse doDelete(DeleteRequestBuilder requestBuilder) throws InterruptedException, ExecutionException {
            throw new InterruptedException("exception");
        }

        @Override
        public void setup(ProcessContext context) {
        }
    };
    runner = TestRunners.newTestRunner(mockDeleteProcessor);
    runner.setValidateExpressionUsage(true);
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.CLUSTER_NAME, "elasticsearch");
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.HOSTS, "127.0.0.1:9300");
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.PING_TIMEOUT, "5s");
    runner.setProperty(AbstractElasticsearch5TransportClientProcessor.SAMPLER_INTERVAL, "5s");
    runner.setProperty(DeleteElasticsearch5.INDEX, INDEX1);
    runner.setProperty(DeleteElasticsearch5.TYPE, TYPE1);
    runner.setProperty(DeleteElasticsearch5.DOCUMENT_ID, "${documentId}");
    runner.assertValid();
    runner.enqueue(new byte[] {}, new HashMap<String, String>() {

        {
            put("documentId", documentId);
        }
    });
    runner.run(1, true, true);
    runner.assertAllFlowFilesTransferred(DeleteElasticsearch5.REL_FAILURE, 1);
    final MockFlowFile out = runner.getFlowFilesForRelationship(DeleteElasticsearch5.REL_FAILURE).get(0);
    assertNotNull(out);
    assertEquals("exception", out.getAttribute(DeleteElasticsearch5.ES_ERROR_MESSAGE));
    out.assertAttributeEquals(DeleteElasticsearch5.ES_REST_STATUS, null);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_FILENAME, documentId);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_INDEX, INDEX1);
    out.assertAttributeEquals(DeleteElasticsearch5.ES_TYPE, TYPE1);
}

Example 65 with ProcessContext

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

the class PutElasticsearchHttp method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final int batchSize = context.getProperty(BATCH_SIZE).evaluateAttributeExpressions().asInteger();
    final List<FlowFile> flowFiles = session.get(batchSize);
    if (flowFiles.isEmpty()) {
        return;
    }
    final String id_attribute = context.getProperty(ID_ATTRIBUTE).getValue();
    // Authentication
    final String username = context.getProperty(USERNAME).evaluateAttributeExpressions().getValue();
    final String password = context.getProperty(PASSWORD).evaluateAttributeExpressions().getValue();
    OkHttpClient okHttpClient = getClient();
    final ComponentLog logger = getLogger();
    // Keep track of the list of flow files that need to be transferred. As they are transferred, remove them from the list.
    List<FlowFile> flowFilesToTransfer = new LinkedList<>(flowFiles);
    final StringBuilder sb = new StringBuilder();
    final String baseUrl = trimToEmpty(context.getProperty(ES_URL).evaluateAttributeExpressions().getValue());
    HttpUrl.Builder urlBuilder = HttpUrl.parse(baseUrl).newBuilder().addPathSegment("_bulk");
    // Find the user-added properties and set them as query parameters on the URL
    for (Map.Entry<PropertyDescriptor, String> property : context.getProperties().entrySet()) {
        PropertyDescriptor pd = property.getKey();
        if (pd.isDynamic()) {
            if (property.getValue() != null) {
                urlBuilder = urlBuilder.addQueryParameter(pd.getName(), context.getProperty(pd).evaluateAttributeExpressions().getValue());
            }
        }
    }
    final URL url = urlBuilder.build().url();
    for (FlowFile file : flowFiles) {
        final String index = context.getProperty(INDEX).evaluateAttributeExpressions(file).getValue();
        final Charset charset = Charset.forName(context.getProperty(CHARSET).evaluateAttributeExpressions(file).getValue());
        if (StringUtils.isEmpty(index)) {
            logger.error("No value for index in for {}, transferring to failure", new Object[] { id_attribute, file });
            flowFilesToTransfer.remove(file);
            session.transfer(file, REL_FAILURE);
            continue;
        }
        final String docType = context.getProperty(TYPE).evaluateAttributeExpressions(file).getValue();
        String indexOp = context.getProperty(INDEX_OP).evaluateAttributeExpressions(file).getValue();
        if (StringUtils.isEmpty(indexOp)) {
            logger.error("No Index operation specified for {}, transferring to failure.", new Object[] { file });
            flowFilesToTransfer.remove(file);
            session.transfer(file, REL_FAILURE);
            continue;
        }
        switch(indexOp.toLowerCase()) {
            case "index":
            case "update":
            case "upsert":
            case "delete":
                break;
            default:
                logger.error("Index operation {} not supported for {}, transferring to failure.", new Object[] { indexOp, file });
                flowFilesToTransfer.remove(file);
                session.transfer(file, REL_FAILURE);
                continue;
        }
        final String id = (id_attribute != null) ? file.getAttribute(id_attribute) : null;
        // a missing ID indicates one is to be auto-generated by Elasticsearch
        if (id == null && !indexOp.equalsIgnoreCase("index")) {
            logger.error("Index operation {} requires a valid identifier value from a flow file attribute, transferring to failure.", new Object[] { indexOp, file });
            flowFilesToTransfer.remove(file);
            session.transfer(file, REL_FAILURE);
            continue;
        }
        final StringBuilder json = new StringBuilder();
        session.read(file, in -> {
            json.append(IOUtils.toString(in, charset).replace("\r\n", " ").replace('\n', ' ').replace('\r', ' '));
        });
        if (indexOp.equalsIgnoreCase("index")) {
            sb.append("{\"index\": { \"_index\": \"");
            sb.append(index);
            sb.append("\", \"_type\": \"");
            sb.append(docType);
            sb.append("\"");
            if (!StringUtils.isEmpty(id)) {
                sb.append(", \"_id\": \"");
                sb.append(id);
                sb.append("\"");
            }
            sb.append("}}\n");
            sb.append(json);
            sb.append("\n");
        } else if (indexOp.equalsIgnoreCase("upsert") || indexOp.equalsIgnoreCase("update")) {
            sb.append("{\"update\": { \"_index\": \"");
            sb.append(index);
            sb.append("\", \"_type\": \"");
            sb.append(docType);
            sb.append("\", \"_id\": \"");
            sb.append(id);
            sb.append("\" }\n");
            sb.append("{\"doc\": ");
            sb.append(json);
            sb.append(", \"doc_as_upsert\": ");
            sb.append(indexOp.equalsIgnoreCase("upsert"));
            sb.append(" }\n");
        } else if (indexOp.equalsIgnoreCase("delete")) {
            sb.append("{\"delete\": { \"_index\": \"");
            sb.append(index);
            sb.append("\", \"_type\": \"");
            sb.append(docType);
            sb.append("\", \"_id\": \"");
            sb.append(id);
            sb.append("\" }\n");
        }
    }
    if (!flowFilesToTransfer.isEmpty()) {
        RequestBody requestBody = RequestBody.create(MediaType.parse("application/json"), sb.toString());
        final Response getResponse;
        try {
            getResponse = sendRequestToElasticsearch(okHttpClient, url, username, password, "PUT", requestBody);
        } catch (final Exception e) {
            logger.error("Routing to {} due to exception: {}", new Object[] { REL_FAILURE.getName(), e }, e);
            flowFilesToTransfer.forEach((flowFileToTransfer) -> {
                flowFileToTransfer = session.penalize(flowFileToTransfer);
                session.transfer(flowFileToTransfer, REL_FAILURE);
            });
            flowFilesToTransfer.clear();
            return;
        }
        final int statusCode = getResponse.code();
        if (isSuccess(statusCode)) {
            ResponseBody responseBody = getResponse.body();
            try {
                final byte[] bodyBytes = responseBody.bytes();
                JsonNode responseJson = parseJsonResponse(new ByteArrayInputStream(bodyBytes));
                boolean errors = responseJson.get("errors").asBoolean(false);
                if (errors) {
                    ArrayNode itemNodeArray = (ArrayNode) responseJson.get("items");
                    if (itemNodeArray.size() > 0) {
                        // All items are returned whether they succeeded or failed, so iterate through the item array
                        // at the same time as the flow file list, moving each to success or failure accordingly,
                        // but only keep the first error for logging
                        String errorReason = null;
                        for (int i = itemNodeArray.size() - 1; i >= 0; i--) {
                            JsonNode itemNode = itemNodeArray.get(i);
                            if (flowFilesToTransfer.size() > i) {
                                FlowFile flowFile = flowFilesToTransfer.remove(i);
                                int status = itemNode.findPath("status").asInt();
                                if (!isSuccess(status)) {
                                    if (errorReason == null) {
                                        // Use "result" if it is present; this happens for status codes like 404 Not Found, which may not have an error/reason
                                        String reason = itemNode.findPath("//result").asText();
                                        if (StringUtils.isEmpty(reason)) {
                                            // If there was no result, we expect an error with a string description in the "reason" field
                                            reason = itemNode.findPath("//error/reason").asText();
                                        }
                                        errorReason = reason;
                                        logger.error("Failed to process {} due to {}, transferring to failure", new Object[] { flowFile, errorReason });
                                    }
                                    flowFile = session.penalize(flowFile);
                                    session.transfer(flowFile, REL_FAILURE);
                                } else {
                                    session.transfer(flowFile, REL_SUCCESS);
                                    // Record provenance event
                                    session.getProvenanceReporter().send(flowFile, url.toString());
                                }
                            }
                        }
                    }
                }
                // Transfer any remaining flowfiles to success
                flowFilesToTransfer.forEach(file -> {
                    session.transfer(file, REL_SUCCESS);
                    // Record provenance event
                    session.getProvenanceReporter().send(file, url.toString());
                });
            } catch (IOException ioe) {
                // Something went wrong when parsing the response, log the error and route to failure
                logger.error("Error parsing Bulk API response: {}", new Object[] { ioe.getMessage() }, ioe);
                session.transfer(flowFilesToTransfer, REL_FAILURE);
                context.yield();
            }
        } else if (statusCode / 100 == 5) {
            // 5xx -> RETRY, but a server error might last a while, so yield
            logger.warn("Elasticsearch returned code {} with message {}, transferring flow file to retry. This is likely a server problem, yielding...", new Object[] { statusCode, getResponse.message() });
            session.transfer(flowFilesToTransfer, REL_RETRY);
            context.yield();
        } else {
            // 1xx, 3xx, 4xx, etc. -> NO RETRY
            logger.warn("Elasticsearch returned code {} with message {}, transferring flow file to failure", new Object[] { statusCode, getResponse.message() });
            session.transfer(flowFilesToTransfer, REL_FAILURE);
        }
        getResponse.close();
    }
}