Examples with OutputStreamCallback org.apache.nifi.processor.io.OutputStreamCallback used on opensource projects

Example 26 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class ExecuteProcess method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    if (proxyOut == null) {
        proxyOut = new ProxyOutputStream(getLogger());
    }
    final Long batchNanos = context.getProperty(BATCH_DURATION).asTimePeriod(TimeUnit.NANOSECONDS);
    final String command = context.getProperty(COMMAND).evaluateAttributeExpressions().getValue();
    final String arguments = context.getProperty(COMMAND_ARGUMENTS).isSet() ? context.getProperty(COMMAND_ARGUMENTS).evaluateAttributeExpressions().getValue() : null;
    final List<String> commandStrings = createCommandStrings(context, command, arguments);
    final String commandString = StringUtils.join(commandStrings, " ");
    if (longRunningProcess == null || longRunningProcess.isDone()) {
        try {
            longRunningProcess = launchProcess(context, commandStrings, batchNanos, proxyOut);
        } catch (final IOException ioe) {
            getLogger().error("Failed to create process due to {}", new Object[] { ioe });
            context.yield();
            return;
        }
    } else {
        getLogger().info("Read from long running process");
    }
    if (!isScheduled()) {
        getLogger().info("User stopped processor; will terminate process immediately");
        longRunningProcess.cancel(true);
        return;
    }
    // Create a FlowFile that we can write to and set the OutputStream for the FlowFile
    // as the delegate for the ProxyOuptutStream, then wait until the process finishes
    // or until the specified amount of time
    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, new OutputStreamCallback() {

        @Override
        public void process(final OutputStream flowFileOut) throws IOException {
            try (final OutputStream out = new BufferedOutputStream(flowFileOut)) {
                proxyOut.setDelegate(out);
                if (batchNanos == null) {
                    // be used to avoid waiting forever.
                    try {
                        longRunningProcess.get();
                    } catch (final InterruptedException ie) {
                    } catch (final ExecutionException ee) {
                        getLogger().error("Process execution failed due to {}", new Object[] { ee.getCause() });
                    }
                } else {
                    // wait the allotted amount of time.
                    try {
                        TimeUnit.NANOSECONDS.sleep(batchNanos);
                    } catch (final InterruptedException ie) {
                    }
                }
                // prevent from writing to this
                proxyOut.setDelegate(null);
            // stream
            }
        }
    });
    if (flowFile.getSize() == 0L) {
        // If no data was written to the file, remove it
        session.remove(flowFile);
    } else if (failure.get()) {
        // If there was a failure processing the output of the Process, remove the FlowFile
        session.remove(flowFile);
        getLogger().error("Failed to read data from Process, so will not generate FlowFile");
    } else {
        // add command and arguments as attribute
        flowFile = session.putAttribute(flowFile, ATTRIBUTE_COMMAND, command);
        if (arguments != null) {
            flowFile = session.putAttribute(flowFile, ATTRIBUTE_COMMAND_ARGS, arguments);
        }
        // All was good. Generate event and transfer FlowFile.
        session.getProvenanceReporter().create(flowFile, "Created from command: " + commandString);
        getLogger().info("Created {} and routed to success", new Object[] { flowFile });
        session.transfer(flowFile, REL_SUCCESS);
    }
    // Commit the session so that the FlowFile is transferred to the next processor
    session.commit();
}

Example 27 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class GetTwitter method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final Event event = eventQueue.poll();
    if (event != null) {
        switch(event.getEventType()) {
            case STOPPED_BY_ERROR:
                getLogger().error("Received error {}: {} due to {}. Will not attempt to reconnect", new Object[] { event.getEventType(), event.getMessage(), event.getUnderlyingException() });
                break;
            case CONNECTION_ERROR:
            case HTTP_ERROR:
                getLogger().error("Received error {}: {}. Will attempt to reconnect", new Object[] { event.getEventType(), event.getMessage() });
                client.reconnect();
                break;
            default:
                break;
        }
    }
    final String tweet = messageQueue.poll();
    if (tweet == null) {
        context.yield();
        return;
    }
    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, new OutputStreamCallback() {

        @Override
        public void process(final OutputStream out) throws IOException {
            out.write(tweet.getBytes(StandardCharsets.UTF_8));
        }
    });
    final Map<String, String> attributes = new HashMap<>();
    attributes.put(CoreAttributes.MIME_TYPE.key(), "application/json");
    attributes.put(CoreAttributes.FILENAME.key(), flowFile.getAttribute(CoreAttributes.FILENAME.key()) + ".json");
    flowFile = session.putAllAttributes(flowFile, attributes);
    session.transfer(flowFile, REL_SUCCESS);
    session.getProvenanceReporter().receive(flowFile, Constants.STREAM_HOST + client.getEndpoint().getURI().toString());
}

Example 28 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class GetSolr method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();
    final AtomicBoolean continuePaging = new AtomicBoolean(true);
    final SolrQuery solrQuery = new SolrQuery();
    try {
        if (id_field == null) {
            id_field = getFieldNameOfUniqueKey();
        }
        final String dateField = context.getProperty(DATE_FIELD).getValue();
        final Map<String, String> stateMap = new HashMap<String, String>();
        stateMap.putAll(context.getStateManager().getState(Scope.CLUSTER).toMap());
        solrQuery.setQuery("*:*");
        final String query = context.getProperty(SOLR_QUERY).getValue();
        if (!StringUtils.isBlank(query) && !query.equals("*:*")) {
            solrQuery.addFilterQuery(query);
        }
        final StringBuilder automatedFilterQuery = (new StringBuilder()).append(dateField).append(":[").append(stateMap.get(STATE_MANAGER_FILTER)).append(" TO *]");
        solrQuery.addFilterQuery(automatedFilterQuery.toString());
        final List<String> fieldList = new ArrayList<String>();
        final String returnFields = context.getProperty(RETURN_FIELDS).getValue();
        if (!StringUtils.isBlank(returnFields)) {
            fieldList.addAll(Arrays.asList(returnFields.trim().split("[,]")));
            if (!fieldList.contains(dateField)) {
                fieldList.add(dateField);
                dateFieldNotInSpecifiedFieldsList.set(true);
            }
            for (String returnField : fieldList) {
                solrQuery.addField(returnField.trim());
            }
        }
        solrQuery.setParam(CursorMarkParams.CURSOR_MARK_PARAM, stateMap.get(STATE_MANAGER_CURSOR_MARK));
        solrQuery.setRows(context.getProperty(BATCH_SIZE).asInteger());
        final StringBuilder sortClause = (new StringBuilder()).append(dateField).append(" asc,").append(id_field).append(" asc");
        solrQuery.setParam("sort", sortClause.toString());
        while (continuePaging.get()) {
            final QueryRequest req = new QueryRequest(solrQuery);
            if (isBasicAuthEnabled()) {
                req.setBasicAuthCredentials(getUsername(), getPassword());
            }
            logger.debug(solrQuery.toQueryString());
            final QueryResponse response = req.process(getSolrClient());
            final SolrDocumentList documentList = response.getResults();
            if (response.getResults().size() > 0) {
                final SolrDocument lastSolrDocument = documentList.get(response.getResults().size() - 1);
                final String latestDateValue = df.format(lastSolrDocument.get(dateField));
                final String newCursorMark = response.getNextCursorMark();
                solrQuery.setParam(CursorMarkParams.CURSOR_MARK_PARAM, newCursorMark);
                stateMap.put(STATE_MANAGER_CURSOR_MARK, newCursorMark);
                stateMap.put(STATE_MANAGER_FILTER, latestDateValue);
                FlowFile flowFile = session.create();
                flowFile = session.putAttribute(flowFile, "solrQuery", solrQuery.toString());
                if (context.getProperty(RETURN_TYPE).getValue().equals(MODE_XML.getValue())) {
                    if (dateFieldNotInSpecifiedFieldsList.get()) {
                        for (SolrDocument doc : response.getResults()) {
                            doc.removeFields(dateField);
                        }
                    }
                    flowFile = session.write(flowFile, SolrUtils.getOutputStreamCallbackToTransformSolrResponseToXml(response));
                    flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), "application/xml");
                } else {
                    final RecordSetWriterFactory writerFactory = context.getProperty(RECORD_WRITER).asControllerService(RecordSetWriterFactory.class);
                    final RecordSchema schema = writerFactory.getSchema(null, null);
                    final RecordSet recordSet = SolrUtils.solrDocumentsToRecordSet(response.getResults(), schema);
                    final StringBuffer mimeType = new StringBuffer();
                    flowFile = session.write(flowFile, new OutputStreamCallback() {

                        @Override
                        public void process(final OutputStream out) throws IOException {
                            try {
                                final RecordSetWriter writer = writerFactory.createWriter(getLogger(), schema, out);
                                writer.write(recordSet);
                                writer.flush();
                                mimeType.append(writer.getMimeType());
                            } catch (SchemaNotFoundException e) {
                                throw new ProcessException("Could not parse Solr response", e);
                            }
                        }
                    });
                    flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), mimeType.toString());
                }
                session.transfer(flowFile, REL_SUCCESS);
            }
            continuePaging.set(response.getResults().size() == Integer.parseInt(context.getProperty(BATCH_SIZE).getValue()));
        }
        context.getStateManager().setState(stateMap, Scope.CLUSTER);
    } catch (SolrServerException | SchemaNotFoundException | IOException e) {
        context.yield();
        session.rollback();
        logger.error("Failed to execute query {} due to {}", new Object[] { solrQuery.toString(), e }, e);
        throw new ProcessException(e);
    } catch (final Throwable t) {
        context.yield();
        session.rollback();
        logger.error("Failed to execute query {} due to {}", new Object[] { solrQuery.toString(), t }, t);
        throw t;
    }
}

Example 29 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class GetSplunk method onTrigger.

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    final long currentTime = System.currentTimeMillis();
    synchronized (isInitialized) {
        if (!isInitialized.get()) {
            splunkService = createSplunkService(context);
            isInitialized.set(true);
        }
    }
    final String query = context.getProperty(QUERY).getValue();
    final String outputMode = context.getProperty(OUTPUT_MODE).getValue();
    final String timeRangeStrategy = context.getProperty(TIME_RANGE_STRATEGY).getValue();
    final String timeZone = context.getProperty(TIME_ZONE).getValue();
    final String timeFieldStrategy = context.getProperty(TIME_FIELD_STRATEGY).getValue();
    final JobExportArgs exportArgs = new JobExportArgs();
    exportArgs.setSearchMode(JobExportArgs.SearchMode.NORMAL);
    exportArgs.setOutputMode(JobExportArgs.OutputMode.valueOf(outputMode));
    String earliestTime = null;
    String latestTime = null;
    if (PROVIDED_VALUE.getValue().equals(timeRangeStrategy)) {
        // for provided we just use the values of the properties
        earliestTime = context.getProperty(EARLIEST_TIME).getValue();
        latestTime = context.getProperty(LATEST_TIME).getValue();
    } else {
        try {
            // not provided so we need to check the previous state
            final TimeRange previousRange = loadState(context.getStateManager());
            final SimpleDateFormat dateFormat = new SimpleDateFormat(DATE_TIME_FORMAT);
            dateFormat.setTimeZone(TimeZone.getTimeZone(timeZone));
            if (previousRange == null) {
                // no previous state so set the earliest time based on the strategy
                if (MANAGED_CURRENT_VALUE.getValue().equals(timeRangeStrategy)) {
                    earliestTime = dateFormat.format(new Date(currentTime));
                }
                // no previous state so set the latest time to the current time
                latestTime = dateFormat.format(new Date(currentTime));
                // initial time saved and next execution will be the first real execution
                if (latestTime.equals(earliestTime)) {
                    saveState(context.getStateManager(), new TimeRange(earliestTime, latestTime));
                    return;
                }
            } else {
                // we have previous state so set earliestTime to (latestTime + 1) of last range
                try {
                    final String previousLastTime = previousRange.getLatestTime();
                    final Date previousLastDate = dateFormat.parse(previousLastTime);
                    earliestTime = dateFormat.format(new Date(previousLastDate.getTime() + 1));
                    latestTime = dateFormat.format(new Date(currentTime));
                } catch (ParseException e) {
                    throw new ProcessException(e);
                }
            }
        } catch (IOException e) {
            getLogger().error("Unable to load data from State Manager due to {}", new Object[] { e.getMessage() }, e);
            context.yield();
            return;
        }
    }
    if (!StringUtils.isBlank(earliestTime)) {
        if (EVENT_TIME_VALUE.getValue().equalsIgnoreCase(timeFieldStrategy)) {
            exportArgs.setEarliestTime(earliestTime);
        } else {
            exportArgs.setIndexEarliest(earliestTime);
        }
    }
    if (!StringUtils.isBlank(latestTime)) {
        if (EVENT_TIME_VALUE.getValue().equalsIgnoreCase(timeFieldStrategy)) {
            exportArgs.setLatestTime(latestTime);
        } else {
            exportArgs.setIndexLatest(latestTime);
        }
    }
    if (EVENT_TIME_VALUE.getValue().equalsIgnoreCase(timeFieldStrategy)) {
        getLogger().debug("Using earliest_time of {} and latest_time of {}", new Object[] { earliestTime, latestTime });
    } else {
        getLogger().debug("Using index_earliest of {} and index_latest of {}", new Object[] { earliestTime, latestTime });
    }
    InputStream export;
    try {
        export = splunkService.export(query, exportArgs);
    // Catch Stale connection exception, reinitialize, and retry
    } catch (com.splunk.HttpException e) {
        getLogger().error("Splunk request status code:" + e.getStatus() + " Retrying the request.");
        splunkService.logout();
        splunkService = createSplunkService(context);
        export = splunkService.export(query, exportArgs);
    }
    final InputStream exportSearch = export;
    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, new OutputStreamCallback() {

        @Override
        public void process(OutputStream rawOut) throws IOException {
            try (BufferedOutputStream out = new BufferedOutputStream(rawOut)) {
                IOUtils.copyLarge(exportSearch, out);
            }
        }
    });
    final Map<String, String> attributes = new HashMap<>(3);
    attributes.put(EARLIEST_TIME_ATTR, earliestTime);
    attributes.put(LATEST_TIME_ATTR, latestTime);
    attributes.put(QUERY_ATTR, query);
    flowFile = session.putAllAttributes(flowFile, attributes);
    session.getProvenanceReporter().receive(flowFile, transitUri);
    session.transfer(flowFile, REL_SUCCESS);
    getLogger().debug("Received {} from Splunk", new Object[] { flowFile });
    // only need to do this for the managed time strategies
    if (!PROVIDED_VALUE.getValue().equals(timeRangeStrategy)) {
        try {
            saveState(context.getStateManager(), new TimeRange(earliestTime, latestTime));
        } catch (IOException e) {
            getLogger().error("Unable to load data from State Manager due to {}", new Object[] { e.getMessage() }, e);
            session.rollback();
            context.yield();
        }
    }
}

Example 30 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class TestEvaluateJsonPath method testNullInput_nullStringRepresentation.

@Test
public void testNullInput_nullStringRepresentation() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new EvaluateJsonPath());
    testRunner.setProperty(EvaluateJsonPath.RETURN_TYPE, EvaluateJsonPath.RETURN_TYPE_JSON);
    testRunner.setProperty(EvaluateJsonPath.DESTINATION, EvaluateJsonPath.DESTINATION_ATTRIBUTE);
    testRunner.setProperty(EvaluateJsonPath.NULL_VALUE_DEFAULT_REPRESENTATION, AbstractJsonPathProcessor.NULL_STRING_OPTION);
    testRunner.setProperty("stringField", "$.stringField");
    testRunner.setProperty("missingField", "$.missingField");
    testRunner.setProperty("nullField", "$.nullField");
    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();
    ff = session.write(ff, new OutputStreamCallback() {

        @Override
        public void process(OutputStream out) throws IOException {
            try (OutputStream outputStream = new BufferedOutputStream(out)) {
                outputStream.write("{\"stringField\": \"String Value\", \"nullField\": null}".getBytes(StandardCharsets.UTF_8));
            }
        }
    });
    testRunner.enqueue(ff);
    testRunner.run();
    testRunner.assertTransferCount(EvaluateJsonPath.REL_MATCH, 1);
    FlowFile output = testRunner.getFlowFilesForRelationship(EvaluateJsonPath.REL_MATCH).get(0);
    String validFieldValue = output.getAttribute("stringField");
    assertEquals("String Value", validFieldValue);
    String missingValue = output.getAttribute("missingField");
    assertEquals("Missing Value", "", missingValue);
    String nullValue = output.getAttribute("nullField");
    assertEquals("Null Value", "null", nullValue);
}