Examples with ProcessSession org.apache.nifi.processor.ProcessSession used on opensource projects

Example 16 with ProcessSession

use of org.apache.nifi.processor.ProcessSession in project nifi by apache.

the class ParseEvtx method processRecordGranularity.

protected void processRecordGranularity(ProcessSession session, ComponentLog componentLog, FlowFile flowFile, String basename, InputStream in) throws IOException {
    FileHeader fileHeader = fileHeaderFactory.create(in, componentLog);
    while (fileHeader.hasNext()) {
        try {
            ChunkHeader chunkHeader = fileHeader.next();
            while (chunkHeader.hasNext()) {
                FlowFile updated = session.create(flowFile);
                AtomicReference<Exception> exceptionReference = new AtomicReference<>(null);
                try {
                    Record record = chunkHeader.next();
                    updated = session.write(updated, out -> {
                        try (RootNodeHandler rootNodeHandler = rootNodeHandlerFactory.create(out)) {
                            try {
                                rootNodeHandler.handle(record.getRootNode());
                            } catch (IOException e) {
                                exceptionReference.set(e);
                            }
                        } catch (IOException e) {
                            exceptionReference.set(e);
                        }
                    });
                    resultProcessor.process(session, componentLog, updated, exceptionReference.get(), getName(basename, chunkHeader.getChunkNumber(), record.getRecordNum(), XML_EXTENSION));
                } catch (Exception e) {
                    exceptionReference.set(e);
                    session.remove(updated);
                }
                if (exceptionReference.get() != null) {
                    malformedChunkHandler.handle(flowFile, session, getName(basename, chunkHeader.getChunkNumber(), null, EVTX_EXTENSION), chunkHeader.getBinaryReader().getBytes());
                }
            }
        } catch (MalformedChunkException e) {
            malformedChunkHandler.handle(flowFile, session, getName(basename, e.getChunkNum(), null, EVTX_EXTENSION), e.getBadChunk());
        }
    }
}

Example 17 with ProcessSession

use of org.apache.nifi.processor.ProcessSession in project nifi by apache.

the class ConsumeAMQP method processResource.

/**
 * Will construct a {@link FlowFile} containing the body of the consumed AMQP message (if {@link GetResponse} returned by {@link AMQPConsumer} is
 * not null) and AMQP properties that came with message which are added to a {@link FlowFile} as attributes, transferring {@link FlowFile} to
 * 'success' {@link Relationship}.
 */
@Override
protected void processResource(final Connection connection, final AMQPConsumer consumer, final ProcessContext context, final ProcessSession session) {
    final GetResponse response = consumer.consume();
    if (response == null) {
        context.yield();
        return;
    }
    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, out -> out.write(response.getBody()));
    final BasicProperties amqpProperties = response.getProps();
    final Map<String, String> attributes = buildAttributes(amqpProperties);
    flowFile = session.putAllAttributes(flowFile, attributes);
    session.getProvenanceReporter().receive(flowFile, connection.toString() + "/" + context.getProperty(QUEUE).getValue());
    session.transfer(flowFile, REL_SUCCESS);
}

Example 18 with ProcessSession

use of org.apache.nifi.processor.ProcessSession in project nifi by apache.

the class MalformedChunkHandlerTest method testHandle.

@Test
public void testHandle() {
    String name = "name";
    byte[] badChunk = { 8 };
    FlowFile original = mock(FlowFile.class);
    FlowFile updated1 = mock(FlowFile.class);
    FlowFile updated2 = mock(FlowFile.class);
    FlowFile updated3 = mock(FlowFile.class);
    FlowFile updated4 = mock(FlowFile.class);
    ProcessSession session = mock(ProcessSession.class);
    when(session.create(original)).thenReturn(updated1);
    when(session.putAttribute(updated1, CoreAttributes.FILENAME.key(), name)).thenReturn(updated2);
    when(session.putAttribute(updated2, CoreAttributes.MIME_TYPE.key(), MediaType.APPLICATION_BINARY.toString())).thenReturn(updated3);
    ByteArrayOutputStream out = new ByteArrayOutputStream();
    when(session.write(eq(updated3), any(OutputStreamCallback.class))).thenAnswer(invocation -> {
        ((OutputStreamCallback) invocation.getArguments()[1]).process(out);
        return updated4;
    });
    malformedChunkHandler.handle(original, session, name, badChunk);
    verify(session).transfer(updated4, badChunkRelationship);
    assertArrayEquals(badChunk, out.toByteArray());
}

Example 19 with ProcessSession

use of org.apache.nifi.processor.ProcessSession in project nifi by apache.

the class ParseEvtxTest method testProcessChunkGranularity.

@Test
public void testProcessChunkGranularity() 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);
    OutputStream out2 = mock(OutputStream.class);
    when(rootNodeHandlerFactory.create(out)).thenReturn(rootNodeHandler1);
    when(rootNodeHandlerFactory.create(out2)).thenReturn(rootNodeHandler2);
    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);
    MalformedChunkException malformedChunkException = new MalformedChunkException("Test", null, offset, chunkNum, badChunk);
    when(session.create(flowFile)).thenReturn(created1).thenReturn(created2).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(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.processChunkGranularity(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).handle(rootNode3);
    verify(rootNodeHandler2).close();
}

Example 20 with ProcessSession

use of org.apache.nifi.processor.ProcessSession in project nifi by apache.

the class ParseEvtxTest method testProcessFileGranularity.

@Test
public void testProcessFileGranularity() throws IOException, MalformedChunkException, XMLStreamException {
    String basename = "basename";
    int chunkNum = 5;
    int offset = 10001;
    byte[] badChunk = { 8 };
    RootNodeHandler rootNodeHandler = mock(RootNodeHandler.class);
    when(rootNodeHandlerFactory.create(out)).thenReturn(rootNodeHandler);
    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);
    AtomicReference<Exception> reference = new AtomicReference<>();
    MalformedChunkException malformedChunkException = new MalformedChunkException("Test", null, offset, chunkNum, badChunk);
    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.processFileGranularity(session, componentLog, flowFile, basename, reference, in, out);
    verify(malformedChunkHandler).handle(flowFile, session, parseEvtx.getName(basename, chunkNum, null, ParseEvtx.EVTX_EXTENSION), badChunk);
    verify(rootNodeHandler).handle(rootNode1);
    verify(rootNodeHandler).handle(rootNode2);
    verify(rootNodeHandler).handle(rootNode3);
    verify(rootNodeHandler).close();
}