Search in sources :

Example 51 with OutputStreamCallback

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

the class ParseEvtxTest method testProcess1RecordGranularity.

@Test
public void testProcess1RecordGranularity() throws IOException, MalformedChunkException, XMLStreamException {
    String basename = "basename";
    int chunkNum = 5;
    int offset = 10001;
    byte[] badChunk = { 8 };
    RootNodeHandler rootNodeHandler1 = mock(RootNodeHandler.class);
    RootNodeHandler rootNodeHandler2 = mock(RootNodeHandler.class);
    RootNodeHandler rootNodeHandler3 = mock(RootNodeHandler.class);
    OutputStream out2 = mock(OutputStream.class);
    OutputStream out3 = mock(OutputStream.class);
    when(rootNodeHandlerFactory.create(out)).thenReturn(rootNodeHandler1);
    when(rootNodeHandlerFactory.create(out2)).thenReturn(rootNodeHandler2);
    when(rootNodeHandlerFactory.create(out3)).thenReturn(rootNodeHandler3);
    ChunkHeader chunkHeader1 = mock(ChunkHeader.class);
    ChunkHeader chunkHeader2 = mock(ChunkHeader.class);
    Record record1 = mock(Record.class);
    Record record2 = mock(Record.class);
    Record record3 = mock(Record.class);
    RootNode rootNode1 = mock(RootNode.class);
    RootNode rootNode2 = mock(RootNode.class);
    RootNode rootNode3 = mock(RootNode.class);
    ProcessSession session = mock(ProcessSession.class);
    FlowFile flowFile = mock(FlowFile.class);
    FlowFile created1 = mock(FlowFile.class);
    FlowFile updated1 = mock(FlowFile.class);
    FlowFile created2 = mock(FlowFile.class);
    FlowFile updated2 = mock(FlowFile.class);
    FlowFile created3 = mock(FlowFile.class);
    FlowFile updated3 = mock(FlowFile.class);
    MalformedChunkException malformedChunkException = new MalformedChunkException("Test", null, offset, chunkNum, badChunk);
    when(session.create(flowFile)).thenReturn(created1).thenReturn(created2).thenReturn(created3).thenReturn(null);
    when(session.write(eq(created1), any(OutputStreamCallback.class))).thenAnswer(invocation -> {
        ((OutputStreamCallback) invocation.getArguments()[1]).process(out);
        return updated1;
    });
    when(session.write(eq(created2), any(OutputStreamCallback.class))).thenAnswer(invocation -> {
        ((OutputStreamCallback) invocation.getArguments()[1]).process(out2);
        return updated2;
    });
    when(session.write(eq(created3), any(OutputStreamCallback.class))).thenAnswer(invocation -> {
        ((OutputStreamCallback) invocation.getArguments()[1]).process(out3);
        return updated3;
    });
    when(record1.getRootNode()).thenReturn(rootNode1);
    when(record2.getRootNode()).thenReturn(rootNode2);
    when(record3.getRootNode()).thenReturn(rootNode3);
    when(fileHeader.hasNext()).thenReturn(true).thenReturn(true).thenReturn(true).thenReturn(false);
    when(fileHeader.next()).thenThrow(malformedChunkException).thenReturn(chunkHeader1).thenReturn(chunkHeader2).thenReturn(null);
    when(chunkHeader1.hasNext()).thenReturn(true).thenReturn(false);
    when(chunkHeader1.next()).thenReturn(record1).thenReturn(null);
    when(chunkHeader2.hasNext()).thenReturn(true).thenReturn(true).thenReturn(false);
    when(chunkHeader2.next()).thenReturn(record2).thenReturn(record3).thenReturn(null);
    parseEvtx.processRecordGranularity(session, componentLog, flowFile, basename, in);
    verify(malformedChunkHandler).handle(flowFile, session, parseEvtx.getName(basename, chunkNum, null, ParseEvtx.EVTX_EXTENSION), badChunk);
    verify(rootNodeHandler1).handle(rootNode1);
    verify(rootNodeHandler1).close();
    verify(rootNodeHandler2).handle(rootNode2);
    verify(rootNodeHandler2).close();
    verify(rootNodeHandler3).handle(rootNode3);
    verify(rootNodeHandler3).close();
}
Also used : ProcessSession(org.apache.nifi.processor.ProcessSession) RootNode(org.apache.nifi.processors.evtx.parser.bxml.RootNode) FlowFile(org.apache.nifi.flowfile.FlowFile) MockFlowFile(org.apache.nifi.util.MockFlowFile) OutputStream(java.io.OutputStream) ChunkHeader(org.apache.nifi.processors.evtx.parser.ChunkHeader) Record(org.apache.nifi.processors.evtx.parser.Record) Mockito.anyString(org.mockito.Mockito.anyString) MalformedChunkException(org.apache.nifi.processors.evtx.parser.MalformedChunkException) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) Test(org.junit.Test)

Example 52 with OutputStreamCallback

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

the class AbstractListenEventBatchingProcessor method getBatches.

/**
 * Batches together up to the batchSize events. Events are grouped together based on a batch key which
 * by default is the sender of the event, but can be override by sub-classes.
 *
 * This method will return when batchSize has been reached, or when no more events are available on the queue.
 *
 * @param session the current session
 * @param totalBatchSize the total number of events to process
 * @param messageDemarcatorBytes the demarcator to put between messages when writing to a FlowFile
 *
 * @return a Map from the batch key to the FlowFile and events for that batch, the size of events in all
 *              the batches will be <= batchSize
 */
protected Map<String, FlowFileEventBatch> getBatches(final ProcessSession session, final int totalBatchSize, final byte[] messageDemarcatorBytes) {
    final Map<String, FlowFileEventBatch> batches = new HashMap<>();
    for (int i = 0; i < totalBatchSize; i++) {
        final E event = getMessage(true, true, session);
        if (event == null) {
            break;
        }
        final String batchKey = getBatchKey(event);
        FlowFileEventBatch batch = batches.get(batchKey);
        // if we don't have a batch for this key then create a new one
        if (batch == null) {
            batch = new FlowFileEventBatch(session.create(), new ArrayList<E>());
            batches.put(batchKey, batch);
        }
        // add the current event to the batch
        batch.getEvents().add(event);
        // append the event's data to the FlowFile, write the demarcator first if not on the first event
        final boolean writeDemarcator = (i > 0);
        try {
            final byte[] rawMessage = event.getData();
            FlowFile appendedFlowFile = session.append(batch.getFlowFile(), new OutputStreamCallback() {

                @Override
                public void process(final OutputStream out) throws IOException {
                    if (writeDemarcator) {
                        out.write(messageDemarcatorBytes);
                    }
                    out.write(rawMessage);
                }
            });
            // update the FlowFile reference in the batch object
            batch.setFlowFile(appendedFlowFile);
        } catch (final Exception e) {
            getLogger().error("Failed to write contents of the message to FlowFile due to {}; will re-queue message and try again", new Object[] { e.getMessage() }, e);
            errorEvents.offer(event);
            break;
        }
    }
    return batches;
}
Also used : FlowFile(org.apache.nifi.flowfile.FlowFile) HashMap(java.util.HashMap) NETWORK_INTF_NAME(org.apache.nifi.processor.util.listen.ListenerProperties.NETWORK_INTF_NAME) OutputStream(java.io.OutputStream) ArrayList(java.util.ArrayList) IOException(java.io.IOException) IOException(java.io.IOException) ProcessException(org.apache.nifi.processor.exception.ProcessException) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback)

Example 53 with OutputStreamCallback

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

the class MonitorActivity method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final long thresholdMillis = context.getProperty(THRESHOLD).asTimePeriod(TimeUnit.MILLISECONDS);
    final long now = System.currentTimeMillis();
    final ComponentLog logger = getLogger();
    final boolean copyAttributes = context.getProperty(COPY_ATTRIBUTES).asBoolean();
    final boolean isClusterScope = isClusterScope(context, false);
    final boolean shouldReportOnlyOnPrimary = shouldReportOnlyOnPrimary(isClusterScope, context);
    final List<FlowFile> flowFiles = session.get(50);
    boolean isInactive = false;
    long updatedLatestSuccessTransfer = -1;
    StateMap clusterState = null;
    if (flowFiles.isEmpty()) {
        final long previousSuccessMillis = latestSuccessTransfer.get();
        boolean sendInactiveMarker = false;
        isInactive = (now >= previousSuccessMillis + thresholdMillis);
        logger.debug("isInactive={}, previousSuccessMillis={}, now={}", new Object[] { isInactive, previousSuccessMillis, now });
        if (isInactive && isClusterScope) {
            // However, if this node is active, we don't have to look at cluster state.
            try {
                clusterState = context.getStateManager().getState(Scope.CLUSTER);
                if (clusterState != null && !StringUtils.isEmpty(clusterState.get(STATE_KEY_LATEST_SUCCESS_TRANSFER))) {
                    final long latestReportedClusterActivity = Long.valueOf(clusterState.get(STATE_KEY_LATEST_SUCCESS_TRANSFER));
                    isInactive = (now >= latestReportedClusterActivity + thresholdMillis);
                    if (!isInactive) {
                        // This node has been inactive, but other node has more recent activity.
                        updatedLatestSuccessTransfer = latestReportedClusterActivity;
                    }
                    logger.debug("isInactive={}, latestReportedClusterActivity={}", new Object[] { isInactive, latestReportedClusterActivity });
                }
            } catch (IOException e) {
                logger.error("Failed to access cluster state. Activity will not be monitored properly until this is addressed.", e);
            }
        }
        if (isInactive) {
            final boolean continual = context.getProperty(CONTINUALLY_SEND_MESSAGES).asBoolean();
            sendInactiveMarker = !inactive.getAndSet(true) || (continual && (now > lastInactiveMessage.get() + thresholdMillis));
        }
        if (sendInactiveMarker && shouldThisNodeReport(isClusterScope, shouldReportOnlyOnPrimary)) {
            lastInactiveMessage.set(System.currentTimeMillis());
            FlowFile inactiveFlowFile = session.create();
            inactiveFlowFile = session.putAttribute(inactiveFlowFile, "inactivityStartMillis", String.valueOf(previousSuccessMillis));
            inactiveFlowFile = session.putAttribute(inactiveFlowFile, "inactivityDurationMillis", String.valueOf(now - previousSuccessMillis));
            final byte[] outBytes = context.getProperty(INACTIVITY_MESSAGE).evaluateAttributeExpressions(inactiveFlowFile).getValue().getBytes(UTF8);
            inactiveFlowFile = session.write(inactiveFlowFile, new OutputStreamCallback() {

                @Override
                public void process(final OutputStream out) throws IOException {
                    out.write(outBytes);
                }
            });
            session.getProvenanceReporter().create(inactiveFlowFile);
            session.transfer(inactiveFlowFile, REL_INACTIVE);
            logger.info("Transferred {} to 'inactive'", new Object[] { inactiveFlowFile });
        } else {
            // no need to dominate CPU checking times; let other processors run for a bit.
            context.yield();
        }
    } else {
        session.transfer(flowFiles, REL_SUCCESS);
        updatedLatestSuccessTransfer = now;
        logger.info("Transferred {} FlowFiles to 'success'", new Object[] { flowFiles.size() });
        final long latestStateReportTimestamp = latestReportedNodeState.get();
        if (isClusterScope && (now - latestStateReportTimestamp) > (thresholdMillis / 3)) {
            // We don't want to hit the state manager every onTrigger(), but often enough to detect activeness.
            try {
                final StateManager stateManager = context.getStateManager();
                final StateMap state = stateManager.getState(Scope.CLUSTER);
                final Map<String, String> newValues = new HashMap<>();
                // Persist attributes so that other nodes can copy it
                if (copyAttributes) {
                    newValues.putAll(flowFiles.get(0).getAttributes());
                }
                newValues.put(STATE_KEY_LATEST_SUCCESS_TRANSFER, String.valueOf(now));
                if (state == null || state.getVersion() == -1) {
                    stateManager.setState(newValues, Scope.CLUSTER);
                } else {
                    final String existingTimestamp = state.get(STATE_KEY_LATEST_SUCCESS_TRANSFER);
                    if (StringUtils.isEmpty(existingTimestamp) || Long.parseLong(existingTimestamp) < now) {
                        // If this returns false due to race condition, it's not a problem since we just need
                        // the latest active timestamp.
                        stateManager.replace(state, newValues, Scope.CLUSTER);
                    } else {
                        logger.debug("Existing state has more recent timestamp, didn't update state.");
                    }
                }
                latestReportedNodeState.set(now);
            } catch (IOException e) {
                logger.error("Failed to access cluster state. Activity will not be monitored properly until this is addressed.", e);
            }
        }
    }
    if (!isInactive) {
        final long inactivityStartMillis = latestSuccessTransfer.get();
        if (updatedLatestSuccessTransfer > -1) {
            latestSuccessTransfer.set(updatedLatestSuccessTransfer);
        }
        if (inactive.getAndSet(false) && shouldThisNodeReport(isClusterScope, shouldReportOnlyOnPrimary)) {
            FlowFile activityRestoredFlowFile = session.create();
            if (copyAttributes) {
                final Map<String, String> attributes = new HashMap<>();
                if (flowFiles.size() > 0) {
                    // copy attributes from the first flow file in the list
                    attributes.putAll(flowFiles.get(0).getAttributes());
                } else if (clusterState != null) {
                    attributes.putAll(clusterState.toMap());
                    attributes.remove(STATE_KEY_LATEST_SUCCESS_TRANSFER);
                }
                // don't copy the UUID
                attributes.remove(CoreAttributes.UUID.key());
                activityRestoredFlowFile = session.putAllAttributes(activityRestoredFlowFile, attributes);
            }
            activityRestoredFlowFile = session.putAttribute(activityRestoredFlowFile, "inactivityStartMillis", String.valueOf(inactivityStartMillis));
            activityRestoredFlowFile = session.putAttribute(activityRestoredFlowFile, "inactivityDurationMillis", String.valueOf(now - inactivityStartMillis));
            final byte[] outBytes = context.getProperty(ACTIVITY_RESTORED_MESSAGE).evaluateAttributeExpressions(activityRestoredFlowFile).getValue().getBytes(UTF8);
            activityRestoredFlowFile = session.write(activityRestoredFlowFile, out -> out.write(outBytes));
            session.getProvenanceReporter().create(activityRestoredFlowFile);
            session.transfer(activityRestoredFlowFile, REL_ACTIVITY_RESTORED);
            logger.info("Transferred {} to 'activity.restored'", new Object[] { activityRestoredFlowFile });
        }
    }
}
Also used : OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) StandardValidators(org.apache.nifi.processor.util.StandardValidators) CapabilityDescription(org.apache.nifi.annotation.documentation.CapabilityDescription) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) HashMap(java.util.HashMap) ComponentLog(org.apache.nifi.logging.ComponentLog) SideEffectFree(org.apache.nifi.annotation.behavior.SideEffectFree) PropertyDescriptor(org.apache.nifi.components.PropertyDescriptor) ArrayList(java.util.ArrayList) HashSet(java.util.HashSet) Charset(java.nio.charset.Charset) WritesAttributes(org.apache.nifi.annotation.behavior.WritesAttributes) Scope(org.apache.nifi.components.state.Scope) Relationship(org.apache.nifi.processor.Relationship) Map(java.util.Map) Requirement(org.apache.nifi.annotation.behavior.InputRequirement.Requirement) TriggerSerially(org.apache.nifi.annotation.behavior.TriggerSerially) ValidationResult(org.apache.nifi.components.ValidationResult) OutputStream(java.io.OutputStream) TriggerWhenEmpty(org.apache.nifi.annotation.behavior.TriggerWhenEmpty) FlowFile(org.apache.nifi.flowfile.FlowFile) StateManager(org.apache.nifi.components.state.StateManager) ProcessContext(org.apache.nifi.processor.ProcessContext) Set(java.util.Set) IOException(java.io.IOException) ProcessSession(org.apache.nifi.processor.ProcessSession) WritesAttribute(org.apache.nifi.annotation.behavior.WritesAttribute) StringUtils(org.apache.nifi.util.StringUtils) AllowableValue(org.apache.nifi.components.AllowableValue) StateMap(org.apache.nifi.components.state.StateMap) TimeUnit(java.util.concurrent.TimeUnit) AtomicLong(java.util.concurrent.atomic.AtomicLong) List(java.util.List) InputRequirement(org.apache.nifi.annotation.behavior.InputRequirement) Stateful(org.apache.nifi.annotation.behavior.Stateful) OnScheduled(org.apache.nifi.annotation.lifecycle.OnScheduled) AbstractProcessor(org.apache.nifi.processor.AbstractProcessor) Tags(org.apache.nifi.annotation.documentation.Tags) CoreAttributes(org.apache.nifi.flowfile.attributes.CoreAttributes) Collections(java.util.Collections) OnStopped(org.apache.nifi.annotation.lifecycle.OnStopped) ProcessorInitializationContext(org.apache.nifi.processor.ProcessorInitializationContext) FlowFile(org.apache.nifi.flowfile.FlowFile) HashMap(java.util.HashMap) StateMap(org.apache.nifi.components.state.StateMap) OutputStream(java.io.OutputStream) IOException(java.io.IOException) ComponentLog(org.apache.nifi.logging.ComponentLog) StateManager(org.apache.nifi.components.state.StateManager) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback)

Example 54 with OutputStreamCallback

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

the class DebugFlow method onTrigger.

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();
    FlowFile ff = session.get();
    try {
        // prevents endless loops in the event of unexpected errors or future changes.)
        for (int pass = 2; pass > 0; pass--) {
            if (ff == null) {
                if (curr_noff_resp.state() == NoFlowFileResponseState.NO_FF_SKIP_RESPONSE) {
                    if (noFlowFileCurrSkip < noFlowFileMaxSkip) {
                        noFlowFileCurrSkip += 1;
                        logger.info("DebugFlow skipping with no flow file");
                        return;
                    } else {
                        noFlowFileCurrSkip = 0;
                        curr_noff_resp.getNextCycle();
                    }
                }
                if (curr_noff_resp.state() == NoFlowFileResponseState.NO_FF_EXCEPTION_RESPONSE) {
                    if (noFlowFileCurrException < noFlowFileMaxException) {
                        noFlowFileCurrException += 1;
                        logger.info("DebugFlow throwing NPE with no flow file");
                        String message = "forced by " + this.getClass().getName();
                        RuntimeException rte;
                        try {
                            rte = noFlowFileExceptionClass.getConstructor(String.class).newInstance(message);
                            throw rte;
                        } catch (InstantiationException | IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {
                            if (logger.isErrorEnabled()) {
                                logger.error("{} unexpected exception throwing DebugFlow exception: {}", new Object[] { this, e });
                            }
                        }
                    } else {
                        noFlowFileCurrException = 0;
                        curr_noff_resp.getNextCycle();
                    }
                }
                if (curr_noff_resp.state() == NoFlowFileResponseState.NO_FF_YIELD_RESPONSE) {
                    if (noFlowFileCurrYield < noFlowFileMaxYield) {
                        noFlowFileCurrYield += 1;
                        logger.info("DebugFlow yielding with no flow file");
                        context.yield();
                        break;
                    } else {
                        noFlowFileCurrYield = 0;
                        curr_noff_resp.getNextCycle();
                    }
                }
                return;
            } else {
                final int writeIterations = context.getProperty(WRITE_ITERATIONS).asInteger();
                if (writeIterations > 0 && pass == 1) {
                    final Random random = new Random();
                    for (int i = 0; i < writeIterations; i++) {
                        final byte[] data = new byte[context.getProperty(CONTENT_SIZE).asDataSize(DataUnit.B).intValue()];
                        random.nextBytes(data);
                        ff = session.write(ff, new OutputStreamCallback() {

                            @Override
                            public void process(final OutputStream out) throws IOException {
                                out.write(data);
                            }
                        });
                    }
                }
                if (curr_ff_resp.state() == FlowFileResponseState.FF_SUCCESS_RESPONSE) {
                    if (flowFileCurrSuccess < flowFileMaxSuccess) {
                        flowFileCurrSuccess += 1;
                        logger.info("DebugFlow transferring to success file={} UUID={}", new Object[] { ff.getAttribute(CoreAttributes.FILENAME.key()), ff.getAttribute(CoreAttributes.UUID.key()) });
                        session.transfer(ff, REL_SUCCESS);
                        break;
                    } else {
                        flowFileCurrSuccess = 0;
                        curr_ff_resp.getNextCycle();
                    }
                }
                if (curr_ff_resp.state() == FlowFileResponseState.FF_FAILURE_RESPONSE) {
                    if (flowFileCurrFailure < flowFileMaxFailure) {
                        flowFileCurrFailure += 1;
                        logger.info("DebugFlow transferring to failure file={} UUID={}", new Object[] { ff.getAttribute(CoreAttributes.FILENAME.key()), ff.getAttribute(CoreAttributes.UUID.key()) });
                        session.transfer(ff, REL_FAILURE);
                        break;
                    } else {
                        flowFileCurrFailure = 0;
                        curr_ff_resp.getNextCycle();
                    }
                }
                if (curr_ff_resp.state() == FlowFileResponseState.FF_ROLLBACK_RESPONSE) {
                    if (flowFileCurrRollback < flowFileMaxRollback) {
                        flowFileCurrRollback += 1;
                        logger.info("DebugFlow rolling back (no penalty) file={} UUID={}", new Object[] { ff.getAttribute(CoreAttributes.FILENAME.key()), ff.getAttribute(CoreAttributes.UUID.key()) });
                        session.rollback();
                        break;
                    } else {
                        flowFileCurrRollback = 0;
                        curr_ff_resp.getNextCycle();
                    }
                }
                if (curr_ff_resp.state() == FlowFileResponseState.FF_YIELD_RESPONSE) {
                    if (flowFileCurrYield < flowFileMaxYield) {
                        flowFileCurrYield += 1;
                        logger.info("DebugFlow yielding file={} UUID={}", new Object[] { ff.getAttribute(CoreAttributes.FILENAME.key()), ff.getAttribute(CoreAttributes.UUID.key()) });
                        session.rollback();
                        context.yield();
                        return;
                    } else {
                        flowFileCurrYield = 0;
                        curr_ff_resp.getNextCycle();
                    }
                }
                if (curr_ff_resp.state() == FlowFileResponseState.FF_PENALTY_RESPONSE) {
                    if (flowFileCurrPenalty < flowFileMaxPenalty) {
                        flowFileCurrPenalty += 1;
                        logger.info("DebugFlow rolling back (with penalty) file={} UUID={}", new Object[] { ff.getAttribute(CoreAttributes.FILENAME.key()), ff.getAttribute(CoreAttributes.UUID.key()) });
                        session.rollback(true);
                        break;
                    } else {
                        flowFileCurrPenalty = 0;
                        curr_ff_resp.getNextCycle();
                    }
                }
                if (curr_ff_resp.state() == FlowFileResponseState.FF_EXCEPTION_RESPONSE) {
                    if (flowFileCurrException < flowFileMaxException) {
                        flowFileCurrException += 1;
                        String message = "forced by " + this.getClass().getName();
                        logger.info("DebugFlow throwing NPE file={} UUID={}", new Object[] { ff.getAttribute(CoreAttributes.FILENAME.key()), ff.getAttribute(CoreAttributes.UUID.key()) });
                        RuntimeException rte;
                        try {
                            rte = flowFileExceptionClass.getConstructor(String.class).newInstance(message);
                            throw rte;
                        } catch (InstantiationException | IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {
                            if (logger.isErrorEnabled()) {
                                logger.error("{} unexpected exception throwing DebugFlow exception: {}", new Object[] { this, e });
                            }
                        }
                    } else {
                        flowFileCurrException = 0;
                        curr_ff_resp.getNextCycle();
                    }
                }
            }
        }
    } finally {
        final long sleepMillis = context.getProperty(ON_TRIGGER_SLEEP_TIME).asTimePeriod(TimeUnit.MILLISECONDS);
        try {
            if (sleepMillis > 0) {
                sleep(sleepMillis, context.getProperty(IGNORE_INTERRUPTS).asBoolean());
                getLogger().info("DebugFlow finishes sleeping at completion of its onTrigger() method");
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }
}
Also used : FlowFile(org.apache.nifi.flowfile.FlowFile) OutputStream(java.io.OutputStream) ComponentLog(org.apache.nifi.logging.ComponentLog) InvocationTargetException(java.lang.reflect.InvocationTargetException) Random(java.util.Random) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback)

Example 55 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback 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);
        }
    }
}
Also used : XPathExpression(javax.xml.xpath.XPathExpression) InputSource(org.xml.sax.InputSource) HashMap(java.util.HashMap) XPathExpressionException(javax.xml.xpath.XPathExpressionException) BufferedOutputStream(org.apache.nifi.stream.io.BufferedOutputStream) ByteArrayOutputStream(java.io.ByteArrayOutputStream) OutputStream(java.io.OutputStream) XPathEvaluator(net.sf.saxon.xpath.XPathEvaluator) Source(javax.xml.transform.Source) InputSource(org.xml.sax.InputSource) SAXSource(javax.xml.transform.sax.SAXSource) SAXException(org.xml.sax.SAXException) XPathFactory(javax.xml.xpath.XPathFactory) BufferedInputStream(org.apache.nifi.stream.io.BufferedInputStream) StringReader(java.io.StringReader) List(java.util.List) ArrayList(java.util.ArrayList) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) BufferedOutputStream(org.apache.nifi.stream.io.BufferedOutputStream) XMLReader(org.xml.sax.XMLReader) TransformerException(javax.xml.transform.TransformerException) FlowFile(org.apache.nifi.flowfile.FlowFile) PropertyDescriptor(org.apache.nifi.components.PropertyDescriptor) QName(javax.xml.namespace.QName) BufferedInputStream(org.apache.nifi.stream.io.BufferedInputStream) InputStream(java.io.InputStream) UnsupportedEncodingException(java.io.UnsupportedEncodingException) AtomicReference(java.util.concurrent.atomic.AtomicReference) EntityResolver(org.xml.sax.EntityResolver) IOException(java.io.IOException) ByteArrayOutputStream(java.io.ByteArrayOutputStream) ComponentLog(org.apache.nifi.logging.ComponentLog) XPathExpressionException(javax.xml.xpath.XPathExpressionException) XPathFactoryConfigurationException(javax.xml.xpath.XPathFactoryConfigurationException) SAXException(org.xml.sax.SAXException) UnsupportedEncodingException(java.io.UnsupportedEncodingException) TransformerException(javax.xml.transform.TransformerException) ProcessException(org.apache.nifi.processor.exception.ProcessException) IOException(java.io.IOException) ProcessException(org.apache.nifi.processor.exception.ProcessException) SAXSource(javax.xml.transform.sax.SAXSource) Relationship(org.apache.nifi.processor.Relationship) InputStreamCallback(org.apache.nifi.processor.io.InputStreamCallback) Map(java.util.Map) HashMap(java.util.HashMap)

Aggregations

FlowFile (org.apache.nifi.flowfile.FlowFile)70 OutputStreamCallback (org.apache.nifi.processor.io.OutputStreamCallback)70 OutputStream (java.io.OutputStream)69 IOException (java.io.IOException)39 ProcessException (org.apache.nifi.processor.exception.ProcessException)27 HashMap (java.util.HashMap)25 InputStream (java.io.InputStream)24 Test (org.junit.Test)24 MockFlowFile (org.apache.nifi.util.MockFlowFile)23 ByteArrayOutputStream (java.io.ByteArrayOutputStream)20 ComponentLog (org.apache.nifi.logging.ComponentLog)17 FileOutputStream (java.io.FileOutputStream)16 FilterOutputStream (java.io.FilterOutputStream)16 InputStreamCallback (org.apache.nifi.processor.io.InputStreamCallback)14 ArrayList (java.util.ArrayList)12 Map (java.util.Map)12 ProcessSession (org.apache.nifi.processor.ProcessSession)12 BufferedOutputStream (org.apache.nifi.stream.io.BufferedOutputStream)10 AtomicReference (java.util.concurrent.atomic.AtomicReference)9 StandardContentClaim (org.apache.nifi.controller.repository.claim.StandardContentClaim)9