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

Example 56 with OutputStreamCallback

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

the class EvaluateXQuery method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final List<FlowFile> flowFileBatch = session.get(50);
    if (flowFileBatch.isEmpty()) {
        return;
    }
    final ComponentLog logger = getLogger();
    final Map<String, XQueryExecutable> attributeToXQueryMap = new HashMap<>();
    final Processor proc = new Processor(false);
    final XMLReader xmlReader;
    try {
        xmlReader = XMLReaderFactory.createXMLReader();
    } catch (SAXException e) {
        logger.error("Error while constructing XMLReader {}", new Object[] { e });
        throw new ProcessException(e.getMessage());
    }
    if (!context.getProperty(VALIDATE_DTD).asBoolean()) {
        xmlReader.setEntityResolver(new EntityResolver() {

            @Override
            public InputSource resolveEntity(String publicId, String systemId) throws SAXException, IOException {
                return new InputSource(new StringReader(""));
            }
        });
    }
    final XQueryCompiler comp = proc.newXQueryCompiler();
    for (final Map.Entry<PropertyDescriptor, String> entry : context.getProperties().entrySet()) {
        if (!entry.getKey().isDynamic()) {
            continue;
        }
        final XQueryExecutable exp;
        try {
            exp = comp.compile(entry.getValue());
            attributeToXQueryMap.put(entry.getKey().getName(), exp);
        } catch (SaxonApiException e) {
            // should not happen because we've already validated the XQuery (in XQueryValidator)
            throw new ProcessException(e);
        }
    }
    final XQueryExecutable slashExpression;
    try {
        slashExpression = comp.compile("/");
    } catch (SaxonApiException e) {
        logger.error("unable to compile XQuery expression due to {}", new Object[] { e });
        session.transfer(flowFileBatch, REL_FAILURE);
        return;
    }
    final String destination = context.getProperty(DESTINATION).getValue();
    flowFileLoop: for (FlowFile flowFile : flowFileBatch) {
        if (!isScheduled()) {
            session.rollback();
            return;
        }
        final AtomicReference<Throwable> error = new AtomicReference<>(null);
        final AtomicReference<XdmNode> sourceRef = new AtomicReference<>(null);
        session.read(flowFile, new InputStreamCallback() {

            @Override
            public void process(final InputStream rawIn) throws IOException {
                try (final InputStream in = new BufferedInputStream(rawIn)) {
                    XQueryEvaluator qe = slashExpression.load();
                    qe.setSource(new SAXSource(xmlReader, new InputSource(in)));
                    DocumentBuilderFactory dfactory = DocumentBuilderFactory.newInstance();
                    dfactory.setNamespaceAware(true);
                    Document dom = dfactory.newDocumentBuilder().newDocument();
                    qe.run(new DOMDestination(dom));
                    XdmNode rootNode = proc.newDocumentBuilder().wrap(dom);
                    sourceRef.set(rootNode);
                } catch (final Exception e) {
                    error.set(e);
                }
            }
        });
        if (error.get() != null) {
            logger.error("unable to evaluate XQuery against {} due to {}; routing to 'failure'", new Object[] { flowFile, error.get() });
            session.transfer(flowFile, REL_FAILURE);
            continue;
        }
        final Map<String, String> xQueryResults = new HashMap<>();
        List<FlowFile> childrenFlowFiles = new ArrayList<>();
        for (final Map.Entry<String, XQueryExecutable> entry : attributeToXQueryMap.entrySet()) {
            try {
                XQueryEvaluator qe = entry.getValue().load();
                qe.setContextItem(sourceRef.get());
                XdmValue result = qe.evaluate();
                if (DESTINATION_ATTRIBUTE.equals(destination)) {
                    int index = 1;
                    for (XdmItem item : result) {
                        String value = formatItem(item, context);
                        String attributeName = entry.getKey();
                        if (result.size() > 1) {
                            attributeName += "." + index++;
                        }
                        xQueryResults.put(attributeName, value);
                    }
                } else {
                    // if (DESTINATION_CONTENT.equals(destination)){
                    if (result.size() == 0) {
                        logger.info("Routing {} to 'unmatched'", new Object[] { flowFile });
                        session.transfer(flowFile, REL_NO_MATCH);
                        continue flowFileLoop;
                    } else if (result.size() == 1) {
                        final XdmItem item = result.itemAt(0);
                        flowFile = session.write(flowFile, new OutputStreamCallback() {

                            @Override
                            public void process(final OutputStream rawOut) throws IOException {
                                try (final OutputStream out = new BufferedOutputStream(rawOut)) {
                                    writeformattedItem(item, context, out);
                                } catch (TransformerFactoryConfigurationError | TransformerException e) {
                                    throw new IOException(e);
                                }
                            }
                        });
                    } else {
                        for (final XdmItem item : result) {
                            FlowFile ff = session.clone(flowFile);
                            ff = session.write(ff, new OutputStreamCallback() {

                                @Override
                                public void process(final OutputStream rawOut) throws IOException {
                                    try (final OutputStream out = new BufferedOutputStream(rawOut)) {
                                        try {
                                            writeformattedItem(item, context, out);
                                        } catch (TransformerFactoryConfigurationError | TransformerException e) {
                                            throw new IOException(e);
                                        }
                                    }
                                }
                            });
                            childrenFlowFiles.add(ff);
                        }
                    }
                }
            } catch (final SaxonApiException e) {
                logger.error("failed to evaluate XQuery for {} for Property {} due to {}; routing to failure", new Object[] { flowFile, entry.getKey(), e });
                session.transfer(flowFile, REL_FAILURE);
                session.remove(childrenFlowFiles);
                continue flowFileLoop;
            } catch (TransformerFactoryConfigurationError | TransformerException | IOException e) {
                logger.error("Failed to write XQuery result for {} due to {}; routing original to 'failure'", new Object[] { flowFile, error.get() });
                session.transfer(flowFile, REL_FAILURE);
                session.remove(childrenFlowFiles);
                continue flowFileLoop;
            }
        }
        if (DESTINATION_ATTRIBUTE.equals(destination)) {
            flowFile = session.putAllAttributes(flowFile, xQueryResults);
            final Relationship destRel = xQueryResults.isEmpty() ? REL_NO_MATCH : REL_MATCH;
            logger.info("Successfully evaluated XQueries against {} and found {} matches; routing to {}", new Object[] { flowFile, xQueryResults.size(), destRel.getName() });
            session.transfer(flowFile, destRel);
            session.getProvenanceReporter().modifyAttributes(flowFile);
        } else {
            // if (DESTINATION_CONTENT.equals(destination)) {
            if (!childrenFlowFiles.isEmpty()) {
                logger.info("Successfully created {} new FlowFiles from {}; routing all to 'matched'", new Object[] { childrenFlowFiles.size(), flowFile });
                session.transfer(childrenFlowFiles, REL_MATCH);
                session.remove(flowFile);
            } else {
                logger.info("Successfully updated content for {}; routing to 'matched'", new Object[] { flowFile });
                session.transfer(flowFile, REL_MATCH);
                session.getProvenanceReporter().modifyContent(flowFile);
            }
        }
    }
// end flowFileLoop
}

Example 57 with OutputStreamCallback

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

the class FetchFileTransfer method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final StopWatch stopWatch = new StopWatch(true);
    final String host = context.getProperty(HOSTNAME).evaluateAttributeExpressions(flowFile).getValue();
    final int port = context.getProperty(UNDEFAULTED_PORT).evaluateAttributeExpressions(flowFile).asInteger();
    final String filename = context.getProperty(REMOTE_FILENAME).evaluateAttributeExpressions(flowFile).getValue();
    // Try to get a FileTransfer object from our cache.
    BlockingQueue<FileTransferIdleWrapper> transferQueue;
    synchronized (fileTransferMap) {
        final Tuple<String, Integer> tuple = new Tuple<>(host, port);
        transferQueue = fileTransferMap.get(tuple);
        if (transferQueue == null) {
            transferQueue = new LinkedBlockingQueue<>();
            fileTransferMap.put(tuple, transferQueue);
        }
        // periodically close idle connections
        if (System.currentTimeMillis() - lastClearTime > IDLE_CONNECTION_MILLIS) {
            closeConnections(false);
            lastClearTime = System.currentTimeMillis();
        }
    }
    // we have a queue of FileTransfer Objects. Get one from the queue or create a new one.
    FileTransfer transfer;
    FileTransferIdleWrapper transferWrapper = transferQueue.poll();
    if (transferWrapper == null) {
        transfer = createFileTransfer(context);
    } else {
        transfer = transferWrapper.getFileTransfer();
    }
    boolean closeConnection = false;
    try {
        // Pull data from remote system.
        final InputStream in;
        try {
            in = transfer.getInputStream(filename, flowFile);
            flowFile = session.write(flowFile, new OutputStreamCallback() {

                @Override
                public void process(final OutputStream out) throws IOException {
                    StreamUtils.copy(in, out);
                }
            });
            if (!transfer.flush(flowFile)) {
                throw new IOException("completePendingCommand returned false, file transfer failed");
            }
        } catch (final FileNotFoundException e) {
            closeConnection = false;
            getLogger().error("Failed to fetch content for {} from filename {} on remote host {} because the file could not be found on the remote system; routing to {}", new Object[] { flowFile, filename, host, REL_NOT_FOUND.getName() });
            session.transfer(session.penalize(flowFile), REL_NOT_FOUND);
            session.getProvenanceReporter().route(flowFile, REL_NOT_FOUND);
            return;
        } catch (final PermissionDeniedException e) {
            closeConnection = false;
            getLogger().error("Failed to fetch content for {} from filename {} on remote host {} due to insufficient permissions; routing to {}", new Object[] { flowFile, filename, host, REL_PERMISSION_DENIED.getName() });
            session.transfer(session.penalize(flowFile), REL_PERMISSION_DENIED);
            session.getProvenanceReporter().route(flowFile, REL_PERMISSION_DENIED);
            return;
        } catch (final ProcessException | IOException e) {
            closeConnection = true;
            getLogger().error("Failed to fetch content for {} from filename {} on remote host {}:{} due to {}; routing to comms.failure", new Object[] { flowFile, filename, host, port, e.toString() }, e);
            session.transfer(session.penalize(flowFile), REL_COMMS_FAILURE);
            return;
        }
        // Add FlowFile attributes
        final String protocolName = transfer.getProtocolName();
        final Map<String, String> attributes = new HashMap<>();
        attributes.put(protocolName + ".remote.host", host);
        attributes.put(protocolName + ".remote.port", String.valueOf(port));
        attributes.put(protocolName + ".remote.filename", filename);
        if (filename.contains("/")) {
            final String path = StringUtils.substringBeforeLast(filename, "/");
            final String filenameOnly = StringUtils.substringAfterLast(filename, "/");
            attributes.put(CoreAttributes.PATH.key(), path);
            attributes.put(CoreAttributes.FILENAME.key(), filenameOnly);
        } else {
            attributes.put(CoreAttributes.FILENAME.key(), filename);
        }
        flowFile = session.putAllAttributes(flowFile, attributes);
        // emit provenance event and transfer FlowFile
        session.getProvenanceReporter().fetch(flowFile, protocolName + "://" + host + ":" + port + "/" + filename, stopWatch.getElapsed(TimeUnit.MILLISECONDS));
        session.transfer(flowFile, REL_SUCCESS);
        // it is critical that we commit the session before moving/deleting the remote file. Otherwise, we could have a situation where
        // we ingest the data, delete/move the remote file, and then NiFi dies/is shut down before the session is committed. This would
        // result in data loss! If we commit the session first, we are safe.
        session.commit();
        final String completionStrategy = context.getProperty(COMPLETION_STRATEGY).getValue();
        if (COMPLETION_DELETE.getValue().equalsIgnoreCase(completionStrategy)) {
            try {
                transfer.deleteFile(flowFile, null, filename);
            } catch (final FileNotFoundException e) {
            // file doesn't exist -- effectively the same as removing it. Move on.
            } catch (final IOException ioe) {
                getLogger().warn("Successfully fetched the content for {} from {}:{}{} but failed to remove the remote file due to {}", new Object[] { flowFile, host, port, filename, ioe }, ioe);
            }
        } else if (COMPLETION_MOVE.getValue().equalsIgnoreCase(completionStrategy)) {
            String targetDir = context.getProperty(MOVE_DESTINATION_DIR).evaluateAttributeExpressions(flowFile).getValue();
            if (!targetDir.endsWith("/")) {
                targetDir = targetDir + "/";
            }
            final String simpleFilename = StringUtils.substringAfterLast(filename, "/");
            final String target = targetDir + simpleFilename;
            try {
                transfer.rename(flowFile, filename, target);
            } catch (final IOException ioe) {
                getLogger().warn("Successfully fetched the content for {} from {}:{}{} but failed to rename the remote file due to {}", new Object[] { flowFile, host, port, filename, ioe }, ioe);
            }
        }
    } finally {
        if (transfer != null) {
            if (closeConnection) {
                getLogger().debug("Closing FileTransfer...");
                try {
                    transfer.close();
                } catch (final IOException e) {
                    getLogger().warn("Failed to close connection to {}:{} due to {}", new Object[] { host, port, e.getMessage() }, e);
                }
            } else {
                getLogger().debug("Returning FileTransfer to pool...");
                transferQueue.offer(new FileTransferIdleWrapper(transfer, System.nanoTime()));
            }
        }
    }
}

Example 58 with OutputStreamCallback

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

the class GenerateFlowFile method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final byte[] data;
    if (context.getProperty(UNIQUE_FLOWFILES).asBoolean()) {
        data = generateData(context);
    } else if (context.getProperty(CUSTOM_TEXT).isSet()) {
        final Charset charset = Charset.forName(context.getProperty(CHARSET).getValue());
        data = context.getProperty(CUSTOM_TEXT).evaluateAttributeExpressions().getValue().getBytes(charset);
    } else {
        data = this.data.get();
    }
    Map<PropertyDescriptor, String> processorProperties = context.getProperties();
    Map<String, String> generatedAttributes = new HashMap<String, String>();
    for (final Map.Entry<PropertyDescriptor, String> entry : processorProperties.entrySet()) {
        PropertyDescriptor property = entry.getKey();
        if (property.isDynamic() && property.isExpressionLanguageSupported()) {
            String dynamicValue = context.getProperty(property).evaluateAttributeExpressions().getValue();
            generatedAttributes.put(property.getName(), dynamicValue);
        }
    }
    for (int i = 0; i < context.getProperty(BATCH_SIZE).asInteger(); i++) {
        FlowFile flowFile = session.create();
        if (data.length > 0) {
            flowFile = session.write(flowFile, new OutputStreamCallback() {

                @Override
                public void process(final OutputStream out) throws IOException {
                    out.write(data);
                }
            });
        }
        flowFile = session.putAllAttributes(flowFile, generatedAttributes);
        session.getProvenanceReporter().create(flowFile);
        session.transfer(flowFile, SUCCESS);
    }
}

Example 59 with OutputStreamCallback

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

the class SpringContextProcessor method receiveFromSpring.

/**
 */
private void receiveFromSpring(ProcessSession processSession) {
    final SpringResponse<?> msgFromSpring = this.exchanger.receive(this.receiveTimeout);
    if (msgFromSpring != null) {
        FlowFile flowFileToProcess = processSession.create();
        flowFileToProcess = processSession.write(flowFileToProcess, new OutputStreamCallback() {

            @Override
            public void process(final OutputStream out) throws IOException {
                Object payload = msgFromSpring.getPayload();
                byte[] payloadBytes = payload instanceof String ? ((String) payload).getBytes() : (byte[]) payload;
                out.write(payloadBytes);
            }
        });
        flowFileToProcess = processSession.putAllAttributes(flowFileToProcess, this.extractFlowFileAttributesFromMessageHeaders(msgFromSpring.getHeaders()));
        processSession.transfer(flowFileToProcess, REL_SUCCESS);
        processSession.getProvenanceReporter().receive(flowFileToProcess, this.applicationContextConfigFileName);
    }
}

Example 60 with OutputStreamCallback

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

the class ConvertExcelToCSVProcessor method handleExcelSheet.

/**
 * Handles an individual Excel sheet from the entire Excel document. Each sheet will result in an individual flowfile.
 *
 * @param session
 *  The NiFi ProcessSession instance for the current invocation.
 */
private void handleExcelSheet(ProcessSession session, FlowFile originalParentFF, final InputStream sheetInputStream, ExcelSheetReadConfig readConfig, CSVFormat csvFormat) throws IOException {
    FlowFile ff = session.create(originalParentFF);
    try {
        final DataFormatter formatter = new DataFormatter();
        final InputSource sheetSource = new InputSource(sheetInputStream);
        final SheetToCSV sheetHandler = new SheetToCSV(readConfig, csvFormat);
        final XMLReader parser = SAXHelper.newXMLReader();
        // If Value Formatting is set to false then don't pass in the styles table.
        // This will cause the XSSF Handler to return the raw value instead of the formatted one.
        final StylesTable sst = readConfig.getFormatValues() ? readConfig.getStyles() : null;
        final XSSFSheetXMLHandler handler = new XSSFSheetXMLHandler(sst, null, readConfig.getSharedStringsTable(), sheetHandler, formatter, false);
        parser.setContentHandler(handler);
        ff = session.write(ff, new OutputStreamCallback() {

            @Override
            public void process(OutputStream out) throws IOException {
                PrintStream outPrint = new PrintStream(out);
                sheetHandler.setOutput(outPrint);
                try {
                    parser.parse(sheetSource);
                    sheetInputStream.close();
                    sheetHandler.close();
                    outPrint.close();
                } catch (SAXException se) {
                    getLogger().error("Error occurred while processing Excel sheet {}", new Object[] { readConfig.getSheetName() }, se);
                }
            }
        });
        ff = session.putAttribute(ff, SHEET_NAME, readConfig.getSheetName());
        ff = session.putAttribute(ff, ROW_NUM, new Long(sheetHandler.getRowCount()).toString());
        if (StringUtils.isNotEmpty(originalParentFF.getAttribute(CoreAttributes.FILENAME.key()))) {
            ff = session.putAttribute(ff, SOURCE_FILE_NAME, originalParentFF.getAttribute(CoreAttributes.FILENAME.key()));
        } else {
            ff = session.putAttribute(ff, SOURCE_FILE_NAME, UNKNOWN_SHEET_NAME);
        }
        // Update the CoreAttributes.FILENAME to have the .csv extension now. Also update MIME.TYPE
        ff = session.putAttribute(ff, CoreAttributes.FILENAME.key(), updateFilenameToCSVExtension(ff.getAttribute(CoreAttributes.UUID.key()), ff.getAttribute(CoreAttributes.FILENAME.key()), readConfig.getSheetName()));
        ff = session.putAttribute(ff, CoreAttributes.MIME_TYPE.key(), CSV_MIME_TYPE);
        session.transfer(ff, SUCCESS);
    } catch (SAXException | ParserConfigurationException saxE) {
        getLogger().error("Failed to create instance of Parser.", saxE);
        ff = session.putAttribute(ff, ConvertExcelToCSVProcessor.class.getName() + ".error", saxE.getMessage());
        session.transfer(ff, FAILURE);
    } finally {
        sheetInputStream.close();
    }
}