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

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

    try {
        Map<String, RecordPath> paths = getRecordPaths(context, input);
        Map<String, String> stats = getStats(input, paths, context, session);

        input = session.putAllAttributes(input, stats);

        session.transfer(input, REL_SUCCESS);

    } catch (Exception ex) {
        getLogger().error("Error processing stats.", ex);
        session.transfer(input, REL_FAILURE);
    }

}
 

protected List<FlowFile> processServiceException(final ProcessSession session, List<FlowFile> flowFiles,
        AmazonServiceException exception) {
    List<FlowFile> failedFlowFiles = new ArrayList<>();
    for (FlowFile flowFile : flowFiles) {
        Map<String,String> attributes = new HashMap<>();
        attributes.put(DYNAMODB_ERROR_EXCEPTION_MESSAGE, exception.getMessage() );
        attributes.put(DYNAMODB_ERROR_CODE, exception.getErrorCode() );
        attributes.put(DYNAMODB_ERROR_MESSAGE, exception.getErrorMessage() );
        attributes.put(DYNAMODB_ERROR_TYPE, exception.getErrorType().name() );
        attributes.put(DYNAMODB_ERROR_SERVICE, exception.getServiceName() );
        attributes.put(DYNAMODB_ERROR_RETRYABLE, Boolean.toString(exception.isRetryable()));
        attributes.put(DYNAMODB_ERROR_REQUEST_ID, exception.getRequestId() );
        attributes.put(DYNAMODB_ERROR_STATUS_CODE, Integer.toString(exception.getStatusCode()) );
        attributes.put(DYNAMODB_ERROR_EXCEPTION_MESSAGE, exception.getMessage() );
        attributes.put(DYNAMODB_ERROR_RETRYABLE, Boolean.toString(exception.isRetryable()));
        flowFile = session.putAllAttributes(flowFile, attributes);
        failedFlowFiles.add(flowFile);
    }
    return failedFlowFiles;
}
 

/**
 * Apply split index, count and other attributes.
 *
 * @param session session
 * @param source source
 * @param splits splits
 * @return generated fragment identifier for the splits
 */
private String finishFragmentAttributes(final ProcessSession session, final FlowFile source, final List<FlowFile> splits) {
    final String originalFilename = source.getAttribute(CoreAttributes.FILENAME.key());

    final String fragmentId = UUID.randomUUID().toString();
    final ArrayList<FlowFile> newList = new ArrayList<>(splits);
    splits.clear();
    for (int i = 1; i <= newList.size(); i++) {
        FlowFile ff = newList.get(i - 1);
        final Map<String, String> attributes = new HashMap<>();
        attributes.put(FRAGMENT_ID, fragmentId);
        attributes.put(FRAGMENT_INDEX, String.valueOf(i));
        attributes.put(FRAGMENT_COUNT, String.valueOf(newList.size()));
        attributes.put(SEGMENT_ORIGINAL_FILENAME, originalFilename);
        FlowFile newFF = session.putAllAttributes(ff, attributes);
        splits.add(newFF);
    }
    return fragmentId;
}
 

protected List<FlowFile> processServiceException(final ProcessSession session, List<FlowFile> flowFiles,
        AmazonServiceException exception) {
    List<FlowFile> failedFlowFiles = new ArrayList<>();
    for (FlowFile flowFile : flowFiles) {
        Map<String,String> attributes = new HashMap<>();
        attributes.put(DYNAMODB_ERROR_EXCEPTION_MESSAGE, exception.getMessage() );
        attributes.put(DYNAMODB_ERROR_CODE, exception.getErrorCode() );
        attributes.put(DYNAMODB_ERROR_MESSAGE, exception.getErrorMessage() );
        attributes.put(DYNAMODB_ERROR_TYPE, exception.getErrorType().name() );
        attributes.put(DYNAMODB_ERROR_SERVICE, exception.getServiceName() );
        attributes.put(DYNAMODB_ERROR_RETRYABLE, Boolean.toString(exception.isRetryable()));
        attributes.put(DYNAMODB_ERROR_REQUEST_ID, exception.getRequestId() );
        attributes.put(DYNAMODB_ERROR_STATUS_CODE, Integer.toString(exception.getStatusCode()) );
        attributes.put(DYNAMODB_ERROR_EXCEPTION_MESSAGE, exception.getMessage() );
        attributes.put(DYNAMODB_ERROR_RETRYABLE, Boolean.toString(exception.isRetryable()));
        flowFile = session.putAllAttributes(flowFile, attributes);
        failedFlowFiles.add(flowFile);
    }
    return failedFlowFiles;
}
 

/**
 * Copies JMS attributes (i.e., headers and properties) as FF attributes.
 * Given that FF attributes mandate that values are of type String, the
 * copied values of JMS attributes will be "stringified" via
 * String.valueOf(attribute).
 */
private FlowFile updateFlowFileAttributesWithJMSAttributes(Map<String, String> jmsAttributes, FlowFile flowFile, ProcessSession processSession) {
    Map<String, String> attributes = new HashMap<>();
    for (Entry<String, String> entry : jmsAttributes.entrySet()) {
        attributes.put(entry.getKey(), entry.getValue());
    }

    flowFile = processSession.putAllAttributes(flowFile, attributes);
    return flowFile;
}
 

/**
 * Populate exception attributes in the flow file
 * @param session process session
 * @param flowFile the flow file
 * @param exception exception thrown during invocation
 * @return FlowFile the updated flow file
 */
private FlowFile populateExceptionAttributes(final ProcessSession session, FlowFile flowFile,
        final AmazonServiceException exception) {
    Map<String,String> attributes = new HashMap<>();
    attributes.put(AWS_LAMBDA_EXCEPTION_MESSAGE, exception.getErrorMessage());
    attributes.put(AWS_LAMBDA_EXCEPTION_ERROR_CODE, exception.getErrorCode());
    attributes.put(AWS_LAMBDA_EXCEPTION_REQUEST_ID, exception.getRequestId());
    attributes.put(AWS_LAMBDA_EXCEPTION_STATUS_CODE, Integer.toString(exception.getStatusCode()));
    if ( exception.getCause() != null )
        attributes.put(AWS_LAMBDA_EXCEPTION_CAUSE, exception.getCause().getMessage());
    attributes.put(AWS_LAMBDA_EXCEPTION_ERROR_TYPE, exception.getErrorType().toString());
    attributes.put(AWS_LAMBDA_EXCEPTION_MESSAGE, exception.getErrorMessage());
    flowFile = session.putAllAttributes(flowFile, attributes);
    return flowFile;
}
 

private void transferQueue(ProcessSession session){
    while (!mqttQueue.isEmpty()) {
        FlowFile messageFlowfile = session.create();
        final MQTTQueueMessage mqttMessage = mqttQueue.peek();

        Map<String, String> attrs = new HashMap<>();
        attrs.put(BROKER_ATTRIBUTE_KEY, broker);
        attrs.put(TOPIC_ATTRIBUTE_KEY, mqttMessage.getTopic());
        attrs.put(QOS_ATTRIBUTE_KEY, String.valueOf(mqttMessage.getQos()));
        attrs.put(IS_DUPLICATE_ATTRIBUTE_KEY, String.valueOf(mqttMessage.isDuplicate()));
        attrs.put(IS_RETAINED_ATTRIBUTE_KEY, String.valueOf(mqttMessage.isRetained()));

        messageFlowfile = session.putAllAttributes(messageFlowfile, attrs);

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

        String transitUri = new StringBuilder(broker).append(mqttMessage.getTopic()).toString();
        session.getProvenanceReporter().receive(messageFlowfile, transitUri);
        session.transfer(messageFlowfile, REL_MESSAGE);
        session.commit();
        if (!mqttQueue.remove(mqttMessage) && logger.isWarnEnabled()) {
            logger.warn(new StringBuilder("FlowFile ")
                    .append(messageFlowfile.getAttribute(CoreAttributes.UUID.key()))
                    .append(" for Mqtt message ")
                    .append(mqttMessage)
                    .append(" had already been removed from queue, possible duplication of flow files")
                    .toString());
        }
    }
}
 

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

@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);
    }
}
 

private void completeFlowFile(final ProcessContext context, final ProcessSession session, final FlowFile flowFile, final RecordSetWriter writer,
        final Relationship relationship, final String details) throws IOException {
    final WriteResult writeResult = writer.finishRecordSet();
    writer.close();

    final String validationDetailsAttributeName = context.getProperty(VALIDATION_DETAILS_ATTRIBUTE_NAME)
            .evaluateAttributeExpressions(flowFile).getValue();

    final Integer maxValidationDetailsLength = context.getProperty(MAX_VALIDATION_DETAILS_LENGTH).evaluateAttributeExpressions(flowFile).asInteger();

    final Map<String, String> attributes = new HashMap<>();
    attributes.putAll(writeResult.getAttributes());
    attributes.put("record.count", String.valueOf(writeResult.getRecordCount()));
    attributes.put(CoreAttributes.MIME_TYPE.key(), writer.getMimeType());

    if(validationDetailsAttributeName != null && details != null && !details.isEmpty()) {
        String truncatedDetails = details;

        //Truncating only when it exceeds the configured maximum
        if (truncatedDetails.length() > maxValidationDetailsLength) {
            truncatedDetails = truncatedDetails.substring(0, maxValidationDetailsLength);
        }

        attributes.put(validationDetailsAttributeName, truncatedDetails);
    }

    session.putAllAttributes(flowFile, attributes);

    session.transfer(flowFile, relationship);
    session.getProvenanceReporter().route(flowFile, relationship, details);
}
 

@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);
    }
}
 

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

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    if (client == null || client.isDone()) {
        connectNewClient();
        if (client.isDone()) {
            context.yield();
            return;
        }
    }
    final Event event = eventQueue.poll();
    if (event != null) {
        switch (event.getEventType()) {
            case STOPPED_BY_ERROR:
                getLogger().error("Received error {}: {} due to {}. Will not attempt to reconnect", new Object[]{event.getEventType(), event.getMessage(), event.getUnderlyingException()});
                break;
            case CONNECTION_ERROR:
            case HTTP_ERROR:
                getLogger().error("Received error {}: {}. Will attempt to reconnect", new Object[]{event.getEventType(), event.getMessage()});
                client.reconnect();
                break;
            default:
                break;
        }
    }

    final String tweet = messageQueue.poll();
    if (tweet == null) {
        context.yield();
        return;
    }

    FlowFile flowFile = session.create();
    flowFile = session.write(flowFile, new OutputStreamCallback() {
        @Override
        public void process(final OutputStream out) throws IOException {
            out.write(tweet.getBytes(StandardCharsets.UTF_8));
        }
    });

    final Map<String, String> attributes = new HashMap<>();
    attributes.put(CoreAttributes.MIME_TYPE.key(), "application/json");
    attributes.put(CoreAttributes.FILENAME.key(), flowFile.getAttribute(CoreAttributes.FILENAME.key()) + ".json");
    flowFile = session.putAllAttributes(flowFile, attributes);

    session.transfer(flowFile, REL_SUCCESS);
    session.getProvenanceReporter().receive(flowFile, Constants.STREAM_HOST + client.getEndpoint().getURI());
}
 

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

    final String sessionId = context.getProperty(PROP_WS_SESSION_ID)
            .evaluateAttributeExpressions(flowfile).getValue();
    final String webSocketServiceId = context.getProperty(PROP_WS_CONTROLLER_SERVICE_ID)
            .evaluateAttributeExpressions(flowfile).getValue();
    final String webSocketServiceEndpoint = context.getProperty(PROP_WS_CONTROLLER_SERVICE_ENDPOINT)
            .evaluateAttributeExpressions(flowfile).getValue();
    final String messageTypeStr = context.getProperty(PROP_WS_MESSAGE_TYPE)
            .evaluateAttributeExpressions(flowfile).getValue();
    final WebSocketMessage.Type messageType = WebSocketMessage.Type.valueOf(messageTypeStr);

    if (StringUtils.isEmpty(sessionId)) {
        getLogger().debug("Specific SessionID not specified. Message will be broadcast to all connected clients.");
    }

    if (StringUtils.isEmpty(webSocketServiceId)
            || StringUtils.isEmpty(webSocketServiceEndpoint)) {
        transferToFailure(processSession, flowfile, "Required WebSocket attribute was not found.");
        return;
    }

    final ControllerService controllerService = context.getControllerServiceLookup().getControllerService(webSocketServiceId);
    if (controllerService == null) {
        transferToFailure(processSession, flowfile, "WebSocket ControllerService was not found.");
        return;
    } else if (!(controllerService instanceof WebSocketService)) {
        transferToFailure(processSession, flowfile, "The ControllerService found was not a WebSocket ControllerService but a "
                + controllerService.getClass().getName());
        return;
    }

    final WebSocketService webSocketService = (WebSocketService)controllerService;
    final byte[] messageContent = new byte[(int) flowfile.getSize()];
    final long startSending = System.currentTimeMillis();

    final AtomicReference<String> transitUri = new AtomicReference<>();
    final Map<String, String> attrs = new HashMap<>();
    attrs.put(ATTR_WS_CS_ID, webSocketService.getIdentifier());

    if (!StringUtils.isEmpty(sessionId)) {
        attrs.put(ATTR_WS_SESSION_ID, sessionId);
    }

    attrs.put(ATTR_WS_ENDPOINT_ID, webSocketServiceEndpoint);
    attrs.put(ATTR_WS_MESSAGE_TYPE, messageTypeStr);

    processSession.read(flowfile, in -> {
        StreamUtils.fillBuffer(in, messageContent, true);
    });

    try {

        webSocketService.sendMessage(webSocketServiceEndpoint, sessionId, sender -> {
            switch (messageType) {
                case TEXT:
                    sender.sendString(new String(messageContent, CHARSET_NAME));
                    break;
                case BINARY:
                    sender.sendBinary(ByteBuffer.wrap(messageContent));
                    break;
            }

            attrs.put(ATTR_WS_LOCAL_ADDRESS, sender.getLocalAddress().toString());
            attrs.put(ATTR_WS_REMOTE_ADDRESS, sender.getRemoteAddress().toString());
            transitUri.set(sender.getTransitUri());
        });

        final FlowFile updatedFlowFile = processSession.putAllAttributes(flowfile, attrs);
        final long transmissionMillis = System.currentTimeMillis() - startSending;
        processSession.getProvenanceReporter().send(updatedFlowFile, transitUri.get(), transmissionMillis);

        processSession.transfer(updatedFlowFile, REL_SUCCESS);

    } catch (WebSocketConfigurationException|IllegalStateException|IOException e) {
        // WebSocketConfigurationException: If the corresponding WebSocketGatewayProcessor has been stopped.
        // IllegalStateException: Session is already closed or not found.
        // IOException: other IO error.
        getLogger().error("Failed to send message via WebSocket due to " + e, e);
        transferToFailure(processSession, flowfile, e.toString());
    }

}
 

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

        if (input == null && context.hasNonLoopConnection()) {
            return;
        }
    }

    try {
        final String query = getQuery(input, context, session);
        final String index = context.getProperty(INDEX).evaluateAttributeExpressions(input).getValue();
        final String type  = context.getProperty(TYPE).isSet()
                ? context.getProperty(TYPE).evaluateAttributeExpressions(input).getValue()
                : null;
        final String queryAttr = context.getProperty(QUERY_ATTRIBUTE).isSet()
                ? context.getProperty(QUERY_ATTRIBUTE).evaluateAttributeExpressions(input).getValue()
                : null;
        DeleteOperationResponse dor = clientService.deleteByQuery(query, index, type);

        if (input == null) {
            input = session.create();
        }

        Map<String, String> attrs = new HashMap<>();
        attrs.put(TOOK_ATTRIBUTE, String.valueOf(dor.getTook()));
        if (!StringUtils.isBlank(queryAttr)) {
            attrs.put(queryAttr, query);
        }

        input = session.putAllAttributes(input, attrs);

        session.transfer(input, REL_SUCCESS);
    } catch (Exception e) {
        if (input != null) {
            input = session.putAttribute(input, ERROR_ATTRIBUTE, e.getMessage());
            session.transfer(input, REL_FAILURE);
        }
        getLogger().error("Error running delete by query: ", e);
        context.yield();
    }
}
 

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

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

        final boolean demarcateFirstRecord;

        BundleTracker tracker = bundleMap.get(bundleInfo);

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

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

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

private void enqueueMessage(final WebSocketMessage incomingMessage){
    final ProcessSession session = processSessionFactory.createSession();
    try {
        FlowFile messageFlowFile = session.create();

        final Map<String, String> attrs = new HashMap<>();
        attrs.put(ATTR_WS_CS_ID, webSocketService.getIdentifier());
        final WebSocketSessionInfo sessionInfo = incomingMessage.getSessionInfo();
        attrs.put(ATTR_WS_SESSION_ID, sessionInfo.getSessionId());
        attrs.put(ATTR_WS_ENDPOINT_ID, endpointId);
        attrs.put(ATTR_WS_LOCAL_ADDRESS, sessionInfo.getLocalAddress().toString());
        attrs.put(ATTR_WS_REMOTE_ADDRESS, sessionInfo.getRemoteAddress().toString());
        final WebSocketMessage.Type messageType = incomingMessage.getType();
        if (messageType != null) {
            attrs.put(ATTR_WS_MESSAGE_TYPE, messageType.name());
        }

        messageFlowFile = session.putAllAttributes(messageFlowFile, attrs);

        final byte[] payload = incomingMessage.getPayload();
        if (payload != null) {
            messageFlowFile = session.write(messageFlowFile, out -> {
                out.write(payload, incomingMessage.getOffset(), incomingMessage.getLength());
            });
        }

        session.getProvenanceReporter().receive(messageFlowFile, getTransitUri(sessionInfo));

        if (incomingMessage instanceof WebSocketConnectedMessage) {
            session.transfer(messageFlowFile, REL_CONNECTED);
        } else {
            switch (messageType) {
                case TEXT:
                    session.transfer(messageFlowFile, REL_MESSAGE_TEXT);
                    break;
                case BINARY:
                    session.transfer(messageFlowFile, REL_MESSAGE_BINARY);
                    break;
            }
        }
        session.commit();

    } catch (Exception e) {
        logger.error("Unable to fully process input due to " + e, e);
        session.rollback();
    }
}
 

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

    synchronized (esClient) {
        if(esClient.get() == null) {
            super.setup(context);
        }
    }

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

    final String index = context.getProperty(INDEX).evaluateAttributeExpressions(flowFile).getValue();
    final String docId = context.getProperty(DOC_ID).evaluateAttributeExpressions(flowFile).getValue();
    final String docType = context.getProperty(TYPE).evaluateAttributeExpressions(flowFile).getValue();
    final Charset charset = Charset.forName(context.getProperty(CHARSET).evaluateAttributeExpressions(flowFile).getValue());

    final ComponentLog logger = getLogger();
    try {

        logger.debug("Fetching {}/{}/{} from Elasticsearch", new Object[]{index, docType, docId});
        GetRequestBuilder getRequestBuilder = esClient.get().prepareGet(index, docType, docId);
        final GetResponse getResponse = getRequestBuilder.execute().actionGet();

        if (getResponse == null || !getResponse.isExists()) {
            logger.debug("Failed to read {}/{}/{} from Elasticsearch: Document not found",
                    new Object[]{index, docType, docId});

            // We couldn't find the document, so penalize it and send it to "not found"
            flowFile = session.penalize(flowFile);
            session.transfer(flowFile, REL_NOT_FOUND);
        } else {
            flowFile = session.putAllAttributes(flowFile, new HashMap<String, String>() {{
                put("filename", docId);
                put("es.index", index);
                put("es.type", docType);
            }});
            flowFile = session.write(flowFile, new OutputStreamCallback() {
                @Override
                public void process(OutputStream out) throws IOException {
                    out.write(getResponse.getSourceAsString().getBytes(charset));
                }
            });
            logger.debug("Elasticsearch document " + docId + " fetched, routing to success");
            // The document is JSON, so update the MIME type of the flow file
            flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), "application/json");
            session.getProvenanceReporter().fetch(flowFile, getResponse.remoteAddress().getAddress());
            session.transfer(flowFile, REL_SUCCESS);
        }
    } catch (NoNodeAvailableException
            | ElasticsearchTimeoutException
            | ReceiveTimeoutTransportException
            | NodeClosedException exceptionToRetry) {
        logger.error("Failed to read into Elasticsearch due to {}, this may indicate an error in configuration "
                        + "(hosts, username/password, etc.), or this issue may be transient. Routing to retry",
                new Object[]{exceptionToRetry.getLocalizedMessage()}, exceptionToRetry);
        session.transfer(flowFile, REL_RETRY);
        context.yield();

    } catch (Exception e) {
        logger.error("Failed to read {} from Elasticsearch due to {}", new Object[]{flowFile, e.getLocalizedMessage()}, e);
        session.transfer(flowFile, REL_FAILURE);
        context.yield();
    }
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    if (!initialized.get()) {
        initializeResolver(context);
        getLogger().warn("Resolver was initialized at onTrigger instead of onScheduled");

    }

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

    final String queryType = context.getProperty(DNS_QUERY_TYPE).getValue();
    final String queryInput = context.getProperty(QUERY_INPUT).evaluateAttributeExpressions(flowFile).getValue();
    final String queryParser = context.getProperty(QUERY_PARSER).getValue();
    final String queryRegex = context.getProperty(QUERY_PARSER_INPUT).getValue();

    boolean found = false;
    try {
        Attributes results = doLookup(queryInput, queryType);
        // NOERROR & NODATA seem to return empty Attributes handled bellow
        // but defaulting to not found in any case
        if (results.size() < 1) {
            found = false;
        } else {
            int recordNumber = 0;
            NamingEnumeration<?> dnsEntryIterator = results.get(queryType).getAll();

            while (dnsEntryIterator.hasMoreElements()) {
                String dnsRecord = dnsEntryIterator.next().toString();
                // While NXDOMAIN is being generated by doLookup catch

                if (dnsRecord != "NXDOMAIN") {
                    // Map<String, String> parsedResults = parseResponse(recordNumber, dnsRecord, queryParser, queryRegex, "dns");
                    Map<String, String> parsedResults = parseResponse(String.valueOf(recordNumber), dnsRecord, queryParser, queryRegex, "dns");
                    flowFile = session.putAllAttributes(flowFile, parsedResults);
                    found = true;
                } else {
                    // Otherwise treat as not found
                    found = false;
                }

                // Increase the counter and iterate over next record....
                recordNumber++;
            }
        }
    } catch (NamingException e) {
        context.yield();
        throw new ProcessException("Unexpected NamingException while processing records. Please review your configuration.", e);

    }

    // Finally prepare to send the data down the pipeline
    if (found) {
        // Sending the resulting flowfile (with attributes) to REL_FOUND
        session.transfer(flowFile, REL_FOUND);
    } else {
        // NXDOMAIN received, accepting the fate but forwarding
        // to REL_NOT_FOUND
        session.transfer(flowFile, REL_NOT_FOUND);
    }
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    if (!initialized.get()) {
        initializeResolver(context);
        getLogger().warn("Resolver was initialized at onTrigger instead of onScheduled");

    }

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

    final String queryType = context.getProperty(DNS_QUERY_TYPE).getValue();
    final String queryInput = context.getProperty(QUERY_INPUT).evaluateAttributeExpressions(flowFile).getValue();
    final String queryParser = context.getProperty(QUERY_PARSER).getValue();
    final String queryRegex = context.getProperty(QUERY_PARSER_INPUT).getValue();

    boolean found = false;
    try {
        Attributes results = doLookup(queryInput, queryType);
        // NOERROR & NODATA seem to return empty Attributes handled bellow
        // but defaulting to not found in any case
        if (results.size() < 1) {
            found = false;
        } else {
            int recordNumber = 0;
            NamingEnumeration<?> dnsEntryIterator = results.get(queryType).getAll();

            while (dnsEntryIterator.hasMoreElements()) {
                String dnsRecord = dnsEntryIterator.next().toString();
                // While NXDOMAIN is being generated by doLookup catch

                if (dnsRecord != "NXDOMAIN") {
                    // Map<String, String> parsedResults = parseResponse(recordNumber, dnsRecord, queryParser, queryRegex, "dns");
                    Map<String, String> parsedResults = parseResponse(String.valueOf(recordNumber), dnsRecord, queryParser, queryRegex, "dns");
                    flowFile = session.putAllAttributes(flowFile, parsedResults);
                    found = true;
                } else {
                    // Otherwise treat as not found
                    found = false;
                }

                // Increase the counter and iterate over next record....
                recordNumber++;
            }
        }
    } catch (NamingException e) {
        context.yield();
        throw new ProcessException("Unexpected NamingException while processing records. Please review your configuration.", e);

    }

    // Finally prepare to send the data down the pipeline
    if (found) {
        // Sending the resulting flowfile (with attributes) to REL_FOUND
        session.transfer(flowFile, REL_FOUND);
    } else {
        // NXDOMAIN received, accepting the fate but forwarding
        // to REL_NOT_FOUND
        session.transfer(flowFile, REL_NOT_FOUND);
    }
}