Examples with BufferedInputStream org.apache.nifi.stream.io.BufferedInputStream used on opensource projects

Example 6 with BufferedInputStream

use of org.apache.nifi.stream.io.BufferedInputStream in project nifi by apache.

the class ExtractTNEFAttachments method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final ComponentLog logger = getLogger();
    final FlowFile originalFlowFile = session.get();
    if (originalFlowFile == null) {
        return;
    }
    final List<FlowFile> attachmentsList = new ArrayList<>();
    final List<FlowFile> invalidFlowFilesList = new ArrayList<>();
    final List<FlowFile> originalFlowFilesList = new ArrayList<>();
    session.read(originalFlowFile, new InputStreamCallback() {

        @Override
        public void process(final InputStream rawIn) throws IOException {
            try (final InputStream in = new BufferedInputStream(rawIn)) {
                Properties props = new Properties();
                HMEFMessage hmefMessage = null;
                // This will trigger an exception in case content is not a TNEF.
                hmefMessage = new HMEFMessage(in);
                // Add otiginal flowfile (may revert later on in case of errors) //
                originalFlowFilesList.add(originalFlowFile);
                if (hmefMessage != null) {
                    // Attachments isn empty, proceeding.
                    if (!hmefMessage.getAttachments().isEmpty()) {
                        final String originalFlowFileName = originalFlowFile.getAttribute(CoreAttributes.FILENAME.key());
                        try {
                            for (final Attachment attachment : hmefMessage.getAttachments()) {
                                FlowFile split = session.create(originalFlowFile);
                                final Map<String, String> attributes = new HashMap<>();
                                if (StringUtils.isNotBlank(attachment.getLongFilename())) {
                                    attributes.put(CoreAttributes.FILENAME.key(), attachment.getFilename());
                                }
                                String parentUuid = originalFlowFile.getAttribute(CoreAttributes.UUID.key());
                                attributes.put(ATTACHMENT_ORIGINAL_UUID, parentUuid);
                                attributes.put(ATTACHMENT_ORIGINAL_FILENAME, originalFlowFileName);
                                // TODO: Extract Mime Type (HMEF doesn't seem to be able to get this info.
                                split = session.append(split, new OutputStreamCallback() {

                                    @Override
                                    public void process(OutputStream out) throws IOException {
                                        out.write(attachment.getContents());
                                    }
                                });
                                split = session.putAllAttributes(split, attributes);
                                attachmentsList.add(split);
                            }
                        } catch (FlowFileHandlingException e) {
                            // Something went wrong
                            // Removing splits that may have been created
                            session.remove(attachmentsList);
                            // Removing the original flow from its list
                            originalFlowFilesList.remove(originalFlowFile);
                            logger.error("Flowfile {} triggered error {} while processing message removing generated FlowFiles from sessions", new Object[] { originalFlowFile, e });
                            invalidFlowFilesList.add(originalFlowFile);
                        }
                    }
                }
            } catch (Exception e) {
                // Another error hit...
                // Removing the original flow from its list
                originalFlowFilesList.remove(originalFlowFile);
                logger.error("Could not parse the flowfile {} as an email, treating as failure", new Object[] { originalFlowFile, e });
                // Message is invalid or triggered an error during parsing
                invalidFlowFilesList.add(originalFlowFile);
            }
        }
    });
    session.transfer(attachmentsList, REL_ATTACHMENTS);
    // As per above code, originalFlowfile may be routed to invalid or
    // original depending on RFC2822 compliance.
    session.transfer(invalidFlowFilesList, REL_FAILURE);
    session.transfer(originalFlowFilesList, REL_ORIGINAL);
    // check if attachments have been extracted
    if (attachmentsList.size() != 0) {
        if (attachmentsList.size() > 10) {
            // If more than 10, summarise log
            logger.info("Split {} into {} files", new Object[] { originalFlowFile, attachmentsList.size() });
        } else {
            // Otherwise be more verbose and list each individual split
            logger.info("Split {} into {} files: {}", new Object[] { originalFlowFile, attachmentsList.size(), attachmentsList });
        }
    }
}

Example 7 with BufferedInputStream

use of org.apache.nifi.stream.io.BufferedInputStream in project nifi by apache.

the class FlowFileEvent method getBody.

@Override
public byte[] getBody() {
    synchronized (bodyLock) {
        if (!bodyLoaded) {
            if (flowFile.getSize() > Integer.MAX_VALUE) {
                throw new RuntimeException("Can't get body of Event because the backing FlowFile is too large (" + flowFile.getSize() + " bytes)");
            }
            final ByteArrayOutputStream baos = new ByteArrayOutputStream((int) flowFile.getSize());
            session.read(flowFile, new InputStreamCallback() {

                @Override
                public void process(InputStream in) throws IOException {
                    try (BufferedInputStream input = new BufferedInputStream(in)) {
                        StreamUtils.copy(input, baos);
                    }
                    baos.close();
                }
            });
            body = baos.toByteArray();
            bodyLoaded = true;
        }
    }
    return body;
}

Example 8 with BufferedInputStream

use of org.apache.nifi.stream.io.BufferedInputStream in project nifi by apache.

the class AbstractJsonPathProcessor method validateAndEstablishJsonContext.

static DocumentContext validateAndEstablishJsonContext(ProcessSession processSession, FlowFile flowFile) {
    // Parse the document once into an associated context to support multiple path evaluations if specified
    final AtomicReference<DocumentContext> contextHolder = new AtomicReference<>(null);
    processSession.read(flowFile, new InputStreamCallback() {

        @Override
        public void process(InputStream in) throws IOException {
            try (BufferedInputStream bufferedInputStream = new BufferedInputStream(in)) {
                DocumentContext ctx = JsonPath.using(STRICT_PROVIDER_CONFIGURATION).parse(bufferedInputStream);
                contextHolder.set(ctx);
            }
        }
    });
    return contextHolder.get();
}

Example 9 with BufferedInputStream

use of org.apache.nifi.stream.io.BufferedInputStream in project nifi by apache.

the class CompressContent method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final ComponentLog logger = getLogger();
    final long sizeBeforeCompression = flowFile.getSize();
    final String compressionMode = context.getProperty(MODE).getValue();
    String compressionFormatValue = context.getProperty(COMPRESSION_FORMAT).getValue();
    if (compressionFormatValue.equals(COMPRESSION_FORMAT_ATTRIBUTE)) {
        final String mimeType = flowFile.getAttribute(CoreAttributes.MIME_TYPE.key());
        if (mimeType == null) {
            logger.error("No {} attribute exists for {}; routing to failure", new Object[] { CoreAttributes.MIME_TYPE.key(), flowFile });
            session.transfer(flowFile, REL_FAILURE);
            return;
        }
        compressionFormatValue = compressionFormatMimeTypeMap.get(mimeType);
        if (compressionFormatValue == null) {
            logger.info("Mime Type of {} is '{}', which does not indicate a supported Compression Format; routing to success without decompressing", new Object[] { flowFile, mimeType });
            session.transfer(flowFile, REL_SUCCESS);
            return;
        }
    }
    final String compressionFormat = compressionFormatValue;
    final AtomicReference<String> mimeTypeRef = new AtomicReference<>(null);
    final StopWatch stopWatch = new StopWatch(true);
    final String fileExtension;
    switch(compressionFormat.toLowerCase()) {
        case COMPRESSION_FORMAT_GZIP:
            fileExtension = ".gz";
            break;
        case COMPRESSION_FORMAT_LZMA:
            fileExtension = ".lzma";
            break;
        case COMPRESSION_FORMAT_XZ_LZMA2:
            fileExtension = ".xz";
            break;
        case COMPRESSION_FORMAT_BZIP2:
            fileExtension = ".bz2";
            break;
        case COMPRESSION_FORMAT_SNAPPY:
            fileExtension = ".snappy";
            break;
        case COMPRESSION_FORMAT_SNAPPY_FRAMED:
            fileExtension = ".sz";
            break;
        default:
            fileExtension = "";
            break;
    }
    try {
        flowFile = session.write(flowFile, new StreamCallback() {

            @Override
            public void process(final InputStream rawIn, final OutputStream rawOut) throws IOException {
                final OutputStream compressionOut;
                final InputStream compressionIn;
                final OutputStream bufferedOut = new BufferedOutputStream(rawOut, 65536);
                final InputStream bufferedIn = new BufferedInputStream(rawIn, 65536);
                try {
                    if (MODE_COMPRESS.equalsIgnoreCase(compressionMode)) {
                        compressionIn = bufferedIn;
                        switch(compressionFormat.toLowerCase()) {
                            case COMPRESSION_FORMAT_GZIP:
                                final int compressionLevel = context.getProperty(COMPRESSION_LEVEL).asInteger();
                                compressionOut = new GZIPOutputStream(bufferedOut, compressionLevel);
                                mimeTypeRef.set("application/gzip");
                                break;
                            case COMPRESSION_FORMAT_LZMA:
                                compressionOut = new LzmaOutputStream.Builder(bufferedOut).build();
                                mimeTypeRef.set("application/x-lzma");
                                break;
                            case COMPRESSION_FORMAT_XZ_LZMA2:
                                compressionOut = new XZOutputStream(bufferedOut, new LZMA2Options());
                                mimeTypeRef.set("application/x-xz");
                                break;
                            case COMPRESSION_FORMAT_SNAPPY:
                                compressionOut = new SnappyOutputStream(bufferedOut);
                                mimeTypeRef.set("application/x-snappy");
                                break;
                            case COMPRESSION_FORMAT_SNAPPY_FRAMED:
                                compressionOut = new SnappyFramedOutputStream(bufferedOut);
                                mimeTypeRef.set("application/x-snappy-framed");
                                break;
                            case COMPRESSION_FORMAT_BZIP2:
                            default:
                                mimeTypeRef.set("application/x-bzip2");
                                compressionOut = new CompressorStreamFactory().createCompressorOutputStream(compressionFormat.toLowerCase(), bufferedOut);
                                break;
                        }
                    } else {
                        compressionOut = bufferedOut;
                        switch(compressionFormat.toLowerCase()) {
                            case COMPRESSION_FORMAT_LZMA:
                                compressionIn = new LzmaInputStream(bufferedIn, new Decoder());
                                break;
                            case COMPRESSION_FORMAT_XZ_LZMA2:
                                compressionIn = new XZInputStream(bufferedIn);
                                break;
                            case COMPRESSION_FORMAT_BZIP2:
                                // need this two-arg constructor to support concatenated streams
                                compressionIn = new BZip2CompressorInputStream(bufferedIn, true);
                                break;
                            case COMPRESSION_FORMAT_GZIP:
                                compressionIn = new GzipCompressorInputStream(bufferedIn, true);
                                break;
                            case COMPRESSION_FORMAT_SNAPPY:
                                compressionIn = new SnappyInputStream(bufferedIn);
                                break;
                            case COMPRESSION_FORMAT_SNAPPY_FRAMED:
                                compressionIn = new SnappyFramedInputStream(bufferedIn);
                                break;
                            default:
                                compressionIn = new CompressorStreamFactory().createCompressorInputStream(compressionFormat.toLowerCase(), bufferedIn);
                        }
                    }
                } catch (final Exception e) {
                    closeQuietly(bufferedOut);
                    throw new IOException(e);
                }
                try (final InputStream in = compressionIn;
                    final OutputStream out = compressionOut) {
                    final byte[] buffer = new byte[8192];
                    int len;
                    while ((len = in.read(buffer)) > 0) {
                        out.write(buffer, 0, len);
                    }
                    out.flush();
                }
            }
        });
        stopWatch.stop();
        final long sizeAfterCompression = flowFile.getSize();
        if (MODE_DECOMPRESS.equalsIgnoreCase(compressionMode)) {
            flowFile = session.removeAttribute(flowFile, CoreAttributes.MIME_TYPE.key());
            if (context.getProperty(UPDATE_FILENAME).asBoolean()) {
                final String filename = flowFile.getAttribute(CoreAttributes.FILENAME.key());
                if (filename.toLowerCase().endsWith(fileExtension)) {
                    flowFile = session.putAttribute(flowFile, CoreAttributes.FILENAME.key(), filename.substring(0, filename.length() - fileExtension.length()));
                }
            }
        } else {
            flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), mimeTypeRef.get());
            if (context.getProperty(UPDATE_FILENAME).asBoolean()) {
                final String filename = flowFile.getAttribute(CoreAttributes.FILENAME.key());
                flowFile = session.putAttribute(flowFile, CoreAttributes.FILENAME.key(), filename + fileExtension);
            }
        }
        logger.info("Successfully {}ed {} using {} compression format; size changed from {} to {} bytes", new Object[] { compressionMode.toLowerCase(), flowFile, compressionFormat, sizeBeforeCompression, sizeAfterCompression });
        session.getProvenanceReporter().modifyContent(flowFile, stopWatch.getDuration(TimeUnit.MILLISECONDS));
        session.transfer(flowFile, REL_SUCCESS);
    } catch (final ProcessException e) {
        logger.error("Unable to {} {} using {} compression format due to {}; routing to failure", new Object[] { compressionMode.toLowerCase(), flowFile, compressionFormat, e });
        session.transfer(flowFile, REL_FAILURE);
    }
}

Example 10 with BufferedInputStream

use of org.apache.nifi.stream.io.BufferedInputStream in project nifi by apache.

the class EvaluateXPath method onTrigger.

@Override
@SuppressWarnings("unchecked")
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final List<FlowFile> flowFiles = session.get(50);
    if (flowFiles.isEmpty()) {
        return;
    }
    final ComponentLog logger = getLogger();
    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 XPathFactory factory = factoryRef.get();
    final XPathEvaluator xpathEvaluator = (XPathEvaluator) factory.newXPath();
    final Map<String, XPathExpression> attributeToXPathMap = new HashMap<>();
    for (final Map.Entry<PropertyDescriptor, String> entry : context.getProperties().entrySet()) {
        if (!entry.getKey().isDynamic()) {
            continue;
        }
        final XPathExpression xpathExpression;
        try {
            xpathExpression = xpathEvaluator.compile(entry.getValue());
            attributeToXPathMap.put(entry.getKey().getName(), xpathExpression);
        } catch (XPathExpressionException e) {
            // should not happen because we've already validated the XPath (in XPathValidator)
            throw new ProcessException(e);
        }
    }
    final XPathExpression slashExpression;
    try {
        slashExpression = xpathEvaluator.compile("/");
    } catch (XPathExpressionException e) {
        logger.error("unable to compile XPath expression due to {}", new Object[] { e });
        session.transfer(flowFiles, REL_FAILURE);
        return;
    }
    final String destination = context.getProperty(DESTINATION).getValue();
    final QName returnType;
    switch(context.getProperty(RETURN_TYPE).getValue()) {
        case RETURN_TYPE_AUTO:
            if (DESTINATION_ATTRIBUTE.equals(destination)) {
                returnType = STRING;
            } else if (DESTINATION_CONTENT.equals(destination)) {
                returnType = NODESET;
            } else {
                throw new IllegalStateException("The only possible destinations should be CONTENT or ATTRIBUTE...");
            }
            break;
        case RETURN_TYPE_NODESET:
            returnType = NODESET;
            break;
        case RETURN_TYPE_STRING:
            returnType = STRING;
            break;
        default:
            throw new IllegalStateException("There are no other return types...");
    }
    flowFileLoop: for (FlowFile flowFile : flowFiles) {
        final AtomicReference<Throwable> error = new AtomicReference<>(null);
        final AtomicReference<Source> 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)) {
                    final List<Source> rootList = (List<Source>) slashExpression.evaluate(new SAXSource(xmlReader, new InputSource(in)), NODESET);
                    sourceRef.set(rootList.get(0));
                } catch (final Exception e) {
                    error.set(e);
                }
            }
        });
        if (error.get() != null) {
            logger.error("unable to evaluate XPath against {} due to {}; routing to 'failure'", new Object[] { flowFile, error.get() });
            session.transfer(flowFile, REL_FAILURE);
            continue;
        }
        final Map<String, String> xpathResults = new HashMap<>();
        for (final Map.Entry<String, XPathExpression> entry : attributeToXPathMap.entrySet()) {
            Object result = null;
            try {
                result = entry.getValue().evaluate(sourceRef.get(), returnType);
                if (result == null) {
                    continue;
                }
            } catch (final XPathExpressionException e) {
                logger.error("failed to evaluate XPath for {} for Property {} due to {}; routing to failure", new Object[] { flowFile, entry.getKey(), e });
                session.transfer(flowFile, REL_FAILURE);
                continue flowFileLoop;
            }
            if (returnType == NODESET) {
                List<Source> nodeList = (List<Source>) result;
                if (nodeList.isEmpty()) {
                    logger.info("Routing {} to 'unmatched'", new Object[] { flowFile });
                    session.transfer(flowFile, REL_NO_MATCH);
                    continue flowFileLoop;
                } else if (nodeList.size() > 1) {
                    logger.error("Routing {} to 'failure' because the XPath evaluated to {} XML nodes", new Object[] { flowFile, nodeList.size() });
                    session.transfer(flowFile, REL_FAILURE);
                    continue flowFileLoop;
                }
                final Source sourceNode = nodeList.get(0);
                if (DESTINATION_ATTRIBUTE.equals(destination)) {
                    try {
                        ByteArrayOutputStream baos = new ByteArrayOutputStream();
                        doTransform(sourceNode, baos);
                        xpathResults.put(entry.getKey(), baos.toString("UTF-8"));
                    } catch (UnsupportedEncodingException e) {
                        throw new ProcessException(e);
                    } catch (TransformerException e) {
                        error.set(e);
                    }
                } else if (DESTINATION_CONTENT.equals(destination)) {
                    flowFile = session.write(flowFile, new OutputStreamCallback() {

                        @Override
                        public void process(final OutputStream rawOut) throws IOException {
                            try (final OutputStream out = new BufferedOutputStream(rawOut)) {
                                doTransform(sourceNode, out);
                            } catch (TransformerException e) {
                                error.set(e);
                            }
                        }
                    });
                }
            } else if (returnType == STRING) {
                final String resultString = (String) result;
                if (DESTINATION_ATTRIBUTE.equals(destination)) {
                    xpathResults.put(entry.getKey(), resultString);
                } else if (DESTINATION_CONTENT.equals(destination)) {
                    flowFile = session.write(flowFile, new OutputStreamCallback() {

                        @Override
                        public void process(final OutputStream rawOut) throws IOException {
                            try (final OutputStream out = new BufferedOutputStream(rawOut)) {
                                out.write(resultString.getBytes("UTF-8"));
                            }
                        }
                    });
                }
            }
        }
        if (error.get() == null) {
            if (DESTINATION_ATTRIBUTE.equals(destination)) {
                flowFile = session.putAllAttributes(flowFile, xpathResults);
                final Relationship destRel = xpathResults.isEmpty() ? REL_NO_MATCH : REL_MATCH;
                logger.info("Successfully evaluated XPaths against {} and found {} matches; routing to {}", new Object[] { flowFile, xpathResults.size(), destRel.getName() });
                session.transfer(flowFile, destRel);
                session.getProvenanceReporter().modifyAttributes(flowFile);
            } else if (DESTINATION_CONTENT.equals(destination)) {
                logger.info("Successfully updated content for {}; routing to 'matched'", new Object[] { flowFile });
                session.transfer(flowFile, REL_MATCH);
                session.getProvenanceReporter().modifyContent(flowFile);
            }
        } else {
            logger.error("Failed to write XPath result for {} due to {}; routing original to 'failure'", new Object[] { flowFile, error.get() });
            session.transfer(flowFile, REL_FAILURE);
        }
    }
}