org.apache.nifi.processor.io.InputStreamCallback Java Examples

/**
 * Will rendezvous with {@link KafkaPublisher} after building
 * {@link PublishingContext} and will produce the resulting {@link FlowFile}.
 * The resulting FlowFile contains all required information to determine
 * if message publishing originated from the provided FlowFile has actually
 * succeeded fully, partially or failed completely (see
 * {@link #isFailedFlowFile(FlowFile)}.
 */
private FlowFile doRendezvousWithKafka(final FlowFile flowFile, final ProcessContext context, final ProcessSession session) {
    final AtomicReference<KafkaPublisherResult> publishResultRef = new AtomicReference<>();
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(InputStream contentStream) throws IOException {
            PublishingContext publishingContext = PutKafka.this.buildPublishingContext(flowFile, context, contentStream);
            KafkaPublisherResult result = null;
            try {
                result = PutKafka.this.kafkaResource.publish(publishingContext);
            } catch (final IllegalArgumentException e) {
                getLogger().error("Failed to publish {}, due to {}", new Object[]{flowFile, e}, e);
                result = new KafkaPublisherResult(0, -1);

            }
            publishResultRef.set(result);
        }
    });

    FlowFile resultFile = publishResultRef.get().isAllAcked()
            ? this.cleanUpFlowFileIfNecessary(flowFile, session)
            : session.putAllAttributes(flowFile, this.buildFailedFlowFileAttributes(publishResultRef.get().getLastMessageAcked(), flowFile, context));

    return resultFile;
}
 

static DocumentContext validateAndEstablishJsonContext(ProcessSession processSession, FlowFile flowFile) {
    // Parse the document once into an associated context to support multiple path evaluations if specified
    final AtomicReference<DocumentContext> contextHolder = new AtomicReference<>(null);
    processSession.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(InputStream in) throws IOException {
            try (BufferedInputStream bufferedInputStream = new BufferedInputStream(in)) {
                DocumentContext ctx = JsonPath.using(STRICT_PROVIDER_CONFIGURATION).parse(bufferedInputStream);
                contextHolder.set(ctx);
            } catch (IllegalArgumentException iae) {
                // The JsonPath.parse() above first parses the json, then creates a context object from the parsed
                // json.  It is possible for the json parsing to complete without error, but produce a null object.
                // In this case the context creation will fail and throw an IllegalArgumentException.  This is in
                // my opinion a bug in the JsonPath library, as it doesn't really throw the correct exception
                // contextually.
                // The general handling in derived classes handles InvalidJsonException.
                // The best thing to do here, is to re-throw with the proper exception, such that the calling logic
                // can route.
                throw new InvalidJsonException(iae);
            }
        }
    });

    return contextHolder.get();
}
 

@Test
public void testReadAfterSessionClosesStream() throws IOException {
    final ContentClaim claim = contentRepo.create(false);
    final FlowFileRecord flowFileRecord = new StandardFlowFileRecord.Builder()
            .contentClaim(claim)
            .addAttribute("uuid", "12345678-1234-1234-1234-123456789012")
            .entryDate(System.currentTimeMillis())
            .build();
    flowFileQueue.put(flowFileRecord);
    final FlowFile flowFile = session.get();
    assertNotNull(flowFile);
    final AtomicReference<InputStream> inputStreamHolder = new AtomicReference<>(null);
    session.read(flowFile, true, new InputStreamCallback() {
        @Override
        public void process(final InputStream inputStream) throws IOException {
            inputStreamHolder.set(inputStream);
        }
    });
    assertDisabled(inputStreamHolder.get());
}
 

@Override
protected PutFlowFile createPut(final ProcessSession session, final ProcessContext context, final FlowFile flowFile) {
    final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions(flowFile).getValue();
    final String row = context.getProperty(ROW_ID).evaluateAttributeExpressions(flowFile).getValue();
    final String columnFamily = context.getProperty(COLUMN_FAMILY).evaluateAttributeExpressions(flowFile).getValue();
    final String columnQualifier = context.getProperty(COLUMN_QUALIFIER).evaluateAttributeExpressions(flowFile).getValue();

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


    final Collection<PutColumn> columns = Collections.singletonList(new PutColumn(columnFamily.getBytes(StandardCharsets.UTF_8),
                                                                        columnQualifier.getBytes(StandardCharsets.UTF_8), buffer));
    byte[] rowKeyBytes = getRow(row,context.getProperty(ROW_ID_ENCODING_STRATEGY).getValue());


    return new PutFlowFile(tableName,rowKeyBytes , columns, flowFile);
}
 

@Override
public void read(FlowFile flowFile, boolean allowSessionStreamManagement, final InputStreamCallback callback) {
    if (callback == null || flowFile == null) {
        throw new IllegalArgumentException("argument cannot be null");
    }

    flowFile = validateState(flowFile);
    if (!(flowFile instanceof StatelessFlowFile)) {
        throw new IllegalArgumentException("Cannot export a flow file that I did not create");
    }

    //allowSessionStreamManagement not used...
    try {
        ((StatelessFlowFile) flowFile).materializeData();
        callback.process(((StatelessFlowFile) flowFile).getDataStream());
    } catch (final IOException e) {
        throw new ProcessException(e.toString(), e);
    }
}
 

/**
 * Parses the Jsoup HTML document from the FlowFile input content.
 *
 * @param inputFlowFile Input FlowFile containing the HTML
 * @param context ProcessContext
 * @param session ProcessSession
 *
 * @return Jsoup Document
 */
protected Document parseHTMLDocumentFromFlowfile(final FlowFile inputFlowFile, final ProcessContext context, final ProcessSession session) {
    final AtomicReference<Document> doc = new AtomicReference<>();
    session.read(inputFlowFile, new InputStreamCallback() {
        @Override
        public void process(InputStream inputStream) throws IOException {
            final String baseUrl = getBaseUrl(inputFlowFile, context);
            if (baseUrl == null || baseUrl.isEmpty()) {
                throw new RuntimeException("Base URL was empty.");
            }
            doc.set(Jsoup.parse(inputStream,
                    context.getProperty(HTML_CHARSET).getValue(),
                    baseUrl));
        }
    });
    return doc.get();
}
 

@Override
public void read(final FlowFile flowFile, boolean allowSessionStreamManagement, final InputStreamCallback callback) {
    if (callback == null || flowFile == null) {
        throw new IllegalArgumentException("argument cannot be null");
    }

    validateState(flowFile);
    if (!(flowFile instanceof MockFlowFile)) {
        throw new IllegalArgumentException("Cannot export a flow file that I did not create");
    }
    final MockFlowFile mock = (MockFlowFile) flowFile;

    final ByteArrayInputStream bais = new ByteArrayInputStream(mock.getData());
    recursionSet.add(flowFile);
    try {
        callback.process(bais);
        if(!allowSessionStreamManagement){
            bais.close();
        }
    } catch (final IOException e) {
        throw new ProcessException(e.toString(), e);
    } finally {
        recursionSet.remove(flowFile);
    }
}
 

/**
 * Infers the Avro schema from the input Flowfile content.
 *
 * @param inputFlowFile
 *  The original input FlowFile containing the JSON content as it entered this processor.
 *
 * @param context
 *  ProcessContext to pull processor configurations.
 *
 * @param session
 *  ProcessSession to transfer FlowFiles
 */
private String inferAvroSchemaFromJSON(final FlowFile inputFlowFile, final ProcessContext context, final ProcessSession session) {

    final AtomicReference<String> avroSchema = new AtomicReference<>();
    session.read(inputFlowFile, new InputStreamCallback() {
        @Override
        public void process(InputStream in) throws IOException {
            Schema as = JsonUtil.inferSchema(
                    in, context.getProperty(RECORD_NAME).evaluateAttributeExpressions(inputFlowFile).getValue(),
                    context.getProperty(NUM_RECORDS_TO_ANALYZE).evaluateAttributeExpressions(inputFlowFile).asInteger());
            avroSchema.set(as.toString(context.getProperty(PRETTY_AVRO_OUTPUT).asBoolean()));

        }
    });

    return avroSchema.get();
}
 

/**
 * Determines the SQL statement that should be executed for the given FlowFile
 *
 * @param session the session that can be used to access the given FlowFile
 * @param flowFile the FlowFile whose SQL statement should be executed
 *
 * @return the SQL that is associated with the given FlowFile
 */
private String getSQL(final ProcessSession session, final FlowFile flowFile) {
    // Read the SQL from the FlowFile's content
    final byte[] buffer = new byte[(int) flowFile.getSize()];
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer);
        }
    });

    // Create the PreparedStatement to use for this FlowFile.
    final String sql = new String(buffer, StandardCharsets.UTF_8);
    return sql;
}
 

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final List<FlowFile> flowFiles = session.get(50);
    if (flowFiles.isEmpty()) {
        return;
    }

    final Schema schema = schemaRef.get();
    final Validator validator = schema.newValidator();
    final ComponentLog logger = getLogger();

    for (FlowFile flowFile : flowFiles) {
        final AtomicBoolean valid = new AtomicBoolean(true);
        final AtomicReference<Exception> exception = new AtomicReference<Exception>(null);

        session.read(flowFile, new InputStreamCallback() {
            @Override
            public void process(final InputStream in) throws IOException {
                try {
                    validator.validate(new StreamSource(in));
                } catch (final IllegalArgumentException | SAXException e) {
                    valid.set(false);
                    exception.set(e);
                }
            }
        });

        if (valid.get()) {
            logger.debug("Successfully validated {} against schema; routing to 'valid'", new Object[]{flowFile});
            session.getProvenanceReporter().route(flowFile, REL_VALID);
            session.transfer(flowFile, REL_VALID);
        } else {
            flowFile = session.putAttribute(flowFile, ERROR_ATTRIBUTE_KEY, exception.get().getLocalizedMessage());
            logger.info("Failed to validate {} against schema due to {}; routing to 'invalid'", new Object[]{flowFile, exception.get().getLocalizedMessage()});
            session.getProvenanceReporter().route(flowFile, REL_INVALID);
            session.transfer(flowFile, REL_INVALID);
        }
    }
}
 

@Test
public void testContentNotFoundExceptionThrownWhenUnableToReadDataOffsetTooLarge() {
    final FlowFileRecord flowFileRecord = new StandardFlowFileRecord.Builder()
            .addAttribute("uuid", "12345678-1234-1234-1234-123456789012")
            .entryDate(System.currentTimeMillis())
        .contentClaim(new StandardContentClaim(resourceClaimManager.newResourceClaim("x", "x", "0", true, false), 0L))
            .build();

    flowFileQueue.put(flowFileRecord);

    FlowFile ff1 = session.get();
    ff1 = session.write(ff1, new OutputStreamCallback() {
        @Override
        public void process(OutputStream out) throws IOException {
        }
    });
    session.transfer(ff1);
    session.commit();

    final FlowFileRecord flowFileRecord2 = new StandardFlowFileRecord.Builder()
            .addAttribute("uuid", "12345678-1234-1234-1234-123456789012")
            .entryDate(System.currentTimeMillis())
        .contentClaim(new StandardContentClaim(resourceClaimManager.newResourceClaim("x", "x", "0", true, false), 0L))
            .contentClaimOffset(1000L).size(1L).build();
    flowFileQueue.put(flowFileRecord2);

    // attempt to read the data.
    try {
        session.get();
        final FlowFile ff2 = session.get();
        session.read(ff2, new InputStreamCallback() {
            @Override
            public void process(InputStream in) throws IOException {
            }
        });
        Assert.fail("Expected MissingFlowFileException");
    } catch (final MissingFlowFileException mffe) {
    }
}
 

/**
 * Send the entire FlowFile as a single message.
 */
private void processSingleMessage(ProcessContext context, ProcessSession session, FlowFile flowFile, ChannelSender sender) {
    // copy the contents of the FlowFile to the ByteArrayOutputStream
    final ByteArrayOutputStream baos = new ByteArrayOutputStream((int)flowFile.getSize() + 1);
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.copy(in, baos);
        }
    });

    // if TCP and we don't end in a new line then add one
    final String protocol = context.getProperty(PROTOCOL).getValue();
    byte[] buf = baos.toByteArray();
    if (protocol.equals(TCP_VALUE.getValue()) && buf[buf.length - 1] != NEW_LINE_CHAR) {
        final byte[] updatedBuf = new byte[buf.length + 1];
        System.arraycopy(buf, 0, updatedBuf, 0, buf.length);
        updatedBuf[updatedBuf.length - 1] = NEW_LINE_CHAR;
        buf = updatedBuf;
    }

    // create a message batch of one message and add to active batches
    final FlowFileMessageBatch messageBatch = new FlowFileMessageBatch(session, flowFile);
    messageBatch.setNumMessages(1);
    activeBatches.add(messageBatch);

    // attempt to send the data and add the appropriate range
    try {
        sender.send(buf);
        messageBatch.addSuccessfulRange(0L, flowFile.getSize());
    } catch (IOException e) {
        messageBatch.addFailedRange(0L, flowFile.getSize(), e);
        context.yield();
    }
}
 

@Override
public byte[] getBody() {
  synchronized (bodyLock) {
    if (!bodyLoaded) {
      if (flowFile.getSize() > Integer.MAX_VALUE) {
        throw new RuntimeException("Can't get body of Event because the backing FlowFile is too large (" + flowFile.getSize() + " bytes)");
      }

      final ByteArrayOutputStream baos = new ByteArrayOutputStream((int) flowFile.getSize());
      session.read(flowFile, new InputStreamCallback() {

        @Override
        public void process(InputStream in) throws IOException {
          try (BufferedInputStream input = new BufferedInputStream(in)) {
            StreamUtils.copy(input, baos);
          }
          baos.close();
        }
      });

      body = baos.toByteArray();
      bodyLoaded = true;
    }
  }

  return body;
}
 

/**
 * Determines the CQL statement that should be executed for the given FlowFile
 *
 * @param session  the session that can be used to access the given FlowFile
 * @param flowFile the FlowFile whose CQL statement should be executed
 * @return the CQL that is associated with the given FlowFile
 */

private String getCQL(final ProcessSession session, final FlowFile flowFile, final Charset charset) {
    // Read the CQL from the FlowFile's content
    final byte[] buffer = new byte[(int) flowFile.getSize()];
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer);
        }
    });

    // Create the PreparedStatement string to use for this FlowFile.
    return new String(buffer, charset);
}
 

/**
 * Extracts contents of the {@link FlowFile} as byte array.
 */
private byte[] extractMessage(FlowFile flowFile, ProcessSession session){
    final byte[] messageContent = new byte[(int) flowFile.getSize()];
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, messageContent, true);
        }
    });
    return messageContent;
}
 

@Test
@Ignore
public void testManyFilesOpened() throws IOException {

    StandardProcessSession[] standardProcessSessions = new StandardProcessSession[100000];
    for (int i = 0; i < 70000; i++) {
        standardProcessSessions[i] = new StandardProcessSession(context);

        FlowFile flowFile = standardProcessSessions[i].create();
        final byte[] buff = new byte["Hello".getBytes().length];

        flowFile = standardProcessSessions[i].append(flowFile, new OutputStreamCallback() {
            @Override
            public void process(OutputStream out) throws IOException {
                out.write("Hello".getBytes());
            }
        });

        try {
            standardProcessSessions[i].read(flowFile, false, new InputStreamCallback() {
                @Override
                public void process(final InputStream in) throws IOException {
                    StreamUtils.fillBuffer(in, buff);
                }
            });
        } catch (Exception e) {
            System.out.println("Failed at file:" + i);
            throw e;
        }
        if (i % 1000 == 0) {
            System.out.println("i:" + i);
        }
    }
}
 

/**
 *
 */
protected FlowFile doTransform(ProcessContext context, ProcessSession session, FlowFile flowFile,  InvocationContextProperties contextProperties) {
    AtomicReference<Map<String, String>> attributeRef = new AtomicReference<Map<String, String>>();
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(InputStream in) throws IOException {
            attributeRef.set(transform(in, null, contextProperties));
        }
    });
    if (attributeRef.get() != null) {
        flowFile = session.putAllAttributes(flowFile, attributeRef.get());
    }
    return flowFile;
}
 

/**
 * Determines the SQL statement that should be executed for the given FlowFile
 *
 * @param session the session that can be used to access the given FlowFile
 * @param flowFile the FlowFile whose SQL statement should be executed
 *
 * @return the SQL that is associated with the given FlowFile
 */
private String getSQL(final ProcessSession session, final FlowFile flowFile) {
    // Read the SQL from the FlowFile's content
    final byte[] buffer = new byte[(int) flowFile.getSize()];
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer);
        }
    });

    // Create the PreparedStatement to use for this FlowFile.
    final String sql = new String(buffer, StandardCharsets.UTF_8);
    return sql;
}
 

/**
 * Infers the Avro schema from the input Flowfile content.
 *
 * @param inputFlowFile
 *  The original input FlowFile containing the JSON content as it entered this processor.
 *
 * @param context
 *  ProcessContext to pull processor configurations.
 *
 * @param session
 *  ProcessSession to transfer FlowFiles
 */
private String inferAvroSchemaFromJSON(final FlowFile inputFlowFile, final ProcessContext context, final ProcessSession session) {

    final AtomicReference<String> avroSchema = new AtomicReference<>();
    session.read(inputFlowFile, new InputStreamCallback() {
        @Override
        public void process(InputStream in) throws IOException {
            Schema as = JsonUtil.inferSchema(
                    in, context.getProperty(RECORD_NAME).evaluateAttributeExpressions(inputFlowFile).getValue(),
                    context.getProperty(NUM_RECORDS_TO_ANALYZE).evaluateAttributeExpressions(inputFlowFile).asInteger());
            avroSchema.set(as.toString(context.getProperty(PRETTY_AVRO_OUTPUT).asBoolean()));

        }
    });

    return avroSchema.get();
}
 

/**
 * Extracts contents of the {@link FlowFile} as byte array.
 */
private byte[] extractMessageBody(FlowFile flowFile, ProcessSession session) {
    final byte[] messageContent = new byte[(int) flowFile.getSize()];
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, messageContent, true);
        }
    });
    return messageContent;
}
 

/**
 * Determines the HiveQL statement that should be executed for the given FlowFile
 *
 * @param session  the session that can be used to access the given FlowFile
 * @param flowFile the FlowFile whose HiveQL statement should be executed
 * @return the HiveQL that is associated with the given FlowFile
 */
protected String getHiveQL(final ProcessSession session, final FlowFile flowFile, final Charset charset) {
    // Read the HiveQL from the FlowFile's content
    final byte[] buffer = new byte[(int) flowFile.getSize()];
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer);
        }
    });

    // Create the PreparedStatement to use for this FlowFile.
    return new String(buffer, charset);
}
 

/**
 * Reads content of the flow file and closes input stream.
 *
 * @param c Closure with one parameter InputStream.
 */
public void read(Closure c) {
    this.read(new InputStreamCallback() {
        public void process(InputStream in) throws IOException {
            c.call(in);
        }
    });
}
 

/**
 * Test BytesMessage to FlowFile conversion
 *
 * @throws java.lang.Exception ex
 */
@Test
public void testMap2FlowFileBytesMessage() throws Exception {

    TestRunner runner = TestRunners.newTestRunner(GetJMSQueue.class);
    BytesMessage bytesMessage = new ActiveMQBytesMessage();

    String sourceString = "Apache NiFi is an easy to use, powerful, and reliable system to process and distribute data.!";
    byte[] payload = sourceString.getBytes("UTF-8");
    bytesMessage.writeBytes(payload);
    bytesMessage.reset();

    ProcessContext context = runner.getProcessContext();
    ProcessSession session = runner.getProcessSessionFactory().createSession();
    ProcessorInitializationContext pic = new MockProcessorInitializationContext(runner.getProcessor(), (MockProcessContext) runner.getProcessContext());

    JmsProcessingSummary summary = JmsConsumer.map2FlowFile(context, session, bytesMessage, true, pic.getLogger());

    assertEquals("BytesMessage content length should equal to FlowFile content size", payload.length, summary.getLastFlowFile().getSize());

    final byte[] buffer = new byte[payload.length];
    runner.clearTransferState();

    session.read(summary.getLastFlowFile(), new InputStreamCallback() {
        @Override
        public void process(InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer, false);
        }
    });

    String contentString = new String(buffer, "UTF-8");
    assertEquals("", sourceString, contentString);
}
 

/**
 * Reads content of the flow file through Reader and closes the stream.
 *
 * @param charset charset to use for Reader
 * @param c       Closure with one parameter Reader.
 */
public void read(String charset, Closure c) {
    this.read(new InputStreamCallback() {
        public void process(InputStream in) throws IOException {
            InputStreamReader r = new InputStreamReader(in, charset);
            c.call(r);
            r.close();
        }
    });
}
 

/**
 * Send the entire FlowFile as a single message.
 */
private void processSingleMessage(ProcessContext context, ProcessSession session, FlowFile flowFile, ChannelSender sender) {
    // copy the contents of the FlowFile to the ByteArrayOutputStream
    final ByteArrayOutputStream baos = new ByteArrayOutputStream((int)flowFile.getSize() + 1);
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.copy(in, baos);
        }
    });

    // if TCP and we don't end in a new line then add one
    final String protocol = context.getProperty(PROTOCOL).getValue();
    byte[] buf = baos.toByteArray();
    if (protocol.equals(TCP_VALUE.getValue()) && buf[buf.length - 1] != NEW_LINE_CHAR) {
        final byte[] updatedBuf = new byte[buf.length + 1];
        System.arraycopy(buf, 0, updatedBuf, 0, buf.length);
        updatedBuf[updatedBuf.length - 1] = NEW_LINE_CHAR;
        buf = updatedBuf;
    }

    // create a message batch of one message and add to active batches
    final FlowFileMessageBatch messageBatch = new FlowFileMessageBatch(session, flowFile);
    messageBatch.setNumMessages(1);
    activeBatches.add(messageBatch);

    // attempt to send the data and add the appropriate range
    try {
        sender.send(buf);
        messageBatch.addSuccessfulRange(0L, flowFile.getSize());
    } catch (IOException e) {
        messageBatch.addFailedRange(0L, flowFile.getSize(), e);
        context.yield();
    }
}
 

/**
 * Determines the HiveQL statement that should be executed for the given FlowFile
 *
 * @param session  the session that can be used to access the given FlowFile
 * @param flowFile the FlowFile whose HiveQL statement should be executed
 * @return the HiveQL that is associated with the given FlowFile
 */
protected String getHiveQL(final ProcessSession session, final FlowFile flowFile, final Charset charset) {
    // Read the HiveQL from the FlowFile's content
    final byte[] buffer = new byte[(int) flowFile.getSize()];
    session.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer);
        }
    });

    // Create the PreparedStatement to use for this FlowFile.
    return new String(buffer, charset);
}
 

/**
 * Extracts contents of the {@link FlowFile} to byte array.
 */
private byte[] extractMessage(FlowFile flowFile, ProcessSession processSession) {
    final byte[] messageContent = new byte[(int) flowFile.getSize()];
    processSession.read(flowFile, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, messageContent, true);
        }
    });
    return messageContent;
}
 

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

    final ComponentLog logger = this.getLogger();
    final AtomicReference<Metadata> value = new AtomicReference<>(null);
    final Integer max = context.getProperty(MAX_NUMBER_OF_ATTRIBUTES).asInteger();

    try {
        session.read(flowfile, new InputStreamCallback() {
            @Override
            public void process(InputStream in) throws IOException {
                try {
                    Metadata imageMetadata = ImageMetadataReader.readMetadata(in);
                    value.set(imageMetadata);
                } catch (ImageProcessingException ex) {
                    throw new ProcessException(ex);
                }
            }
        });

        Metadata metadata = value.get();
        Map<String, String> results = getTags(max, metadata);

        // Write the results to an attribute
        if (!results.isEmpty()) {
            flowfile = session.putAllAttributes(flowfile, results);
        }

        session.transfer(flowfile, SUCCESS);
    } catch (ProcessException e) {
        logger.error("Failed to extract image metadata from {} due to {}", new Object[]{flowfile, e});
        session.transfer(flowfile, FAILURE);
    }
}
 

@Override
public byte[] getBody() {
  synchronized (bodyLock) {
    if (!bodyLoaded) {
      if (flowFile.getSize() > Integer.MAX_VALUE) {
        throw new RuntimeException("Can't get body of Event because the backing FlowFile is too large (" + flowFile.getSize() + " bytes)");
      }

      final ByteArrayOutputStream baos = new ByteArrayOutputStream((int) flowFile.getSize());
      session.read(flowFile, new InputStreamCallback() {

        @Override
        public void process(InputStream in) throws IOException {
          try (BufferedInputStream input = new BufferedInputStream(in)) {
            StreamUtils.copy(input, baos);
          }
          baos.close();
        }
      });

      body = baos.toByteArray();
      bodyLoaded = true;
    }
  }

  return body;
}
 

/**
 * Test BytesMessage to FlowFile conversion
 *
 * @throws java.lang.Exception ex
 */
@Test
public void testMap2FlowFileBytesMessage() throws Exception {

    TestRunner runner = TestRunners.newTestRunner(GetJMSQueue.class);
    BytesMessage bytesMessage = new ActiveMQBytesMessage();

    String sourceString = "Apache NiFi is an easy to use, powerful, and reliable system to process and distribute data.!";
    byte[] payload = sourceString.getBytes("UTF-8");
    bytesMessage.writeBytes(payload);
    bytesMessage.reset();

    ProcessContext context = runner.getProcessContext();
    ProcessSession session = runner.getProcessSessionFactory().createSession();
    ProcessorInitializationContext pic = new MockProcessorInitializationContext(runner.getProcessor(), (MockProcessContext) runner.getProcessContext());

    JmsProcessingSummary summary = JmsConsumer.map2FlowFile(context, session, bytesMessage, true, pic.getLogger());

    assertEquals("BytesMessage content length should equal to FlowFile content size", payload.length, summary.getLastFlowFile().getSize());

    final byte[] buffer = new byte[payload.length];
    runner.clearTransferState();

    session.read(summary.getLastFlowFile(), new InputStreamCallback() {
        @Override
        public void process(InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, buffer, false);
        }
    });

    String contentString = new String(buffer, "UTF-8");
    assertEquals("", sourceString, contentString);
}