mirror of https://github.com/apache/nifi.git
NIFI-1156
This commit is contained in:
Jeremy Dyer
parent
8d37af07b9
commit
8966643d48
|
|
@ -0,0 +1,234 @@
|
|||
/*
|
||||
* 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.processors.kite;
|
||||
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.BufferedReader;
|
||||
import java.io.OutputStream;
|
||||
import java.io.InputStream;
|
||||
import java.io.InputStreamReader;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
import java.util.HashSet;
|
||||
import java.util.Collections;
|
||||
import java.util.ArrayList;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
import org.apache.nifi.annotation.behavior.InputRequirement;
|
||||
import org.apache.nifi.annotation.behavior.InputRequirement.Requirement;
|
||||
import org.apache.nifi.annotation.documentation.CapabilityDescription;
|
||||
import org.apache.nifi.annotation.documentation.SeeAlso;
|
||||
import org.apache.nifi.annotation.documentation.Tags;
|
||||
import org.apache.nifi.components.PropertyDescriptor;
|
||||
import org.apache.nifi.flowfile.FlowFile;
|
||||
import org.apache.nifi.processor.ProcessContext;
|
||||
import org.apache.nifi.processor.ProcessSession;
|
||||
import org.apache.nifi.processor.ProcessorInitializationContext;
|
||||
import org.apache.nifi.processor.Relationship;
|
||||
import org.apache.nifi.processor.exception.ProcessException;
|
||||
import org.apache.nifi.processor.io.InputStreamCallback;
|
||||
import org.apache.nifi.processor.io.OutputStreamCallback;
|
||||
import org.apache.nifi.processor.util.StandardValidators;
|
||||
import org.kitesdk.data.spi.filesystem.CSVProperties;
|
||||
import org.kitesdk.data.spi.filesystem.CSVUtil;
|
||||
import org.kitesdk.shaded.com.google.common.collect.ImmutableSet;
|
||||
|
||||
@Tags({"kite", "csv", "avro", "infer", "schema"})
|
||||
@SeeAlso({InferAvroSchemaFromCSV.class})
|
||||
@InputRequirement(Requirement.INPUT_REQUIRED)
|
||||
@CapabilityDescription("Creates an Avro schema from a CSV file header. The header line definition can either be provided" +
|
||||
"as a property to the processor OR present in the first line of CSV in the incoming FlowFile content. If a header" +
|
||||
" property is specified for this processor no attempt will be made to use the header line that may be present" +
|
||||
" in the incoming CSV FlowFile content.")
|
||||
public class InferAvroSchemaFromCSV
|
||||
extends AbstractKiteProcessor {
|
||||
|
||||
public static final String CSV_DELIMITER = ",";
|
||||
|
||||
public static final PropertyDescriptor HEADER_LINE = new PropertyDescriptor.Builder()
|
||||
.name("CSV Header Line")
|
||||
.description("Comma separated string defining the column names expected in the CSV data. " +
|
||||
"EX: \"fname,lname,zip,address\"")
|
||||
.addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
|
||||
.build();
|
||||
|
||||
public static final PropertyDescriptor HEADER_LINE_SKIP_COUNT = new PropertyDescriptor.Builder()
|
||||
.name("CSV Header Line Skip Count")
|
||||
.description("Specifies the number of header lines that should be skipped when reading the CSV data. If the " +
|
||||
" first line of the CSV data is a header line and you specify the \"CSV Header Line\" property " +
|
||||
"you need to set this vlaue to 1 otherwise the header line will be treated as actual data.")
|
||||
.required(true)
|
||||
.defaultValue("0")
|
||||
.addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
|
||||
.build();
|
||||
|
||||
public static final PropertyDescriptor ESCAPE_STRING = new PropertyDescriptor.Builder()
|
||||
.name("CSV escape string")
|
||||
.description("String that represents an escape sequence in the CSV FlowFile content data.")
|
||||
.required(true)
|
||||
.defaultValue("\\")
|
||||
.addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
|
||||
.build();
|
||||
|
||||
public static final PropertyDescriptor QUOTE_STRING = new PropertyDescriptor.Builder()
|
||||
.name("CSV quote string")
|
||||
.description("String that represents a literal quote character in the CSV FlowFile content data.")
|
||||
.required(true)
|
||||
.defaultValue("'")
|
||||
.addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
|
||||
.build();
|
||||
|
||||
public static final PropertyDescriptor RECORD_NAME = new PropertyDescriptor.Builder()
|
||||
.name("Avro Record Name")
|
||||
.description("Value to be placed in the Avro record schema \"name\" field.")
|
||||
.required(true)
|
||||
.addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
|
||||
.build();
|
||||
|
||||
public static final PropertyDescriptor CHARSET = new PropertyDescriptor.Builder()
|
||||
.name("Charset")
|
||||
.description("Character encoding of CSV data.")
|
||||
.required(true)
|
||||
.defaultValue("UTF-8")
|
||||
.addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
|
||||
.build();
|
||||
|
||||
public static final PropertyDescriptor PRETTY_AVRO_OUTPUT = new PropertyDescriptor.Builder()
|
||||
.name("Pretty Avro Output")
|
||||
.description("If true the Avro output will be formatted.")
|
||||
.required(true)
|
||||
.defaultValue("true")
|
||||
.allowableValues("true", "false")
|
||||
.addValidator(StandardValidators.BOOLEAN_VALIDATOR)
|
||||
.build();
|
||||
|
||||
|
||||
public static final Relationship REL_SUCCESS = new Relationship.Builder().name("success")
|
||||
.description("Successfully created Avro schema for CSV data.").build();
|
||||
|
||||
public static final Relationship REL_ORIGINAL = new Relationship.Builder().name("original")
|
||||
.description("Original incoming FlowFile CSV data").build();
|
||||
|
||||
public static final Relationship REL_FAILURE = new Relationship.Builder().name("failure")
|
||||
.description("Failed to create Avro schema for CSV data.").build();
|
||||
|
||||
private List<PropertyDescriptor> properties;
|
||||
private Set<Relationship> relationships;
|
||||
|
||||
@Override
|
||||
protected void init(final ProcessorInitializationContext context) {
|
||||
final List<PropertyDescriptor> properties = new ArrayList<>();
|
||||
properties.add(HEADER_LINE);
|
||||
properties.add(HEADER_LINE_SKIP_COUNT);
|
||||
properties.add(ESCAPE_STRING);
|
||||
properties.add(QUOTE_STRING);
|
||||
properties.add(PRETTY_AVRO_OUTPUT);
|
||||
properties.add(RECORD_NAME);
|
||||
properties.add(CHARSET);
|
||||
this.properties = Collections.unmodifiableList(properties);
|
||||
|
||||
final Set<Relationship> relationships = new HashSet<>();
|
||||
relationships.add(REL_SUCCESS);
|
||||
relationships.add(REL_FAILURE);
|
||||
relationships.add(REL_ORIGINAL);
|
||||
this.relationships = Collections.unmodifiableSet(relationships);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
|
||||
return properties;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Set<Relationship> getRelationships() {
|
||||
return relationships;
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
public void onTrigger(final ProcessContext context, ProcessSession session) throws ProcessException {
|
||||
final FlowFile original = session.get();
|
||||
if (original == null) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
|
||||
//Determines the header line either from the property input or the first line of the delimited file.
|
||||
final AtomicReference<String> header = new AtomicReference<>();
|
||||
final AtomicReference<Boolean> hasHeader = new AtomicReference<>();
|
||||
|
||||
if (context.getProperty(HEADER_LINE).isSet()) {
|
||||
header.set(context.getProperty(HEADER_LINE).getValue());
|
||||
hasHeader.set(Boolean.FALSE);
|
||||
} else {
|
||||
//Read the first line of the file to get the header value.
|
||||
session.read(original, new InputStreamCallback() {
|
||||
@Override
|
||||
public void process(InputStream in) throws IOException {
|
||||
BufferedReader br = new BufferedReader(new InputStreamReader(in));
|
||||
header.set(br.readLine());
|
||||
hasHeader.set(Boolean.TRUE);
|
||||
br.close();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
//Prepares the CSVProperties for kite
|
||||
final CSVProperties props = new CSVProperties.Builder()
|
||||
.delimiter(CSV_DELIMITER)
|
||||
.escape(context.getProperty(ESCAPE_STRING).getValue())
|
||||
.quote(context.getProperty(QUOTE_STRING).getValue())
|
||||
.header(header.get())
|
||||
.hasHeader(hasHeader.get())
|
||||
.linesToSkip(context.getProperty(HEADER_LINE_SKIP_COUNT).asInteger())
|
||||
.charset(context.getProperty(CHARSET).getValue())
|
||||
.build();
|
||||
|
||||
final Set<String> required = ImmutableSet.of();
|
||||
final AtomicReference<String> avroSchema = new AtomicReference<>();
|
||||
|
||||
session.read(original, new InputStreamCallback() {
|
||||
@Override
|
||||
public void process(InputStream in) throws IOException {
|
||||
avroSchema.set(CSVUtil
|
||||
.inferNullableSchema(
|
||||
context.getProperty(RECORD_NAME).getValue(), in, props, required)
|
||||
.toString(context.getProperty(PRETTY_AVRO_OUTPUT).asBoolean()));
|
||||
}
|
||||
});
|
||||
|
||||
FlowFile avroSchemaFF = session.write(session.create(), new OutputStreamCallback() {
|
||||
@Override
|
||||
public void process(OutputStream out) throws IOException {
|
||||
out.write(avroSchema.get().getBytes());
|
||||
}
|
||||
});
|
||||
|
||||
//Transfer the sessions.
|
||||
session.transfer(original, REL_ORIGINAL);
|
||||
session.transfer(avroSchemaFF, REL_SUCCESS);
|
||||
|
||||
} catch (Exception ex) {
|
||||
getLogger().error(ex.getMessage());
|
||||
session.transfer(original, REL_FAILURE);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,164 @@
|
|||
/*
|
||||
* 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.processors.kite;
|
||||
|
||||
import org.apache.avro.Schema;
|
||||
import org.apache.nifi.annotation.behavior.InputRequirement;
|
||||
import org.apache.nifi.annotation.documentation.CapabilityDescription;
|
||||
import org.apache.nifi.annotation.documentation.SeeAlso;
|
||||
import org.apache.nifi.annotation.documentation.Tags;
|
||||
import org.apache.nifi.components.PropertyDescriptor;
|
||||
import org.apache.nifi.flowfile.FlowFile;
|
||||
import org.apache.nifi.processor.ProcessContext;
|
||||
import org.apache.nifi.processor.ProcessSession;
|
||||
import org.apache.nifi.processor.ProcessorInitializationContext;
|
||||
import org.apache.nifi.processor.Relationship;
|
||||
import org.apache.nifi.processor.exception.ProcessException;
|
||||
import org.apache.nifi.processor.io.InputStreamCallback;
|
||||
import org.apache.nifi.processor.io.OutputStreamCallback;
|
||||
import org.apache.nifi.processor.util.StandardValidators;
|
||||
import org.kitesdk.data.spi.JsonUtil;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.io.OutputStream;
|
||||
import java.util.List;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Set;
|
||||
import java.util.HashSet;
|
||||
import java.util.Collections;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
@Tags({"kite", "json", "avro", "infer", "schema"})
|
||||
@SeeAlso({InferAvroSchemaFromJSON.class})
|
||||
@InputRequirement(InputRequirement.Requirement.INPUT_REQUIRED)
|
||||
@CapabilityDescription("Creates an Avro schema from JSON data. The Avro schema is inferred by examining the fields " +
|
||||
"in the JSON input. The Avro schema generated by kite will use the same names present in the incoming JSON payload")
|
||||
public class InferAvroSchemaFromJSON
|
||||
extends AbstractKiteProcessor {
|
||||
|
||||
public static final PropertyDescriptor RECORD_NAME = new PropertyDescriptor.Builder()
|
||||
.name("Avro Record Name")
|
||||
.description("Value to be placed in the Avro record schema \"name\" field.")
|
||||
.required(true)
|
||||
.addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
|
||||
.build();
|
||||
|
||||
public static final PropertyDescriptor NUM_RECORDS_TO_ANALYZE = new PropertyDescriptor.Builder()
|
||||
.name("Number of records to analyze")
|
||||
.description("Number of records that should be analyzed by kite to infer the Avro schema")
|
||||
.required(true)
|
||||
.defaultValue("10")
|
||||
.addValidator(StandardValidators.INTEGER_VALIDATOR)
|
||||
.build();
|
||||
|
||||
public static final PropertyDescriptor CHARSET = new PropertyDescriptor.Builder()
|
||||
.name("Charset")
|
||||
.description("Character encoding of CSV data.")
|
||||
.required(true)
|
||||
.defaultValue("UTF-8")
|
||||
.addValidator(StandardValidators.NON_EMPTY_VALIDATOR)
|
||||
.build();
|
||||
|
||||
public static final PropertyDescriptor PRETTY_AVRO_OUTPUT = new PropertyDescriptor.Builder()
|
||||
.name("Pretty Avro Output")
|
||||
.description("If true the Avro output will be formatted.")
|
||||
.required(true)
|
||||
.defaultValue("true")
|
||||
.allowableValues("true", "false")
|
||||
.addValidator(StandardValidators.BOOLEAN_VALIDATOR)
|
||||
.build();
|
||||
|
||||
|
||||
public static final Relationship REL_SUCCESS = new Relationship.Builder().name("success")
|
||||
.description("Successfully created Avro schema for JSON data.").build();
|
||||
|
||||
public static final Relationship REL_ORIGINAL = new Relationship.Builder().name("original")
|
||||
.description("Original incoming FlowFile JSON data").build();
|
||||
|
||||
public static final Relationship REL_FAILURE = new Relationship.Builder().name("failure")
|
||||
.description("Failed to create Avro schema for JSON data.").build();
|
||||
|
||||
private List<PropertyDescriptor> properties;
|
||||
private Set<Relationship> relationships;
|
||||
|
||||
@Override
|
||||
protected void init(final ProcessorInitializationContext context) {
|
||||
final List<PropertyDescriptor> properties = new ArrayList<>();
|
||||
properties.add(CHARSET);
|
||||
properties.add(PRETTY_AVRO_OUTPUT);
|
||||
properties.add(RECORD_NAME);
|
||||
properties.add(NUM_RECORDS_TO_ANALYZE);
|
||||
this.properties = Collections.unmodifiableList(properties);
|
||||
|
||||
final Set<Relationship> relationships = new HashSet<>();
|
||||
relationships.add(REL_SUCCESS);
|
||||
relationships.add(REL_FAILURE);
|
||||
relationships.add(REL_ORIGINAL);
|
||||
this.relationships = Collections.unmodifiableSet(relationships);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
|
||||
return properties;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Set<Relationship> getRelationships() {
|
||||
return relationships;
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
public void onTrigger(final ProcessContext context, ProcessSession session) throws ProcessException {
|
||||
final FlowFile original = session.get();
|
||||
if (original == null) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
|
||||
final AtomicReference<String> avroSchema = new AtomicReference<>();
|
||||
session.read(original, new InputStreamCallback() {
|
||||
@Override
|
||||
public void process(InputStream in) throws IOException {
|
||||
Schema as = JsonUtil.inferSchema(
|
||||
in, context.getProperty(RECORD_NAME).getValue(), context.getProperty(NUM_RECORDS_TO_ANALYZE).asInteger());
|
||||
avroSchema.set(as.toString(context.getProperty(PRETTY_AVRO_OUTPUT).asBoolean()));
|
||||
|
||||
}
|
||||
});
|
||||
|
||||
FlowFile avroSchemaFF = session.write(session.create(), new OutputStreamCallback() {
|
||||
@Override
|
||||
public void process(OutputStream out) throws IOException {
|
||||
out.write(avroSchema.get().getBytes());
|
||||
}
|
||||
});
|
||||
|
||||
//Transfer the FlowFiles
|
||||
session.transfer(original, REL_ORIGINAL);
|
||||
session.transfer(avroSchemaFF, REL_SUCCESS);
|
||||
|
||||
} catch (Exception ex) {
|
||||
getLogger().error(ex.getMessage());
|
||||
session.transfer(original, REL_FAILURE);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -16,3 +16,5 @@ org.apache.nifi.processors.kite.StoreInKiteDataset
|
|||
org.apache.nifi.processors.kite.ConvertCSVToAvro
|
||||
org.apache.nifi.processors.kite.ConvertJSONToAvro
|
||||
org.apache.nifi.processors.kite.ConvertAvroSchema
|
||||
org.apache.nifi.processors.kite.InferAvroSchemaFromCSV
|
||||
org.apache.nifi.processors.kite.InferAvroSchemaFromJSON
|
||||
|
|
|
|||
|
|
@ -0,0 +1,129 @@
|
|||
/*
|
||||
* 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.processors.kite;
|
||||
;
|
||||
import org.apache.nifi.flowfile.FlowFile;
|
||||
import org.apache.nifi.processor.ProcessSession;
|
||||
import org.apache.nifi.processor.io.OutputStreamCallback;
|
||||
import org.apache.nifi.processor.io.StreamCallback;
|
||||
import org.apache.nifi.util.TestRunner;
|
||||
import org.apache.nifi.util.TestRunners;
|
||||
import org.junit.Test;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.io.OutputStream;
|
||||
|
||||
public class TestInferAvroSchemaFromCSV {
|
||||
|
||||
private final String CSV_HEADER_LINE = "fname,lname,age,zip";
|
||||
|
||||
@Test
|
||||
public void inferSchemaFromHeaderLineOfCSV() throws Exception {
|
||||
TestRunner runner = TestRunners.newTestRunner(InferAvroSchemaFromCSV.class);
|
||||
|
||||
runner.assertNotValid();
|
||||
runner.setProperty(InferAvroSchemaFromCSV.HEADER_LINE_SKIP_COUNT, "0");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.ESCAPE_STRING, "\\");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.QUOTE_STRING, "'");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.RECORD_NAME, "contact");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.CHARSET, "UTF-8");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.PRETTY_AVRO_OUTPUT, "true");
|
||||
|
||||
runner.assertValid();
|
||||
|
||||
ProcessSession session = runner.getProcessSessionFactory().createSession();
|
||||
FlowFile ff = session.write(session.create(), new OutputStreamCallback() {
|
||||
@Override
|
||||
public void process(OutputStream out) throws IOException {
|
||||
out.write((CSV_HEADER_LINE + "\nJeremy,Dyer,29,55555").getBytes());
|
||||
}
|
||||
});
|
||||
|
||||
//Enqueue the empty FlowFile
|
||||
runner.enqueue(ff);
|
||||
runner.run();
|
||||
runner.assertTransferCount(InferAvroSchemaFromCSV.REL_FAILURE, 0);
|
||||
runner.assertTransferCount(InferAvroSchemaFromCSV.REL_ORIGINAL, 1);
|
||||
runner.assertTransferCount(InferAvroSchemaFromCSV.REL_SUCCESS, 1);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void inferSchemaFormHeaderLinePropertyOfProcessor() throws Exception {
|
||||
TestRunner runner = TestRunners.newTestRunner(InferAvroSchemaFromCSV.class);
|
||||
|
||||
runner.assertNotValid();
|
||||
runner.setProperty(InferAvroSchemaFromCSV.HEADER_LINE, CSV_HEADER_LINE);
|
||||
runner.setProperty(InferAvroSchemaFromCSV.HEADER_LINE_SKIP_COUNT, "1");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.ESCAPE_STRING, "\\");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.QUOTE_STRING, "'");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.RECORD_NAME, "contact");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.CHARSET, "UTF-8");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.PRETTY_AVRO_OUTPUT, "true");
|
||||
|
||||
runner.assertValid();
|
||||
|
||||
ProcessSession session = runner.getProcessSessionFactory().createSession();
|
||||
FlowFile ff = session.write(session.create(), new StreamCallback() {
|
||||
@Override
|
||||
public void process(InputStream in, OutputStream out) throws IOException {
|
||||
out.write((CSV_HEADER_LINE + "\nJeremy,Dyer,29,55555").getBytes());
|
||||
}
|
||||
});
|
||||
|
||||
//Enqueue the empty FlowFile
|
||||
runner.enqueue(ff);
|
||||
runner.run();
|
||||
runner.assertTransferCount(InferAvroSchemaFromCSV.REL_FAILURE, 0);
|
||||
runner.assertTransferCount(InferAvroSchemaFromCSV.REL_ORIGINAL, 1);
|
||||
runner.assertTransferCount(InferAvroSchemaFromCSV.REL_SUCCESS, 1);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void inferSchemaFromEmptyContent() throws Exception {
|
||||
TestRunner runner = TestRunners.newTestRunner(InferAvroSchemaFromCSV.class);
|
||||
|
||||
runner.assertNotValid();
|
||||
runner.setProperty(InferAvroSchemaFromCSV.HEADER_LINE, CSV_HEADER_LINE);
|
||||
runner.setProperty(InferAvroSchemaFromCSV.HEADER_LINE_SKIP_COUNT, "1");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.ESCAPE_STRING, "\\");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.QUOTE_STRING, "'");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.RECORD_NAME, "contact");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.CHARSET, "UTF-8");
|
||||
runner.setProperty(InferAvroSchemaFromCSV.PRETTY_AVRO_OUTPUT, "true");
|
||||
|
||||
runner.assertValid();
|
||||
|
||||
ProcessSession session = runner.getProcessSessionFactory().createSession();
|
||||
FlowFile ff = session.write(session.create(), new StreamCallback() {
|
||||
@Override
|
||||
public void process(InputStream in, OutputStream out) throws IOException {
|
||||
out.write("".getBytes());
|
||||
}
|
||||
});
|
||||
|
||||
//Enqueue the empty FlowFile
|
||||
runner.enqueue(ff);
|
||||
runner.run();
|
||||
runner.assertTransferCount(InferAvroSchemaFromCSV.REL_FAILURE, 1);
|
||||
runner.assertTransferCount(InferAvroSchemaFromCSV.REL_ORIGINAL, 0);
|
||||
runner.assertTransferCount(InferAvroSchemaFromCSV.REL_SUCCESS, 0);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,215 @@
|
|||
/*
|
||||
* 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.processors.kite;
|
||||
|
||||
import com.google.common.collect.Lists;
|
||||
import org.apache.avro.Schema;
|
||||
import org.apache.avro.SchemaBuilder;
|
||||
import org.apache.avro.file.DataFileStream;
|
||||
import org.apache.avro.generic.GenericData;
|
||||
import org.apache.avro.generic.GenericDatumReader;
|
||||
import org.apache.nifi.util.MockFlowFile;
|
||||
import org.apache.nifi.util.TestRunner;
|
||||
import org.apache.nifi.util.TestRunners;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.IOException;
|
||||
import java.util.List;
|
||||
|
||||
import static org.apache.nifi.processors.kite.TestUtil.streamFor;
|
||||
|
||||
public class TestInferAvroSchemaFromJSON {
|
||||
|
||||
public static final Schema INPUT_SCHEMA = SchemaBuilder.record("InputTest")
|
||||
.fields().requiredString("id").requiredString("primaryColor")
|
||||
.optionalString("secondaryColor").optionalString("price")
|
||||
.endRecord();
|
||||
|
||||
public static final Schema OUTPUT_SCHEMA = SchemaBuilder.record("Test")
|
||||
.fields().requiredLong("id").requiredString("color")
|
||||
.optionalDouble("price").endRecord();
|
||||
|
||||
public static final String MAPPING = "[{\"source\":\"primaryColor\", \"target\":\"color\"}]";
|
||||
|
||||
public static final String FAILURE_SUMMARY = "Cannot convert free to double";
|
||||
|
||||
@Test
|
||||
public void testBasicConversion() throws IOException {
|
||||
TestRunner runner = TestRunners.newTestRunner(ConvertAvroSchema.class);
|
||||
runner.assertNotValid();
|
||||
runner.setProperty(ConvertAvroSchema.INPUT_SCHEMA,
|
||||
INPUT_SCHEMA.toString());
|
||||
runner.setProperty(ConvertAvroSchema.OUTPUT_SCHEMA,
|
||||
OUTPUT_SCHEMA.toString());
|
||||
runner.setProperty("primaryColor", "color");
|
||||
runner.assertValid();
|
||||
|
||||
// Two valid rows, and one invalid because "free" is not a double.
|
||||
GenericData.Record goodRecord1 = dataBasic("1", "blue", null, null);
|
||||
GenericData.Record goodRecord2 = dataBasic("2", "red", "yellow", "5.5");
|
||||
GenericData.Record badRecord = dataBasic("3", "red", "yellow", "free");
|
||||
List<GenericData.Record> input = Lists.newArrayList(goodRecord1, goodRecord2,
|
||||
badRecord);
|
||||
|
||||
runner.enqueue(streamFor(input));
|
||||
runner.run();
|
||||
|
||||
long converted = runner.getCounterValue("Converted records");
|
||||
long errors = runner.getCounterValue("Conversion errors");
|
||||
Assert.assertEquals("Should convert 2 rows", 2, converted);
|
||||
Assert.assertEquals("Should reject 1 rows", 1, errors);
|
||||
|
||||
runner.assertTransferCount("success", 1);
|
||||
runner.assertTransferCount("failure", 1);
|
||||
|
||||
MockFlowFile incompatible = runner.getFlowFilesForRelationship(
|
||||
"failure").get(0);
|
||||
GenericDatumReader<GenericData.Record> reader = new GenericDatumReader<GenericData.Record>(
|
||||
INPUT_SCHEMA);
|
||||
DataFileStream<GenericData.Record> stream = new DataFileStream<GenericData.Record>(
|
||||
new ByteArrayInputStream(
|
||||
runner.getContentAsByteArray(incompatible)), reader);
|
||||
int count = 0;
|
||||
for (GenericData.Record r : stream) {
|
||||
Assert.assertEquals(badRecord, r);
|
||||
count++;
|
||||
}
|
||||
stream.close();
|
||||
Assert.assertEquals(1, count);
|
||||
Assert.assertEquals("Should accumulate error messages",
|
||||
FAILURE_SUMMARY, incompatible.getAttribute("errors"));
|
||||
|
||||
GenericDatumReader<GenericData.Record> successReader = new GenericDatumReader<GenericData.Record>(
|
||||
OUTPUT_SCHEMA);
|
||||
DataFileStream<GenericData.Record> successStream = new DataFileStream<GenericData.Record>(
|
||||
new ByteArrayInputStream(runner.getContentAsByteArray(runner
|
||||
.getFlowFilesForRelationship("success").get(0))),
|
||||
successReader);
|
||||
count = 0;
|
||||
for (GenericData.Record r : successStream) {
|
||||
if (count == 0) {
|
||||
Assert.assertEquals(convertBasic(goodRecord1), r);
|
||||
} else {
|
||||
Assert.assertEquals(convertBasic(goodRecord2), r);
|
||||
}
|
||||
count++;
|
||||
}
|
||||
successStream.close();
|
||||
Assert.assertEquals(2, count);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testNestedConversion() throws IOException {
|
||||
TestRunner runner = TestRunners.newTestRunner(ConvertAvroSchema.class);
|
||||
runner.assertNotValid();
|
||||
runner.setProperty(ConvertAvroSchema.INPUT_SCHEMA,
|
||||
TestAvroRecordConverter.NESTED_RECORD_SCHEMA.toString());
|
||||
runner.setProperty(ConvertAvroSchema.OUTPUT_SCHEMA,
|
||||
TestAvroRecordConverter.UNNESTED_OUTPUT_SCHEMA.toString());
|
||||
runner.setProperty("parent.id", "parentId");
|
||||
runner.assertValid();
|
||||
|
||||
// Two valid rows
|
||||
GenericData.Record goodRecord1 = dataNested(1L, "200", null, null);
|
||||
GenericData.Record goodRecord2 = dataNested(2L, "300", 5L, "ParentCompany");
|
||||
List<GenericData.Record> input = Lists.newArrayList(goodRecord1, goodRecord2);
|
||||
|
||||
runner.enqueue(streamFor(input));
|
||||
runner.run();
|
||||
|
||||
long converted = runner.getCounterValue("Converted records");
|
||||
long errors = runner.getCounterValue("Conversion errors");
|
||||
Assert.assertEquals("Should convert 2 rows", 2, converted);
|
||||
Assert.assertEquals("Should reject 0 rows", 0, errors);
|
||||
|
||||
runner.assertTransferCount("success", 1);
|
||||
runner.assertTransferCount("failure", 0);
|
||||
|
||||
GenericDatumReader<GenericData.Record> successReader = new GenericDatumReader<GenericData.Record>(
|
||||
TestAvroRecordConverter.UNNESTED_OUTPUT_SCHEMA);
|
||||
DataFileStream<GenericData.Record> successStream = new DataFileStream<GenericData.Record>(
|
||||
new ByteArrayInputStream(runner.getContentAsByteArray(runner
|
||||
.getFlowFilesForRelationship("success").get(0))),
|
||||
successReader);
|
||||
int count = 0;
|
||||
for (GenericData.Record r : successStream) {
|
||||
if (count == 0) {
|
||||
Assert.assertEquals(convertNested(goodRecord1), r);
|
||||
} else {
|
||||
Assert.assertEquals(convertNested(goodRecord2), r);
|
||||
}
|
||||
count++;
|
||||
}
|
||||
successStream.close();
|
||||
Assert.assertEquals(2, count);
|
||||
}
|
||||
|
||||
private GenericData.Record convertBasic(GenericData.Record inputRecord) {
|
||||
GenericData.Record result = new GenericData.Record(OUTPUT_SCHEMA);
|
||||
result.put("id", Long.parseLong(inputRecord.get("id").toString()));
|
||||
result.put("color", inputRecord.get("primaryColor").toString());
|
||||
if (inputRecord.get("price") == null) {
|
||||
result.put("price", null);
|
||||
} else {
|
||||
result.put("price",
|
||||
Double.parseDouble(inputRecord.get("price").toString()));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private GenericData.Record dataBasic(String id, String primaryColor,
|
||||
String secondaryColor, String price) {
|
||||
GenericData.Record result = new GenericData.Record(INPUT_SCHEMA);
|
||||
result.put("id", id);
|
||||
result.put("primaryColor", primaryColor);
|
||||
result.put("secondaryColor", secondaryColor);
|
||||
result.put("price", price);
|
||||
return result;
|
||||
}
|
||||
|
||||
private GenericData.Record convertNested(GenericData.Record inputRecord) {
|
||||
GenericData.Record result = new GenericData.Record(
|
||||
TestAvroRecordConverter.UNNESTED_OUTPUT_SCHEMA);
|
||||
result.put("l1", inputRecord.get("l1"));
|
||||
result.put("s1", Long.parseLong(inputRecord.get("s1").toString()));
|
||||
if (inputRecord.get("parent") != null) {
|
||||
// output schema doesn't have parent name.
|
||||
result.put("parentId",
|
||||
((GenericData.Record) inputRecord.get("parent")).get("id"));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private GenericData.Record dataNested(long id, String companyName, Long parentId,
|
||||
String parentName) {
|
||||
GenericData.Record result = new GenericData.Record(TestAvroRecordConverter.NESTED_RECORD_SCHEMA);
|
||||
result.put("l1", id);
|
||||
result.put("s1", companyName);
|
||||
if (parentId != null || parentName != null) {
|
||||
GenericData.Record parent = new GenericData.Record(
|
||||
TestAvroRecordConverter.NESTED_PARENT_SCHEMA);
|
||||
parent.put("id", parentId);
|
||||
parent.put("name", parentName);
|
||||
result.put("parent", parent);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue