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

private void appendLine(final ProcessSession session, final Map<Relationship, Map<Group, FlowFile>> flowFileMap, final Relationship relationship,
    final FlowFile original, final String line, final Charset charset, final Group group) {

    Map<Group, FlowFile> groupToFlowFileMap = flowFileMap.get(relationship);
    if (groupToFlowFileMap == null) {
        groupToFlowFileMap = new HashMap<>();
        flowFileMap.put(relationship, groupToFlowFileMap);
    }

    FlowFile flowFile = groupToFlowFileMap.get(group);
    if (flowFile == null) {
        flowFile = session.create(original);
    }

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

    groupToFlowFileMap.put(group, flowFile);
}
 

/**
 * Will concatenate the contents of the provided array of {@link FlowFile}s
 * into a single {@link FlowFile}. While this operation is as general as it
 * is described in the previous sentence, in the context of this processor
 * there can only be two {@link FlowFile}s with the first {@link FlowFile}
 * representing the header content of the split and the second
 * {@link FlowFile} represents the split itself.
 */
private FlowFile concatenateContents(FlowFile sourceFlowFile, ProcessSession session, FlowFile... flowFiles) {
    FlowFile mergedFlowFile = session.create(sourceFlowFile);
    for (FlowFile flowFile : flowFiles) {
        mergedFlowFile = session.append(mergedFlowFile, new OutputStreamCallback() {
            @Override
            public void process(OutputStream out) throws IOException {
                try (InputStream is = session.read(flowFile)) {
                    IOUtils.copy(is, out);
                }
            }
        });
    }
    session.remove(flowFiles[1]); // in current usage we always have 2 files
    return mergedFlowFile;
}
 

private void appendLine(final ProcessSession session, final Map<Relationship, Map<Group, FlowFile>> flowFileMap, final Relationship relationship,
    final FlowFile original, final String line, final Charset charset, final Group group) {

    final Map<Group, FlowFile> groupToFlowFileMap = flowFileMap.computeIfAbsent(relationship, k -> new HashMap<>());

    FlowFile flowFile = groupToFlowFileMap.get(group);
    if (flowFile == null) {
        flowFile = session.create(original);
    }

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

    groupToFlowFileMap.put(group, flowFile);
}
 

/**
 * Will concatenate the contents of the provided array of {@link FlowFile}s
 * into a single {@link FlowFile}. While this operation is as general as it
 * is described in the previous sentence, in the context of this processor
 * there can only be two {@link FlowFile}s with the first {@link FlowFile}
 * representing the header content of the split and the second
 * {@link FlowFile} represents the split itself.
 */
private FlowFile concatenateContents(FlowFile sourceFlowFile, ProcessSession session, FlowFile... flowFiles) {
    FlowFile mergedFlowFile = session.create(sourceFlowFile);
    for (FlowFile flowFile : flowFiles) {
        mergedFlowFile = session.append(mergedFlowFile, new OutputStreamCallback() {
            @Override
            public void process(OutputStream out) throws IOException {
                try (InputStream is = session.read(flowFile)) {
                    IOUtils.copy(is, out);
                }
            }
        });
    }
    session.remove(flowFiles[1]); // in current usage we always have 2 files
    return mergedFlowFile;
}
 

private void finalizeFlowFile(final ProcessSession session, final HBaseClientService hBaseClientService,
        FlowFile flowFile, final String tableName, Long rowsPulled, Exception e) {
    Relationship rel = REL_SUCCESS;
    flowFile = session.putAttribute(flowFile, HBASE_ROWS_COUNT_ATTR, rowsPulled.toString());

    final AtomicReference<IOException> ioe = new AtomicReference<>(null);
    flowFile = session.append(flowFile, (out) -> {
        try{
            out.write("]".getBytes());
        }catch(IOException ei){
            ioe.set(ei);
        }
    });
    if (e != null || ioe.get() != null) {
        flowFile = session.putAttribute(flowFile, "scanhbase.error", (e==null?e:ioe.get()).toString());
        rel = REL_FAILURE;
    } else {
        session.getProvenanceReporter().receive(flowFile, hBaseClientService.toTransitUri(tableName, "{ids}"));
    }
    session.transfer(flowFile, rel);
}
 

private void writeData(final ProcessSession session, final List<ConsumerRecord<byte[], byte[]>> records, final TopicPartition topicPartition) {
    final ConsumerRecord<byte[], byte[]> firstRecord = records.get(0);
    final boolean demarcateFirstRecord;
    BundleTracker tracker = bundleMap.get(topicPartition);
    FlowFile flowFile;
    if (tracker == null) {
        tracker = new BundleTracker(firstRecord, topicPartition, keyEncoding);
        flowFile = session.create();
        tracker.updateFlowFile(flowFile);
        demarcateFirstRecord = false; //have not yet written records for this topic/partition in this lease
    } else {
        demarcateFirstRecord = true; //have already been writing records for this topic/partition in this lease
    }
    flowFile = tracker.flowFile;
    tracker.incrementRecordCount(records.size());
    flowFile = session.append(flowFile, out -> {
        boolean useDemarcator = demarcateFirstRecord;
        for (final ConsumerRecord<byte[], byte[]> record : records) {
            if (useDemarcator) {
                out.write(demarcatorBytes);
            }
            out.write(record.value());
            useDemarcator = true;
        }
    });
    tracker.updateFlowFile(flowFile);
    bundleMap.put(topicPartition, tracker);
}
 

private void writeData(final ProcessSession session, final List<ConsumerRecord<byte[], byte[]>> records, final TopicPartition topicPartition) {
    final ConsumerRecord<byte[], byte[]> firstRecord = records.get(0);
    final boolean demarcateFirstRecord;
    BundleTracker tracker = bundleMap.get(topicPartition);
    FlowFile flowFile;
    if (tracker == null) {
        tracker = new BundleTracker(firstRecord, topicPartition, keyEncoding);
        flowFile = session.create();
        tracker.updateFlowFile(flowFile);
        demarcateFirstRecord = false; //have not yet written records for this topic/partition in this lease
    } else {
        demarcateFirstRecord = true; //have already been writing records for this topic/partition in this lease
    }
    flowFile = tracker.flowFile;
    tracker.incrementRecordCount(records.size());
    flowFile = session.append(flowFile, out -> {
        boolean useDemarcator = demarcateFirstRecord;
        for (final ConsumerRecord<byte[], byte[]> record : records) {
            if (useDemarcator) {
                out.write(demarcatorBytes);
            }
            out.write(record.value());
            useDemarcator = true;
        }
    });
    tracker.updateFlowFile(flowFile);
    bundleMap.put(topicPartition, tracker);
}
 

private FlowFile addAsContent(ExecutionContext executionContext, ProcessSession session, HDFSObjectInfoDetails o, FlowFile ff) {
    if (executionContext.nrOfWaitingHDFSObjects > 0) {
        ff = session.append(ff, (out) -> out.write(("\n").getBytes()));
    }

    return session.append(ff, (out) -> out.write((o.toJsonString()).getBytes()));
}
 

@Test
public void testAppendToFlowFileWhereResourceClaimHasMultipleContentClaims() throws IOException {
    final Relationship relationship = new Relationship.Builder().name("A").build();

    FlowFile ffa = session.create();
    ffa = session.write(ffa, (out) -> out.write('A'));
    session.transfer(ffa, relationship);

    FlowFile ffb = session.create();
    ffb = session.write(ffb, (out) -> out.write('B'));
    session.transfer(ffb, relationship);
    session.commit();

    final ProcessSession newSession = new StandardProcessSession(context, () -> false);
    FlowFile toUpdate = newSession.get();
    newSession.append(toUpdate, out -> out.write('C'));

    // Read the content back and ensure that it is correct
    final byte[] buff;
    try (final ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
        newSession.read(toUpdate, in -> StreamUtils.copy(in, baos));
        buff = baos.toByteArray();
    }

    final String output = new String(buff, StandardCharsets.UTF_8);
    assertEquals("AC", output);
    newSession.transfer(toUpdate);
    newSession.commit();
}
 

private void writeDemarcatedData(final ProcessSession session, final List<ConsumerRecord<byte[], byte[]>> records, final TopicPartition topicPartition) {
    final ConsumerRecord<byte[], byte[]> firstRecord = records.get(0);
    final boolean demarcateFirstRecord;
    final BundleInformation bundleInfo = new BundleInformation(topicPartition, null);
    BundleTracker tracker = bundleMap.get(bundleInfo);
    FlowFile flowFile;
    if (tracker == null) {
        tracker = new BundleTracker(firstRecord, topicPartition, keyEncoding);
        flowFile = session.create();
        tracker.updateFlowFile(flowFile);
        demarcateFirstRecord = false; //have not yet written records for this topic/partition in this lease
    } else {
        demarcateFirstRecord = true; //have already been writing records for this topic/partition in this lease
    }
    flowFile = tracker.flowFile;

    tracker.incrementRecordCount(records.size());
    flowFile = session.append(flowFile, out -> {
        boolean useDemarcator = demarcateFirstRecord;
        for (final ConsumerRecord<byte[], byte[]> record : records) {
            if (useDemarcator) {
                out.write(demarcatorBytes);
            }
            final byte[] value = record.value();
            if (value != null) {
                out.write(record.value());
            }
            useDemarcator = true;
        }
    });
    tracker.updateFlowFile(flowFile);
    bundleMap.put(bundleInfo, tracker);
}
 

private void writeData(final ProcessSession session, final List<ConsumerRecord<byte[], byte[]>> records, final TopicPartition topicPartition) {
    final ConsumerRecord<byte[], byte[]> firstRecord = records.get(0);
    final boolean demarcateFirstRecord;
    BundleTracker tracker = bundleMap.get(topicPartition);
    FlowFile flowFile;
    if (tracker == null) {
        tracker = new BundleTracker(firstRecord, topicPartition, keyEncoding);
        flowFile = session.create();
        tracker.updateFlowFile(flowFile);
        demarcateFirstRecord = false; //have not yet written records for this topic/partition in this lease
    } else {
        demarcateFirstRecord = true; //have already been writing records for this topic/partition in this lease
    }
    flowFile = tracker.flowFile;
    tracker.incrementRecordCount(records.size());
    flowFile = session.append(flowFile, out -> {
        boolean useDemarcator = demarcateFirstRecord;
        for (final ConsumerRecord<byte[], byte[]> record : records) {
            if (useDemarcator) {
                out.write(demarcatorBytes);
            }
            final byte[] value = record.value();
            if (value != null) {
                out.write(record.value());
            }
            useDemarcator = true;
        }
    });
    tracker.updateFlowFile(flowFile);
    bundleMap.put(topicPartition, tracker);
}
 

/**
 * Batches together up to the batchSize events. Events are grouped together based on a batch key which
 * by default is the sender of the event, but can be override by sub-classes.
 *
 * This method will return when batchSize has been reached, or when no more events are available on the queue.
 *
 * @param session the current session
 * @param totalBatchSize the total number of events to process
 * @param messageDemarcatorBytes the demarcator to put between messages when writing to a FlowFile
 *
 * @return a Map from the batch key to the FlowFile and events for that batch, the size of events in all
 *              the batches will be <= batchSize
 */
protected Map<String,FlowFileEventBatch> getBatches(final ProcessSession session, final int totalBatchSize,
                                                    final byte[] messageDemarcatorBytes) {

    final Map<String,FlowFileEventBatch> batches = new HashMap<>();
    for (int i=0; i < totalBatchSize; i++) {
        final E event = getMessage(true, true, session);
        if (event == null) {
            break;
        }

        final String batchKey = getBatchKey(event);
        FlowFileEventBatch batch = batches.get(batchKey);

        // if we don't have a batch for this key then create a new one
        if (batch == null) {
            batch = new FlowFileEventBatch(session.create(), new ArrayList<E>());
            batches.put(batchKey, batch);
        }

        // add the current event to the batch
        batch.getEvents().add(event);

        // append the event's data to the FlowFile, write the demarcator first if not on the first event
        final boolean writeDemarcator = (i > 0);
        try {
            final byte[] rawMessage = event.getData();
            FlowFile appendedFlowFile = session.append(batch.getFlowFile(), new OutputStreamCallback() {
                @Override
                public void process(final OutputStream out) throws IOException {
                    if (writeDemarcator) {
                        out.write(messageDemarcatorBytes);
                    }

                    out.write(rawMessage);
                }
            });

            // update the FlowFile reference in the batch object
            batch.setFlowFile(appendedFlowFile);

        } catch (final Exception e) {
            getLogger().error("Failed to write contents of the message to FlowFile due to {}; will re-queue message and try again",
                    new Object[] {e.getMessage()}, e);
            errorEvents.offer(event);
            break;
        }
    }

    return batches;
}
 

/**
 * Batches together up to the batchSize events. Events are grouped together based on a batch key which
 * by default is the sender of the event, but can be override by sub-classes.
 *
 * This method will return when batchSize has been reached, or when no more events are available on the queue.
 *
 * @param session the current session
 * @param totalBatchSize the total number of events to process
 * @param messageDemarcatorBytes the demarcator to put between messages when writing to a FlowFile
 *
 * @return a Map from the batch key to the FlowFile and events for that batch, the size of events in all
 *              the batches will be <= batchSize
 */
protected Map<String,FlowFileEventBatch> getBatches(final ProcessSession session, final int totalBatchSize,
                                                    final byte[] messageDemarcatorBytes) {

    final Map<String,FlowFileEventBatch> batches = new HashMap<>();
    for (int i=0; i < totalBatchSize; i++) {
        final E event = getMessage(true, true, session);
        if (event == null) {
            break;
        }

        final String batchKey = getBatchKey(event);
        FlowFileEventBatch batch = batches.get(batchKey);

        // if we don't have a batch for this key then create a new one
        if (batch == null) {
            batch = new FlowFileEventBatch(session.create(), new ArrayList<E>());
            batches.put(batchKey, batch);
        }

        // add the current event to the batch
        batch.getEvents().add(event);

        // append the event's data to the FlowFile, write the demarcator first if not on the first event
        final boolean writeDemarcator = (i > 0);
        try {
            final byte[] rawMessage = event.getData();
            FlowFile appendedFlowFile = session.append(batch.getFlowFile(), new OutputStreamCallback() {
                @Override
                public void process(final OutputStream out) throws IOException {
                    if (writeDemarcator) {
                        out.write(messageDemarcatorBytes);
                    }

                    out.write(rawMessage);
                }
            });

            // update the FlowFile reference in the batch object
            batch.setFlowFile(appendedFlowFile);

        } catch (final Exception e) {
            getLogger().error("Failed to write contents of the message to FlowFile due to {}; will re-queue message and try again",
                    new Object[] {e.getMessage()}, e);
            errorEvents.offer(event);
            break;
        }
    }

    return batches;
}
 

private void writeDemarcatedData(final ProcessSession session, final List<ConsumerRecord<byte[], byte[]>> records, final TopicPartition topicPartition) {
    // Group the Records by their BundleInformation
    final Map<BundleInformation, List<ConsumerRecord<byte[], byte[]>>> map = records.stream()
        .collect(Collectors.groupingBy(rec -> new BundleInformation(topicPartition, null, getAttributes(rec))));

    for (final Map.Entry<BundleInformation, List<ConsumerRecord<byte[], byte[]>>> entry : map.entrySet()) {
        final BundleInformation bundleInfo = entry.getKey();
        final List<ConsumerRecord<byte[], byte[]>> recordList = entry.getValue();

        final boolean demarcateFirstRecord;

        BundleTracker tracker = bundleMap.get(bundleInfo);

        FlowFile flowFile;
        if (tracker == null) {
            tracker = new BundleTracker(recordList.get(0), topicPartition, keyEncoding);
            flowFile = session.create();
            flowFile = session.putAllAttributes(flowFile, bundleInfo.attributes);
            tracker.updateFlowFile(flowFile);
            demarcateFirstRecord = false; //have not yet written records for this topic/partition in this lease
        } else {
            demarcateFirstRecord = true; //have already been writing records for this topic/partition in this lease
        }
        flowFile = tracker.flowFile;

        tracker.incrementRecordCount(recordList.size());
        flowFile = session.append(flowFile, out -> {
            boolean useDemarcator = demarcateFirstRecord;
            for (final ConsumerRecord<byte[], byte[]> record : recordList) {
                if (useDemarcator) {
                    out.write(demarcatorBytes);
                }
                final byte[] value = record.value();
                if (value != null) {
                    out.write(record.value());
                }
                useDemarcator = true;
            }
        });

        tracker.updateFlowFile(flowFile);
        bundleMap.put(bundleInfo, tracker);
    }
}
 

private void writeDemarcatedData(final ProcessSession session, final List<ConsumerRecord<byte[], byte[]>> records, final TopicPartition topicPartition) {
    // Group the Records by their BundleInformation
    final Map<BundleInformation, List<ConsumerRecord<byte[], byte[]>>> map = records.stream()
        .collect(Collectors.groupingBy(rec -> new BundleInformation(topicPartition, null, getAttributes(rec))));

    for (final Map.Entry<BundleInformation, List<ConsumerRecord<byte[], byte[]>>> entry : map.entrySet()) {
        final BundleInformation bundleInfo = entry.getKey();
        final List<ConsumerRecord<byte[], byte[]>> recordList = entry.getValue();

        final boolean demarcateFirstRecord;

        BundleTracker tracker = bundleMap.get(bundleInfo);

        FlowFile flowFile;
        if (tracker == null) {
            tracker = new BundleTracker(recordList.get(0), topicPartition, keyEncoding);
            flowFile = session.create();
            flowFile = session.putAllAttributes(flowFile, bundleInfo.attributes);
            tracker.updateFlowFile(flowFile);
            demarcateFirstRecord = false; //have not yet written records for this topic/partition in this lease
        } else {
            demarcateFirstRecord = true; //have already been writing records for this topic/partition in this lease
        }
        flowFile = tracker.flowFile;

        tracker.incrementRecordCount(recordList.size());
        flowFile = session.append(flowFile, out -> {
            boolean useDemarcator = demarcateFirstRecord;
            for (final ConsumerRecord<byte[], byte[]> record : recordList) {
                if (useDemarcator) {
                    out.write(demarcatorBytes);
                }
                final byte[] value = record.value();
                if (value != null) {
                    out.write(record.value());
                }
                useDemarcator = true;
            }
        });

        tracker.updateFlowFile(flowFile);
        bundleMap.put(bundleInfo, tracker);
    }
}
 

private void writeDemarcatedData(final ProcessSession session, final List<ConsumerRecord<byte[], byte[]>> records, final TopicPartition topicPartition) {
    // Group the Records by their BundleInformation
    final Map<BundleInformation, List<ConsumerRecord<byte[], byte[]>>> map = records.stream()
        .collect(Collectors.groupingBy(rec -> new BundleInformation(topicPartition, null, getAttributes(rec))));

    for (final Map.Entry<BundleInformation, List<ConsumerRecord<byte[], byte[]>>> entry : map.entrySet()) {
        final BundleInformation bundleInfo = entry.getKey();
        final List<ConsumerRecord<byte[], byte[]>> recordList = entry.getValue();

        final boolean demarcateFirstRecord;

        BundleTracker tracker = bundleMap.get(bundleInfo);

        FlowFile flowFile;
        if (tracker == null) {
            tracker = new BundleTracker(recordList.get(0), topicPartition, keyEncoding);
            flowFile = session.create();
            flowFile = session.putAllAttributes(flowFile, bundleInfo.attributes);
            tracker.updateFlowFile(flowFile);
            demarcateFirstRecord = false; //have not yet written records for this topic/partition in this lease
        } else {
            demarcateFirstRecord = true; //have already been writing records for this topic/partition in this lease
        }
        flowFile = tracker.flowFile;

        tracker.incrementRecordCount(recordList.size());
        flowFile = session.append(flowFile, out -> {
            boolean useDemarcator = demarcateFirstRecord;
            for (final ConsumerRecord<byte[], byte[]> record : recordList) {
                if (useDemarcator) {
                    out.write(demarcatorBytes);
                }
                final byte[] value = record.value();
                if (value != null) {
                    out.write(record.value());
                }
                useDemarcator = true;
            }
        });

        tracker.updateFlowFile(flowFile);
        bundleMap.put(bundleInfo, tracker);
    }
}