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

Example 21 with ProcessSession

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

the class ResultProcessorTest method testProcessResultFileSuccess.

@Test
public void testProcessResultFileSuccess() {
    ProcessSession processSession = mock(ProcessSession.class);
    ComponentLog componentLog = mock(ComponentLog.class);
    FlowFile flowFile = mock(FlowFile.class);
    Exception exception = null;
    String name = "basename";
    when(processSession.putAttribute(eq(flowFile), anyString(), anyString())).thenReturn(flowFile);
    resultProcessor.process(processSession, componentLog, flowFile, exception, name);
    verify(processSession).putAttribute(flowFile, CoreAttributes.FILENAME.key(), name);
    verify(processSession).putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), MediaType.APPLICATION_XML_UTF_8.toString());
    verify(processSession).transfer(flowFile, successRelationship);
    verifyNoMoreInteractions(componentLog);
}

Example 22 with ProcessSession

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

the class BinFiles method processBins.

private int processBins(final ProcessContext context) {
    final ComponentLog logger = getLogger();
    int processedBins = 0;
    Bin bin;
    while ((bin = readyBins.poll()) != null) {
        boolean binAlreadyCommitted;
        try {
            binAlreadyCommitted = this.processBin(bin, context);
        } catch (final ProcessException e) {
            logger.error("Failed to process bundle of {} files due to {}", new Object[] { bin.getContents().size(), e });
            final ProcessSession binSession = bin.getSession();
            for (final FlowFile flowFile : bin.getContents()) {
                binSession.transfer(flowFile, REL_FAILURE);
            }
            binSession.commit();
            continue;
        } catch (final Exception e) {
            logger.error("Failed to process bundle of {} files due to {}; rolling back sessions", new Object[] { bin.getContents().size(), e });
            bin.getSession().rollback();
            continue;
        }
        // If this bin's session has been committed, move on.
        if (!binAlreadyCommitted) {
            final ProcessSession binSession = bin.getSession();
            binSession.transfer(bin.getContents(), REL_ORIGINAL);
            binSession.commit();
        }
        processedBins++;
    }
    return processedBins;
}

Example 23 with ProcessSession

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

the class BinFiles method binFlowFiles.

private int binFlowFiles(final ProcessContext context, final ProcessSessionFactory sessionFactory) {
    int flowFilesBinned = 0;
    while (binManager.getBinCount() <= context.getProperty(MAX_BIN_COUNT).asInteger()) {
        if (!isScheduled()) {
            break;
        }
        final ProcessSession session = sessionFactory.createSession();
        final List<FlowFile> flowFiles = session.get(1000);
        if (flowFiles.isEmpty()) {
            break;
        }
        final Map<String, List<FlowFile>> flowFileGroups = new HashMap<>();
        for (FlowFile flowFile : flowFiles) {
            flowFile = this.preprocessFlowFile(context, session, flowFile);
            try {
                final String groupingIdentifier = getGroupId(context, flowFile, session);
                flowFileGroups.computeIfAbsent(groupingIdentifier, id -> new ArrayList<>()).add(flowFile);
            } catch (final Exception e) {
                getLogger().error("Could not determine which Bin to add {} to; will route to failure", new Object[] { flowFile }, e);
                session.transfer(flowFile, REL_FAILURE);
                continue;
            }
        }
        for (final Map.Entry<String, List<FlowFile>> entry : flowFileGroups.entrySet()) {
            final Set<FlowFile> unbinned = binManager.offer(entry.getKey(), entry.getValue(), session, sessionFactory);
            for (final FlowFile flowFile : unbinned) {
                Bin bin = new Bin(sessionFactory.createSession(), 0, Long.MAX_VALUE, 0, Integer.MAX_VALUE, null);
                bin.offer(flowFile, session);
                this.readyBins.add(bin);
            }
            flowFilesBinned += entry.getValue().size();
        }
    }
    return flowFilesBinned;
}

Example 24 with ProcessSession

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

the class AbstractMQTTProcessor method onTrigger.

@Override
public final void onTrigger(final ProcessContext context, final ProcessSessionFactory sessionFactory) throws ProcessException {
    if (processSessionFactory == null) {
        processSessionFactory = sessionFactory;
    }
    ProcessSession session = sessionFactory.createSession();
    try {
        onTrigger(context, session);
        session.commit();
    } catch (final Throwable t) {
        getLogger().error("{} failed to process due to {}; rolling back session", new Object[] { this, t });
        session.rollback(true);
        throw t;
    }
}

Example 25 with ProcessSession

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

the class TestConsumeMQTT method testMessageNotConsumedOnCommitFail.

/**
 * If the session.commit() fails, we should not remove the unprocessed message
 */
@Test
public void testMessageNotConsumedOnCommitFail() throws NoSuchFieldException, IllegalAccessException, NoSuchMethodException, InvocationTargetException {
    testRunner.run(1, false);
    ConsumeMQTT processor = (ConsumeMQTT) testRunner.getProcessor();
    MQTTQueueMessage mock = mock(MQTTQueueMessage.class);
    when(mock.getPayload()).thenReturn(new byte[0]);
    when(mock.getTopic()).thenReturn("testTopic");
    BlockingQueue<MQTTQueueMessage> mqttQueue = getMqttQueue(processor);
    mqttQueue.add(mock);
    try {
        ProcessSession session = testRunner.getProcessSessionFactory().createSession();
        transferQueue(processor, (ProcessSession) Proxy.newProxyInstance(getClass().getClassLoader(), new Class[] { ProcessSession.class }, (proxy, method, args) -> {
            if (method.getName().equals("commit")) {
                throw new RuntimeException();
            } else {
                return method.invoke(session, args);
            }
        }));
        fail("Expected runtime exception");
    } catch (InvocationTargetException e) {
        assertTrue("Expected generic runtime exception, not " + e, e.getCause() instanceof RuntimeException);
    }
    assertTrue("Expected mqttQueue to contain uncommitted message.", mqttQueue.contains(mock));
}