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

The following examples show how to use org.apache.nifi.processor.ProcessSession#create() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: PutSlackTest.java    From nifi with Apache License 2.0 6 votes vote down vote up
@Test
public void testInvalidDynamicPropertiesWithExpressionLanguage() {
    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();
    Map<String, String> props = new HashMap<>();
    props.put("foo", "\"\"bar\"");
    props.put("ping", "\"pong");
    ff = session.putAllAttributes(ff, props);

    testRunner.setProperty(PutSlack.WEBHOOK_URL, server.getUrl());
    testRunner.setProperty(PutSlack.WEBHOOK_TEXT, WEBHOOK_TEST_TEXT);
    PropertyDescriptor dynamicProp = new PropertyDescriptor.Builder()
            .dynamic(true)
            .name("foo")
            .build();
    testRunner.setProperty(dynamicProp, "{\"foo\": ${foo}, \"ping\":\"${ping}\"}");

    testRunner.enqueue(ff);
    testRunner.run(1);
    testRunner.assertTransferCount(PutSlack.REL_SUCCESS, 0);
    testRunner.assertTransferCount(PutSlack.REL_FAILURE, 1);
}
 
Example 2
Source File: DataGeneratorTestProcessor.java    From localization_nifi with Apache License 2.0 6 votes vote down vote up
@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();
}
 
Example 3
Source File: TestSplitJson.java    From localization_nifi with Apache License 2.0 6 votes vote down vote up
@Test
public void testSplit_pathToNullValue() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new SplitJson());
    testRunner.setProperty(SplitJson.ARRAY_JSON_PATH_EXPRESSION, "$.nullField");

    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();

    ff = session.write(ff, new OutputStreamCallback() {
        @Override
        public void process(OutputStream out) throws IOException {
            try (OutputStream outputStream = new BufferedOutputStream(out)) {
                outputStream.write("{\"stringField\": \"String Value\", \"nullField\": null}".getBytes(StandardCharsets.UTF_8));
            }
        }
    });

    testRunner.enqueue(ff);
    testRunner.run();

    testRunner.assertTransferCount(SplitJson.REL_FAILURE, 1);
}
 
Example 4
Source File: SpringContextProcessor.java    From nifi with Apache License 2.0 6 votes vote down vote up
/**
 *
 */
private void receiveFromSpring(ProcessSession processSession) {
    final SpringResponse<?> msgFromSpring = this.exchanger.receive(this.receiveTimeout);
    if (msgFromSpring != null) {
        FlowFile flowFileToProcess = processSession.create();
        flowFileToProcess = processSession.write(flowFileToProcess, new OutputStreamCallback() {
            @Override
            public void process(final OutputStream out) throws IOException {
                Object payload = msgFromSpring.getPayload();
                byte[] payloadBytes = payload instanceof String ? ((String) payload).getBytes() : (byte[]) payload;
                out.write(payloadBytes);
            }
        });
        flowFileToProcess = processSession.putAllAttributes(flowFileToProcess,
                this.extractFlowFileAttributesFromMessageHeaders(msgFromSpring.getHeaders()));
        processSession.transfer(flowFileToProcess, REL_SUCCESS);
        processSession.getProvenanceReporter().receive(flowFileToProcess, this.applicationContextConfigFileName);
    }
}
 
Example 5
Source File: TestAttributesToJSON.java    From localization_nifi with Apache License 2.0 6 votes vote down vote up
@Test
public void testAttribute_singleUserDefinedAttributeWithWhiteSpace() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new AttributesToJSON());
    testRunner.setProperty(AttributesToJSON.ATTRIBUTES_LIST, " " + TEST_ATTRIBUTE_KEY + " ");
    testRunner.setProperty(AttributesToJSON.DESTINATION, AttributesToJSON.DESTINATION_ATTRIBUTE);

    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();
    ff = session.putAttribute(ff, TEST_ATTRIBUTE_KEY, TEST_ATTRIBUTE_VALUE);

    testRunner.enqueue(ff);
    testRunner.run();

    testRunner.getFlowFilesForRelationship(AttributesToJSON.REL_SUCCESS).get(0)
            .assertAttributeExists(AttributesToJSON.JSON_ATTRIBUTE_NAME);
    testRunner.assertTransferCount(AttributesToJSON.REL_SUCCESS, 1);
    testRunner.assertTransferCount(AttributesToJSON.REL_FAILURE, 0);

    String json = testRunner.getFlowFilesForRelationship(AttributesToJSON.REL_SUCCESS)
            .get(0).getAttribute(AttributesToJSON.JSON_ATTRIBUTE_NAME);

    ObjectMapper mapper = new ObjectMapper();
    Map<String, String> val = mapper.readValue(json, HashMap.class);
    assertTrue(val.get(TEST_ATTRIBUTE_KEY).equals(TEST_ATTRIBUTE_VALUE));
    assertTrue(val.size() == 1);
}
 
Example 6
Source File: TestStandardProcessorTestRunner.java    From localization_nifi with Apache License 2.0 5 votes vote down vote up
@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile ff = session.create();
    if (counter % 2 == 0) {
        ff = session.putAttribute(ff, KEY, "value");
        session.transfer(ff, REL_SUCCESS);
    } else {
        session.transfer(ff, REL_FAILURE);
    }
    counter++;
}
 
Example 7
Source File: TestAttributesToJSON.java    From localization_nifi with Apache License 2.0 5 votes vote down vote up
@Test
public void testEmptyStringValueForEmptyAttribute() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new AttributesToJSON());
    testRunner.setProperty(AttributesToJSON.DESTINATION, AttributesToJSON.DESTINATION_ATTRIBUTE);
    final String NON_PRESENT_ATTRIBUTE_KEY = "NonExistingAttributeKey";
    testRunner.setProperty(AttributesToJSON.ATTRIBUTES_LIST, NON_PRESENT_ATTRIBUTE_KEY);
    testRunner.setProperty(AttributesToJSON.NULL_VALUE_FOR_EMPTY_STRING, "false");

    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();

    testRunner.enqueue(ff);
    testRunner.run();

    //Expecting success transition because Jackson is taking care of escaping the bad JSON characters
    testRunner.getFlowFilesForRelationship(AttributesToJSON.REL_SUCCESS).get(0)
            .assertAttributeExists(AttributesToJSON.JSON_ATTRIBUTE_NAME);
    testRunner.assertTransferCount(AttributesToJSON.REL_SUCCESS, 1);
    testRunner.assertTransferCount(AttributesToJSON.REL_FAILURE, 0);

    //Make sure that the value is a true JSON null for the non existing attribute
    String json = testRunner.getFlowFilesForRelationship(AttributesToJSON.REL_SUCCESS)
            .get(0).getAttribute(AttributesToJSON.JSON_ATTRIBUTE_NAME);

    ObjectMapper mapper = new ObjectMapper();
    Map<String, String> val = mapper.readValue(json, HashMap.class);

    assertEquals(val.get(NON_PRESENT_ATTRIBUTE_KEY), "");
}
 
Example 8
Source File: AbstractBinlogEventWriter.java    From nifi with Apache License 2.0 5 votes vote down vote up
@Override
public long writeEvent(ProcessSession session, String transitUri, T eventInfo, long currentSequenceId, Relationship relationship) {
    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, (outputStream) -> {
        super.startJson(outputStream, eventInfo);
        writeJson(eventInfo);
        // Nothing in the body
        super.endJson();
    });
    flowFile = session.putAllAttributes(flowFile, getCommonAttributes(currentSequenceId, eventInfo));
    session.transfer(flowFile, relationship);
    session.getProvenanceReporter().receive(flowFile, transitUri);
    return currentSequenceId + 1;
}
 
Example 9
Source File: TestAttributesToJSON.java    From localization_nifi with Apache License 2.0 5 votes vote down vote up
@Test(expected = AssertionError.class)
public void testInvalidIncludeCoreAttributesProperty() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new AttributesToJSON());
    testRunner.setProperty(AttributesToJSON.ATTRIBUTES_LIST, "val1,val2");
    testRunner.setProperty(AttributesToJSON.DESTINATION, AttributesToJSON.DESTINATION_ATTRIBUTE);
    testRunner.setProperty(AttributesToJSON.INCLUDE_CORE_ATTRIBUTES, "maybe");

    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();

    testRunner.enqueue(ff);
    testRunner.run();
}
 
Example 10
Source File: TestEvaluateJsonPath.java    From nifi with Apache License 2.0 5 votes vote down vote up
@Test
public void testNullInput_nullStringRepresentation() throws Exception {
    final TestRunner testRunner = TestRunners.newTestRunner(new EvaluateJsonPath());
    testRunner.setProperty(EvaluateJsonPath.RETURN_TYPE, EvaluateJsonPath.RETURN_TYPE_JSON);
    testRunner.setProperty(EvaluateJsonPath.DESTINATION, EvaluateJsonPath.DESTINATION_ATTRIBUTE);
    testRunner.setProperty(EvaluateJsonPath.NULL_VALUE_DEFAULT_REPRESENTATION, AbstractJsonPathProcessor.NULL_STRING_OPTION);
    testRunner.setProperty("stringField", "$.stringField");
    testRunner.setProperty("missingField", "$.missingField");
    testRunner.setProperty("nullField", "$.nullField");

    ProcessSession session = testRunner.getProcessSessionFactory().createSession();
    FlowFile ff = session.create();

    ff = session.write(ff, new OutputStreamCallback() {
        @Override
        public void process(OutputStream out) throws IOException {
            try (OutputStream outputStream = new BufferedOutputStream(out)) {
                outputStream.write("{\"stringField\": \"String Value\", \"nullField\": null}".getBytes(StandardCharsets.UTF_8));
            }
        }
    });

    testRunner.enqueue(ff);
    testRunner.run();

    testRunner.assertTransferCount(EvaluateJsonPath.REL_MATCH, 1);

    FlowFile output = testRunner.getFlowFilesForRelationship(EvaluateJsonPath.REL_MATCH).get(0);

    String validFieldValue = output.getAttribute("stringField");
    assertEquals("String Value", validFieldValue);

    String missingValue = output.getAttribute("missingField");
    assertEquals("Missing Value", "", missingValue);

    String nullValue = output.getAttribute("nullField");
    assertEquals("Null Value", "null", nullValue);
}
 
Example 11
Source File: AbstractBinlogTableEventWriter.java    From nifi with Apache License 2.0 5 votes vote down vote up
@Override
public long writeEvent(ProcessSession session, String transitUri, T eventInfo, long currentSequenceId, Relationship relationship) {
    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, (outputStream) -> {
        super.startJson(outputStream, eventInfo);
        writeJson(eventInfo);
        // Nothing in the body
        super.endJson();
    });
    flowFile = session.putAllAttributes(flowFile, getCommonAttributes(currentSequenceId, eventInfo));
    session.transfer(flowFile, relationship);
    session.getProvenanceReporter().receive(flowFile, transitUri);
    return currentSequenceId + 1;
}
 
Example 12
Source File: GenerateProcessor.java    From nifi with Apache License 2.0 5 votes vote down vote up
@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    for (int i=0; i < context.getProperty(COUNT).asInteger(); i++) {
        FlowFile flowFile = session.create();

        final int size = context.getProperty(CONTENT_SIZE).asDataSize(DataUnit.B).intValue();
        final byte[] data = new byte[size];
        session.write(flowFile, out -> out.write(data));
        session.transfer(flowFile, REL_SUCCESS);
    }
}
 
Example 13
Source File: ConsumerLease.java    From nifi with Apache License 2.0 5 votes vote down vote up
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);
        });
    }
    tracker.updateFlowFile(flowFile);
    populateAttributes(tracker);
    session.transfer(tracker.flowFile, REL_SUCCESS);
}
 
Example 14
Source File: ConsumerLease.java    From nifi with Apache License 2.0 5 votes vote down vote up
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 15
Source File: ConsumerLease.java    From nifi with Apache License 2.0 4 votes vote down vote up
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);
    }
}
 
Example 16
Source File: GetHDFS.java    From nifi with Apache License 2.0 4 votes vote down vote up
protected void processBatchOfFiles(final List<Path> files, final ProcessContext context, final ProcessSession session) {
    // process the batch of files
    InputStream stream = null;
    CompressionCodec codec = null;
    Configuration conf = getConfiguration();
    FileSystem hdfs = getFileSystem();
    final boolean keepSourceFiles = context.getProperty(KEEP_SOURCE_FILE).asBoolean();
    final Double bufferSizeProp = context.getProperty(BUFFER_SIZE).asDataSize(DataUnit.B);
    int bufferSize = bufferSizeProp != null ? bufferSizeProp.intValue() : conf.getInt(BUFFER_SIZE_KEY,
            BUFFER_SIZE_DEFAULT);
    final Path rootDir = new Path(context.getProperty(DIRECTORY).evaluateAttributeExpressions().getValue());

    final CompressionType compressionType = CompressionType.valueOf(context.getProperty(COMPRESSION_CODEC).toString());
    final boolean inferCompressionCodec = compressionType == CompressionType.AUTOMATIC;
    if (inferCompressionCodec || compressionType != CompressionType.NONE) {
        codec = getCompressionCodec(context, getConfiguration());
    }
    final CompressionCodecFactory compressionCodecFactory = new CompressionCodecFactory(conf);
    for (final Path file : files) {
        try {
            if (!getUserGroupInformation().doAs((PrivilegedExceptionAction<Boolean>) () -> hdfs.exists(file))) {
                continue; // if file is no longer there then move on
            }
            final String originalFilename = file.getName();
            final String relativePath = getPathDifference(rootDir, file);

            stream = getUserGroupInformation().doAs((PrivilegedExceptionAction<FSDataInputStream>) () -> hdfs.open(file, bufferSize));

            final String outputFilename;
            // Check if we should infer compression codec
            if (inferCompressionCodec) {
                codec = compressionCodecFactory.getCodec(file);
            }
            // Check if compression codec is defined (inferred or otherwise)
            if (codec != null) {
                stream = codec.createInputStream(stream);
                outputFilename = StringUtils.removeEnd(originalFilename, codec.getDefaultExtension());
            } else {
                outputFilename = originalFilename;
            }

            FlowFile flowFile = session.create();

            final StopWatch stopWatch = new StopWatch(true);
            flowFile = session.importFrom(stream, flowFile);
            stopWatch.stop();
            final String dataRate = stopWatch.calculateDataRate(flowFile.getSize());
            final long millis = stopWatch.getDuration(TimeUnit.MILLISECONDS);

            flowFile = session.putAttribute(flowFile, CoreAttributes.PATH.key(), relativePath.isEmpty() ? "." : relativePath);
            flowFile = session.putAttribute(flowFile, CoreAttributes.FILENAME.key(), outputFilename);

            if (!keepSourceFiles && !getUserGroupInformation().doAs((PrivilegedExceptionAction<Boolean>) () -> hdfs.delete(file, false))) {
                getLogger().warn("Could not remove {} from HDFS. Not ingesting this file ...",
                        new Object[]{file});
                session.remove(flowFile);
                continue;
            }

            session.getProvenanceReporter().receive(flowFile, file.toString());
            session.transfer(flowFile, REL_SUCCESS);
            getLogger().info("retrieved {} from HDFS {} in {} milliseconds at a rate of {}",
                    new Object[]{flowFile, file, millis, dataRate});
            session.commit();
        } catch (final Throwable t) {
            getLogger().error("Error retrieving file {} from HDFS due to {}", new Object[]{file, t});
            session.rollback();
            context.yield();
        } finally {
            IOUtils.closeQuietly(stream);
            stream = null;
        }
    }
}
 
Example 17
Source File: GetMongo.java    From localization_nifi with Apache License 2.0 4 votes vote down vote up
@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();

    final Document query = context.getProperty(QUERY).isSet() ? Document.parse(context.getProperty(QUERY).getValue()) : null;
    final Document projection = context.getProperty(PROJECTION).isSet() ? Document.parse(context.getProperty(PROJECTION).getValue()) : null;
    final Document sort = context.getProperty(SORT).isSet() ? Document.parse(context.getProperty(SORT).getValue()) : null;

    final MongoCollection<Document> collection = getCollection(context);

    try {
        final FindIterable<Document> it = query != null ? collection.find(query) : collection.find();
        if (projection != null) {
            it.projection(projection);
        }
        if (sort != null) {
            it.sort(sort);
        }
        if (context.getProperty(LIMIT).isSet()) {
            it.limit(context.getProperty(LIMIT).asInteger());
        }
        if (context.getProperty(BATCH_SIZE).isSet()) {
            it.batchSize(context.getProperty(BATCH_SIZE).asInteger());
        }

        final MongoCursor<Document> cursor = it.iterator();
        try {
            FlowFile flowFile = null;
            while (cursor.hasNext()) {
                flowFile = session.create();
                flowFile = session.write(flowFile, new OutputStreamCallback() {
                    @Override
                    public void process(OutputStream out) throws IOException {
                        IOUtils.write(cursor.next().toJson(), out);
                    }
                });

                session.getProvenanceReporter().receive(flowFile, context.getProperty(URI).getValue());
                session.transfer(flowFile, REL_SUCCESS);
            }

            session.commit();

        } finally {
            cursor.close();
        }

    } catch (final RuntimeException e) {
        context.yield();
        session.rollback();
        logger.error("Failed to execute query {} due to {}", new Object[] { query, e }, e);
    }
}
 
Example 18
Source File: ConsumeGCPubSub.java    From nifi with Apache License 2.0 4 votes vote down vote up
@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    if (subscriber == null) {

        if (storedException.get() != null) {
            getLogger().error("Failed to create Google Cloud PubSub subscriber due to {}", new Object[]{storedException.get()});
        } else {
            getLogger().error("Google Cloud PubSub Subscriber was not properly created. Yielding the processor...");
        }

        context.yield();
        return;
    }

    final PullResponse pullResponse = subscriber.pullCallable().call(pullRequest);
    final List<String> ackIds = new ArrayList<>();

    for (ReceivedMessage message : pullResponse.getReceivedMessagesList()) {
        if (message.hasMessage()) {
            FlowFile flowFile = session.create();

            final Map<String, String> attributes = new HashMap<>();
            ackIds.add(message.getAckId());

            attributes.put(ACK_ID_ATTRIBUTE, message.getAckId());
            attributes.put(SERIALIZED_SIZE_ATTRIBUTE, String.valueOf(message.getSerializedSize()));
            attributes.put(MESSAGE_ID_ATTRIBUTE, message.getMessage().getMessageId());
            attributes.put(MSG_ATTRIBUTES_COUNT_ATTRIBUTE, String.valueOf(message.getMessage().getAttributesCount()));
            attributes.put(MSG_PUBLISH_TIME_ATTRIBUTE, String.valueOf(message.getMessage().getPublishTime().getSeconds()));
            attributes.putAll(message.getMessage().getAttributesMap());

            flowFile = session.putAllAttributes(flowFile, attributes);
            flowFile = session.write(flowFile, out -> out.write(message.getMessage().getData().toByteArray()));

            session.transfer(flowFile, REL_SUCCESS);
            session.getProvenanceReporter().receive(flowFile, getSubscriptionName(context));
        }
    }

    if (!ackIds.isEmpty()) {
        AcknowledgeRequest acknowledgeRequest = AcknowledgeRequest.newBuilder()
                .addAllAckIds(ackIds)
                .setSubscription(getSubscriptionName(context))
                .build();
        subscriber.acknowledgeCallable().call(acknowledgeRequest);
    }
}
 
Example 19
Source File: MergeContent.java    From nifi with Apache License 2.0 4 votes vote down vote up
@Override
public FlowFile merge(final Bin bin, final ProcessContext context) {
    final ProcessSession session = bin.getSession();
    final List<FlowFile> contents = bin.getContents();

    FlowFile bundle = session.create(contents);

    try {
        bundle = session.write(bundle, new OutputStreamCallback() {
            @Override
            public void process(final OutputStream rawOut) throws IOException {
                try (final OutputStream bufferedOut = new BufferedOutputStream(rawOut)) {
                    // we don't want the packager closing the stream. V1 creates a TAR Output Stream, which then gets
                    // closed, which in turn closes the underlying OutputStream, and we want to protect ourselves against that.
                    final OutputStream out = new NonCloseableOutputStream(bufferedOut);

                    for (final FlowFile flowFile : contents) {
                        bin.getSession().read(flowFile, false, new InputStreamCallback() {
                            @Override
                            public void process(final InputStream rawIn) throws IOException {
                                try (final InputStream in = new BufferedInputStream(rawIn)) {
                                    final Map<String, String> attributes = new HashMap<>(flowFile.getAttributes());

                                    // for backward compatibility purposes, we add the "legacy" NiFi attributes
                                    attributes.put("nf.file.name", attributes.get(CoreAttributes.FILENAME.key()));
                                    attributes.put("nf.file.path", attributes.get(CoreAttributes.PATH.key()));
                                    if (attributes.containsKey(CoreAttributes.MIME_TYPE.key())) {
                                        attributes.put("content-type", attributes.get(CoreAttributes.MIME_TYPE.key()));
                                    }
                                    packager.packageFlowFile(in, out, attributes, flowFile.getSize());
                                }
                            }
                        });
                    }
                }
            }
        });
    } catch (final Exception e) {
        session.remove(bundle);
        throw e;
    }

    bundle = session.putAttribute(bundle, CoreAttributes.FILENAME.key(), createFilename(contents) + ".pkg");
    session.getProvenanceReporter().join(contents, bundle);
    return bundle;
}
 
Example 20
Source File: GetHTMLElement.java    From localization_nifi with Apache License 2.0 4 votes vote down vote up
@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final FlowFile flowFile = session.get();
    if ( flowFile == null ) {
        return;
    }

    final Document doc;
    final Elements eles;

    try {
        doc = parseHTMLDocumentFromFlowfile(flowFile, context, session);
        eles = doc.select(context.getProperty(CSS_SELECTOR).evaluateAttributeExpressions(flowFile).getValue());
    } catch (final Exception ex) {
        getLogger().error("Failed to extract HTML from {} due to {}; routing to {}", new Object[] {flowFile, ex, REL_INVALID_HTML}, ex);
        session.transfer(flowFile, REL_INVALID_HTML);
        return;
    }

    final String prependValue = context.getProperty(PREPEND_ELEMENT_VALUE).evaluateAttributeExpressions(flowFile).getValue();
    final String appendValue = context.getProperty(APPEND_ELEMENT_VALUE).evaluateAttributeExpressions(flowFile).getValue();
    final String outputType = context.getProperty(OUTPUT_TYPE).getValue();
    final String attributeKey = context.getProperty(ATTRIBUTE_KEY).evaluateAttributeExpressions(flowFile).getValue();

    if (eles == null || eles.isEmpty()) {
        // No element found
        session.transfer(flowFile, REL_NOT_FOUND);
    } else {
        // Create a new FlowFile for each matching element.
        for (final Element ele : eles) {
            final String extractedElementValue = extractElementValue(prependValue, outputType, appendValue, ele, attributeKey);

            final FlowFile ff = session.create(flowFile);
            FlowFile updatedFF = ff;

            switch (context.getProperty(DESTINATION).getValue()) {
                case DESTINATION_ATTRIBUTE:
                    updatedFF = session.putAttribute(ff, HTML_ELEMENT_ATTRIBUTE_NAME, extractedElementValue);
                    break;
                case DESTINATION_CONTENT:
                    updatedFF = session.write(ff, new StreamCallback() {
                        @Override
                        public void process(final InputStream inputStream, final OutputStream outputStream) throws IOException {
                            outputStream.write(extractedElementValue.getBytes(StandardCharsets.UTF_8));
                        }
                    });

                    break;
            }

            session.transfer(updatedFF, REL_SUCCESS);
        }

        // Transfer the original HTML
        session.transfer(flowFile, REL_ORIGINAL);
    }
}