Examples with RecordSetWriterFactory org.apache.nifi.serialization.RecordSetWriterFactory used on opensource projects

Example 21 with RecordSetWriterFactory

use of org.apache.nifi.serialization.RecordSetWriterFactory in project nifi by apache.

the class ScriptedRecordSetWriter method reloadScript.

/**
 * Reloads the script RecordSetWriterFactory. This must be called within the lock.
 *
 * @param scriptBody An input stream associated with the script content
 * @return Whether the script was successfully reloaded
 */
@Override
protected boolean reloadScript(final String scriptBody) {
    // note we are starting here with a fresh listing of validation
    // results since we are (re)loading a new/updated script. any
    // existing validation results are not relevant
    final Collection<ValidationResult> results = new HashSet<>();
    try {
        // get the engine and ensure its invocable
        if (scriptEngine instanceof Invocable) {
            final Invocable invocable = (Invocable) scriptEngine;
            // Find a custom configurator and invoke their eval() method
            ScriptEngineConfigurator configurator = scriptingComponentHelper.scriptEngineConfiguratorMap.get(scriptingComponentHelper.getScriptEngineName().toLowerCase());
            if (configurator != null) {
                configurator.eval(scriptEngine, scriptBody, scriptingComponentHelper.getModules());
            } else {
                // evaluate the script
                scriptEngine.eval(scriptBody);
            }
            // get configured processor from the script (if it exists)
            final Object obj = scriptEngine.get("writer");
            if (obj != null) {
                final ComponentLog logger = getLogger();
                try {
                    // set the logger if the processor wants it
                    invocable.invokeMethod(obj, "setLogger", logger);
                } catch (final NoSuchMethodException nsme) {
                    if (logger.isDebugEnabled()) {
                        logger.debug("Configured script RecordSetWriterFactory does not contain a setLogger method.");
                    }
                }
                if (configurationContext != null) {
                    try {
                        // set the logger if the processor wants it
                        invocable.invokeMethod(obj, "setConfigurationContext", configurationContext);
                    } catch (final NoSuchMethodException nsme) {
                        if (logger.isDebugEnabled()) {
                            logger.debug("Configured script RecordSetWriterFactory does not contain a setConfigurationContext method.");
                        }
                    }
                }
                // record the processor for use later
                final RecordSetWriterFactory scriptedWriter = invocable.getInterface(obj, RecordSetWriterFactory.class);
                recordFactory.set(scriptedWriter);
            } else {
                throw new ScriptException("No RecordSetWriterFactory was defined by the script.");
            }
        }
    } catch (final Exception ex) {
        final ComponentLog logger = getLogger();
        final String message = "Unable to load script: " + ex.getLocalizedMessage();
        logger.error(message, ex);
        results.add(new ValidationResult.Builder().subject("ScriptValidation").valid(false).explanation("Unable to load script due to " + ex.getLocalizedMessage()).input(scriptingComponentHelper.getScriptPath()).build());
    }
    // store the updated validation results
    validationResults.set(results);
    // return whether there was any issues loading the configured script
    return results.isEmpty();
}

Example 22 with RecordSetWriterFactory

use of org.apache.nifi.serialization.RecordSetWriterFactory in project nifi by apache.

the class ConsumeKafkaRecord_0_10 method createConsumerPool.

protected ConsumerPool createConsumerPool(final ProcessContext context, final ComponentLog log) {
    final int maxLeases = context.getMaxConcurrentTasks();
    final long maxUncommittedTime = context.getProperty(MAX_UNCOMMITTED_TIME).asTimePeriod(TimeUnit.MILLISECONDS);
    final Map<String, Object> props = new HashMap<>();
    KafkaProcessorUtils.buildCommonKafkaProperties(context, ConsumerConfig.class, props);
    props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
    props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
    props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
    final String topicListing = context.getProperty(ConsumeKafkaRecord_0_10.TOPICS).evaluateAttributeExpressions().getValue();
    final String topicType = context.getProperty(ConsumeKafkaRecord_0_10.TOPIC_TYPE).evaluateAttributeExpressions().getValue();
    final List<String> topics = new ArrayList<>();
    final String securityProtocol = context.getProperty(KafkaProcessorUtils.SECURITY_PROTOCOL).getValue();
    final String bootstrapServers = context.getProperty(KafkaProcessorUtils.BOOTSTRAP_SERVERS).evaluateAttributeExpressions().getValue();
    final RecordReaderFactory readerFactory = context.getProperty(RECORD_READER).asControllerService(RecordReaderFactory.class);
    final RecordSetWriterFactory writerFactory = context.getProperty(RECORD_WRITER).asControllerService(RecordSetWriterFactory.class);
    if (topicType.equals(TOPIC_NAME.getValue())) {
        for (final String topic : topicListing.split(",", 100)) {
            final String trimmedName = topic.trim();
            if (!trimmedName.isEmpty()) {
                topics.add(trimmedName);
            }
        }
        return new ConsumerPool(maxLeases, readerFactory, writerFactory, props, topics, maxUncommittedTime, securityProtocol, bootstrapServers, log);
    } else if (topicType.equals(TOPIC_PATTERN.getValue())) {
        final Pattern topicPattern = Pattern.compile(topicListing.trim());
        return new ConsumerPool(maxLeases, readerFactory, writerFactory, props, topicPattern, maxUncommittedTime, securityProtocol, bootstrapServers, log);
    } else {
        getLogger().error("Subscription type has an unknown value {}", new Object[] { topicType });
        return null;
    }
}

Example 23 with RecordSetWriterFactory

use of org.apache.nifi.serialization.RecordSetWriterFactory in project nifi by apache.

the class TestConsumeKafkaRecord_0_10 method setup.

@Before
public void setup() throws InitializationException {
    mockLease = mock(ConsumerLease.class);
    mockConsumerPool = mock(ConsumerPool.class);
    ConsumeKafkaRecord_0_10 proc = new ConsumeKafkaRecord_0_10() {

        @Override
        protected ConsumerPool createConsumerPool(final ProcessContext context, final ComponentLog log) {
            return mockConsumerPool;
        }
    };
    runner = TestRunners.newTestRunner(proc);
    runner.setProperty(KafkaProcessorUtils.BOOTSTRAP_SERVERS, "okeydokey:1234");
    final String readerId = "record-reader";
    final MockRecordParser readerService = new MockRecordParser();
    readerService.addSchemaField("name", RecordFieldType.STRING);
    readerService.addSchemaField("age", RecordFieldType.INT);
    runner.addControllerService(readerId, readerService);
    runner.enableControllerService(readerService);
    final String writerId = "record-writer";
    final RecordSetWriterFactory writerService = new MockRecordWriter("name, age");
    runner.addControllerService(writerId, writerService);
    runner.enableControllerService(writerService);
    runner.setProperty(ConsumeKafkaRecord_0_10.RECORD_READER, readerId);
    runner.setProperty(ConsumeKafkaRecord_0_10.RECORD_WRITER, writerId);
}

Example 24 with RecordSetWriterFactory

use of org.apache.nifi.serialization.RecordSetWriterFactory in project nifi by apache.

the class ConsumeKafkaRecord_0_11 method createConsumerPool.

protected ConsumerPool createConsumerPool(final ProcessContext context, final ComponentLog log) {
    final int maxLeases = context.getMaxConcurrentTasks();
    final long maxUncommittedTime = context.getProperty(MAX_UNCOMMITTED_TIME).asTimePeriod(TimeUnit.MILLISECONDS);
    final Map<String, Object> props = new HashMap<>();
    KafkaProcessorUtils.buildCommonKafkaProperties(context, ConsumerConfig.class, props);
    props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
    props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
    props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
    final String topicListing = context.getProperty(ConsumeKafkaRecord_0_11.TOPICS).evaluateAttributeExpressions().getValue();
    final String topicType = context.getProperty(ConsumeKafkaRecord_0_11.TOPIC_TYPE).evaluateAttributeExpressions().getValue();
    final List<String> topics = new ArrayList<>();
    final String securityProtocol = context.getProperty(KafkaProcessorUtils.SECURITY_PROTOCOL).getValue();
    final String bootstrapServers = context.getProperty(KafkaProcessorUtils.BOOTSTRAP_SERVERS).evaluateAttributeExpressions().getValue();
    final RecordReaderFactory readerFactory = context.getProperty(RECORD_READER).asControllerService(RecordReaderFactory.class);
    final RecordSetWriterFactory writerFactory = context.getProperty(RECORD_WRITER).asControllerService(RecordSetWriterFactory.class);
    final boolean honorTransactions = context.getProperty(HONOR_TRANSACTIONS).asBoolean();
    final String charsetName = context.getProperty(MESSAGE_HEADER_ENCODING).evaluateAttributeExpressions().getValue();
    final Charset charset = Charset.forName(charsetName);
    final String headerNameRegex = context.getProperty(HEADER_NAME_REGEX).getValue();
    final Pattern headerNamePattern = headerNameRegex == null ? null : Pattern.compile(headerNameRegex);
    if (topicType.equals(TOPIC_NAME.getValue())) {
        for (final String topic : topicListing.split(",", 100)) {
            final String trimmedName = topic.trim();
            if (!trimmedName.isEmpty()) {
                topics.add(trimmedName);
            }
        }
        return new ConsumerPool(maxLeases, readerFactory, writerFactory, props, topics, maxUncommittedTime, securityProtocol, bootstrapServers, log, honorTransactions, charset, headerNamePattern);
    } else if (topicType.equals(TOPIC_PATTERN.getValue())) {
        final Pattern topicPattern = Pattern.compile(topicListing.trim());
        return new ConsumerPool(maxLeases, readerFactory, writerFactory, props, topicPattern, maxUncommittedTime, securityProtocol, bootstrapServers, log, honorTransactions, charset, headerNamePattern);
    } else {
        getLogger().error("Subscription type has an unknown value {}", new Object[] { topicType });
        return null;
    }
}

Example 25 with RecordSetWriterFactory

use of org.apache.nifi.serialization.RecordSetWriterFactory in project nifi by apache.

the class PublishKafkaRecord_0_11 method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final List<FlowFile> flowFiles = session.get(FlowFileFilters.newSizeBasedFilter(1, DataUnit.MB, 500));
    if (flowFiles.isEmpty()) {
        return;
    }
    final PublisherPool pool = getPublisherPool(context);
    if (pool == null) {
        context.yield();
        return;
    }
    final String securityProtocol = context.getProperty(KafkaProcessorUtils.SECURITY_PROTOCOL).getValue();
    final String bootstrapServers = context.getProperty(KafkaProcessorUtils.BOOTSTRAP_SERVERS).evaluateAttributeExpressions().getValue();
    final RecordSetWriterFactory writerFactory = context.getProperty(RECORD_WRITER).asControllerService(RecordSetWriterFactory.class);
    final RecordReaderFactory readerFactory = context.getProperty(RECORD_READER).asControllerService(RecordReaderFactory.class);
    final boolean useTransactions = context.getProperty(USE_TRANSACTIONS).asBoolean();
    final long startTime = System.nanoTime();
    try (final PublisherLease lease = pool.obtainPublisher()) {
        if (useTransactions) {
            lease.beginTransaction();
        }
        // Send each FlowFile to Kafka asynchronously.
        final Iterator<FlowFile> itr = flowFiles.iterator();
        while (itr.hasNext()) {
            final FlowFile flowFile = itr.next();
            if (!isScheduled()) {
                // If stopped, re-queue FlowFile instead of sending it
                if (useTransactions) {
                    session.rollback();
                    lease.rollback();
                    return;
                }
                session.transfer(flowFile);
                itr.remove();
                continue;
            }
            final String topic = context.getProperty(TOPIC).evaluateAttributeExpressions(flowFile).getValue();
            final String messageKeyField = context.getProperty(MESSAGE_KEY_FIELD).evaluateAttributeExpressions(flowFile).getValue();
            try {
                session.read(flowFile, new InputStreamCallback() {

                    @Override
                    public void process(final InputStream rawIn) throws IOException {
                        try (final InputStream in = new BufferedInputStream(rawIn)) {
                            final RecordReader reader = readerFactory.createRecordReader(flowFile, in, getLogger());
                            final RecordSet recordSet = reader.createRecordSet();
                            final RecordSchema schema = writerFactory.getSchema(flowFile.getAttributes(), recordSet.getSchema());
                            lease.publish(flowFile, recordSet, writerFactory, schema, messageKeyField, topic);
                        } catch (final SchemaNotFoundException | MalformedRecordException e) {
                            throw new ProcessException(e);
                        }
                    }
                });
            } catch (final Exception e) {
                // The FlowFile will be obtained and the error logged below, when calling publishResult.getFailedFlowFiles()
                lease.fail(flowFile, e);
                continue;
            }
        }
        // Complete the send
        final PublishResult publishResult = lease.complete();
        if (publishResult.isFailure()) {
            getLogger().info("Failed to send FlowFile to kafka; transferring to failure");
            session.transfer(flowFiles, REL_FAILURE);
            return;
        }
        // Transfer any successful FlowFiles.
        final long transmissionMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime);
        for (FlowFile success : flowFiles) {
            final String topic = context.getProperty(TOPIC).evaluateAttributeExpressions(success).getValue();
            final int msgCount = publishResult.getSuccessfulMessageCount(success);
            success = session.putAttribute(success, MSG_COUNT, String.valueOf(msgCount));
            session.adjustCounter("Messages Sent", msgCount, true);
            final String transitUri = KafkaProcessorUtils.buildTransitURI(securityProtocol, bootstrapServers, topic);
            session.getProvenanceReporter().send(success, transitUri, "Sent " + msgCount + " messages", transmissionMillis);
            session.transfer(success, REL_SUCCESS);
        }
    }
}