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

Example 31 with ProcessSession

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

the class ListDatabaseTables method onTrigger.

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();
    final DBCPService dbcpService = context.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class);
    final String catalog = context.getProperty(CATALOG).getValue();
    final String schemaPattern = context.getProperty(SCHEMA_PATTERN).getValue();
    final String tableNamePattern = context.getProperty(TABLE_NAME_PATTERN).getValue();
    final String[] tableTypes = context.getProperty(TABLE_TYPES).isSet() ? context.getProperty(TABLE_TYPES).getValue().split("\\s*,\\s*") : null;
    final boolean includeCount = context.getProperty(INCLUDE_COUNT).asBoolean();
    final long refreshInterval = context.getProperty(REFRESH_INTERVAL).asTimePeriod(TimeUnit.MILLISECONDS);
    final StateManager stateManager = context.getStateManager();
    final StateMap stateMap;
    final Map<String, String> stateMapProperties;
    try {
        stateMap = stateManager.getState(Scope.CLUSTER);
        stateMapProperties = new HashMap<>(stateMap.toMap());
    } catch (IOException ioe) {
        throw new ProcessException(ioe);
    }
    try (final Connection con = dbcpService.getConnection()) {
        DatabaseMetaData dbMetaData = con.getMetaData();
        ResultSet rs = dbMetaData.getTables(catalog, schemaPattern, tableNamePattern, tableTypes);
        while (rs.next()) {
            final String tableCatalog = rs.getString(1);
            final String tableSchema = rs.getString(2);
            final String tableName = rs.getString(3);
            final String tableType = rs.getString(4);
            final String tableRemarks = rs.getString(5);
            // Build fully-qualified name
            String fqn = Stream.of(tableCatalog, tableSchema, tableName).filter(segment -> !StringUtils.isEmpty(segment)).collect(Collectors.joining("."));
            String lastTimestampForTable = stateMapProperties.get(fqn);
            boolean refreshTable = true;
            try {
                // Refresh state if the interval has elapsed
                long lastRefreshed = -1;
                final long currentTime = System.currentTimeMillis();
                if (!StringUtils.isEmpty(lastTimestampForTable)) {
                    lastRefreshed = Long.parseLong(lastTimestampForTable);
                }
                if (lastRefreshed == -1 || (refreshInterval > 0 && currentTime >= (lastRefreshed + refreshInterval))) {
                    stateMapProperties.remove(lastTimestampForTable);
                } else {
                    refreshTable = false;
                }
            } catch (final NumberFormatException nfe) {
                getLogger().error("Failed to retrieve observed last table fetches from the State Manager. Will not perform " + "query until this is accomplished.", nfe);
                context.yield();
                return;
            }
            if (refreshTable) {
                FlowFile flowFile = session.create();
                logger.info("Found {}: {}", new Object[] { tableType, fqn });
                if (includeCount) {
                    try (Statement st = con.createStatement()) {
                        final String countQuery = "SELECT COUNT(1) FROM " + fqn;
                        logger.debug("Executing query: {}", new Object[] { countQuery });
                        ResultSet countResult = st.executeQuery(countQuery);
                        if (countResult.next()) {
                            flowFile = session.putAttribute(flowFile, DB_TABLE_COUNT, Long.toString(countResult.getLong(1)));
                        }
                    } catch (SQLException se) {
                        logger.error("Couldn't get row count for {}", new Object[] { fqn });
                        session.remove(flowFile);
                        continue;
                    }
                }
                if (tableCatalog != null) {
                    flowFile = session.putAttribute(flowFile, DB_TABLE_CATALOG, tableCatalog);
                }
                if (tableSchema != null) {
                    flowFile = session.putAttribute(flowFile, DB_TABLE_SCHEMA, tableSchema);
                }
                flowFile = session.putAttribute(flowFile, DB_TABLE_NAME, tableName);
                flowFile = session.putAttribute(flowFile, DB_TABLE_FULLNAME, fqn);
                flowFile = session.putAttribute(flowFile, DB_TABLE_TYPE, tableType);
                if (tableRemarks != null) {
                    flowFile = session.putAttribute(flowFile, DB_TABLE_REMARKS, tableRemarks);
                }
                String transitUri;
                try {
                    transitUri = dbMetaData.getURL();
                } catch (SQLException sqle) {
                    transitUri = "<unknown>";
                }
                session.getProvenanceReporter().receive(flowFile, transitUri);
                session.transfer(flowFile, REL_SUCCESS);
                stateMapProperties.put(fqn, Long.toString(System.currentTimeMillis()));
            }
        }
        // Update the timestamps for listed tables
        if (stateMap.getVersion() == -1) {
            stateManager.setState(stateMapProperties, Scope.CLUSTER);
        } else {
            stateManager.replace(stateMap, stateMapProperties, Scope.CLUSTER);
        }
    } catch (final SQLException | IOException e) {
        throw new ProcessException(e);
    }
}

Example 32 with ProcessSession

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

the class ListenUDPRecord method handleParseFailure.

private void handleParseFailure(final StandardEvent event, final ProcessSession session, final Exception cause, final String message) {
    // If we are unable to parse the data, we need to transfer it to 'parse failure' relationship
    final Map<String, String> attributes = getAttributes(event.getSender());
    FlowFile failureFlowFile = session.create();
    failureFlowFile = session.write(failureFlowFile, out -> out.write(event.getData()));
    failureFlowFile = session.putAllAttributes(failureFlowFile, attributes);
    final String transitUri = getTransitUri(event.getSender());
    session.getProvenanceReporter().receive(failureFlowFile, transitUri);
    session.transfer(failureFlowFile, REL_PARSE_FAILURE);
    if (cause == null) {
        getLogger().error(message);
    } else {
        getLogger().error(message, cause);
    }
    session.adjustCounter("Parse Failures", 1, false);
}

Example 33 with ProcessSession

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

the class MergeContent method processBin.

@Override
protected boolean processBin(final Bin bin, final ProcessContext context) throws ProcessException {
    final String mergeFormat = context.getProperty(MERGE_FORMAT).getValue();
    MergeBin merger;
    switch(mergeFormat) {
        case MERGE_FORMAT_TAR_VALUE:
            merger = new TarMerge();
            break;
        case MERGE_FORMAT_ZIP_VALUE:
            merger = new ZipMerge(context.getProperty(COMPRESSION_LEVEL).asInteger());
            break;
        case MERGE_FORMAT_FLOWFILE_STREAM_V3_VALUE:
            merger = new FlowFileStreamMerger(new FlowFilePackagerV3(), "application/flowfile-v3");
            break;
        case MERGE_FORMAT_FLOWFILE_STREAM_V2_VALUE:
            merger = new FlowFileStreamMerger(new FlowFilePackagerV2(), "application/flowfile-v2");
            break;
        case MERGE_FORMAT_FLOWFILE_TAR_V1_VALUE:
            merger = new FlowFileStreamMerger(new FlowFilePackagerV1(), "application/flowfile-v1");
            break;
        case MERGE_FORMAT_CONCAT_VALUE:
            merger = new BinaryConcatenationMerge();
            break;
        case MERGE_FORMAT_AVRO_VALUE:
            merger = new AvroMerge();
            break;
        default:
            throw new AssertionError();
    }
    final AttributeStrategy attributeStrategy = AttributeStrategyUtil.strategyFor(context);
    final List<FlowFile> contents = bin.getContents();
    final ProcessSession binSession = bin.getSession();
    if (MERGE_STRATEGY_DEFRAGMENT.equals(context.getProperty(MERGE_STRATEGY).getValue())) {
        final String error = getDefragmentValidationError(bin.getContents());
        // Fail the flow files and commit them
        if (error != null) {
            final String binDescription = contents.size() <= 10 ? contents.toString() : contents.size() + " FlowFiles";
            getLogger().error(error + "; routing {} to failure", new Object[] { binDescription });
            binSession.transfer(contents, REL_FAILURE);
            binSession.commit();
            return true;
        }
        Collections.sort(contents, new FragmentComparator());
    }
    FlowFile bundle = merger.merge(bin, context);
    // keep the filename, as it is added to the bundle.
    final String filename = bundle.getAttribute(CoreAttributes.FILENAME.key());
    // merge all of the attributes
    final Map<String, String> bundleAttributes = attributeStrategy.getMergedAttributes(contents);
    bundleAttributes.put(CoreAttributes.MIME_TYPE.key(), merger.getMergedContentType());
    // restore the filename of the bundle
    bundleAttributes.put(CoreAttributes.FILENAME.key(), filename);
    bundleAttributes.put(MERGE_COUNT_ATTRIBUTE, Integer.toString(contents.size()));
    bundleAttributes.put(MERGE_BIN_AGE_ATTRIBUTE, Long.toString(bin.getBinAge()));
    bundle = binSession.putAllAttributes(bundle, bundleAttributes);
    final String inputDescription = contents.size() < 10 ? contents.toString() : contents.size() + " FlowFiles";
    getLogger().info("Merged {} into {}", new Object[] { inputDescription, bundle });
    binSession.transfer(bundle, REL_MERGED);
    for (final FlowFile unmerged : merger.getUnmergedFlowFiles()) {
        final FlowFile unmergedCopy = binSession.clone(unmerged);
        binSession.transfer(unmergedCopy, REL_FAILURE);
    }
    // We haven't committed anything, parent will take care of it
    return false;
}

Example 34 with ProcessSession

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

the class MergeRecord method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSessionFactory sessionFactory) throws ProcessException {
    RecordBinManager manager = binManager.get();
    while (manager == null) {
        manager = new RecordBinManager(context, sessionFactory, getLogger());
        manager.setMaxBinAge(context.getProperty(MAX_BIN_AGE).asTimePeriod(TimeUnit.NANOSECONDS), TimeUnit.NANOSECONDS);
        final boolean updated = binManager.compareAndSet(null, manager);
        if (!updated) {
            manager = binManager.get();
        }
    }
    final ProcessSession session = sessionFactory.createSession();
    final List<FlowFile> flowFiles = session.get(FlowFileFilters.newSizeBasedFilter(250, DataUnit.KB, 250));
    if (getLogger().isDebugEnabled()) {
        final List<String> ids = flowFiles.stream().map(ff -> "id=" + ff.getId()).collect(Collectors.toList());
        getLogger().debug("Pulled {} FlowFiles from queue: {}", new Object[] { ids.size(), ids });
    }
    final String mergeStrategy = context.getProperty(MERGE_STRATEGY).getValue();
    final boolean block;
    if (MERGE_STRATEGY_DEFRAGMENT.equals(mergeStrategy)) {
        block = true;
    } else if (context.getProperty(CORRELATION_ATTRIBUTE_NAME).isSet()) {
        block = true;
    } else {
        block = false;
    }
    try {
        for (final FlowFile flowFile : flowFiles) {
            try {
                binFlowFile(context, flowFile, session, manager, block);
            } catch (final Exception e) {
                getLogger().error("Failed to bin {} due to {}", new Object[] { flowFile, e });
                session.transfer(flowFile, REL_FAILURE);
            }
        }
    } finally {
        session.commit();
    }
    try {
        manager.completeExpiredBins();
    } catch (final Exception e) {
        getLogger().error("Failed to merge FlowFiles to create new bin due to " + e, e);
    }
    if (flowFiles.isEmpty()) {
        getLogger().debug("No FlowFiles to bin; will yield");
        context.yield();
    }
}

Example 35 with ProcessSession

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

the class Notify method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();
    final PropertyValue signalIdProperty = context.getProperty(RELEASE_SIGNAL_IDENTIFIER);
    final PropertyValue counterNameProperty = context.getProperty(SIGNAL_COUNTER_NAME);
    final PropertyValue deltaProperty = context.getProperty(SIGNAL_COUNTER_DELTA);
    final String attributeCacheRegex = context.getProperty(ATTRIBUTE_CACHE_REGEX).getValue();
    final Integer bufferCount = context.getProperty(SIGNAL_BUFFER_COUNT).asInteger();
    // the cache client used to interact with the distributed cache.
    final AtomicDistributedMapCacheClient cache = context.getProperty(DISTRIBUTED_CACHE_SERVICE).asControllerService(AtomicDistributedMapCacheClient.class);
    final WaitNotifyProtocol protocol = new WaitNotifyProtocol(cache);
    final Map<String, SignalBuffer> signalBuffers = new HashMap<>();
    for (int i = 0; i < bufferCount; i++) {
        final FlowFile flowFile = session.get();
        if (flowFile == null) {
            break;
        }
        // Signal id is computed from attribute 'RELEASE_SIGNAL_IDENTIFIER' with expression language support
        final String signalId = signalIdProperty.evaluateAttributeExpressions(flowFile).getValue();
        // if the computed value is null, or empty, we transfer the flow file to failure relationship
        if (StringUtils.isBlank(signalId)) {
            logger.error("FlowFile {} has no attribute for given Release Signal Identifier", new Object[] { flowFile });
            // set 'notified' attribute
            session.transfer(session.putAttribute(flowFile, NOTIFIED_ATTRIBUTE_NAME, String.valueOf(false)), REL_FAILURE);
            continue;
        }
        String counterName = counterNameProperty.evaluateAttributeExpressions(flowFile).getValue();
        if (StringUtils.isEmpty(counterName)) {
            counterName = WaitNotifyProtocol.DEFAULT_COUNT_NAME;
        }
        int delta = 1;
        if (deltaProperty.isSet()) {
            final String deltaStr = deltaProperty.evaluateAttributeExpressions(flowFile).getValue();
            try {
                delta = Integer.parseInt(deltaStr);
            } catch (final NumberFormatException e) {
                logger.error("Failed to calculate delta for FlowFile {} due to {}", new Object[] { flowFile, e }, e);
                session.transfer(session.putAttribute(flowFile, NOTIFIED_ATTRIBUTE_NAME, String.valueOf(false)), REL_FAILURE);
                continue;
            }
        }
        if (!signalBuffers.containsKey(signalId)) {
            signalBuffers.put(signalId, new SignalBuffer());
        }
        final SignalBuffer signalBuffer = signalBuffers.get(signalId);
        if (StringUtils.isNotEmpty(attributeCacheRegex)) {
            flowFile.getAttributes().entrySet().stream().filter(e -> (!e.getKey().equals("uuid") && e.getKey().matches(attributeCacheRegex))).forEach(e -> signalBuffer.attributesToCache.put(e.getKey(), e.getValue()));
        }
        signalBuffer.incrementDelta(counterName, delta);
        signalBuffer.flowFiles.add(flowFile);
        if (logger.isDebugEnabled()) {
            logger.debug("Cached release signal identifier {} counterName {} from FlowFile {}", new Object[] { signalId, counterName, flowFile });
        }
    }
    signalBuffers.forEach((signalId, signalBuffer) -> {
        // retry after yielding for a while.
        try {
            protocol.notify(signalId, signalBuffer.deltas, signalBuffer.attributesToCache);
            signalBuffer.flowFiles.forEach(flowFile -> session.transfer(session.putAttribute(flowFile, NOTIFIED_ATTRIBUTE_NAME, String.valueOf(true)), REL_SUCCESS));
        } catch (IOException e) {
            throw new RuntimeException(String.format("Unable to communicate with cache when processing %s due to %s", signalId, e), e);
        }
    });
}