Examples with LivySessionService org.apache.nifi.controller.api.livy.LivySessionService used on opensource projects

Example 1 with LivySessionService

use of org.apache.nifi.controller.api.livy.LivySessionService in project nifi by apache.

the class ExecuteSparkInteractive method onTrigger.

@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 = livySessionService.getSession();
    if (livyController == null || livyController.isEmpty()) {
        log.debug("No Spark session available (yet), routing flowfile to wait");
        session.transfer(flowFile, REL_WAIT);
        return;
    }
    final long statusCheckInterval = context.getProperty(STATUS_CHECK_INTERVAL).evaluateAttributeExpressions(flowFile).asTimePeriod(TimeUnit.MILLISECONDS);
    Charset charset;
    try {
        charset = Charset.forName(context.getProperty(CHARSET).evaluateAttributeExpressions(flowFile).getValue());
    } catch (Exception e) {
        log.warn("Illegal character set name specified, defaulting to UTF-8");
        charset = StandardCharsets.UTF_8;
    }
    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.escapeJavaScript(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()));
                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()));
                flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), LivySessionService.APPLICATION_JSON);
                flowFile = session.penalize(flowFile);
                session.transfer(flowFile, REL_FAILURE);
            }
        }
    } catch (IOException ioe) {
        log.error("Failure processing flowfile {} due to {}, penalizing and routing to failure", new Object[] { flowFile, ioe.getMessage() }, ioe);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
    }
}