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

Example 26 with ProcessSession

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

the class ConsumerLease method writeData.

private void writeData(final ProcessSession session, ConsumerRecord<byte[], byte[]> record, final TopicPartition topicPartition) {
    FlowFile flowFile = session.create();
    final BundleTracker tracker = new BundleTracker(record, topicPartition, keyEncoding);
    tracker.incrementRecordCount(1);
    final byte[] value = record.value();
    if (value != null) {
        flowFile = session.write(flowFile, out -> {
            out.write(value);
        });
    }
    flowFile = session.putAllAttributes(flowFile, getAttributes(record));
    tracker.updateFlowFile(flowFile);
    populateAttributes(tracker);
    session.transfer(tracker.flowFile, REL_SUCCESS);
}

Example 27 with ProcessSession

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

the class PutJMS method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();
    final List<FlowFile> flowFiles = session.get(context.getProperty(BATCH_SIZE).asInteger().intValue());
    if (flowFiles.isEmpty()) {
        return;
    }
    WrappedMessageProducer wrappedProducer = producerQueue.poll();
    if (wrappedProducer == null) {
        try {
            wrappedProducer = JmsFactory.createMessageProducer(context, true);
            logger.info("Connected to JMS server {}", new Object[] { context.getProperty(URL).getValue() });
        } catch (final JMSException e) {
            logger.error("Failed to connect to JMS Server due to {}", new Object[] { e });
            session.transfer(flowFiles, REL_FAILURE);
            context.yield();
            return;
        }
    }
    final Session jmsSession = wrappedProducer.getSession();
    final MessageProducer producer = wrappedProducer.getProducer();
    final int maxBufferSize = context.getProperty(MAX_BUFFER_SIZE).asDataSize(DataUnit.B).intValue();
    try {
        final Set<FlowFile> successfulFlowFiles = new HashSet<>();
        for (FlowFile flowFile : flowFiles) {
            if (flowFile.getSize() > maxBufferSize) {
                session.transfer(flowFile, REL_FAILURE);
                logger.warn("Routing {} to failure because its size exceeds the configured max", new Object[] { flowFile });
                continue;
            }
            // Read the contents of the FlowFile into a byte array
            final byte[] messageContent = new byte[(int) flowFile.getSize()];
            session.read(flowFile, new InputStreamCallback() {

                @Override
                public void process(final InputStream in) throws IOException {
                    StreamUtils.fillBuffer(in, messageContent, true);
                }
            });
            final Long ttl = context.getProperty(MESSAGE_TTL).asTimePeriod(TimeUnit.MILLISECONDS);
            final String replyToQueueName = context.getProperty(REPLY_TO_QUEUE).evaluateAttributeExpressions(flowFile).getValue();
            final Destination replyToQueue = replyToQueueName == null ? null : JmsFactory.createQueue(context, replyToQueueName);
            int priority = DEFAULT_MESSAGE_PRIORITY;
            try {
                final Integer priorityInt = context.getProperty(MESSAGE_PRIORITY).evaluateAttributeExpressions(flowFile).asInteger();
                priority = priorityInt == null ? priority : priorityInt;
            } catch (final NumberFormatException e) {
                logger.warn("Invalid value for JMS Message Priority: {}; defaulting to priority of {}", new Object[] { context.getProperty(MESSAGE_PRIORITY).evaluateAttributeExpressions(flowFile).getValue(), DEFAULT_MESSAGE_PRIORITY });
            }
            try {
                final Message message = createMessage(jmsSession, context, messageContent, flowFile, replyToQueue, priority);
                if (ttl == null) {
                    producer.setTimeToLive(0L);
                } else {
                    producer.setTimeToLive(ttl);
                }
                producer.send(message);
            } catch (final JMSException e) {
                logger.error("Failed to send {} to JMS Server due to {}", new Object[] { flowFile, e });
                session.transfer(flowFiles, REL_FAILURE);
                context.yield();
                try {
                    jmsSession.rollback();
                } catch (final JMSException jmse) {
                    logger.warn("Unable to roll back JMS Session due to {}", new Object[] { jmse });
                }
                wrappedProducer.close(logger);
                return;
            }
            successfulFlowFiles.add(flowFile);
            session.getProvenanceReporter().send(flowFile, context.getProperty(URL).getValue());
        }
        try {
            jmsSession.commit();
            session.transfer(successfulFlowFiles, REL_SUCCESS);
            final String flowFileDescription = successfulFlowFiles.size() > 10 ? successfulFlowFiles.size() + " FlowFiles" : successfulFlowFiles.toString();
            logger.info("Sent {} to JMS Server and transferred to 'success'", new Object[] { flowFileDescription });
        } catch (JMSException e) {
            logger.error("Failed to commit JMS Session due to {} and transferred to 'failure'", new Object[] { e });
            session.transfer(flowFiles, REL_FAILURE);
            context.yield();
            wrappedProducer.close(logger);
        }
    } finally {
        if (!wrappedProducer.isClosed()) {
            producerQueue.offer(wrappedProducer);
        }
    }
}

Example 28 with ProcessSession

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

the class PutTCP method onTrigger.

/**
 * event handler method to handle the FlowFile being forwarded to the Processor by the framework. The FlowFile contents is sent out over a TCP connection using an acquired ChannelSender object. If
 * the FlowFile contents was sent out successfully then the FlowFile is forwarded to the success relationship. If an error occurred then the FlowFile is forwarded to the failure relationship.
 *
 * @param context
 *            - the current process context.
 *
 * @param sessionFactory
 *            - a factory object to obtain a process session.
 */
@Override
public void onTrigger(final ProcessContext context, final ProcessSessionFactory sessionFactory) throws ProcessException {
    final ProcessSession session = sessionFactory.createSession();
    final FlowFile flowFile = session.get();
    if (flowFile == null) {
        final PruneResult result = pruneIdleSenders(context.getProperty(IDLE_EXPIRATION).asTimePeriod(TimeUnit.MILLISECONDS).longValue());
        // yield if we closed an idle connection, or if there were no connections in the first place
        if (result.getNumClosed() > 0 || (result.getNumClosed() == 0 && result.getNumConsidered() == 0)) {
            context.yield();
        }
        return;
    }
    ChannelSender sender = acquireSender(context, session, flowFile);
    if (sender == null) {
        return;
    }
    // can cast to a SocketChannelSender later in order to obtain the OutputStream
    if (!(sender instanceof SocketChannelSender)) {
        getLogger().error("Processor can only be used with a SocketChannelSender, but obtained: " + sender.getClass().getCanonicalName());
        context.yield();
        return;
    }
    boolean closeSender = isConnectionPerFlowFile(context);
    try {
        // We might keep the connection open across invocations of the processor so don't auto-close this
        final OutputStream out = ((SocketChannelSender) sender).getOutputStream();
        final String delimiter = getOutgoingMessageDelimiter(context, flowFile);
        final StopWatch stopWatch = new StopWatch(true);
        try (final InputStream rawIn = session.read(flowFile);
            final BufferedInputStream in = new BufferedInputStream(rawIn)) {
            IOUtils.copy(in, out);
            if (delimiter != null) {
                final Charset charSet = Charset.forName(context.getProperty(CHARSET).getValue());
                out.write(delimiter.getBytes(charSet), 0, delimiter.length());
            }
            out.flush();
        } catch (final Exception e) {
            closeSender = true;
            throw e;
        }
        session.getProvenanceReporter().send(flowFile, transitUri, stopWatch.getElapsed(TimeUnit.MILLISECONDS));
        session.transfer(flowFile, REL_SUCCESS);
        session.commit();
    } catch (Exception e) {
        onFailure(context, session, flowFile);
        getLogger().error("Exception while handling a process session, transferring {} to failure.", new Object[] { flowFile }, e);
    } finally {
        if (closeSender) {
            getLogger().debug("Closing sender");
            sender.close();
        } else {
            getLogger().debug("Relinquishing sender");
            relinquishSender(sender);
        }
    }
}

Example 29 with ProcessSession

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

the class QueryDatabaseTable method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSessionFactory sessionFactory) throws ProcessException {
    // Fetch the column/table info once
    if (!setupComplete.get()) {
        super.setup(context);
    }
    ProcessSession session = sessionFactory.createSession();
    final List<FlowFile> resultSetFlowFiles = new ArrayList<>();
    final ComponentLog logger = getLogger();
    final DBCPService dbcpService = context.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class);
    final DatabaseAdapter dbAdapter = dbAdapters.get(context.getProperty(DB_TYPE).getValue());
    final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions().getValue();
    final String columnNames = context.getProperty(COLUMN_NAMES).evaluateAttributeExpressions().getValue();
    final String maxValueColumnNames = context.getProperty(MAX_VALUE_COLUMN_NAMES).evaluateAttributeExpressions().getValue();
    final String customWhereClause = context.getProperty(WHERE_CLAUSE).evaluateAttributeExpressions().getValue();
    final Integer fetchSize = context.getProperty(FETCH_SIZE).evaluateAttributeExpressions().asInteger();
    final Integer maxRowsPerFlowFile = context.getProperty(MAX_ROWS_PER_FLOW_FILE).evaluateAttributeExpressions().asInteger();
    final Integer outputBatchSizeField = context.getProperty(OUTPUT_BATCH_SIZE).evaluateAttributeExpressions().asInteger();
    final int outputBatchSize = outputBatchSizeField == null ? 0 : outputBatchSizeField;
    final Integer maxFragments = context.getProperty(MAX_FRAGMENTS).isSet() ? context.getProperty(MAX_FRAGMENTS).evaluateAttributeExpressions().asInteger() : 0;
    final JdbcCommon.AvroConversionOptions options = JdbcCommon.AvroConversionOptions.builder().recordName(tableName).maxRows(maxRowsPerFlowFile).convertNames(context.getProperty(NORMALIZE_NAMES_FOR_AVRO).asBoolean()).useLogicalTypes(context.getProperty(USE_AVRO_LOGICAL_TYPES).asBoolean()).defaultPrecision(context.getProperty(DEFAULT_PRECISION).evaluateAttributeExpressions().asInteger()).defaultScale(context.getProperty(DEFAULT_SCALE).evaluateAttributeExpressions().asInteger()).build();
    final StateManager stateManager = context.getStateManager();
    final StateMap stateMap;
    try {
        stateMap = stateManager.getState(Scope.CLUSTER);
    } catch (final IOException ioe) {
        getLogger().error("Failed to retrieve observed maximum values from the State Manager. Will not perform " + "query until this is accomplished.", ioe);
        context.yield();
        return;
    }
    // Make a mutable copy of the current state property map. This will be updated by the result row callback, and eventually
    // set as the current state map (after the session has been committed)
    final Map<String, String> statePropertyMap = new HashMap<>(stateMap.toMap());
    // If an initial max value for column(s) has been specified using properties, and this column is not in the state manager, sync them to the state property map
    for (final Map.Entry<String, String> maxProp : maxValueProperties.entrySet()) {
        String maxPropKey = maxProp.getKey().toLowerCase();
        String fullyQualifiedMaxPropKey = getStateKey(tableName, maxPropKey);
        if (!statePropertyMap.containsKey(fullyQualifiedMaxPropKey)) {
            String newMaxPropValue;
            // but store the new initial max value under the fully-qualified key.
            if (statePropertyMap.containsKey(maxPropKey)) {
                newMaxPropValue = statePropertyMap.get(maxPropKey);
            } else {
                newMaxPropValue = maxProp.getValue();
            }
            statePropertyMap.put(fullyQualifiedMaxPropKey, newMaxPropValue);
        }
    }
    List<String> maxValueColumnNameList = StringUtils.isEmpty(maxValueColumnNames) ? null : Arrays.asList(maxValueColumnNames.split("\\s*,\\s*"));
    final String selectQuery = getQuery(dbAdapter, tableName, columnNames, maxValueColumnNameList, customWhereClause, statePropertyMap);
    final StopWatch stopWatch = new StopWatch(true);
    final String fragmentIdentifier = UUID.randomUUID().toString();
    try (final Connection con = dbcpService.getConnection();
        final Statement st = con.createStatement()) {
        if (fetchSize != null && fetchSize > 0) {
            try {
                st.setFetchSize(fetchSize);
            } catch (SQLException se) {
                // Not all drivers support this, just log the error (at debug level) and move on
                logger.debug("Cannot set fetch size to {} due to {}", new Object[] { fetchSize, se.getLocalizedMessage() }, se);
            }
        }
        String jdbcURL = "DBCPService";
        try {
            DatabaseMetaData databaseMetaData = con.getMetaData();
            if (databaseMetaData != null) {
                jdbcURL = databaseMetaData.getURL();
            }
        } catch (SQLException se) {
        // Ignore and use default JDBC URL. This shouldn't happen unless the driver doesn't implement getMetaData() properly
        }
        final Integer queryTimeout = context.getProperty(QUERY_TIMEOUT).evaluateAttributeExpressions().asTimePeriod(TimeUnit.SECONDS).intValue();
        // timeout in seconds
        st.setQueryTimeout(queryTimeout);
        try {
            logger.debug("Executing query {}", new Object[] { selectQuery });
            final ResultSet resultSet = st.executeQuery(selectQuery);
            int fragmentIndex = 0;
            while (true) {
                final AtomicLong nrOfRows = new AtomicLong(0L);
                FlowFile fileToProcess = session.create();
                try {
                    fileToProcess = session.write(fileToProcess, out -> {
                        // Max values will be updated in the state property map by the callback
                        final MaxValueResultSetRowCollector maxValCollector = new MaxValueResultSetRowCollector(tableName, statePropertyMap, dbAdapter);
                        try {
                            nrOfRows.set(JdbcCommon.convertToAvroStream(resultSet, out, options, maxValCollector));
                        } catch (SQLException | RuntimeException e) {
                            throw new ProcessException("Error during database query or conversion of records to Avro.", e);
                        }
                    });
                } catch (ProcessException e) {
                    // Add flowfile to results before rethrowing so it will be removed from session in outer catch
                    resultSetFlowFiles.add(fileToProcess);
                    throw e;
                }
                if (nrOfRows.get() > 0) {
                    // set attribute how many rows were selected
                    fileToProcess = session.putAttribute(fileToProcess, RESULT_ROW_COUNT, String.valueOf(nrOfRows.get()));
                    fileToProcess = session.putAttribute(fileToProcess, RESULT_TABLENAME, tableName);
                    fileToProcess = session.putAttribute(fileToProcess, CoreAttributes.MIME_TYPE.key(), JdbcCommon.MIME_TYPE_AVRO_BINARY);
                    if (maxRowsPerFlowFile > 0) {
                        fileToProcess = session.putAttribute(fileToProcess, "fragment.identifier", fragmentIdentifier);
                        fileToProcess = session.putAttribute(fileToProcess, "fragment.index", String.valueOf(fragmentIndex));
                    }
                    logger.info("{} contains {} Avro records; transferring to 'success'", new Object[] { fileToProcess, nrOfRows.get() });
                    session.getProvenanceReporter().receive(fileToProcess, jdbcURL, stopWatch.getElapsed(TimeUnit.MILLISECONDS));
                    resultSetFlowFiles.add(fileToProcess);
                    // If we've reached the batch size, send out the flow files
                    if (outputBatchSize > 0 && resultSetFlowFiles.size() >= outputBatchSize) {
                        session.transfer(resultSetFlowFiles, REL_SUCCESS);
                        session.commit();
                        resultSetFlowFiles.clear();
                    }
                } else {
                    // If there were no rows returned, don't send the flowfile
                    session.remove(fileToProcess);
                    context.yield();
                    break;
                }
                fragmentIndex++;
                if (maxFragments > 0 && fragmentIndex >= maxFragments) {
                    break;
                }
            }
            // Even though the maximum value and total count are known at this point, to maintain consistent behavior if Output Batch Size is set, do not store the attributes
            if (outputBatchSize == 0) {
                for (int i = 0; i < resultSetFlowFiles.size(); i++) {
                    // Add maximum values as attributes
                    for (Map.Entry<String, String> entry : statePropertyMap.entrySet()) {
                        // Get just the column name from the key
                        String key = entry.getKey();
                        String colName = key.substring(key.lastIndexOf(NAMESPACE_DELIMITER) + NAMESPACE_DELIMITER.length());
                        resultSetFlowFiles.set(i, session.putAttribute(resultSetFlowFiles.get(i), "maxvalue." + colName, entry.getValue()));
                    }
                    // set count on all FlowFiles
                    if (maxRowsPerFlowFile > 0) {
                        resultSetFlowFiles.set(i, session.putAttribute(resultSetFlowFiles.get(i), "fragment.count", Integer.toString(fragmentIndex)));
                    }
                }
            }
        } catch (final SQLException e) {
            throw e;
        }
        session.transfer(resultSetFlowFiles, REL_SUCCESS);
    } catch (final ProcessException | SQLException e) {
        logger.error("Unable to execute SQL select query {} due to {}", new Object[] { selectQuery, e });
        if (!resultSetFlowFiles.isEmpty()) {
            session.remove(resultSetFlowFiles);
        }
        context.yield();
    } finally {
        session.commit();
        try {
            // Update the state
            stateManager.setState(statePropertyMap, Scope.CLUSTER);
        } catch (IOException ioe) {
            getLogger().error("{} failed to update State Manager, maximum observed values will not be recorded", new Object[] { this, ioe });
        }
    }
}

Example 30 with ProcessSession

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

the class SplitJson method onTrigger.

@Override
public void onTrigger(final ProcessContext processContext, final ProcessSession processSession) {
    FlowFile original = processSession.get();
    if (original == null) {
        return;
    }
    final ComponentLog logger = getLogger();
    DocumentContext documentContext;
    try {
        documentContext = validateAndEstablishJsonContext(processSession, original);
    } catch (InvalidJsonException e) {
        logger.error("FlowFile {} did not have valid JSON content.", new Object[] { original });
        processSession.transfer(original, REL_FAILURE);
        return;
    }
    final JsonPath jsonPath = JSON_PATH_REF.get();
    Object jsonPathResult;
    try {
        jsonPathResult = documentContext.read(jsonPath);
    } catch (PathNotFoundException e) {
        logger.warn("JsonPath {} could not be found for FlowFile {}", new Object[] { jsonPath.getPath(), original });
        processSession.transfer(original, REL_FAILURE);
        return;
    }
    if (!(jsonPathResult instanceof List)) {
        logger.error("The evaluated value {} of {} was not a JSON Array compatible type and cannot be split.", new Object[] { jsonPathResult, jsonPath.getPath() });
        processSession.transfer(original, REL_FAILURE);
        return;
    }
    List resultList = (List) jsonPathResult;
    Map<String, String> attributes = new HashMap<>();
    final String fragmentId = UUID.randomUUID().toString();
    attributes.put(FRAGMENT_ID.key(), fragmentId);
    attributes.put(FRAGMENT_COUNT.key(), Integer.toString(resultList.size()));
    for (int i = 0; i < resultList.size(); i++) {
        Object resultSegment = resultList.get(i);
        FlowFile split = processSession.create(original);
        split = processSession.write(split, (out) -> {
            String resultSegmentContent = getResultRepresentation(resultSegment, nullDefaultValue);
            out.write(resultSegmentContent.getBytes(StandardCharsets.UTF_8));
        });
        attributes.put(SEGMENT_ORIGINAL_FILENAME.key(), split.getAttribute(CoreAttributes.FILENAME.key()));
        attributes.put(FRAGMENT_INDEX.key(), Integer.toString(i));
        processSession.transfer(processSession.putAllAttributes(split, attributes), REL_SPLIT);
    }
    original = copyAttributesToOriginal(processSession, original, fragmentId, resultList.size());
    processSession.transfer(original, REL_ORIGINAL);
    logger.info("Split {} into {} FlowFiles", new Object[] { original, resultList.size() });
}