Examples with SoftLimitBoundedByteArrayOutputStream org.apache.nifi.processors.standard.util.SoftLimitBoundedByteArrayOutputStream used on opensource projects

Example 1 with SoftLimitBoundedByteArrayOutputStream

use of org.apache.nifi.processors.standard.util.SoftLimitBoundedByteArrayOutputStream in project nifi by apache.

the class InvokeHTTP method onTrigger.

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    OkHttpClient okHttpClient = okHttpClientAtomicReference.get();
    FlowFile requestFlowFile = session.get();
    // Checking to see if the property to put the body of the response in an attribute was set
    boolean putToAttribute = context.getProperty(PROP_PUT_OUTPUT_IN_ATTRIBUTE).isSet();
    if (requestFlowFile == null) {
        if (context.hasNonLoopConnection()) {
            return;
        }
        String request = context.getProperty(PROP_METHOD).evaluateAttributeExpressions().getValue().toUpperCase();
        if ("POST".equals(request) || "PUT".equals(request) || "PATCH".equals(request)) {
            return;
        } else if (putToAttribute) {
            requestFlowFile = session.create();
        }
    }
    // Setting some initial variables
    final int maxAttributeSize = context.getProperty(PROP_PUT_ATTRIBUTE_MAX_LENGTH).asInteger();
    final ComponentLog logger = getLogger();
    // log ETag cache metrics
    final boolean eTagEnabled = context.getProperty(PROP_USE_ETAG).asBoolean();
    if (eTagEnabled && logger.isDebugEnabled()) {
        final Cache cache = okHttpClient.cache();
        logger.debug("OkHttp ETag cache metrics :: Request Count: {} | Network Count: {} | Hit Count: {}", new Object[] { cache.requestCount(), cache.networkCount(), cache.hitCount() });
    }
    // Every request/response cycle has a unique transaction id which will be stored as a flowfile attribute.
    final UUID txId = UUID.randomUUID();
    FlowFile responseFlowFile = null;
    try {
        // read the url property from the context
        final String urlstr = trimToEmpty(context.getProperty(PROP_URL).evaluateAttributeExpressions(requestFlowFile).getValue());
        final URL url = new URL(urlstr);
        Request httpRequest = configureRequest(context, session, requestFlowFile, url);
        // log request
        logRequest(logger, httpRequest);
        // emit send provenance event if successfully sent to the server
        if (httpRequest.body() != null) {
            session.getProvenanceReporter().send(requestFlowFile, url.toExternalForm(), true);
        }
        final long startNanos = System.nanoTime();
        Response responseHttp = okHttpClient.newCall(httpRequest).execute();
        // output the raw response headers (DEBUG level only)
        logResponse(logger, url, responseHttp);
        // store the status code and message
        int statusCode = responseHttp.code();
        String statusMessage = responseHttp.message();
        if (statusCode == 0) {
            throw new IllegalStateException("Status code unknown, connection hasn't been attempted.");
        }
        // Create a map of the status attributes that are always written to the request and response FlowFiles
        Map<String, String> statusAttributes = new HashMap<>();
        statusAttributes.put(STATUS_CODE, String.valueOf(statusCode));
        statusAttributes.put(STATUS_MESSAGE, statusMessage);
        statusAttributes.put(REQUEST_URL, url.toExternalForm());
        statusAttributes.put(TRANSACTION_ID, txId.toString());
        if (requestFlowFile != null) {
            requestFlowFile = session.putAllAttributes(requestFlowFile, statusAttributes);
        }
        // If the property to add the response headers to the request flowfile is true then add them
        if (context.getProperty(PROP_ADD_HEADERS_TO_REQUEST).asBoolean() && requestFlowFile != null) {
            // write the response headers as attributes
            // this will overwrite any existing flowfile attributes
            requestFlowFile = session.putAllAttributes(requestFlowFile, convertAttributesFromHeaders(url, responseHttp));
        }
        boolean outputBodyToRequestAttribute = (!isSuccess(statusCode) || putToAttribute) && requestFlowFile != null;
        boolean outputBodyToResponseContent = (isSuccess(statusCode) && !putToAttribute) || context.getProperty(PROP_OUTPUT_RESPONSE_REGARDLESS).asBoolean();
        ResponseBody responseBody = responseHttp.body();
        boolean bodyExists = responseBody != null;
        InputStream responseBodyStream = null;
        SoftLimitBoundedByteArrayOutputStream outputStreamToRequestAttribute = null;
        TeeInputStream teeInputStream = null;
        try {
            responseBodyStream = bodyExists ? responseBody.byteStream() : null;
            if (responseBodyStream != null && outputBodyToRequestAttribute && outputBodyToResponseContent) {
                outputStreamToRequestAttribute = new SoftLimitBoundedByteArrayOutputStream(maxAttributeSize);
                teeInputStream = new TeeInputStream(responseBodyStream, outputStreamToRequestAttribute);
            }
            if (outputBodyToResponseContent) {
                // clone the flowfile to capture the response
                if (requestFlowFile != null) {
                    responseFlowFile = session.create(requestFlowFile);
                } else {
                    responseFlowFile = session.create();
                }
                // write attributes to response flowfile
                responseFlowFile = session.putAllAttributes(responseFlowFile, statusAttributes);
                // write the response headers as attributes
                // this will overwrite any existing flowfile attributes
                responseFlowFile = session.putAllAttributes(responseFlowFile, convertAttributesFromHeaders(url, responseHttp));
                // can potentially be null in edge cases
                if (bodyExists) {
                    // write content type attribute to response flowfile if it is available
                    if (responseBody.contentType() != null) {
                        responseFlowFile = session.putAttribute(responseFlowFile, CoreAttributes.MIME_TYPE.key(), responseBody.contentType().toString());
                    }
                    if (teeInputStream != null) {
                        responseFlowFile = session.importFrom(teeInputStream, responseFlowFile);
                    } else {
                        responseFlowFile = session.importFrom(responseBodyStream, responseFlowFile);
                    }
                    // emit provenance event
                    final long millis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
                    if (requestFlowFile != null) {
                        session.getProvenanceReporter().fetch(responseFlowFile, url.toExternalForm(), millis);
                    } else {
                        session.getProvenanceReporter().receive(responseFlowFile, url.toExternalForm(), millis);
                    }
                }
            }
            // if not successful and request flowfile is not null, store the response body into a flowfile attribute
            if (outputBodyToRequestAttribute && bodyExists) {
                String attributeKey = context.getProperty(PROP_PUT_OUTPUT_IN_ATTRIBUTE).evaluateAttributeExpressions(requestFlowFile).getValue();
                if (attributeKey == null) {
                    attributeKey = RESPONSE_BODY;
                }
                byte[] outputBuffer;
                int size;
                if (outputStreamToRequestAttribute != null) {
                    outputBuffer = outputStreamToRequestAttribute.getBuffer();
                    size = outputStreamToRequestAttribute.size();
                } else {
                    outputBuffer = new byte[maxAttributeSize];
                    size = StreamUtils.fillBuffer(responseBodyStream, outputBuffer, false);
                }
                String bodyString = new String(outputBuffer, 0, size, getCharsetFromMediaType(responseBody.contentType()));
                requestFlowFile = session.putAttribute(requestFlowFile, attributeKey, bodyString);
                final long millis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
                session.getProvenanceReporter().modifyAttributes(requestFlowFile, "The " + attributeKey + " has been added. The value of which is the body of a http call to " + url.toExternalForm() + ". It took " + millis + "millis,");
            }
        } finally {
            if (outputStreamToRequestAttribute != null) {
                outputStreamToRequestAttribute.close();
                outputStreamToRequestAttribute = null;
            }
            if (teeInputStream != null) {
                teeInputStream.close();
                teeInputStream = null;
            } else if (responseBodyStream != null) {
                responseBodyStream.close();
                responseBodyStream = null;
            }
        }
        route(requestFlowFile, responseFlowFile, session, context, statusCode);
    } catch (final Exception e) {
        // penalize or yield
        if (requestFlowFile != null) {
            logger.error("Routing to {} due to exception: {}", new Object[] { REL_FAILURE.getName(), e }, e);
            requestFlowFile = session.penalize(requestFlowFile);
            requestFlowFile = session.putAttribute(requestFlowFile, EXCEPTION_CLASS, e.getClass().getName());
            requestFlowFile = session.putAttribute(requestFlowFile, EXCEPTION_MESSAGE, e.getMessage());
            // transfer original to failure
            session.transfer(requestFlowFile, REL_FAILURE);
        } else {
            logger.error("Yielding processor due to exception encountered as a source processor: {}", e);
            context.yield();
        }
        // cleanup response flowfile, if applicable
        try {
            if (responseFlowFile != null) {
                session.remove(responseFlowFile);
            }
        } catch (final Exception e1) {
            logger.error("Could not cleanup response flowfile due to exception: {}", new Object[] { e1 }, e1);
        }
    }
}