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

/**
 * 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();
}
 

private String getMessage(final FlowFile flowFile, final ProcessContext context, final ProcessSession session) {
    String messageText = "";

    if(context.getProperty(CONTENT_AS_MESSAGE).evaluateAttributeExpressions(flowFile).asBoolean()) {
        // reading all the content of the input flow file
        final byte[] byteBuffer = new byte[(int) flowFile.getSize()];
        session.read(flowFile, new InputStreamCallback() {
            @Override
            public void process(InputStream in) throws IOException {
                StreamUtils.fillBuffer(in, byteBuffer, false);
            }
        });

        messageText = new String(byteBuffer, 0, byteBuffer.length, Charset.forName("UTF-8"));
    } else if (context.getProperty(MESSAGE).isSet()) {
        messageText = context.getProperty(MESSAGE).evaluateAttributeExpressions(flowFile).getValue();
    }

    if (context.getProperty(INCLUDE_ALL_ATTRIBUTES).asBoolean()) {
        return formatAttributes(flowFile, messageText);
    }

    return messageText;
}
 

/**
 * 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} 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;
}
 

/**
 * 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 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) 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);
    }
}
 

/**
 * 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);
}
 

@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 (final FlowFile flowFile : flowFiles) {
        final AtomicBoolean valid = new AtomicBoolean(true);
        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);
                    logger.debug("Failed to validate {} against schema due to {}", new Object[]{flowFile, e});
                }
            }
        });

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

/**
 * 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;
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
  final ProcessorLog log = this.getLogger();
  final AtomicReference<List<Map.Entry<File, Boolean>>> value = new AtomicReference<>();
  final File tempDir = new File(context.getProperty(TEMP_DIR).getValue());
  System.getProperties().setProperty("jna.library.path", context.getProperty(JNI_PATH).getValue());
  FlowFile flowfile = session.get();
  session.read(flowfile, in -> {
    try {
      value.set(convert(in, tempDir));
    }
    catch(Exception e) {
      log.error("Unable to convert: " + e.getMessage(), e);
    }
  });
  if(value.get() != null) {
    for(Map.Entry<File, Boolean> kv : value.get()) {
      final File convertedFile = kv.getKey();
      try {
        final int pageNumber = getPageNumber(convertedFile.getName());
        if(kv.getValue()) {
          FlowFile ff = session.clone(flowfile);
          ff = session.putAttribute(ff, "pageNumber", "" + pageNumber);
          ff = session.write(ff, out -> IOUtils.copy(new BufferedInputStream(new FileInputStream(convertedFile)), out));
          session.transfer(ff, SUCCESS);
        }
      }
      finally {
        if(convertedFile != null && convertedFile.exists()) {
          convertedFile.delete();
        }
      }
    }
  }
  session.transfer(flowfile, RAW);
}
 

@Override
public void unpack(final ProcessSession session, final FlowFile source, final List<FlowFile> unpacked) {
    final String fragmentId = UUID.randomUUID().toString();
    session.read(source, new InputStreamCallback() {
        @Override
        public void process(final InputStream in) throws IOException {
            int fragmentCount = 0;
            try (final ZipArchiveInputStream zipIn = new ZipArchiveInputStream(new BufferedInputStream(in))) {
                ArchiveEntry zipEntry;
                while ((zipEntry = zipIn.getNextEntry()) != null) {
                    if (zipEntry.isDirectory() || !fileMatches(zipEntry)) {
                        continue;
                    }
                    final File file = new File(zipEntry.getName());
                    final String parentDirectory = (file.getParent() == null) ? "/" : file.getParent();
                    final Path absPath = file.toPath().toAbsolutePath();
                    final String absPathString = absPath.getParent().toString() + "/";

                    FlowFile unpackedFile = session.create(source);
                    try {
                        final Map<String, String> attributes = new HashMap<>();
                        attributes.put(CoreAttributes.FILENAME.key(), file.getName());
                        attributes.put(CoreAttributes.PATH.key(), parentDirectory);
                        attributes.put(CoreAttributes.ABSOLUTE_PATH.key(), absPathString);
                        attributes.put(CoreAttributes.MIME_TYPE.key(), OCTET_STREAM);

                        attributes.put(FRAGMENT_ID, fragmentId);
                        attributes.put(FRAGMENT_INDEX, String.valueOf(++fragmentCount));

                        unpackedFile = session.putAllAttributes(unpackedFile, attributes);
                        unpackedFile = session.write(unpackedFile, new OutputStreamCallback() {
                            @Override
                            public void process(final OutputStream out) throws IOException {
                                StreamUtils.copy(zipIn, out);
                            }
                        });
                    } finally {
                        unpacked.add(unpackedFile);
                    }
                }
            }
        }
    });
}
 

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

    final Charset charset = Charset.forName(context.getProperty(CHARACTER_SET).evaluateAttributeExpressions(flowFile).getValue());
    final Boolean useSegmentNames = context.getProperty(USE_SEGMENT_NAMES).asBoolean();
    final Boolean parseSegmentFields = context.getProperty(PARSE_SEGMENT_FIELDS).asBoolean();
    final Boolean skipValidation = context.getProperty(SKIP_VALIDATION).asBoolean();
    final String inputVersion = context.getProperty(HL7_INPUT_VERSION).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);
        }
    });

    @SuppressWarnings("resource")
    final HapiContext hapiContext = new DefaultHapiContext();
    if (!inputVersion.equals("autodetect")) {
        hapiContext.setModelClassFactory(new CanonicalModelClassFactory(inputVersion));
    }
    if (skipValidation) {
        hapiContext.setValidationContext((ValidationContext) ValidationContextFactory.noValidation());
    }

    final PipeParser parser = hapiContext.getPipeParser();
    final String hl7Text = new String(buffer, charset);
    try {
        final Message message = parser.parse(hl7Text);
        final Map<String, String> attributes = getAttributes(message, useSegmentNames, parseSegmentFields);
        flowFile = session.putAllAttributes(flowFile, attributes);
        getLogger().debug("Added the following attributes for {}: {}", new Object[]{flowFile, attributes});
    } catch (final HL7Exception e) {
        getLogger().error("Failed to extract attributes from {} due to {}", new Object[]{flowFile, e});
        session.transfer(flowFile, REL_FAILURE);
        return;
    }

    session.transfer(flowFile, REL_SUCCESS);
}
 

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

    if(mqttClient == null || !mqttClient.isConnected()){
        logger.info("Was disconnected from client or was never connected, attempting to connect.");
        try {
            reconnect();
        } catch (MqttException e) {
            context.yield();
            session.transfer(flowfile, REL_FAILURE);
            logger.error("MQTT client is disconnected and re-connecting failed. Transferring FlowFile to fail and yielding", e);
            return;
        }
    }

    // get the MQTT topic
    String topic = context.getProperty(PROP_TOPIC).evaluateAttributeExpressions(flowfile).getValue();

    if (topic == null || topic.isEmpty()) {
        logger.warn("Evaluation of the topic property returned null or evaluated to be empty, routing to failure");
        session.transfer(flowfile, REL_FAILURE);
        return;
    }

    // do the read
    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);
        }
    });

    int qos = context.getProperty(PROP_QOS).evaluateAttributeExpressions(flowfile).asInteger();
    final MqttMessage mqttMessage = new MqttMessage(messageContent);
    mqttMessage.setQos(qos);
    mqttMessage.setPayload(messageContent);
    mqttMessage.setRetained(context.getProperty(PROP_RETAIN).evaluateAttributeExpressions(flowfile).asBoolean());

    try {
        mqttClientConnectLock.readLock().lock();
        final StopWatch stopWatch = new StopWatch(true);
        try {
            /*
             * Underlying method waits for the message to publish (according to set QoS), so it executes synchronously:
             *     MqttClient.java:361 aClient.publish(topic, message, null, null).waitForCompletion(getTimeToWait());
             */
            mqttClient.publish(topic, mqttMessage);
        } finally {
            mqttClientConnectLock.readLock().unlock();
        }

        session.getProvenanceReporter().send(flowfile, broker, stopWatch.getElapsed(TimeUnit.MILLISECONDS));
        session.transfer(flowfile, REL_SUCCESS);
    } catch(MqttException me) {
        logger.error("Failed to publish message.", me);
        session.transfer(flowfile, REL_FAILURE);
    }
}
 

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

    // Determine the algorithm to use
    final String algorithmName = context.getProperty(HASH_ALGORITHM).getValue();
    logger.debug("Using algorithm {}", new Object[]{algorithmName});
    HashAlgorithm algorithm = HashAlgorithm.fromName(algorithmName);

    if (flowFile.getSize() == 0) {
        if (context.getProperty(FAIL_WHEN_EMPTY).asBoolean()) {
            logger.info("Routing {} to 'failure' because content is empty (and FAIL_WHEN_EMPTY is true)");
            session.transfer(flowFile, REL_FAILURE);
            return;
        } else {
            logger.debug("Flowfile content is empty; hashing with {} anyway", new Object[]{algorithmName});
        }
    }

    // Generate a hash with the configured algorithm for the content
    // and create a new attribute with the configured name
    logger.debug("Generating {} hash of content", new Object[]{algorithmName});
    final AtomicReference<String> hashValueHolder = new AtomicReference<>(null);

    try {
        // Read the flowfile content via a lambda InputStreamCallback and hash the content
        session.read(flowFile, in -> hashValueHolder.set(HashService.hashValueStreaming(algorithm, in)));

        // Determine the destination attribute name
        final String attributeName = "content_" + algorithmName;
        logger.debug("Writing {} hash to attribute '{}'", new Object[]{algorithmName, attributeName});

        // Write the attribute
        flowFile = session.putAttribute(flowFile, attributeName, hashValueHolder.get());
        logger.info("Successfully added attribute '{}' to {} with a value of {}; routing to success", new Object[]{attributeName, flowFile, hashValueHolder.get()});

        // Update provenance and route to success
        session.getProvenanceReporter().modifyAttributes(flowFile);
        session.transfer(flowFile, REL_SUCCESS);
    } catch (ProcessException e) {
        logger.error("Failed to process {} due to {}; routing to failure", new Object[]{flowFile, e});
        session.transfer(flowFile, REL_FAILURE);
    }
}
 

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

    final Charset charset = Charset.forName(context.getProperty(CHARACTER_SET).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);
        }
    });

    @SuppressWarnings("resource")
    final HapiContext hapiContext = new DefaultHapiContext();
    hapiContext.setValidationContext((ca.uhn.hl7v2.validation.ValidationContext) ValidationContextFactory.noValidation());

    final PipeParser parser = hapiContext.getPipeParser();
    final String hl7Text = new String(buffer, charset);
    final HL7Message message;
    try {
        final Message hapiMessage = parser.parse(hl7Text);
        message = new HapiMessage(hapiMessage);
    } catch (final Exception e) {
        getLogger().error("Failed to parse {} as HL7 due to {}; routing to failure", new Object[]{flowFile, e});
        session.transfer(flowFile, REL_FAILURE);
        return;
    }

    final Set<String> matchingRels = new HashSet<>();
    final Map<Relationship, HL7Query> queryMap = queries;
    for (final Map.Entry<Relationship, HL7Query> entry : queryMap.entrySet()) {
        final Relationship relationship = entry.getKey();
        final HL7Query query = entry.getValue();

        final QueryResult result = query.evaluate(message);
        if (result.isMatch()) {
            FlowFile clone = session.clone(flowFile);
            clone = session.putAttribute(clone, "RouteHL7.Route", relationship.getName());
            session.transfer(clone, relationship);
            session.getProvenanceReporter().route(clone, relationship);
            matchingRels.add(relationship.getName());
        }
    }

    session.transfer(flowFile, REL_ORIGINAL);
    getLogger().info("Routed a copy of {} to {} relationships: {}", new Object[]{flowFile, matchingRels.size(), matchingRels});
}
 

@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, final ProcessSession session) throws ProcessException {

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

    final ComponentLog log = getLogger();
    final LivySessionService livySessionService = context.getProperty(LIVY_CONTROLLER_SERVICE).asControllerService(LivySessionService.class);
    final Map<String, String> livyController;
    try {
        livyController = livySessionService.getSession();
        if (livyController == null || livyController.isEmpty()) {
            log.debug("No Spark session available (yet), routing flowfile to wait");
            session.transfer(flowFile, REL_WAIT);
            context.yield();
            return;
        }
    } catch (SessionManagerException sme) {
        log.error("Error opening spark session, routing flowfile to wait", sme);
        session.transfer(flowFile, REL_WAIT);
        context.yield();
        return;
    }
    final long statusCheckInterval = context.getProperty(STATUS_CHECK_INTERVAL).evaluateAttributeExpressions(flowFile).asTimePeriod(TimeUnit.MILLISECONDS);
    Charset charset = Charset.forName(context.getProperty(CHARSET).evaluateAttributeExpressions(flowFile).getValue());

    String sessionId = livyController.get("sessionId");
    String livyUrl = livyController.get("livyUrl");
    String code = context.getProperty(CODE).evaluateAttributeExpressions(flowFile).getValue();
    if (StringUtils.isEmpty(code)) {
        try (InputStream inputStream = session.read(flowFile)) {
            // If no code was provided, assume it is in the content of the incoming flow file
            code = IOUtils.toString(inputStream, charset);
        } catch (IOException ioe) {
            log.error("Error reading input flowfile, penalizing and routing to failure", new Object[]{flowFile, ioe.getMessage()}, ioe);
            flowFile = session.penalize(flowFile);
            session.transfer(flowFile, REL_FAILURE);
            return;
        }
    }

    code = StringEscapeUtils.escapeJson(code);
    String payload = "{\"code\":\"" + code + "\"}";
    try {
        final JSONObject result = submitAndHandleJob(livyUrl, livySessionService, sessionId, payload, statusCheckInterval);
        log.debug("ExecuteSparkInteractive Result of Job Submit: " + result);
        if (result == null) {
            session.transfer(flowFile, REL_FAILURE);
        } else {
            try {
                final JSONObject output = result.getJSONObject("data");
                flowFile = session.write(flowFile, out -> out.write(output.toString().getBytes(charset)));
                flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), LivySessionService.APPLICATION_JSON);
                session.transfer(flowFile, REL_SUCCESS);
            } catch (JSONException je) {
                // The result doesn't contain the data, just send the output object as the flow file content to failure (after penalizing)
                log.error("Spark Session returned an error, sending the output JSON object as the flow file content to failure (after penalizing)");
                flowFile = session.write(flowFile, out -> out.write(result.toString().getBytes(charset)));
                flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), LivySessionService.APPLICATION_JSON);
                flowFile = session.penalize(flowFile);
                session.transfer(flowFile, REL_FAILURE);
            }
        }
    } catch (IOException | SessionManagerException e) {
        log.error("Failure processing flowfile {} due to {}, penalizing and routing to failure", new Object[]{flowFile, e.getMessage()}, e);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
    }
}
 

@Override
public List<FlowFile> split(final ProcessSession session, final FlowFile originalFlowFile, final SplitWriter splitWriter) {
    final List<FlowFile> childFlowFiles = new ArrayList<>();
    final AtomicReference<GenericRecord> recordHolder = new AtomicReference<>(null);

    session.read(originalFlowFile, new InputStreamCallback() {
        @Override
        public void process(InputStream rawIn) throws IOException {
            try (final InputStream in = new BufferedInputStream(rawIn);
                 final DataFileStream<GenericRecord> reader = new DataFileStream<>(in, new GenericDatumReader<GenericRecord>())) {

                final AtomicReference<String> codec = new AtomicReference<>(reader.getMetaString(DataFileConstants.CODEC));
                if (codec.get() == null) {
                    codec.set(DataFileConstants.NULL_CODEC);
                }

                // while records are left, start a new split by spawning a FlowFile
                final AtomicReference<Boolean> hasNextHolder = new AtomicReference<Boolean>(reader.hasNext());
                while (hasNextHolder.get()) {
                    FlowFile childFlowFile = session.create(originalFlowFile);
                    childFlowFile = session.write(childFlowFile, new OutputStreamCallback() {
                        @Override
                        public void process(OutputStream rawOut) throws IOException {
                            try (final BufferedOutputStream out = new BufferedOutputStream(rawOut)) {
                                splitWriter.init(reader, codec.get(), out);

                                // append to the current FlowFile until no more records, or splitSize is reached
                                int recordCount = 0;
                                while (hasNextHolder.get() && recordCount < splitSize) {
                                    recordHolder.set(reader.next(recordHolder.get()));
                                    splitWriter.write(recordHolder.get());
                                    recordCount++;
                                    hasNextHolder.set(reader.hasNext());
                                }

                                splitWriter.flush();
                            } finally {
                                splitWriter.close();
                            }
                        }
                    });

                    // would prefer this to be part of the SplitWriter, but putting the metadata in FlowFile attributes
                    // can't be done inside of an OutputStream callback which is where the splitWriter is used
                    if (splitWriter instanceof BareRecordSplitWriter && transferMetadata) {
                        final Map<String,String> metadata = new HashMap<>();
                        for (String metaKey : reader.getMetaKeys()) {
                            metadata.put(metaKey, reader.getMetaString(metaKey));
                        }
                        childFlowFile = session.putAllAttributes(childFlowFile, metadata);
                    }

                    childFlowFiles.add(childFlowFile);
                }
            }
        }
    });

    return childFlowFiles;
}
 

@Override
protected void doDelete(ProcessContext context, ProcessSession session) throws Exception {
    FlowFile input = session.get();
    if (input == null) {
        return;
    }

    final String separator = context.getProperty(SEPARATOR).evaluateAttributeExpressions(input).getValue();
    final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions(input).getValue();
    List<String> rowKeys = new ArrayList<>();
    int lineNum = 1;
    try (InputStream is = session.read(input)) {
        Scanner scanner = new Scanner(is);
        List<DeleteRequest> deletes = new ArrayList<>();
        while (scanner.hasNextLine()) {
            String line = scanner.nextLine().trim();
            if (line.equals("")) {
                continue;
            }
            String[] parts = line.split(separator);
            if (parts.length < 3 || parts.length > 4) {
                final String msg = String.format("Invalid line length. It must have 3 or 4 components. It had %d.", parts.length);
                is.close();
                input = writeErrorAttributes(lineNum, msg, input, session);
                session.transfer(input, REL_FAILURE);
                getLogger().error(msg);
                return;
            }
            String rowId = parts[0];
            String family = parts[1];
            String column = parts[2];
            String visibility = parts.length == 4 ? parts[3] : null;

            DeleteRequest request = new DeleteRequest(rowId.getBytes(), family.getBytes(), column.getBytes(), visibility);
            deletes.add(request);
            if (!rowKeys.contains(rowId)) {
                rowKeys.add(rowId);
            }

            if (getLogger().isDebugEnabled()) {
                logCell(rowId, family, column, visibility);
            }

            lineNum++;
        }
        is.close();
        clientService.deleteCells(tableName, deletes);
        for (int index = 0; index < rowKeys.size(); index++) { //Could be many row keys in one flowfile.
            session.getProvenanceReporter().invokeRemoteProcess(input, clientService.toTransitUri(tableName, rowKeys.get(index)));
        }

        session.transfer(input, REL_SUCCESS);
    } catch (Exception ex) {
        input = writeErrorAttributes(lineNum, ex.getMessage(), input, session);
        session.transfer(input, REL_FAILURE);
    }
}