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

Example 51 with InputStreamCallback

use of org.apache.nifi.processor.io.InputStreamCallback 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 52 with InputStreamCallback

use of org.apache.nifi.processor.io.InputStreamCallback 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);
        }
    }
}

Example 53 with InputStreamCallback

use of org.apache.nifi.processor.io.InputStreamCallback 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 54 with InputStreamCallback

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

the class ExtractText 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 Charset charset = Charset.forName(context.getProperty(CHARACTER_SET).getValue());
    final int maxCaptureGroupLength = context.getProperty(MAX_CAPTURE_GROUP_LENGTH).asInteger();
    final String contentString;
    byte[] buffer = bufferQueue.poll();
    if (buffer == null) {
        final int maxBufferSize = context.getProperty(MAX_BUFFER_SIZE).asDataSize(DataUnit.B).intValue();
        buffer = new byte[maxBufferSize];
    }
    try {
        final byte[] byteBuffer = buffer;
        session.read(flowFile, new InputStreamCallback() {

            @Override
            public void process(InputStream in) throws IOException {
                StreamUtils.fillBuffer(in, byteBuffer, false);
            }
        });
        final long len = Math.min(byteBuffer.length, flowFile.getSize());
        contentString = new String(byteBuffer, 0, (int) len, charset);
    } finally {
        bufferQueue.offer(buffer);
    }
    final Map<String, String> regexResults = new HashMap<>();
    final Map<String, Pattern> patternMap = compiledPattersMapRef.get();
    final int startGroupIdx = context.getProperty(INCLUDE_CAPTURE_GROUP_ZERO).asBoolean() ? 0 : 1;
    for (final Map.Entry<String, Pattern> entry : patternMap.entrySet()) {
        final Matcher matcher = entry.getValue().matcher(contentString);
        int j = 0;
        while (matcher.find()) {
            final String baseKey = entry.getKey();
            int start = j == 0 ? startGroupIdx : 1;
            for (int i = start; i <= matcher.groupCount(); i++) {
                final String key = new StringBuilder(baseKey).append(".").append(i + j).toString();
                String value = matcher.group(i);
                if (value != null && !value.isEmpty()) {
                    if (value.length() > maxCaptureGroupLength) {
                        value = value.substring(0, maxCaptureGroupLength);
                    }
                    regexResults.put(key, value);
                    if (i == 1 && j == 0) {
                        regexResults.put(baseKey, value);
                    }
                }
            }
            j += matcher.groupCount();
            if (!context.getProperty(ENABLE_REPEATING_CAPTURE_GROUP).asBoolean()) {
                break;
            }
        }
    }
    if (!regexResults.isEmpty()) {
        flowFile = session.putAllAttributes(flowFile, regexResults);
        session.getProvenanceReporter().modifyAttributes(flowFile);
        session.transfer(flowFile, REL_MATCH);
        logger.info("Matched {} Regular Expressions and added attributes to FlowFile {}", new Object[] { regexResults.size(), flowFile });
    } else {
        session.transfer(flowFile, REL_NO_MATCH);
        logger.info("Did not match any Regular Expressions for  FlowFile {}", new Object[] { flowFile });
    }
}

Example 55 with InputStreamCallback

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

the class PutSplunk method processSingleMessage.

/**
 * Send the entire FlowFile as a single message.
 */
private void processSingleMessage(ProcessContext context, ProcessSession session, FlowFile flowFile, ChannelSender sender) {
    // copy the contents of the FlowFile to the ByteArrayOutputStream
    final ByteArrayOutputStream baos = new ByteArrayOutputStream((int) flowFile.getSize() + 1);
    session.read(flowFile, new InputStreamCallback() {

        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.copy(in, baos);
        }
    });
    // if TCP and we don't end in a new line then add one
    final String protocol = context.getProperty(PROTOCOL).getValue();
    byte[] buf = baos.toByteArray();
    if (protocol.equals(TCP_VALUE.getValue()) && buf[buf.length - 1] != NEW_LINE_CHAR) {
        final byte[] updatedBuf = new byte[buf.length + 1];
        System.arraycopy(buf, 0, updatedBuf, 0, buf.length);
        updatedBuf[updatedBuf.length - 1] = NEW_LINE_CHAR;
        buf = updatedBuf;
    }
    // create a message batch of one message and add to active batches
    final FlowFileMessageBatch messageBatch = new FlowFileMessageBatch(session, flowFile);
    messageBatch.setNumMessages(1);
    activeBatches.add(messageBatch);
    // attempt to send the data and add the appropriate range
    try {
        sender.send(buf);
        messageBatch.addSuccessfulRange(0L, flowFile.getSize());
    } catch (IOException e) {
        messageBatch.addFailedRange(0L, flowFile.getSize(), e);
        context.yield();
    }
}