Java Code Examples for org.apache.nifi.processor.ProcessSession#transfer()

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final Map<String, Pattern> patterns = regexMapRef.get();
    final ComponentLog logger = getLogger();

    final SortedMap<String, String> attributes = getRelevantAttributes(flowFile, patterns);
    if (attributes.size() != patterns.size()) {
        final Set<String> wantedKeys = patterns.keySet();
        final Set<String> foundKeys = attributes.keySet();
        final StringBuilder missingKeys = new StringBuilder();
        for (final String wantedKey : wantedKeys) {
            if (!foundKeys.contains(wantedKey)) {
                missingKeys.append(wantedKey).append(" ");
            }
        }

        logger.error("routing {} to 'failure' because of missing attributes: {}", new Object[]{flowFile, missingKeys.toString()});
        session.transfer(flowFile, REL_FAILURE);
    } else {
        // create single string of attribute key/value pairs to use for group ID hash
        final StringBuilder hashableValue = new StringBuilder();
        for (final Map.Entry<String, String> entry : attributes.entrySet()) {
            hashableValue.append(entry.getKey());
            if (StringUtils.isBlank(entry.getValue())) {
                hashableValue.append("EMPTY");
            } else {
                hashableValue.append(entry.getValue());
            }
        }

        // create group ID
        final String hashValue = DigestUtils.md5Hex(hashableValue.toString());

        logger.info("adding Hash Value {} to attributes for {} and routing to success", new Object[]{hashValue, flowFile});
        flowFile = session.putAttribute(flowFile, context.getProperty(HASH_VALUE_ATTRIBUTE).getValue(), hashValue);
        session.getProvenanceReporter().modifyAttributes(flowFile);
        session.transfer(flowFile, REL_SUCCESS);
    }
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    FlowFile toRemove = session.get();
    if (toRemove != null) {
        LOG.warn("Removing flow file");
        session.remove(toRemove);
    }

    FlowFile flowFile = session.create();
    final Random random = new Random();
    final byte[] data = new byte[4096];
    random.nextBytes(data);

    flowFile = session.write(flowFile, new OutputStreamCallback() {
        @Override
        public void process(final OutputStream out) throws IOException {
            out.write(data);
        }
    });

    LOG.info("{} transferring {} to success", new Object[]{this, flowFile});
    session.transfer(flowFile, REL_SUCCESS);
    session.commit();
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final StopWatch stopWatch = new StopWatch(true);
        final byte[] buffer = new byte[(int) flowFile.getSize()];
        session.read(flowFile, in -> StreamUtils.fillBuffer(in, buffer));

        try {
            sendMessage(buffer);
        } catch (final ProcessException processException) {
            getLogger().error("Failed to send {} to EventHub due to {}; routing to failure", new Object[]{flowFile, processException}, processException);
            session.transfer(session.penalize(flowFile), REL_FAILURE);
            return;
        }

        final String namespace = context.getProperty(NAMESPACE).getValue();
        final String eventHubName = context.getProperty(EVENT_HUB_NAME).getValue();
        session.getProvenanceReporter().send(flowFile, "amqps://" + namespace + ".servicebus.windows.net" + "/" + eventHubName, stopWatch.getElapsed(TimeUnit.MILLISECONDS));
        session.transfer(flowFile, REL_SUCCESS);

}
 

@Override
public long writeEvent(ProcessSession session, String transitUri, DDLEventInfo eventInfo, long currentSequenceId, Relationship relationship) {
    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, (outputStream) -> {
        super.startJson(outputStream, eventInfo);
        super.writeJson(eventInfo);
        jsonGenerator.writeStringField("query", eventInfo.getQuery());
        super.endJson();
    });
    flowFile = session.putAllAttributes(flowFile, getCommonAttributes(currentSequenceId, eventInfo));
    session.transfer(flowFile, relationship);
    session.getProvenanceReporter().receive(flowFile, transitUri);
    return currentSequenceId + 1;
}
 

private void createFlowFiles(final Set<FileStatus> fileStatuses, final ProcessSession session) {
    for (final FileStatus status : fileStatuses) {
        final Map<String, String> attributes = createAttributes(status);
        FlowFile flowFile = session.create();
        flowFile = session.putAllAttributes(flowFile, attributes);
        session.transfer(flowFile, REL_SUCCESS);
    }
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final long startNanos = System.nanoTime();

    final String bucket = context.getProperty(BUCKET)
                               .evaluateAttributeExpressions(flowFile)
                               .getValue();
    final String key = context.getProperty(KEY)
                               .evaluateAttributeExpressions(flowFile)
                               .getValue();

    final Long generation = context.getProperty(GENERATION)
            .evaluateAttributeExpressions(flowFile)
            .asLong();


    final Storage storage = getCloudService();

    // Deletes a key on Google Cloud
    try {
        storage.delete(BlobId.of(bucket, key, generation));
    } catch (Exception e) {
        getLogger().error(e.getMessage(), e);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
        return;
    }

    session.transfer(flowFile, REL_SUCCESS);
    final long millis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
    getLogger().info("Successfully deleted GCS Object for {} in {} millis; routing to success", new Object[]{flowFile, millis});
}
 

private void finishProcessing(ExecutionContext executionContext, ProcessSession session) {
    executionContext.nrOfWaitingHDFSObjects++;

    if (req.groupping == Groupping.NONE && req.isDestContent && executionContext.nrOfWaitingHDFSObjects < req.batchSize) {
        return;
    }

    session.transfer(executionContext.flowfile, REL_SUCCESS);

    executionContext.reset();
}
 

private void handleFile(GridFSBucket bucket, ProcessSession session, ProcessContext context, FlowFile parent, GridFSFile input, String query) {
    Map<String, String> attrs = new HashMap<>();
    attrs.put(METADATA_ATTRIBUTE, input.getMetadata() != null ? input.getMetadata().toJson() : "{}");
    if (context.getProperty(QUERY_ATTRIBUTE).isSet()) {
        String key = context.getProperty(QUERY_ATTRIBUTE).evaluateAttributeExpressions(parent).getValue();
        attrs.put(key, query);
    }
    attrs.put(CoreAttributes.FILENAME.key(), input.getFilename());
    FlowFile output = parent != null ? session.create(parent) : session.create();
    output = session.write(output, out -> bucket.downloadToStream(input.getObjectId(), out));
    output = session.putAllAttributes(output, attrs);
    session.transfer(output, REL_SUCCESS);
    session.getProvenanceReporter().receive(output, getTransitUri(input.getObjectId(), output, context));
}
 

/**
 * Will release flow file. Releasing of the flow file in the context of this
 * operation implies the following:
 *
 * If Empty then remove from session and return
 * If has something then transfer to REL_SUCCESS
 */
private void releaseFlowFile(FlowFile flowFile, ProcessSession session, Map<String, String> attributes, long start, String topic, int msgCount){
    if (flowFile.getSize() == 0L) {
        session.remove(flowFile);
    } else {
        flowFile = session.putAllAttributes(flowFile, attributes);
        final long millis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start);
        session.getProvenanceReporter().receive(flowFile, "kafka://" + topic, "Received " + msgCount + " Kafka messages", millis);
        getLogger().info("Successfully received {} from Kafka with {} messages in {} millis", new Object[]{flowFile, msgCount, millis});
        session.transfer(flowFile, REL_SUCCESS);
    }
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final String logLevelValue = context.getProperty(LOG_LEVEL).getValue().toLowerCase();

    final DebugLevels logLevel;
    try {
        logLevel = DebugLevels.valueOf(logLevelValue);
    } catch (Exception e) {
        throw new ProcessException(e);
    }

    final ComponentLog LOG = getLogger();
    boolean isLogLevelEnabled = false;
    switch (logLevel) {
        case trace:
            isLogLevelEnabled = LOG.isTraceEnabled();
            break;
        case debug:
            isLogLevelEnabled = LOG.isDebugEnabled();
            break;
        case info:
            isLogLevelEnabled = LOG.isInfoEnabled();
            break;
        case warn:
            isLogLevelEnabled = LOG.isWarnEnabled();
            break;
        case error:
            isLogLevelEnabled = LOG.isErrorEnabled();
            break;
    }

    if (!isLogLevelEnabled) {
        transferChunk(session);
        return;
    }

    final FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    processFlowFile(LOG, logLevel, flowFile, session, context);
    session.transfer(flowFile, REL_SUCCESS);
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    try {
        restoreState(context);
    } catch (IOException ioe) {
        getLogger().error("Failed to restore processor state; yielding", ioe);
        context.yield();
        return;
    }

    final long startNanos = System.nanoTime();
    final String bucket = context.getProperty(BUCKET).evaluateAttributeExpressions().getValue();

    final AmazonS3 client = getClient();
    int listCount = 0;
    long maxTimestamp = 0L;
    String delimiter = context.getProperty(DELIMITER).getValue();
    String prefix = context.getProperty(PREFIX).evaluateAttributeExpressions().getValue();

    boolean useVersions = context.getProperty(USE_VERSIONS).asBoolean();

    S3BucketLister bucketLister = useVersions
            ? new S3VersionBucketLister(client)
            : new S3ObjectBucketLister(client);

    bucketLister.setBucketName(bucket);

    if (delimiter != null && !delimiter.isEmpty()) {
        bucketLister.setDelimiter(delimiter);
    }
    if (prefix != null && !prefix.isEmpty()) {
        bucketLister.setPrefix(prefix);
    }

    VersionListing versionListing;
    do {
        versionListing = bucketLister.listVersions();
        for (S3VersionSummary versionSummary : versionListing.getVersionSummaries()) {
            long lastModified = versionSummary.getLastModified().getTime();
            if (lastModified < currentTimestamp
                    || lastModified == currentTimestamp && currentKeys.contains(versionSummary.getKey())) {
                continue;
            }

            // Create the attributes
            final Map<String, String> attributes = new HashMap<>();
            attributes.put(CoreAttributes.FILENAME.key(), versionSummary.getKey());
            attributes.put("s3.bucket", versionSummary.getBucketName());
            if (versionSummary.getOwner() != null) { // We may not have permission to read the owner
                attributes.put("s3.owner", versionSummary.getOwner().getId());
            }
            attributes.put("s3.etag", versionSummary.getETag());
            attributes.put("s3.lastModified", String.valueOf(lastModified));
            attributes.put("s3.length", String.valueOf(versionSummary.getSize()));
            attributes.put("s3.storeClass", versionSummary.getStorageClass());
            attributes.put("s3.isLatest", String.valueOf(versionSummary.isLatest()));
            if (versionSummary.getVersionId() != null) {
                attributes.put("s3.version", versionSummary.getVersionId());
            }

            // Create the flowfile
            FlowFile flowFile = session.create();
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.transfer(flowFile, REL_SUCCESS);

            // Update state
            if (lastModified > maxTimestamp) {
                maxTimestamp = lastModified;
                currentKeys.clear();
            }
            if (lastModified == maxTimestamp) {
                currentKeys.add(versionSummary.getKey());
            }
            listCount++;
        }
        bucketLister.setNextMarker();

        commit(context, session, listCount);
        listCount = 0;
    } while (bucketLister.isTruncated());
    currentTimestamp = maxTimestamp;

    final long listMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
    getLogger().info("Successfully listed S3 bucket {} in {} millis", new Object[]{bucket, listMillis});

    if (!commit(context, session, listCount)) {
        if (currentTimestamp > 0) {
            persistState(context);
        }
        getLogger().debug("No new objects in S3 bucket {} to list. Yielding.", new Object[]{bucket});
        context.yield();
    }
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final StopWatch stopWatch = new StopWatch(true);
    final String key = context.getProperty(KEY).getValue();
    final String sourceLanguage = context.getProperty(SOURCE_LANGUAGE).evaluateAttributeExpressions(flowFile).getValue();
    final String targetLanguage = context.getProperty(TARGET_LANGUAGE).evaluateAttributeExpressions(flowFile).getValue();
    final String encoding = context.getProperty(CHARACTER_SET).evaluateAttributeExpressions(flowFile).getValue();

    final List<String> attributeNames = new ArrayList<>();
    final List<String> textValues = new ArrayList<>();
    for (final PropertyDescriptor descriptor : context.getProperties().keySet()) {
        if (descriptor.isDynamic()) {
            attributeNames.add(descriptor.getName()); // add to list so that we know the order when the translations come back.
            textValues.add(context.getProperty(descriptor).evaluateAttributeExpressions(flowFile).getValue());
        }
    }

    if (context.getProperty(TRANSLATE_CONTENT).asBoolean()) {
        final byte[] buff = new byte[(int) flowFile.getSize()];
        session.read(flowFile, new InputStreamCallback() {
            @Override
            public void process(final InputStream in) throws IOException {
                StreamUtils.fillBuffer(in, buff);
            }
        });
        final String content = new String(buff, Charset.forName(encoding));
        textValues.add(content);
    }

    final WebResource.Builder builder = prepareResource(key, textValues, sourceLanguage, targetLanguage);

    final ClientResponse response;
    try {
        response = builder.post(ClientResponse.class);
    } catch (final Exception e) {
        getLogger().error("Failed to make request to Yandex to transate text for {} due to {}; routing to comms.failure", new Object[]{flowFile, e});
        session.transfer(flowFile, REL_COMMS_FAILURE);
        return;
    }

    if (response.getStatus() != Status.OK.getStatusCode()) {
        getLogger().error("Failed to translate text using Yandex for {}; response was {}: {}; routing to {}", new Object[]{
            flowFile, response.getStatus(), response.getStatusInfo().getReasonPhrase(), REL_TRANSLATION_FAILED.getName()});
        flowFile = session.putAttribute(flowFile, "yandex.translate.failure.reason", response.getStatusInfo().getReasonPhrase());
        session.transfer(flowFile, REL_TRANSLATION_FAILED);
        return;
    }

    final Map<String, String> newAttributes = new HashMap<>();
    final Translation translation = response.getEntity(Translation.class);
    final List<String> texts = translation.getText();
    for (int i = 0; i < texts.size(); i++) {
        final String text = texts.get(i);
        if (i < attributeNames.size()) {
            final String attributeName = attributeNames.get(i);
            newAttributes.put(attributeName, text);
        } else {
            flowFile = session.write(flowFile, new OutputStreamCallback() {
                @Override
                public void process(final OutputStream out) throws IOException {
                    out.write(text.getBytes(encoding));
                }
            });

            newAttributes.put("language", targetLanguage);
        }
    }

    if (!newAttributes.isEmpty()) {
        flowFile = session.putAllAttributes(flowFile, newAttributes);
    }

    stopWatch.stop();
    session.transfer(flowFile, REL_SUCCESS);
    getLogger().info("Successfully translated {} items for {} from {} to {} in {}; routing to success",
            new Object[]{texts.size(), flowFile, sourceLanguage, targetLanguage, stopWatch.getDuration()});
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final String projectId = context.getProperty(PROJECT_ID).evaluateAttributeExpressions().getValue();
    final String dataset = context.getProperty(DATASET).evaluateAttributeExpressions(flowFile).getValue();
    final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions(flowFile).getValue();

    final TableId tableId;
    if (StringUtils.isEmpty(projectId)) {
        tableId = TableId.of(dataset, tableName);
    } else {
        tableId = TableId.of(projectId, dataset, tableName);
    }

    try {

        InsertAllRequest.Builder request = InsertAllRequest.newBuilder(tableId);
        int nbrecord = 0;

        try (final InputStream in = session.read(flowFile)) {
            final RecordReaderFactory readerFactory = context.getProperty(RECORD_READER).asControllerService(RecordReaderFactory.class);
            try (final RecordReader reader = readerFactory.createRecordReader(flowFile, in, getLogger());) {
                Record currentRecord;
                while ((currentRecord = reader.nextRecord()) != null) {
                    request.addRow(convertMapRecord(currentRecord.toMap()));
                    nbrecord++;
                }
            }
        }

        request.setIgnoreUnknownValues(context.getProperty(IGNORE_UNKNOWN).evaluateAttributeExpressions(flowFile).asBoolean());
        request.setSkipInvalidRows(context.getProperty(SKIP_INVALID_ROWS).evaluateAttributeExpressions(flowFile).asBoolean());

        InsertAllResponse response = getCloudService().insertAll(request.build());

        final Map<String, String> attributes = new HashMap<>();

        if (response.hasErrors()) {
            getLogger().log(LogLevel.WARN, "Failed to insert {} of {} records into BigQuery {} table.", new Object[] { response.getInsertErrors().size(), nbrecord, tableName });
            if (getLogger().isDebugEnabled()) {
                for (long index : response.getInsertErrors().keySet()) {
                    for (BigQueryError e : response.getInsertErrors().get(index)) {
                        getLogger().log(LogLevel.DEBUG, "Failed to insert record #{}: {}", new Object[] { index, e.getMessage() });
                    }
                }
            }

            attributes.put(BigQueryAttributes.JOB_NB_RECORDS_ATTR, Long.toString(nbrecord - response.getInsertErrors().size()));

            flowFile = session.penalize(flowFile);
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.transfer(flowFile, REL_FAILURE);
        } else {
            attributes.put(BigQueryAttributes.JOB_NB_RECORDS_ATTR, Long.toString(nbrecord));
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.transfer(flowFile, REL_SUCCESS);
        }

    } catch (Exception ex) {
        getLogger().log(LogLevel.ERROR, ex.getMessage(), ex);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
    }
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    List<FlowFile> flowFiles = session.get(new ThrottleFilter(MAX_FLOW_FILES_PER_BATCH));
    if (flowFiles.isEmpty()) {
        context.yield();
        return;
    }

    // Periodically clear any Throttle that has not been used in more than 2 throttling periods
    final long lastClearTime = lastThrottleClearTime.get();
    final long throttleExpirationMillis = System.currentTimeMillis() - 2 * context.getProperty(TIME_PERIOD).asTimePeriod(TimeUnit.MILLISECONDS);
    if (lastClearTime < throttleExpirationMillis) {
        if (lastThrottleClearTime.compareAndSet(lastClearTime, System.currentTimeMillis())) {
            final Iterator<Map.Entry<String, Throttle>> itr = throttleMap.entrySet().iterator();
            while (itr.hasNext()) {
                final Map.Entry<String, Throttle> entry = itr.next();
                final Throttle throttle = entry.getValue();
                if (throttle.tryLock()) {
                    try {
                        if (throttle.lastUpdateTime() < lastClearTime) {
                            itr.remove();
                        }
                    } finally {
                        throttle.unlock();
                    }
                }
            }
        }
    }

    final ComponentLog logger = getLogger();
    for (FlowFile flowFile : flowFiles) {
        // call this to capture potential error
        final long accrualAmount = getFlowFileAccrual(flowFile);
        if (accrualAmount < 0) {
            logger.error("Routing {} to 'failure' due to missing or invalid attribute", new Object[]{flowFile});
            session.transfer(flowFile, REL_FAILURE);
        } else {
            logger.info("transferring {} to 'success'", new Object[]{flowFile});
            session.transfer(flowFile, REL_SUCCESS);
        }
    }
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {

    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final String functionName = context.getProperty(AWS_LAMBDA_FUNCTION_NAME).getValue();

    final String qualifier = context.getProperty(AWS_LAMBDA_FUNCTION_QUALIFIER).getValue();

    // Max size of message is 6 MB
    if ( flowFile.getSize() > MAX_REQUEST_SIZE) {
        getLogger().error("Max size for request body is 6mb but was {} for flow file {} for function {}",
            new Object[]{flowFile.getSize(), flowFile, functionName});
        session.transfer(flowFile, REL_FAILURE);
        return;
    }

    final AWSLambdaClient client = getClient();

    try {
        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
        session.exportTo(flowFile, baos);

        InvokeRequest invokeRequest = new InvokeRequest()
            .withFunctionName(functionName)
            .withLogType(LogType.Tail).withInvocationType(InvocationType.RequestResponse)
            .withPayload(ByteBuffer.wrap(baos.toByteArray()))
            .withQualifier(qualifier);
        long startTime = System.nanoTime();

        InvokeResult result = client.invoke(invokeRequest);

        flowFile = session.putAttribute(flowFile, AWS_LAMBDA_RESULT_STATUS_CODE, result.getStatusCode().toString());

        if ( !StringUtils.isBlank(result.getLogResult() )) {
            flowFile = session.putAttribute(flowFile, AWS_LAMBDA_RESULT_LOG, new String(Base64.decode(result.getLogResult()),Charset.defaultCharset()));
        }

        if ( result.getPayload() != null ) {
            flowFile = session.putAttribute(flowFile, AWS_LAMBDA_RESULT_PAYLOAD, new String(result.getPayload().array(),Charset.defaultCharset()));
        }

        if ( ! StringUtils.isBlank(result.getFunctionError()) ){
            flowFile = session.putAttribute(flowFile, AWS_LAMBDA_RESULT_FUNCTION_ERROR, result.getFunctionError());
            session.transfer(flowFile, REL_FAILURE);
        } else {
            session.transfer(flowFile, REL_SUCCESS);
            final long totalTimeMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime);
            session.getProvenanceReporter().send(flowFile, functionName, totalTimeMillis);
        }
    } catch (final InvalidRequestContentException
        | InvalidParameterValueException
        | RequestTooLargeException
        | ResourceNotFoundException
        | UnsupportedMediaTypeException unrecoverableException) {
            getLogger().error("Failed to invoke lambda {} with unrecoverable exception {} for flow file {}",
                new Object[]{functionName, unrecoverableException, flowFile});
            flowFile = populateExceptionAttributes(session, flowFile, unrecoverableException);
            session.transfer(flowFile, REL_FAILURE);
    } catch (final TooManyRequestsException retryableServiceException) {
        getLogger().error("Failed to invoke lambda {} with exception {} for flow file {}, therefore penalizing flowfile",
            new Object[]{functionName, retryableServiceException, flowFile});
        flowFile = populateExceptionAttributes(session, flowFile, retryableServiceException);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
        context.yield();
    } catch (final AmazonServiceException unrecoverableServiceException) {
        getLogger().error("Failed to invoke lambda {} with exception {} for flow file {} sending to fail",
            new Object[]{functionName, unrecoverableServiceException, flowFile});
        flowFile = populateExceptionAttributes(session, flowFile, unrecoverableServiceException);
        session.transfer(flowFile, REL_FAILURE);
        context.yield();
    } catch (final Exception exception) {
        getLogger().error("Failed to invoke lambda {} with exception {} for flow file {}",
            new Object[]{functionName, exception, flowFile});
        session.transfer(flowFile, REL_FAILURE);
        context.yield();
    }
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final FlowFile original = session.get();
    if (original == null) {
        return;
    }

    final int depth = context.getProperty(SPLIT_DEPTH).asInteger();
    final ComponentLog logger = getLogger();

    final List<FlowFile> splits = new ArrayList<>();
    final String fragmentIdentifier = UUID.randomUUID().toString();
    final AtomicInteger numberOfRecords = new AtomicInteger(0);
    final XmlSplitterSaxParser parser = new XmlSplitterSaxParser(xmlTree -> {
        FlowFile split = session.create(original);
        split = session.write(split, out -> out.write(xmlTree.getBytes("UTF-8")));
        split = session.putAttribute(split, FRAGMENT_ID.key(), fragmentIdentifier);
        split = session.putAttribute(split, FRAGMENT_INDEX.key(), Integer.toString(numberOfRecords.getAndIncrement()));
        split = session.putAttribute(split, SEGMENT_ORIGINAL_FILENAME.key(), split.getAttribute(CoreAttributes.FILENAME.key()));
        splits.add(split);
    }, depth);

    final AtomicBoolean failed = new AtomicBoolean(false);
    session.read(original, rawIn -> {
        try (final InputStream in = new BufferedInputStream(rawIn)) {
            SAXParser saxParser = null;
            try {
                saxParser = saxParserFactory.newSAXParser();
                final XMLReader reader = saxParser.getXMLReader();
                reader.setContentHandler(parser);
                reader.parse(new InputSource(in));
            } catch (final ParserConfigurationException | SAXException e) {
                logger.error("Unable to parse {} due to {}", new Object[]{original, e});
                failed.set(true);
            }
        }
    });

    if (failed.get()) {
        session.transfer(original, REL_FAILURE);
        session.remove(splits);
    } else {
        splits.forEach((split) -> {
            split = session.putAttribute(split, FRAGMENT_COUNT.key(), Integer.toString(numberOfRecords.get()));
            session.transfer(split, REL_SPLIT);
        });

        final FlowFile originalToTransfer = copyAttributesToOriginal(session, original, fragmentIdentifier, numberOfRecords.get());
        session.transfer(originalToTransfer, REL_ORIGINAL);
        logger.info("Split {} into {} FlowFiles", new Object[]{originalToTransfer, splits.size()});
    }
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final ComponentLog logger = getLogger();
    final String cacheKey = context.getProperty(PROP_CACHE_ENTRY_IDENTIFIER).evaluateAttributeExpressions(flowFile).getValue();
    // This block retains the previous behavior when only one Cache Entry Identifier was allowed, so as not to change the expected error message
    if (StringUtils.isBlank(cacheKey)) {
        logger.error("FlowFile {} has no attribute for given Cache Entry Identifier", new Object[]{flowFile});
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
        return;
    }
    List<String> cacheKeys = Arrays.stream(cacheKey.split(",")).filter(path -> !StringUtils.isEmpty(path)).map(String::trim).collect(Collectors.toList());
    for (int i = 0; i < cacheKeys.size(); i++) {
        if (StringUtils.isBlank(cacheKeys.get(i))) {
            // Log first missing identifier, route to failure, and return
            logger.error("FlowFile {} has no attribute for Cache Entry Identifier in position {}", new Object[]{flowFile, i});
            flowFile = session.penalize(flowFile);
            session.transfer(flowFile, REL_FAILURE);
            return;
        }
    }

    final DistributedMapCacheClient cache = context.getProperty(PROP_DISTRIBUTED_CACHE_SERVICE).asControllerService(DistributedMapCacheClient.class);

    try {
        final Map<String, byte[]> cacheValues;
        final boolean singleKey = cacheKeys.size() == 1;
        if (singleKey) {
            cacheValues = new HashMap<>(1);
            cacheValues.put(cacheKeys.get(0), cache.get(cacheKey, keySerializer, valueDeserializer));
        } else {
            cacheValues = cache.subMap(new HashSet<>(cacheKeys), keySerializer, valueDeserializer);
        }
        boolean notFound = false;
        for(Map.Entry<String,byte[]> cacheValueEntry : cacheValues.entrySet()) {
            final byte[] cacheValue = cacheValueEntry.getValue();

            if (cacheValue == null) {
                logger.info("Could not find an entry in cache for {}; routing to not-found", new Object[]{flowFile});
                notFound = true;
                break;
            } else {
                boolean putInAttribute = context.getProperty(PROP_PUT_CACHE_VALUE_IN_ATTRIBUTE).isSet();
                if (putInAttribute) {
                    String attributeName = context.getProperty(PROP_PUT_CACHE_VALUE_IN_ATTRIBUTE).evaluateAttributeExpressions(flowFile).getValue();
                    if (!singleKey) {
                        // Append key to attribute name if multiple keys
                        attributeName += "." + cacheValueEntry.getKey();
                    }
                    String attributeValue = new String(cacheValue, context.getProperty(PROP_CHARACTER_SET).getValue());

                    int maxLength = context.getProperty(PROP_PUT_ATTRIBUTE_MAX_LENGTH).asInteger();
                    if (maxLength < attributeValue.length()) {
                        attributeValue = attributeValue.substring(0, maxLength);
                    }

                    flowFile = session.putAttribute(flowFile, attributeName, attributeValue);

                } else if (cacheKeys.size() > 1) {
                    throw new IOException("Multiple Cache Value Identifiers specified without Put Cache Value In Attribute set");
                } else {
                    // Write single value to content
                    flowFile = session.write(flowFile, out -> out.write(cacheValue));
                }

                if (putInAttribute) {
                    logger.info("Found a cache key of {} and added an attribute to {} with it's value.", new Object[]{cacheKey, flowFile});
                } else {
                    logger.info("Found a cache key of {} and replaced the contents of {} with it's value.", new Object[]{cacheKey, flowFile});
                }
            }
        }
        // If the loop was exited because a cache entry was not found, route to REL_NOT_FOUND; otherwise route to REL_SUCCESS
        if (notFound) {
            session.transfer(flowFile, REL_NOT_FOUND);
        } else {
            session.transfer(flowFile, REL_SUCCESS);
        }
    } catch (final IOException e) {
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
        logger.error("Unable to communicate with cache when processing {} due to {}", new Object[]{flowFile, e});
    }
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final String segmentId = UUID.randomUUID().toString();
    final long segmentSize = context.getProperty(SIZE).asDataSize(DataUnit.B).longValue();

    final String originalFileName = flowFile.getAttribute(CoreAttributes.FILENAME.key());

    if (flowFile.getSize() <= segmentSize) {
        flowFile = session.putAttribute(flowFile, SEGMENT_ID, segmentId);
        flowFile = session.putAttribute(flowFile, SEGMENT_INDEX, "1");
        flowFile = session.putAttribute(flowFile, SEGMENT_COUNT, "1");
        flowFile = session.putAttribute(flowFile, SEGMENT_ORIGINAL_FILENAME, originalFileName);

        flowFile = session.putAttribute(flowFile, FRAGMENT_ID, segmentId);
        flowFile = session.putAttribute(flowFile, FRAGMENT_INDEX, "1");
        flowFile = session.putAttribute(flowFile, FRAGMENT_COUNT, "1");

        FlowFile clone = session.clone(flowFile);
        session.transfer(flowFile, REL_ORIGINAL);
        session.transfer(clone, REL_SEGMENTS);
        return;
    }

    int totalSegments = (int) (flowFile.getSize() / segmentSize);
    if (totalSegments * segmentSize < flowFile.getSize()) {
        totalSegments++;
    }

    final Map<String, String> segmentAttributes = new HashMap<>();
    segmentAttributes.put(SEGMENT_ID, segmentId);
    segmentAttributes.put(SEGMENT_COUNT, String.valueOf(totalSegments));
    segmentAttributes.put(SEGMENT_ORIGINAL_FILENAME, originalFileName);

    segmentAttributes.put(FRAGMENT_ID, segmentId);
    segmentAttributes.put(FRAGMENT_COUNT, String.valueOf(totalSegments));

    final Set<FlowFile> segmentSet = new HashSet<>();
    for (int i = 1; i <= totalSegments; i++) {
        final long segmentOffset = segmentSize * (i - 1);
        FlowFile segment = session.clone(flowFile, segmentOffset, Math.min(segmentSize, flowFile.getSize() - segmentOffset));
        segmentAttributes.put(SEGMENT_INDEX, String.valueOf(i));
        segmentAttributes.put(FRAGMENT_INDEX, String.valueOf(i));
        segment = session.putAllAttributes(segment, segmentAttributes);
        segmentSet.add(segment);
    }

    session.transfer(segmentSet, REL_SEGMENTS);
    flowFile = FragmentAttributes.copyAttributesToOriginal(session, flowFile, segmentId, totalSegments);
    session.transfer(flowFile, REL_ORIGINAL);

    if (totalSegments <= 10) {
        getLogger().info("Segmented {} into {} segments: {}", new Object[]{flowFile, totalSegments, segmentSet});
    } else {
        getLogger().info("Segmented {} into {} segments", new Object[]{flowFile, totalSegments});
    }
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final byte[] content = new byte[(int) flowFile.getSize()];
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, content, true);
        }
    });

    String docId = flowFile.getAttribute(CoreAttributes.UUID.key());
    if (!StringUtils.isEmpty(context.getProperty(DOC_ID).getValue())) {
        docId = context.getProperty(DOC_ID).evaluateAttributeExpressions(flowFile).getValue();
    }

    try {
        Document<?> doc = null;
        final DocumentType documentType = DocumentType.valueOf(context.getProperty(DOCUMENT_TYPE).getValue());
        switch (documentType) {
            case Json: {
                doc = RawJsonDocument.create(docId, new String(content, StandardCharsets.UTF_8));
                break;
            }
            case Binary: {
                final ByteBuf buf = Unpooled.copiedBuffer(content);
                doc = BinaryDocument.create(docId, buf);
                break;
            }
        }

        final PersistTo persistTo = PersistTo.valueOf(context.getProperty(PERSIST_TO).getValue());
        final ReplicateTo replicateTo = ReplicateTo.valueOf(context.getProperty(REPLICATE_TO).getValue());
        doc = openBucket(context).upsert(doc, persistTo, replicateTo);

        final Map<String, String> updatedAttrs = new HashMap<>();
        updatedAttrs.put(CouchbaseAttributes.Cluster.key(), context.getProperty(COUCHBASE_CLUSTER_SERVICE).getValue());
        updatedAttrs.put(CouchbaseAttributes.Bucket.key(), context.getProperty(BUCKET_NAME).getValue());
        updatedAttrs.put(CouchbaseAttributes.DocId.key(), docId);
        updatedAttrs.put(CouchbaseAttributes.Cas.key(), String.valueOf(doc.cas()));
        updatedAttrs.put(CouchbaseAttributes.Expiry.key(), String.valueOf(doc.expiry()));

        flowFile = session.putAllAttributes(flowFile, updatedAttrs);
        session.getProvenanceReporter().send(flowFile, getTransitUrl(context, docId));
        session.transfer(flowFile, REL_SUCCESS);
    } catch (final CouchbaseException e) {
        String errMsg = String.format("Writing document %s to Couchbase Server using %s failed due to %s", docId, flowFile, e);
        handleCouchbaseException(context, session, logger, flowFile, e, errMsg);
    }
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {

    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    String mimeType = flowFile.getAttribute(CoreAttributes.MIME_TYPE.key());
    String packagingFormat = NOT_PACKAGED;
    if (null != mimeType) {
        switch (mimeType.toLowerCase()) {
            case "application/tar":
                packagingFormat = TAR_FORMAT;
                break;
            case "application/zip":
                packagingFormat = ZIP_FORMAT;
                break;
            case "application/flowfile-v3":
                packagingFormat = FLOWFILE_STREAM_FORMAT_V3;
                break;
            default:
                getLogger().warn(
                        "Cannot unpack {} because its mime.type attribute is set to '{}', which is not a format that can be unpacked",
                        new Object[]{flowFile, mimeType});
        }
    }
    final SequenceFileWriter sequenceFileWriter;
    switch (packagingFormat) {
        case TAR_FORMAT:
            sequenceFileWriter = new TarUnpackerSequenceFileWriter();
            break;
        case ZIP_FORMAT:
            sequenceFileWriter = new ZipUnpackerSequenceFileWriter();
            break;
        case FLOWFILE_STREAM_FORMAT_V3:
            sequenceFileWriter = new FlowFileStreamUnpackerSequenceFileWriter();
            break;
        default:
            sequenceFileWriter = new SequenceFileWriterImpl();
    }

    final Configuration configuration = getConfiguration();
    if (configuration == null) {
        getLogger().error("HDFS not configured properly");
        session.transfer(flowFile, RELATIONSHIP_FAILURE);
        context.yield();
        return;
    }

    final CompressionCodec codec = getCompressionCodec(context, configuration);

    final String value = context.getProperty(COMPRESSION_TYPE).getValue();
    final SequenceFile.CompressionType compressionType = value == null
        ? SequenceFile.CompressionType.valueOf(DEFAULT_COMPRESSION_TYPE) : SequenceFile.CompressionType.valueOf(value);

    final String fileName = flowFile.getAttribute(CoreAttributes.FILENAME.key()) + ".sf";
    flowFile = session.putAttribute(flowFile, CoreAttributes.FILENAME.key(), fileName);

    try {
        StopWatch stopWatch = new StopWatch(true);
        flowFile = sequenceFileWriter.writeSequenceFile(flowFile, session, configuration, compressionType, codec);
        session.getProvenanceReporter().modifyContent(flowFile, stopWatch.getElapsed(TimeUnit.MILLISECONDS));
        session.transfer(flowFile, RELATIONSHIP_SUCCESS);
        getLogger().info("Transferred flowfile {} to {}", new Object[]{flowFile, RELATIONSHIP_SUCCESS});
    } catch (ProcessException e) {
        getLogger().error("Failed to create Sequence File. Transferring {} to 'failure'", new Object[]{flowFile}, e);
        session.transfer(flowFile, RELATIONSHIP_FAILURE);
    }

}