Search in sources :

Example 21 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class ListenSyslog method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    // poll the queue with a small timeout to avoid unnecessarily yielding below
    RawSyslogEvent rawSyslogEvent = getMessage(true, true, session);
    // throttling even when no data is available
    if (rawSyslogEvent == null) {
        return;
    }
    final int maxBatchSize = context.getProperty(MAX_BATCH_SIZE).asInteger();
    final String port = context.getProperty(PORT).evaluateAttributeExpressions().getValue();
    final String protocol = context.getProperty(PROTOCOL).getValue();
    final Map<String, String> defaultAttributes = new HashMap<>(4);
    defaultAttributes.put(SyslogAttributes.PROTOCOL.key(), protocol);
    defaultAttributes.put(SyslogAttributes.PORT.key(), port);
    defaultAttributes.put(CoreAttributes.MIME_TYPE.key(), "text/plain");
    final int numAttributes = SyslogAttributes.values().length + 2;
    final boolean shouldParse = context.getProperty(PARSE_MESSAGES).asBoolean();
    final Map<String, FlowFile> flowFilePerSender = new HashMap<>();
    final SyslogParser parser = getParser();
    for (int i = 0; i < maxBatchSize; i++) {
        SyslogEvent event = null;
        // If this is our first iteration, we have already polled our queues. Otherwise, poll on each iteration.
        if (i > 0) {
            rawSyslogEvent = getMessage(true, false, session);
            if (rawSyslogEvent == null) {
                break;
            }
        }
        final String sender = rawSyslogEvent.getSender();
        FlowFile flowFile = flowFilePerSender.computeIfAbsent(sender, k -> session.create());
        if (shouldParse) {
            boolean valid = true;
            try {
                event = parser.parseEvent(rawSyslogEvent.getData(), sender);
            } catch (final ProcessException pe) {
                getLogger().warn("Failed to parse Syslog event; routing to invalid");
                valid = false;
            }
            // because the 'flowFile' object may already have data written to it.
            if (!valid || event == null || !event.isValid()) {
                FlowFile invalidFlowFile = session.create();
                invalidFlowFile = session.putAllAttributes(invalidFlowFile, defaultAttributes);
                if (sender != null) {
                    invalidFlowFile = session.putAttribute(invalidFlowFile, SyslogAttributes.SENDER.key(), sender);
                }
                try {
                    final byte[] rawBytes = rawSyslogEvent.getData();
                    invalidFlowFile = session.write(invalidFlowFile, new OutputStreamCallback() {

                        @Override
                        public void process(final OutputStream out) throws IOException {
                            out.write(rawBytes);
                        }
                    });
                } catch (final Exception e) {
                    getLogger().error("Failed to write contents of Syslog message to FlowFile due to {}; will re-queue message and try again", e);
                    errorEvents.offer(rawSyslogEvent);
                    session.remove(invalidFlowFile);
                    break;
                }
                session.transfer(invalidFlowFile, REL_INVALID);
                break;
            }
            getLogger().trace(event.getFullMessage());
            final Map<String, String> attributes = new HashMap<>(numAttributes);
            attributes.put(SyslogAttributes.PRIORITY.key(), event.getPriority());
            attributes.put(SyslogAttributes.SEVERITY.key(), event.getSeverity());
            attributes.put(SyslogAttributes.FACILITY.key(), event.getFacility());
            attributes.put(SyslogAttributes.VERSION.key(), event.getVersion());
            attributes.put(SyslogAttributes.TIMESTAMP.key(), event.getTimeStamp());
            attributes.put(SyslogAttributes.HOSTNAME.key(), event.getHostName());
            attributes.put(SyslogAttributes.BODY.key(), event.getMsgBody());
            attributes.put(SyslogAttributes.VALID.key(), String.valueOf(event.isValid()));
            flowFile = session.putAllAttributes(flowFile, attributes);
        }
        // figure out if we should write the bytes from the raw event or parsed event
        final boolean writeDemarcator = (i > 0);
        try {
            // write the raw bytes of the message as the FlowFile content
            final byte[] rawMessage = (event == null) ? rawSyslogEvent.getData() : event.getRawMessage();
            flowFile = session.append(flowFile, new OutputStreamCallback() {

                @Override
                public void process(final OutputStream out) throws IOException {
                    if (writeDemarcator) {
                        out.write(messageDemarcatorBytes);
                    }
                    out.write(rawMessage);
                }
            });
        } catch (final Exception e) {
            getLogger().error("Failed to write contents of Syslog message to FlowFile due to {}; will re-queue message and try again", e);
            errorEvents.offer(rawSyslogEvent);
            break;
        }
        flowFilePerSender.put(sender, flowFile);
    }
    for (final Map.Entry<String, FlowFile> entry : flowFilePerSender.entrySet()) {
        final String sender = entry.getKey();
        FlowFile flowFile = entry.getValue();
        if (flowFile.getSize() == 0L) {
            session.remove(flowFile);
            getLogger().debug("No data written to FlowFile from Sender {}; removing FlowFile", new Object[] { sender });
            continue;
        }
        final Map<String, String> newAttributes = new HashMap<>(defaultAttributes.size() + 1);
        newAttributes.putAll(defaultAttributes);
        newAttributes.put(SyslogAttributes.SENDER.key(), sender);
        flowFile = session.putAllAttributes(flowFile, newAttributes);
        getLogger().debug("Transferring {} to success", new Object[] { flowFile });
        session.transfer(flowFile, REL_SUCCESS);
        session.adjustCounter("FlowFiles Transferred to Success", 1L, false);
        final String senderHost = sender.startsWith("/") && sender.length() > 1 ? sender.substring(1) : sender;
        final String transitUri = new StringBuilder().append(protocol.toLowerCase()).append("://").append(senderHost).append(":").append(port).toString();
        session.getProvenanceReporter().receive(flowFile, transitUri);
    }
}
Also used : FlowFile(org.apache.nifi.flowfile.FlowFile) HashMap(java.util.HashMap) OutputStream(java.io.OutputStream) ProcessException(org.apache.nifi.processor.exception.ProcessException) IOException(java.io.IOException) SyslogEvent(org.apache.nifi.processors.standard.syslog.SyslogEvent) ProcessException(org.apache.nifi.processor.exception.ProcessException) SyslogParser(org.apache.nifi.processors.standard.syslog.SyslogParser) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) Map(java.util.Map) HashMap(java.util.HashMap)

Example 22 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class ModifyBytes method onTrigger.

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile ff = session.get();
    if (null == ff) {
        return;
    }
    final ComponentLog logger = getLogger();
    final long startOffset = context.getProperty(START_OFFSET).evaluateAttributeExpressions(ff).asDataSize(DataUnit.B).longValue();
    final long endOffset = context.getProperty(END_OFFSET).evaluateAttributeExpressions(ff).asDataSize(DataUnit.B).longValue();
    final boolean removeAll = context.getProperty(REMOVE_ALL).asBoolean();
    final long newFileSize = removeAll ? 0L : ff.getSize() - startOffset - endOffset;
    final StopWatch stopWatch = new StopWatch(true);
    if (newFileSize <= 0) {
        ff = session.write(ff, new OutputStreamCallback() {

            @Override
            public void process(final OutputStream out) throws IOException {
                out.write(new byte[0]);
            }
        });
    } else {
        ff = session.write(ff, new StreamCallback() {

            @Override
            public void process(final InputStream in, final OutputStream out) throws IOException {
                in.skip(startOffset);
                StreamUtils.copy(in, out, newFileSize);
            }
        });
    }
    logger.info("Transferred {} to 'success'", new Object[] { ff });
    session.getProvenanceReporter().modifyContent(ff, stopWatch.getElapsed(TimeUnit.MILLISECONDS));
    session.transfer(ff, REL_SUCCESS);
}
Also used : FlowFile(org.apache.nifi.flowfile.FlowFile) InputStream(java.io.InputStream) OutputStream(java.io.OutputStream) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) ComponentLog(org.apache.nifi.logging.ComponentLog) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) StreamCallback(org.apache.nifi.processor.io.StreamCallback) StopWatch(org.apache.nifi.util.StopWatch)

Example 23 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class ParseCEF method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final CEFParser parser = new CEFParser(validator);
    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);
        }
    });
    CommonEvent event;
    try {
        // parcefoneLocale defaults to en_US, so this should not fail. But we force failure in case the custom
        // validator failed to identify an invalid Locale
        final Locale parcefoneLocale = Locale.forLanguageTag(context.getProperty(DATETIME_REPRESENTATION).getValue());
        event = parser.parse(buffer, true, parcefoneLocale);
    } catch (Exception e) {
        // This should never trigger but adding in here as a fencing mechanism to
        // address possible ParCEFone bugs.
        getLogger().error("Parser returned unexpected Exception {} while processing {}; routing to failure", new Object[] { e, flowFile });
        session.transfer(flowFile, REL_FAILURE);
        return;
    }
    // event, so we test
    if (event == null) {
        getLogger().error("Failed to parse {} as a CEF message: it does not conform to the CEF standard; routing to failure", new Object[] { flowFile });
        session.transfer(flowFile, REL_FAILURE);
        return;
    }
    try {
        final String destination = context.getProperty(FIELDS_DESTINATION).getValue();
        switch(destination) {
            case DESTINATION_ATTRIBUTES:
                final Map<String, String> attributes = new HashMap<>();
                // Process KVs of the Header field
                for (Map.Entry<String, Object> entry : event.getHeader().entrySet()) {
                    attributes.put("cef.header." + entry.getKey(), prettyResult(entry.getValue(), tzId));
                }
                // Process KVs composing the Extension field
                for (Map.Entry<String, Object> entry : event.getExtension(true).entrySet()) {
                    attributes.put("cef.extension." + entry.getKey(), prettyResult(entry.getValue(), tzId));
                    flowFile = session.putAllAttributes(flowFile, attributes);
                }
                break;
            case DESTINATION_CONTENT:
                ObjectNode results = mapper.createObjectNode();
                // Add two JSON objects containing one CEF field each
                results.set("header", mapper.valueToTree(event.getHeader()));
                results.set("extension", mapper.valueToTree(event.getExtension(true)));
                // to the resulting JSON
                if (context.getProperty(APPEND_RAW_MESSAGE_TO_JSON).asBoolean()) {
                    results.set("_raw", mapper.valueToTree(new String(buffer)));
                }
                flowFile = session.write(flowFile, new OutputStreamCallback() {

                    @Override
                    public void process(OutputStream out) throws IOException {
                        try (OutputStream outputStream = new BufferedOutputStream(out)) {
                            outputStream.write(mapper.writeValueAsBytes(results));
                        }
                    }
                });
                // Adjust the FlowFile mime.type attribute
                flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), "application/json");
                // Update the provenance for good measure
                session.getProvenanceReporter().modifyContent(flowFile, "Replaced content with parsed CEF fields and values");
                break;
        }
        // whatever the parsing stratgy, ready to transfer to success and commit
        session.transfer(flowFile, REL_SUCCESS);
        session.commit();
    } catch (CEFHandlingException e) {
        // The flowfile has failed parsing & validation, routing to failure and committing
        getLogger().error("Failed to parse {} as a CEF message due to {}; routing to failure", new Object[] { flowFile, e });
        // Create a provenance event recording the routing to failure
        session.getProvenanceReporter().route(flowFile, REL_FAILURE);
        session.transfer(flowFile, REL_FAILURE);
        session.commit();
        return;
    } finally {
        session.rollback();
    }
}
Also used : Locale(java.util.Locale) FlowFile(org.apache.nifi.flowfile.FlowFile) ObjectNode(com.fasterxml.jackson.databind.node.ObjectNode) HashMap(java.util.HashMap) InputStream(java.io.InputStream) BufferedOutputStream(org.apache.nifi.stream.io.BufferedOutputStream) OutputStream(java.io.OutputStream) IOException(java.io.IOException) CommonEvent(com.fluenda.parcefone.event.CommonEvent) CEFHandlingException(com.fluenda.parcefone.event.CEFHandlingException) ProcessException(org.apache.nifi.processor.exception.ProcessException) JsonProcessingException(com.fasterxml.jackson.core.JsonProcessingException) IOException(java.io.IOException) InputStreamCallback(org.apache.nifi.processor.io.InputStreamCallback) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) CEFParser(com.fluenda.parcefone.parser.CEFParser) Map(java.util.Map) HashMap(java.util.HashMap) BufferedOutputStream(org.apache.nifi.stream.io.BufferedOutputStream) CEFHandlingException(com.fluenda.parcefone.event.CEFHandlingException)

Example 24 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class ListenHTTPServlet method doPost.

@Override
protected void doPost(final HttpServletRequest request, final HttpServletResponse response) throws ServletException, IOException {
    final ProcessContext context = processContext;
    ProcessSessionFactory sessionFactory;
    do {
        sessionFactory = sessionFactoryHolder.get();
        if (sessionFactory == null) {
            try {
                Thread.sleep(10);
            } catch (final InterruptedException e) {
            }
        }
    } while (sessionFactory == null);
    final ProcessSession session = sessionFactory.createSession();
    FlowFile flowFile = null;
    String holdUuid = null;
    String foundSubject = null;
    try {
        final long n = filesReceived.getAndIncrement() % FILES_BEFORE_CHECKING_DESTINATION_SPACE;
        if (n == 0 || !spaceAvailable.get()) {
            if (context.getAvailableRelationships().isEmpty()) {
                spaceAvailable.set(false);
                if (logger.isDebugEnabled()) {
                    logger.debug("Received request from " + request.getRemoteHost() + " but no space available; Indicating Service Unavailable");
                }
                response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE);
                return;
            } else {
                spaceAvailable.set(true);
            }
        }
        response.setHeader("Content-Type", MediaType.TEXT_PLAIN);
        final boolean contentGzipped = Boolean.parseBoolean(request.getHeader(GZIPPED_HEADER));
        final X509Certificate[] certs = (X509Certificate[]) request.getAttribute("javax.servlet.request.X509Certificate");
        foundSubject = DEFAULT_FOUND_SUBJECT;
        if (certs != null && certs.length > 0) {
            for (final X509Certificate cert : certs) {
                foundSubject = cert.getSubjectDN().getName();
                if (authorizedPattern.matcher(foundSubject).matches()) {
                    break;
                } else {
                    logger.warn("Rejecting transfer attempt from " + foundSubject + " because the DN is not authorized, host=" + request.getRemoteHost());
                    response.sendError(HttpServletResponse.SC_FORBIDDEN, "not allowed based on dn");
                    return;
                }
            }
        }
        final String destinationVersion = request.getHeader(PROTOCOL_VERSION_HEADER);
        Integer protocolVersion = null;
        if (destinationVersion != null) {
            try {
                protocolVersion = Integer.valueOf(destinationVersion);
            } catch (final NumberFormatException e) {
            // Value was invalid. Treat as if the header were missing.
            }
        }
        final boolean destinationIsLegacyNiFi = (protocolVersion == null);
        final boolean createHold = Boolean.parseBoolean(request.getHeader(FLOWFILE_CONFIRMATION_HEADER));
        final String contentType = request.getContentType();
        final InputStream unthrottled = contentGzipped ? new GZIPInputStream(request.getInputStream()) : request.getInputStream();
        final InputStream in = (streamThrottler == null) ? unthrottled : streamThrottler.newThrottledInputStream(unthrottled);
        if (logger.isDebugEnabled()) {
            logger.debug("Received request from " + request.getRemoteHost() + ", createHold=" + createHold + ", content-type=" + contentType + ", gzip=" + contentGzipped);
        }
        final AtomicBoolean hasMoreData = new AtomicBoolean(false);
        final FlowFileUnpackager unpackager;
        if (APPLICATION_FLOW_FILE_V3.equals(contentType)) {
            unpackager = new FlowFileUnpackagerV3();
        } else if (APPLICATION_FLOW_FILE_V2.equals(contentType)) {
            unpackager = new FlowFileUnpackagerV2();
        } else if (APPLICATION_FLOW_FILE_V1.equals(contentType)) {
            unpackager = new FlowFileUnpackagerV1();
        } else {
            unpackager = null;
        }
        final Set<FlowFile> flowFileSet = new HashSet<>();
        do {
            final long startNanos = System.nanoTime();
            final Map<String, String> attributes = new HashMap<>();
            flowFile = session.create();
            flowFile = session.write(flowFile, new OutputStreamCallback() {

                @Override
                public void process(final OutputStream rawOut) throws IOException {
                    try (final BufferedOutputStream bos = new BufferedOutputStream(rawOut, 65536)) {
                        if (unpackager == null) {
                            IOUtils.copy(in, bos);
                            hasMoreData.set(false);
                        } else {
                            attributes.putAll(unpackager.unpackageFlowFile(in, bos));
                            if (destinationIsLegacyNiFi) {
                                if (attributes.containsKey("nf.file.name")) {
                                    // for backward compatibility with old nifi...
                                    attributes.put(CoreAttributes.FILENAME.key(), attributes.remove("nf.file.name"));
                                }
                                if (attributes.containsKey("nf.file.path")) {
                                    attributes.put(CoreAttributes.PATH.key(), attributes.remove("nf.file.path"));
                                }
                            }
                            hasMoreData.set(unpackager.hasMoreData());
                        }
                    }
                }
            });
            final long transferNanos = System.nanoTime() - startNanos;
            final long transferMillis = TimeUnit.MILLISECONDS.convert(transferNanos, TimeUnit.NANOSECONDS);
            // put metadata on flowfile
            final String nameVal = request.getHeader(CoreAttributes.FILENAME.key());
            if (StringUtils.isNotBlank(nameVal)) {
                attributes.put(CoreAttributes.FILENAME.key(), nameVal);
            }
            // put arbitrary headers on flow file
            for (Enumeration<String> headerEnum = request.getHeaderNames(); headerEnum.hasMoreElements(); ) {
                String headerName = headerEnum.nextElement();
                if (headerPattern != null && headerPattern.matcher(headerName).matches()) {
                    String headerValue = request.getHeader(headerName);
                    attributes.put(headerName, headerValue);
                }
            }
            String sourceSystemFlowFileIdentifier = attributes.get(CoreAttributes.UUID.key());
            if (sourceSystemFlowFileIdentifier != null) {
                sourceSystemFlowFileIdentifier = "urn:nifi:" + sourceSystemFlowFileIdentifier;
                // If we receveied a UUID, we want to give the FlowFile a new UUID and register the sending system's
                // identifier as the SourceSystemFlowFileIdentifier field in the Provenance RECEIVE event
                attributes.put(CoreAttributes.UUID.key(), UUID.randomUUID().toString());
            }
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.getProvenanceReporter().receive(flowFile, request.getRequestURL().toString(), sourceSystemFlowFileIdentifier, "Remote DN=" + foundSubject, transferMillis);
            flowFile = session.putAttribute(flowFile, "restlistener.remote.source.host", request.getRemoteHost());
            flowFile = session.putAttribute(flowFile, "restlistener.request.uri", request.getRequestURI());
            flowFile = session.putAttribute(flowFile, "restlistener.remote.user.dn", foundSubject);
            flowFileSet.add(flowFile);
            if (holdUuid == null) {
                holdUuid = flowFile.getAttribute(CoreAttributes.UUID.key());
            }
        } while (hasMoreData.get());
        if (createHold) {
            String uuid = (holdUuid == null) ? UUID.randomUUID().toString() : holdUuid;
            if (flowFileMap.containsKey(uuid)) {
                uuid = UUID.randomUUID().toString();
            }
            final FlowFileEntryTimeWrapper wrapper = new FlowFileEntryTimeWrapper(session, flowFileSet, System.currentTimeMillis(), request.getRemoteHost());
            FlowFileEntryTimeWrapper previousWrapper;
            do {
                previousWrapper = flowFileMap.putIfAbsent(uuid, wrapper);
                if (previousWrapper != null) {
                    uuid = UUID.randomUUID().toString();
                }
            } while (previousWrapper != null);
            response.setStatus(HttpServletResponse.SC_SEE_OTHER);
            final String ackUri = "/" + basePath + "/holds/" + uuid;
            response.addHeader(LOCATION_HEADER_NAME, ackUri);
            response.addHeader(LOCATION_URI_INTENT_NAME, LOCATION_URI_INTENT_VALUE);
            response.getOutputStream().write(ackUri.getBytes("UTF-8"));
            if (logger.isDebugEnabled()) {
                logger.debug("Ingested {} from Remote Host: [{}] Port [{}] SubjectDN [{}]; placed hold on these {} files with ID {}", new Object[] { flowFileSet, request.getRemoteHost(), request.getRemotePort(), foundSubject, flowFileSet.size(), uuid });
            }
        } else {
            response.setStatus(this.returnCode);
            logger.info("Received from Remote Host: [{}] Port [{}] SubjectDN [{}]; transferring to 'success' {}", new Object[] { request.getRemoteHost(), request.getRemotePort(), foundSubject, flowFile });
            session.transfer(flowFileSet, ListenHTTP.RELATIONSHIP_SUCCESS);
            session.commit();
        }
    } catch (final Throwable t) {
        session.rollback();
        if (flowFile == null) {
            logger.error("Unable to receive file from Remote Host: [{}] SubjectDN [{}] due to {}", new Object[] { request.getRemoteHost(), foundSubject, t });
        } else {
            logger.error("Unable to receive file {} from Remote Host: [{}] SubjectDN [{}] due to {}", new Object[] { flowFile, request.getRemoteHost(), foundSubject, t });
        }
        response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, t.toString());
    }
}
Also used : ProcessSession(org.apache.nifi.processor.ProcessSession) HashMap(java.util.HashMap) OutputStream(java.io.OutputStream) BufferedOutputStream(org.apache.nifi.stream.io.BufferedOutputStream) ProcessContext(org.apache.nifi.processor.ProcessContext) GZIPInputStream(java.util.zip.GZIPInputStream) FlowFileUnpackagerV3(org.apache.nifi.util.FlowFileUnpackagerV3) FlowFileUnpackagerV2(org.apache.nifi.util.FlowFileUnpackagerV2) FlowFileUnpackagerV1(org.apache.nifi.util.FlowFileUnpackagerV1) ProcessSessionFactory(org.apache.nifi.processor.ProcessSessionFactory) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) BufferedOutputStream(org.apache.nifi.stream.io.BufferedOutputStream) HashSet(java.util.HashSet) FlowFile(org.apache.nifi.flowfile.FlowFile) GZIPInputStream(java.util.zip.GZIPInputStream) InputStream(java.io.InputStream) FlowFileUnpackager(org.apache.nifi.util.FlowFileUnpackager) X509Certificate(java.security.cert.X509Certificate) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) FlowFileEntryTimeWrapper(org.apache.nifi.processors.standard.ListenHTTP.FlowFileEntryTimeWrapper)

Example 25 with OutputStreamCallback

use of org.apache.nifi.processor.io.OutputStreamCallback in project nifi by apache.

the class ConvertJSONToSQL method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final boolean translateFieldNames = context.getProperty(TRANSLATE_FIELD_NAMES).asBoolean();
    final boolean ignoreUnmappedFields = IGNORE_UNMATCHED_FIELD.getValue().equalsIgnoreCase(context.getProperty(UNMATCHED_FIELD_BEHAVIOR).getValue());
    final String statementType = context.getProperty(STATEMENT_TYPE).getValue();
    final String updateKeys = context.getProperty(UPDATE_KEY).evaluateAttributeExpressions(flowFile).getValue();
    final String catalog = context.getProperty(CATALOG_NAME).evaluateAttributeExpressions(flowFile).getValue();
    final String schemaName = context.getProperty(SCHEMA_NAME).evaluateAttributeExpressions(flowFile).getValue();
    final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions(flowFile).getValue();
    final SchemaKey schemaKey = new SchemaKey(catalog, tableName);
    final boolean includePrimaryKeys = UPDATE_TYPE.equals(statementType) && updateKeys == null;
    // Is the unmatched column behaviour fail or warning?
    final boolean failUnmappedColumns = FAIL_UNMATCHED_COLUMN.getValue().equalsIgnoreCase(context.getProperty(UNMATCHED_COLUMN_BEHAVIOR).getValue());
    final boolean warningUnmappedColumns = WARNING_UNMATCHED_COLUMN.getValue().equalsIgnoreCase(context.getProperty(UNMATCHED_COLUMN_BEHAVIOR).getValue());
    // Escape column names?
    final boolean escapeColumnNames = context.getProperty(QUOTED_IDENTIFIERS).asBoolean();
    // Quote table name?
    final boolean quoteTableName = context.getProperty(QUOTED_TABLE_IDENTIFIER).asBoolean();
    // Attribute prefix
    final String attributePrefix = context.getProperty(SQL_PARAM_ATTR_PREFIX).evaluateAttributeExpressions(flowFile).getValue();
    // get the database schema from the cache, if one exists. We do this in a synchronized block, rather than
    // using a ConcurrentMap because the Map that we are using is a LinkedHashMap with a capacity such that if
    // the Map grows beyond this capacity, old elements are evicted. We do this in order to avoid filling the
    // Java Heap if there are a lot of different SQL statements being generated that reference different tables.
    TableSchema schema;
    synchronized (this) {
        schema = schemaCache.get(schemaKey);
        if (schema == null) {
            // No schema exists for this table yet. Query the database to determine the schema and put it into the cache.
            final DBCPService dbcpService = context.getProperty(CONNECTION_POOL).asControllerService(DBCPService.class);
            try (final Connection conn = dbcpService.getConnection()) {
                schema = TableSchema.from(conn, catalog, schemaName, tableName, translateFieldNames, includePrimaryKeys);
                schemaCache.put(schemaKey, schema);
            } catch (final SQLException e) {
                getLogger().error("Failed to convert {} into a SQL statement due to {}; routing to failure", new Object[] { flowFile, e.toString() }, e);
                session.transfer(flowFile, REL_FAILURE);
                return;
            }
        }
    }
    // Parse the JSON document
    final ObjectMapper mapper = new ObjectMapper();
    final AtomicReference<JsonNode> rootNodeRef = new AtomicReference<>(null);
    try {
        session.read(flowFile, new InputStreamCallback() {

            @Override
            public void process(final InputStream in) throws IOException {
                try (final InputStream bufferedIn = new BufferedInputStream(in)) {
                    rootNodeRef.set(mapper.readTree(bufferedIn));
                }
            }
        });
    } catch (final ProcessException pe) {
        getLogger().error("Failed to parse {} as JSON due to {}; routing to failure", new Object[] { flowFile, pe.toString() }, pe);
        session.transfer(flowFile, REL_FAILURE);
        return;
    }
    final JsonNode rootNode = rootNodeRef.get();
    // The node may or may not be a Json Array. If it isn't, we will create an
    // ArrayNode and add just the root node to it. We do this so that we can easily iterate
    // over the array node, rather than duplicating the logic or creating another function that takes many variables
    // in order to implement the logic.
    final ArrayNode arrayNode;
    if (rootNode.isArray()) {
        arrayNode = (ArrayNode) rootNode;
    } else {
        final JsonNodeFactory nodeFactory = JsonNodeFactory.instance;
        arrayNode = new ArrayNode(nodeFactory);
        arrayNode.add(rootNode);
    }
    final String fragmentIdentifier = UUID.randomUUID().toString();
    final Set<FlowFile> created = new HashSet<>();
    for (int i = 0; i < arrayNode.size(); i++) {
        final JsonNode jsonNode = arrayNode.get(i);
        final String sql;
        final Map<String, String> attributes = new HashMap<>();
        try {
            // build the fully qualified table name
            final StringBuilder tableNameBuilder = new StringBuilder();
            if (catalog != null) {
                tableNameBuilder.append(catalog).append(".");
            }
            if (schemaName != null) {
                tableNameBuilder.append(schemaName).append(".");
            }
            tableNameBuilder.append(tableName);
            final String fqTableName = tableNameBuilder.toString();
            if (INSERT_TYPE.equals(statementType)) {
                sql = generateInsert(jsonNode, attributes, fqTableName, schema, translateFieldNames, ignoreUnmappedFields, failUnmappedColumns, warningUnmappedColumns, escapeColumnNames, quoteTableName, attributePrefix);
            } else if (UPDATE_TYPE.equals(statementType)) {
                sql = generateUpdate(jsonNode, attributes, fqTableName, updateKeys, schema, translateFieldNames, ignoreUnmappedFields, failUnmappedColumns, warningUnmappedColumns, escapeColumnNames, quoteTableName, attributePrefix);
            } else {
                sql = generateDelete(jsonNode, attributes, fqTableName, schema, translateFieldNames, ignoreUnmappedFields, failUnmappedColumns, warningUnmappedColumns, escapeColumnNames, quoteTableName, attributePrefix);
            }
        } catch (final ProcessException pe) {
            getLogger().error("Failed to convert {} to a SQL {} statement due to {}; routing to failure", new Object[] { flowFile, statementType, pe.toString() }, pe);
            session.remove(created);
            session.transfer(flowFile, REL_FAILURE);
            return;
        }
        FlowFile sqlFlowFile = session.create(flowFile);
        created.add(sqlFlowFile);
        sqlFlowFile = session.write(sqlFlowFile, new OutputStreamCallback() {

            @Override
            public void process(final OutputStream out) throws IOException {
                out.write(sql.getBytes(StandardCharsets.UTF_8));
            }
        });
        attributes.put(CoreAttributes.MIME_TYPE.key(), "text/plain");
        attributes.put(attributePrefix + ".table", tableName);
        attributes.put(FRAGMENT_ID.key(), fragmentIdentifier);
        attributes.put(FRAGMENT_COUNT.key(), String.valueOf(arrayNode.size()));
        attributes.put(FRAGMENT_INDEX.key(), String.valueOf(i));
        if (catalog != null) {
            attributes.put(attributePrefix + ".catalog", catalog);
        }
        sqlFlowFile = session.putAllAttributes(sqlFlowFile, attributes);
        session.transfer(sqlFlowFile, REL_SQL);
    }
    flowFile = copyAttributesToOriginal(session, flowFile, fragmentIdentifier, arrayNode.size());
    session.transfer(flowFile, REL_ORIGINAL);
}
Also used : SQLException(java.sql.SQLException) HashMap(java.util.HashMap) LinkedHashMap(java.util.LinkedHashMap) OutputStream(java.io.OutputStream) JsonNode(org.codehaus.jackson.JsonNode) BufferedInputStream(java.io.BufferedInputStream) ArrayNode(org.codehaus.jackson.node.ArrayNode) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback) ObjectMapper(org.codehaus.jackson.map.ObjectMapper) HashSet(java.util.HashSet) FlowFile(org.apache.nifi.flowfile.FlowFile) BufferedInputStream(java.io.BufferedInputStream) InputStream(java.io.InputStream) Connection(java.sql.Connection) AtomicReference(java.util.concurrent.atomic.AtomicReference) IOException(java.io.IOException) JsonNodeFactory(org.codehaus.jackson.node.JsonNodeFactory) ProcessException(org.apache.nifi.processor.exception.ProcessException) DBCPService(org.apache.nifi.dbcp.DBCPService) InputStreamCallback(org.apache.nifi.processor.io.InputStreamCallback)

Aggregations

FlowFile (org.apache.nifi.flowfile.FlowFile)70 OutputStreamCallback (org.apache.nifi.processor.io.OutputStreamCallback)70 OutputStream (java.io.OutputStream)69 IOException (java.io.IOException)39 ProcessException (org.apache.nifi.processor.exception.ProcessException)27 HashMap (java.util.HashMap)25 InputStream (java.io.InputStream)24 Test (org.junit.Test)24 MockFlowFile (org.apache.nifi.util.MockFlowFile)23 ByteArrayOutputStream (java.io.ByteArrayOutputStream)20 ComponentLog (org.apache.nifi.logging.ComponentLog)17 FileOutputStream (java.io.FileOutputStream)16 FilterOutputStream (java.io.FilterOutputStream)16 InputStreamCallback (org.apache.nifi.processor.io.InputStreamCallback)14 ArrayList (java.util.ArrayList)12 Map (java.util.Map)12 ProcessSession (org.apache.nifi.processor.ProcessSession)12 BufferedOutputStream (org.apache.nifi.stream.io.BufferedOutputStream)10 AtomicReference (java.util.concurrent.atomic.AtomicReference)9 StandardContentClaim (org.apache.nifi.controller.repository.claim.StandardContentClaim)9