Examples with StreamCallback org.apache.nifi.processor.io.StreamCallback used on opensource projects

Example 1 with StreamCallback

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

the class TestStandardProcessSession method testCreateThenRollbackRemovesContent.

@Test
public void testCreateThenRollbackRemovesContent() throws IOException {
    final FlowFileRecord flowFileRecord = new StandardFlowFileRecord.Builder().addAttribute("uuid", "12345678-1234-1234-1234-123456789012").entryDate(System.currentTimeMillis()).build();
    flowFileQueue.put(flowFileRecord);
    final StreamCallback nop = new StreamCallback() {

        @Override
        public void process(InputStream in, OutputStream out) throws IOException {
        }
    };
    session.create();
    FlowFile flowFile = session.create(flowFileRecord);
    flowFile = session.write(flowFile, nop);
    FlowFile flowFile2 = session.create(flowFileRecord);
    flowFile2 = session.write(flowFile2, nop);
    session.write(flowFile2, nop);
    final FlowFile flowFile3 = session.create();
    session.write(flowFile3, nop);
    session.rollback();
    assertEquals(4, contentRepo.getClaimsRemoved());
}

Example 2 with StreamCallback

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

the class ConvertAvroToJSON method onTrigger.

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final String containerOption = context.getProperty(CONTAINER_OPTIONS).getValue();
    final boolean useContainer = containerOption.equals(CONTAINER_ARRAY);
    // Wrap a single record (inclusive of no records) only when a container is being used
    final boolean wrapSingleRecord = context.getProperty(WRAP_SINGLE_RECORD).asBoolean() && useContainer;
    final String stringSchema = context.getProperty(SCHEMA).getValue();
    final boolean schemaLess = stringSchema != null;
    try {
        flowFile = session.write(flowFile, new StreamCallback() {

            @Override
            public void process(final InputStream rawIn, final OutputStream rawOut) throws IOException {
                final GenericData genericData = GenericData.get();
                if (schemaLess) {
                    if (schema == null) {
                        schema = new Schema.Parser().parse(stringSchema);
                    }
                    try (final InputStream in = new BufferedInputStream(rawIn);
                        final OutputStream out = new BufferedOutputStream(rawOut)) {
                        final DatumReader<GenericRecord> reader = new GenericDatumReader<GenericRecord>(schema);
                        final BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(in, null);
                        final GenericRecord record = reader.read(null, decoder);
                        // need to be true before we wrap it with an array
                        if (useContainer && wrapSingleRecord) {
                            out.write('[');
                        }
                        final byte[] outputBytes = (record == null) ? EMPTY_JSON_OBJECT : genericData.toString(record).getBytes(StandardCharsets.UTF_8);
                        out.write(outputBytes);
                        if (useContainer && wrapSingleRecord) {
                            out.write(']');
                        }
                    }
                } else {
                    try (final InputStream in = new BufferedInputStream(rawIn);
                        final OutputStream out = new BufferedOutputStream(rawOut);
                        final DataFileStream<GenericRecord> reader = new DataFileStream<>(in, new GenericDatumReader<GenericRecord>())) {
                        int recordCount = 0;
                        GenericRecord currRecord = null;
                        if (reader.hasNext()) {
                            currRecord = reader.next();
                            recordCount++;
                        }
                        // if configured to wrap single record
                        if (reader.hasNext() && useContainer || wrapSingleRecord) {
                            out.write('[');
                        }
                        // Determine the initial output record, inclusive if we should have an empty set of Avro records
                        final byte[] outputBytes = (currRecord == null) ? EMPTY_JSON_OBJECT : genericData.toString(currRecord).getBytes(StandardCharsets.UTF_8);
                        out.write(outputBytes);
                        while (reader.hasNext()) {
                            if (useContainer) {
                                out.write(',');
                            } else {
                                out.write('\n');
                            }
                            currRecord = reader.next(currRecord);
                            out.write(genericData.toString(currRecord).getBytes(StandardCharsets.UTF_8));
                            recordCount++;
                        }
                        // configured to wrap a single record
                        if (recordCount > 1 && useContainer || wrapSingleRecord) {
                            out.write(']');
                        }
                    }
                }
            }
        });
    } catch (final ProcessException pe) {
        getLogger().error("Failed to convert {} from Avro to JSON due to {}; transferring to failure", new Object[] { flowFile, pe });
        session.transfer(flowFile, REL_FAILURE);
        return;
    }
    flowFile = session.putAttribute(flowFile, CoreAttributes.MIME_TYPE.key(), "application/json");
    session.transfer(flowFile, REL_SUCCESS);
}

Example 3 with StreamCallback

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

the class TransformXml method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final FlowFile original = session.get();
    if (original == null) {
        return;
    }
    final ComponentLog logger = getLogger();
    final StopWatch stopWatch = new StopWatch(true);
    final String xsltFileName = context.getProperty(XSLT_FILE_NAME).evaluateAttributeExpressions(original).getValue();
    final Boolean indentOutput = context.getProperty(INDENT_OUTPUT).asBoolean();
    try {
        FlowFile transformed = session.write(original, new StreamCallback() {

            @Override
            public void process(final InputStream rawIn, final OutputStream out) throws IOException {
                try (final InputStream in = new BufferedInputStream(rawIn)) {
                    final Templates templates;
                    if (cache != null) {
                        templates = cache.get(xsltFileName);
                    } else {
                        templates = newTemplates(context, xsltFileName);
                    }
                    final Transformer transformer = templates.newTransformer();
                    transformer.setOutputProperty(OutputKeys.INDENT, (indentOutput ? "yes" : "no"));
                    // pass all dynamic properties to the transformer
                    for (final Map.Entry<PropertyDescriptor, String> entry : context.getProperties().entrySet()) {
                        if (entry.getKey().isDynamic()) {
                            String value = context.newPropertyValue(entry.getValue()).evaluateAttributeExpressions(original).getValue();
                            transformer.setParameter(entry.getKey().getName(), value);
                        }
                    }
                    // use a StreamSource with Saxon
                    StreamSource source = new StreamSource(in);
                    StreamResult result = new StreamResult(out);
                    transformer.transform(source, result);
                } catch (final Exception e) {
                    throw new IOException(e);
                }
            }
        });
        session.transfer(transformed, REL_SUCCESS);
        session.getProvenanceReporter().modifyContent(transformed, stopWatch.getElapsed(TimeUnit.MILLISECONDS));
        logger.info("Transformed {}", new Object[] { original });
    } catch (ProcessException e) {
        logger.error("Unable to transform {} due to {}", new Object[] { original, e });
        session.transfer(original, REL_FAILURE);
    }
}

Example 4 with StreamCallback

use of org.apache.nifi.processor.io.StreamCallback 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);
}

Example 5 with StreamCallback

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

the class AbstractRecordProcessor method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final RecordReaderFactory readerFactory = context.getProperty(RECORD_READER).asControllerService(RecordReaderFactory.class);
    final RecordSetWriterFactory writerFactory = context.getProperty(RECORD_WRITER).asControllerService(RecordSetWriterFactory.class);
    final Map<String, String> attributes = new HashMap<>();
    final AtomicInteger recordCount = new AtomicInteger();
    final FlowFile original = flowFile;
    final Map<String, String> originalAttributes = flowFile.getAttributes();
    try {
        flowFile = session.write(flowFile, new StreamCallback() {

            @Override
            public void process(final InputStream in, final OutputStream out) throws IOException {
                try (final RecordReader reader = readerFactory.createRecordReader(originalAttributes, in, getLogger())) {
                    final RecordSchema writeSchema = writerFactory.getSchema(originalAttributes, reader.getSchema());
                    try (final RecordSetWriter writer = writerFactory.createWriter(getLogger(), writeSchema, out)) {
                        writer.beginRecordSet();
                        Record record;
                        while ((record = reader.nextRecord()) != null) {
                            final Record processed = AbstractRecordProcessor.this.process(record, writeSchema, original, context);
                            writer.write(processed);
                        }
                        final WriteResult writeResult = writer.finishRecordSet();
                        attributes.put("record.count", String.valueOf(writeResult.getRecordCount()));
                        attributes.put(CoreAttributes.MIME_TYPE.key(), writer.getMimeType());
                        attributes.putAll(writeResult.getAttributes());
                        recordCount.set(writeResult.getRecordCount());
                    }
                } catch (final SchemaNotFoundException e) {
                    throw new ProcessException(e.getLocalizedMessage(), e);
                } catch (final MalformedRecordException e) {
                    throw new ProcessException("Could not parse incoming data", e);
                }
            }
        });
    } catch (final Exception e) {
        getLogger().error("Failed to process {}; will route to failure", new Object[] { flowFile, e });
        session.transfer(flowFile, REL_FAILURE);
        return;
    }
    flowFile = session.putAllAttributes(flowFile, attributes);
    session.transfer(flowFile, REL_SUCCESS);
    final int count = recordCount.get();
    session.adjustCounter("Records Processed", count, false);
    getLogger().info("Successfully converted {} records for {}", new Object[] { count, flowFile });
}