NIFI-8224: Add LoggingRecordSink controller service

Signed-off-by: Pierre Villard <pierre.villard.fr@gmail.com>

This closes #4820.
This commit is contained in:
Matthew Burgess 2021-02-10 19:04:07 -05:00 committed by Pierre Villard
parent 6ed496c714
commit b08d6071dd
No known key found for this signature in database
GPG Key ID: F92A93B30C07C6D5
6 changed files with 205 additions and 1 deletions

View File

@ -59,5 +59,10 @@
<artifactId>nifi-mock</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.nifi</groupId>
<artifactId>nifi-mock-record-utils</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View File

@ -0,0 +1,103 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.nifi.record.sink;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.Tags;
import org.apache.nifi.annotation.lifecycle.OnEnabled;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.controller.AbstractControllerService;
import org.apache.nifi.controller.ConfigurationContext;
import org.apache.nifi.controller.ControllerServiceInitializationContext;
import org.apache.nifi.logging.ComponentLog;
import org.apache.nifi.logging.LogLevel;
import org.apache.nifi.reporting.InitializationException;
import org.apache.nifi.schema.access.SchemaNotFoundException;
import org.apache.nifi.serialization.RecordSetWriter;
import org.apache.nifi.serialization.RecordSetWriterFactory;
import org.apache.nifi.serialization.WriteResult;
import org.apache.nifi.serialization.record.Record;
import org.apache.nifi.serialization.record.RecordSet;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
@Tags({"record", "sink", "log"})
@CapabilityDescription("Provides a RecordSinkService that can be used to log records to the application log (nifi-app.log, e.g.) using the specified writer for formatting.")
public class LoggingRecordSink extends AbstractControllerService implements RecordSinkService {
private List<PropertyDescriptor> properties;
private volatile RecordSetWriterFactory writerFactory;
private volatile LogLevel logLevel;
public static final PropertyDescriptor LOG_LEVEL = new PropertyDescriptor.Builder()
.name("logsink-log-level")
.displayName("Log Level")
.required(true)
.description("The Log Level at which to log records (INFO, DEBUG, e.g.)")
.allowableValues(LogLevel.values())
.defaultValue(LogLevel.INFO.name())
.build();
@Override
protected void init(final ControllerServiceInitializationContext context) throws InitializationException {
final List<PropertyDescriptor> properties = new ArrayList<>();
properties.add(RecordSinkService.RECORD_WRITER_FACTORY);
properties.add(LOG_LEVEL);
this.properties = Collections.unmodifiableList(properties);
}
@Override
protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
return properties;
}
@OnEnabled
public void onEnabled(final ConfigurationContext context) throws InitializationException {
writerFactory = context.getProperty(RecordSinkService.RECORD_WRITER_FACTORY).asControllerService(RecordSetWriterFactory.class);
logLevel = LogLevel.valueOf(context.getProperty(LOG_LEVEL).getValue());
}
@Override
public WriteResult sendData(RecordSet recordSet, Map<String, String> attributes, boolean sendZeroResults) throws IOException {
WriteResult writeResult;
try (final ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
ComponentLog log = getLogger();
try (final RecordSetWriter writer = writerFactory.createWriter(getLogger(), recordSet.getSchema(), baos, attributes)) {
writer.beginRecordSet();
Record r;
while ((r = recordSet.next()) != null) {
baos.reset();
writer.write(r);
writer.flush();
log.log(logLevel, baos.toString());
}
writeResult = writer.finishRecordSet();
writer.flush();
}
} catch (SchemaNotFoundException e) {
throw new IOException(e);
}
return writeResult;
}
}

View File

@ -13,3 +13,4 @@
# See the License for the specific language governing permissions and
# limitations under the License.
org.apache.nifi.record.sink.lookup.RecordSinkServiceLookup
org.apache.nifi.record.sink.LoggingRecordSink

View File

@ -0,0 +1,94 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.nifi.record.sink;
import org.apache.nifi.reporting.InitializationException;
import org.apache.nifi.serialization.SimpleRecordSchema;
import org.apache.nifi.serialization.WriteResult;
import org.apache.nifi.serialization.record.MapRecord;
import org.apache.nifi.serialization.record.MockRecordWriter;
import org.apache.nifi.serialization.record.Record;
import org.apache.nifi.serialization.record.RecordField;
import org.apache.nifi.serialization.record.RecordFieldType;
import org.apache.nifi.serialization.record.RecordSchema;
import org.apache.nifi.serialization.record.RecordSet;
import org.apache.nifi.util.TestRunner;
import org.apache.nifi.util.TestRunners;
import org.junit.Before;
import org.junit.Test;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.fail;
public class TestLoggingRecordSink {
private LoggingRecordSink recordSink;
private RecordSet recordSet;
@Before
public void setup() throws InitializationException {
TestRunner runner = TestRunners.newTestRunner(TestProcessor.class);
recordSink = new LoggingRecordSink();
runner.addControllerService("log", recordSink);
MockRecordWriter writerFactory = new MockRecordWriter();
runner.addControllerService("writer", writerFactory);
runner.setProperty(recordSink, LoggingRecordSink.RECORD_WRITER_FACTORY, "writer");
runner.enableControllerService(writerFactory);
runner.enableControllerService(recordSink);
final List<RecordField> fields = new ArrayList<>();
fields.add(new RecordField("a", RecordFieldType.STRING.getDataType()));
fields.add(new RecordField("b", RecordFieldType.BOOLEAN.getDataType()));
final RecordSchema schema = new SimpleRecordSchema(fields);
final Map<String, Object> valueMap1 = new HashMap<>();
valueMap1.put("a", "Hello");
valueMap1.put("b", true);
final Record record1 = new MapRecord(schema, valueMap1);
final Map<String, Object> valueMap2 = new HashMap<>();
valueMap2.put("a", "World");
valueMap2.put("b", false);
final Record record2 = new MapRecord(schema, valueMap2);
recordSet = RecordSet.of(schema, record1, record2);
}
@Test
public void testLogging() {
try {
final WriteResult writeResult = recordSink.sendData(recordSet, Collections.emptyMap(), false);
assertNotNull(writeResult);
assertEquals(2, writeResult.getRecordCount());
} catch (IOException ioe) {
fail("Should have completed successfully");
}
}
}

View File

@ -14,7 +14,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.nifi.record.sink.lookup;
package org.apache.nifi.record.sink;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.processor.AbstractProcessor;

View File

@ -17,6 +17,7 @@
package org.apache.nifi.record.sink.lookup;
import org.apache.nifi.record.sink.TestProcessor;
import org.apache.nifi.reporting.InitializationException;
import org.apache.nifi.serialization.SimpleRecordSchema;
import org.apache.nifi.serialization.WriteResult;