http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/main/java/org/apache/nifi/serialization/record/util/DataTypeUtils.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/main/java/org/apache/nifi/serialization/record/util/DataTypeUtils.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/main/java/org/apache/nifi/serialization/record/util/DataTypeUtils.java index 1cdefb8..05b3157 100644 --- a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/main/java/org/apache/nifi/serialization/record/util/DataTypeUtils.java +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/main/java/org/apache/nifi/serialization/record/util/DataTypeUtils.java @@ -42,44 +42,46 @@ public class DataTypeUtils { private static final TimeZone gmt = TimeZone.getTimeZone("gmt"); - public static Object convertType(final Object value, final DataType dataType) { - return convertType(value, dataType, RecordFieldType.DATE.getDefaultFormat(), RecordFieldType.TIME.getDefaultFormat(), RecordFieldType.TIMESTAMP.getDefaultFormat()); + public static Object convertType(final Object value, final DataType dataType, final String fieldName) { + return convertType(value, dataType, RecordFieldType.DATE.getDefaultFormat(), RecordFieldType.TIME.getDefaultFormat(), RecordFieldType.TIMESTAMP.getDefaultFormat(), fieldName); } - public static Object convertType(final Object value, final DataType dataType, final String dateFormat, final String timeFormat, final String timestampFormat) { + public static Object convertType(final Object value, final DataType dataType, final String dateFormat, final String timeFormat, final String timestampFormat, final String fieldName) { switch (dataType.getFieldType()) { case BIGINT: - return toBigInt(value); + return toBigInt(value, fieldName); case BOOLEAN: - return toBoolean(value); + return toBoolean(value, fieldName); case BYTE: - return toByte(value); + return toByte(value, fieldName); case CHAR: - return toCharacter(value); + return toCharacter(value, fieldName); case DATE: - return toDate(value, dateFormat); + return toDate(value, dateFormat, fieldName); case DOUBLE: - return toDouble(value); + return toDouble(value, fieldName); case FLOAT: - return toFloat(value); + return toFloat(value, fieldName); case INT: - return toInteger(value); + return toInteger(value, fieldName); case LONG: - return toLong(value); + return toLong(value, fieldName); case SHORT: - return toShort(value); + return toShort(value, fieldName); case STRING: return toString(value, dateFormat, timeFormat, timestampFormat); case TIME: - return toTime(value, timeFormat); + return toTime(value, timeFormat, fieldName); case TIMESTAMP: - return toTimestamp(value, timestampFormat); + return toTimestamp(value, timestampFormat, fieldName); case ARRAY: - return toArray(value); + return toArray(value, fieldName); + case MAP: + return toMap(value, fieldName); case RECORD: final RecordDataType recordType = (RecordDataType) dataType; final RecordSchema childSchema = recordType.getChildSchema(); - return toRecord(value, childSchema); + return toRecord(value, childSchema, fieldName); case CHOICE: { if (value == null) { return null; @@ -89,10 +91,10 @@ public class DataTypeUtils { final DataType chosenDataType = chooseDataType(value, choiceDataType); if (chosenDataType == null) { throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() - + " to any of the following available Sub-Types for a Choice: " + choiceDataType.getPossibleSubTypes()); + + " for field " + fieldName + " to any of the following available Sub-Types for a Choice: " + choiceDataType.getPossibleSubTypes()); } - return convertType(value, chosenDataType); + return convertType(value, chosenDataType, fieldName); } } @@ -132,6 +134,8 @@ public class DataTypeUtils { return isTimestampTypeCompatible(value, dataType.getFormat()); case STRING: return isStringTypeCompatible(value); + case MAP: + return isMapTypeCompatible(value); case CHOICE: { final DataType chosenDataType = chooseDataType(value, (ChoiceDataType) dataType); return chosenDataType != null; @@ -151,7 +155,7 @@ public class DataTypeUtils { return null; } - public static Record toRecord(final Object value, final RecordSchema recordSchema) { + public static Record toRecord(final Object value, final RecordSchema recordSchema, final String fieldName) { if (value == null) { return null; } @@ -163,7 +167,7 @@ public class DataTypeUtils { if (value instanceof Map) { if (recordSchema == null) { throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() - + " to Record because the value is a Map but no Record Schema was provided"); + + " to Record for field " + fieldName + " because the value is a Map but no Record Schema was provided"); } final Map<?, ?> map = (Map<?, ?>) value; @@ -182,21 +186,21 @@ public class DataTypeUtils { } final Object rawValue = entry.getValue(); - final Object coercedValue = convertType(rawValue, desiredTypeOption.get()); + final Object coercedValue = convertType(rawValue, desiredTypeOption.get(), fieldName); coercedValues.put(key, coercedValue); } return new MapRecord(recordSchema, coercedValues); } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Record"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Record for field " + fieldName); } public static boolean isRecordTypeCompatible(final Object value) { return value != null && value instanceof Record; } - public static Object[] toArray(final Object value) { + public static Object[] toArray(final Object value, final String fieldName) { if (value == null) { return null; } @@ -205,13 +209,70 @@ public class DataTypeUtils { return (Object[]) value; } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Object Array"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Object Array for field " + fieldName); } public static boolean isArrayTypeCompatible(final Object value) { return value != null && value instanceof Object[]; } + @SuppressWarnings("unchecked") + public static Map<String, Object> toMap(final Object value, final String fieldName) { + if (value == null) { + return null; + } + + if (value instanceof Map) { + final Map<?, ?> original = (Map<?, ?>) value; + + boolean keysAreStrings = true; + for (final Object key : original.keySet()) { + if (!(key instanceof String)) { + keysAreStrings = false; + } + } + + if (keysAreStrings) { + return (Map<String, Object>) value; + } + + final Map<String, Object> transformed = new HashMap<>(); + for (final Map.Entry<?, ?> entry : original.entrySet()) { + final Object key = entry.getKey(); + if (key == null) { + transformed.put(null, entry.getValue()); + } else { + transformed.put(key.toString(), entry.getValue()); + } + } + + return transformed; + } + + if (value instanceof Record) { + final Record record = (Record) value; + final RecordSchema recordSchema = record.getSchema(); + if (recordSchema == null) { + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type Record to Map for field " + fieldName + + " because Record does not have an associated Schema"); + } + + final Map<String, Object> map = new HashMap<>(); + for (final String recordFieldName : recordSchema.getFieldNames()) { + map.put(recordFieldName, record.getValue(recordFieldName)); + } + + return map; + } + + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Map for field " + fieldName); + } + + public static boolean isMapTypeCompatible(final Object value) { + return value != null && value instanceof Map; + } + + public static String toString(final Object value, final String dateFormat, final String timeFormat, final String timestampFormat) { if (value == null) { return null; @@ -238,10 +299,10 @@ public class DataTypeUtils { } public static boolean isStringTypeCompatible(final Object value) { - return value != null && (value instanceof String || value instanceof java.util.Date); + return value != null; } - public static java.sql.Date toDate(final Object value, final String format) { + public static java.sql.Date toDate(final Object value, final String format, final String fieldName) { if (value == null) { return null; } @@ -261,11 +322,11 @@ public class DataTypeUtils { return new Date(utilDate.getTime()); } catch (final ParseException e) { throw new IllegalTypeConversionException("Could not convert value [" + value - + "] of type java.lang.String to Date because the value is not in the expected date format: " + format); + + "] of type java.lang.String to Date because the value is not in the expected date format: " + format + " for field " + fieldName); } } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Date"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Date for field " + fieldName); } public static boolean isDateTypeCompatible(final Object value, final String format) { @@ -289,7 +350,7 @@ public class DataTypeUtils { return false; } - public static Time toTime(final Object value, final String format) { + public static Time toTime(final Object value, final String format, final String fieldName) { if (value == null) { return null; } @@ -309,11 +370,11 @@ public class DataTypeUtils { return new Time(utilDate.getTime()); } catch (final ParseException e) { throw new IllegalTypeConversionException("Could not convert value [" + value - + "] of type java.lang.String to Time because the value is not in the expected date format: " + format); + + "] of type java.lang.String to Time for field " + fieldName + " because the value is not in the expected date format: " + format); } } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Time"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Time for field " + fieldName); } private static DateFormat getDateFormat(final String format) { @@ -326,7 +387,7 @@ public class DataTypeUtils { return isDateTypeCompatible(value, format); } - public static Timestamp toTimestamp(final Object value, final String format) { + public static Timestamp toTimestamp(final Object value, final String format, final String fieldName) { if (value == null) { return null; } @@ -346,11 +407,11 @@ public class DataTypeUtils { return new Timestamp(utilDate.getTime()); } catch (final ParseException e) { throw new IllegalTypeConversionException("Could not convert value [" + value - + "] of type java.lang.String to Timestamp because the value is not in the expected date format: " + format); + + "] of type java.lang.String to Timestamp for field " + fieldName + " because the value is not in the expected date format: " + format); } } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Timestamp"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Timestamp for field " + fieldName); } public static boolean isTimestampTypeCompatible(final Object value, final String format) { @@ -358,7 +419,7 @@ public class DataTypeUtils { } - public static BigInteger toBigInt(final Object value) { + public static BigInteger toBigInt(final Object value, final String fieldName) { if (value == null) { return null; } @@ -370,14 +431,14 @@ public class DataTypeUtils { return BigInteger.valueOf((Long) value); } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to BigInteger"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to BigInteger for field " + fieldName); } public static boolean isBigIntTypeCompatible(final Object value) { return value == null && (value instanceof BigInteger || value instanceof Long); } - public static Boolean toBoolean(final Object value) { + public static Boolean toBoolean(final Object value, final String fieldName) { if (value == null) { return null; } @@ -394,7 +455,7 @@ public class DataTypeUtils { } } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Boolean"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Boolean for field " + fieldName); } public static boolean isBooleanTypeCompatible(final Object value) { @@ -411,7 +472,7 @@ public class DataTypeUtils { return false; } - public static Double toDouble(final Object value) { + public static Double toDouble(final Object value, final String fieldName) { if (value == null) { return null; } @@ -424,7 +485,7 @@ public class DataTypeUtils { return Double.parseDouble((String) value); } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Double"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Double for field " + fieldName); } public static boolean isDoubleTypeCompatible(final Object value) { @@ -452,7 +513,7 @@ public class DataTypeUtils { return false; } - public static Float toFloat(final Object value) { + public static Float toFloat(final Object value, final String fieldName) { if (value == null) { return null; } @@ -465,14 +526,14 @@ public class DataTypeUtils { return Float.parseFloat((String) value); } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Float"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Float for field " + fieldName); } public static boolean isFloatTypeCompatible(final Object value) { return isNumberTypeCompatible(value, s -> Float.parseFloat(s)); } - public static Long toLong(final Object value) { + public static Long toLong(final Object value, final String fieldName) { if (value == null) { return null; } @@ -489,7 +550,7 @@ public class DataTypeUtils { return ((java.util.Date) value).getTime(); } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Long"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Long for field " + fieldName); } public static boolean isLongTypeCompatible(final Object value) { @@ -518,7 +579,7 @@ public class DataTypeUtils { } - public static Integer toInteger(final Object value) { + public static Integer toInteger(final Object value, final String fieldName) { if (value == null) { return null; } @@ -531,7 +592,7 @@ public class DataTypeUtils { return Integer.parseInt((String) value); } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Integer"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Integer for field " + fieldName); } public static boolean isIntegerTypeCompatible(final Object value) { @@ -539,7 +600,7 @@ public class DataTypeUtils { } - public static Short toShort(final Object value) { + public static Short toShort(final Object value, final String fieldName) { if (value == null) { return null; } @@ -552,14 +613,14 @@ public class DataTypeUtils { return Short.parseShort((String) value); } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Short"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Short for field " + fieldName); } public static boolean isShortTypeCompatible(final Object value) { return isNumberTypeCompatible(value, s -> Short.parseShort(s)); } - public static Byte toByte(final Object value) { + public static Byte toByte(final Object value, final String fieldName) { if (value == null) { return null; } @@ -572,7 +633,7 @@ public class DataTypeUtils { return Byte.parseByte((String) value); } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Byte"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Byte for field " + fieldName); } public static boolean isByteTypeCompatible(final Object value) { @@ -580,7 +641,7 @@ public class DataTypeUtils { } - public static Character toCharacter(final Object value) { + public static Character toCharacter(final Object value, final String fieldName) { if (value == null) { return null; } @@ -592,13 +653,14 @@ public class DataTypeUtils { if (value instanceof CharSequence) { final CharSequence charSeq = (CharSequence) value; if (charSeq.length() == 0) { - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Character because it has a length of 0"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + + " to Character because it has a length of 0 for field " + fieldName); } return charSeq.charAt(0); } - throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Character"); + throw new IllegalTypeConversionException("Cannot convert value [" + value + "] of type " + value.getClass() + " to Character for field " + fieldName); } public static boolean isCharacterTypeCompatible(final Object value) {
http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/test/java/org/apache/nifi/serialization/TestSimpleRecordSchema.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/test/java/org/apache/nifi/serialization/TestSimpleRecordSchema.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/test/java/org/apache/nifi/serialization/TestSimpleRecordSchema.java new file mode 100644 index 0000000..5a61275 --- /dev/null +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/test/java/org/apache/nifi/serialization/TestSimpleRecordSchema.java @@ -0,0 +1,79 @@ +/* + * 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.serialization; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import org.apache.nifi.serialization.record.RecordField; +import org.apache.nifi.serialization.record.RecordFieldType; +import org.junit.Assert; +import org.junit.Test; + +public class TestSimpleRecordSchema { + + @Test + public void testPreventsTwoFieldsWithSameAlias() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("hello", RecordFieldType.STRING.getDataType(), null, set("foo", "bar"))); + fields.add(new RecordField("goodbye", RecordFieldType.STRING.getDataType(), null, set("baz", "bar"))); + + try { + new SimpleRecordSchema(fields); + Assert.fail("Was able to create two fields with same alias"); + } catch (final IllegalArgumentException expected) { + } + } + + @Test + public void testPreventsTwoFieldsWithSameName() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("hello", RecordFieldType.STRING.getDataType(), null, set("foo", "bar"))); + fields.add(new RecordField("hello", RecordFieldType.STRING.getDataType())); + + try { + new SimpleRecordSchema(fields); + Assert.fail("Was able to create two fields with same name"); + } catch (final IllegalArgumentException expected) { + } + } + + @Test + public void testPreventsTwoFieldsWithConflictingNamesAliases() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("hello", RecordFieldType.STRING.getDataType(), null, set("foo", "bar"))); + fields.add(new RecordField("bar", RecordFieldType.STRING.getDataType())); + + try { + new SimpleRecordSchema(fields); + Assert.fail("Was able to create two fields with conflicting names/aliases"); + } catch (final IllegalArgumentException expected) { + } + } + + private Set<String> set(final String... values) { + final Set<String> set = new HashSet<>(); + for (final String value : values) { + set.add(value); + } + return set; + } + +} http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/test/java/org/apache/nifi/serialization/record/TestMapRecord.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/test/java/org/apache/nifi/serialization/record/TestMapRecord.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/test/java/org/apache/nifi/serialization/record/TestMapRecord.java new file mode 100644 index 0000000..82e20a6 --- /dev/null +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-service-api/src/test/java/org/apache/nifi/serialization/record/TestMapRecord.java @@ -0,0 +1,188 @@ +/* + * 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.serialization.record; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import org.apache.nifi.serialization.SimpleRecordSchema; +import org.junit.Assert; +import org.junit.Test; + +public class TestMapRecord { + + @Test + public void testDefaultValue() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("noDefault", RecordFieldType.STRING.getDataType())); + fields.add(new RecordField("defaultOfHello", RecordFieldType.STRING.getDataType(), "hello")); + + final RecordSchema schema = new SimpleRecordSchema(fields); + final Map<String, Object> values = new HashMap<>(); + final Record record = new MapRecord(schema, values); + + assertNull(record.getValue("noDefault")); + assertEquals("hello", record.getValue("defaultOfHello")); + } + + @Test + public void testDefaultValueInGivenField() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("noDefault", RecordFieldType.STRING.getDataType())); + fields.add(new RecordField("defaultOfHello", RecordFieldType.STRING.getDataType(), "hello")); + + final RecordSchema schema = new SimpleRecordSchema(fields); + final Map<String, Object> values = new HashMap<>(); + final Record record = new MapRecord(schema, values); + + assertNull(record.getValue("noDefault")); + assertEquals("hello", record.getValue("defaultOfHello")); + + final RecordField newField = new RecordField("noDefault", RecordFieldType.STRING.getDataType(), "new"); + assertEquals("new", record.getValue(newField)); + } + + @Test + public void testIllegalDefaultValue() { + new RecordField("hello", RecordFieldType.STRING.getDataType(), 84); + new RecordField("hello", RecordFieldType.STRING.getDataType(), (Object) null); + new RecordField("hello", RecordFieldType.INT.getDataType(), 84); + new RecordField("hello", RecordFieldType.INT.getDataType(), (Object) null); + + try { + new RecordField("hello", RecordFieldType.INT.getDataType(), "foo"); + Assert.fail("Was able to set a default value of \"foo\" for INT type"); + } catch (final IllegalArgumentException expected) { + // expected + } + } + + private Set<String> set(final String... values) { + final Set<String> set = new HashSet<>(); + for (final String value : values) { + set.add(value); + } + return set; + } + + @Test + public void testAliasOneValue() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("foo", RecordFieldType.STRING.getDataType(), null, set("bar", "baz"))); + + final RecordSchema schema = new SimpleRecordSchema(fields); + final Map<String, Object> values = new HashMap<>(); + values.put("bar", 1); + + final Record record = new MapRecord(schema, values); + assertEquals(1, record.getValue("foo")); + assertEquals(1, record.getValue("bar")); + assertEquals(1, record.getValue("baz")); + } + + @Test + public void testAliasConflictingValues() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("foo", RecordFieldType.STRING.getDataType(), null, set("bar", "baz"))); + + final RecordSchema schema = new SimpleRecordSchema(fields); + final Map<String, Object> values = new HashMap<>(); + values.put("bar", 1); + values.put("foo", null); + + final Record record = new MapRecord(schema, values); + assertEquals(1, record.getValue("foo")); + assertEquals(1, record.getValue("bar")); + assertEquals(1, record.getValue("baz")); + } + + @Test + public void testAliasConflictingAliasValues() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("foo", RecordFieldType.STRING.getDataType(), null, set("bar", "baz"))); + + final RecordSchema schema = new SimpleRecordSchema(fields); + final Map<String, Object> values = new HashMap<>(); + values.put("baz", 1); + values.put("bar", 33); + + final Record record = new MapRecord(schema, values); + assertEquals(33, record.getValue("foo")); + assertEquals(33, record.getValue("bar")); + assertEquals(33, record.getValue("baz")); + } + + @Test + public void testAliasInGivenField() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("foo", RecordFieldType.STRING.getDataType(), null, set("bar", "baz"))); + + final RecordSchema schema = new SimpleRecordSchema(fields); + final Map<String, Object> values = new HashMap<>(); + values.put("bar", 33); + + final Record record = new MapRecord(schema, values); + assertEquals(33, record.getValue("foo")); + assertEquals(33, record.getValue("bar")); + assertEquals(33, record.getValue("baz")); + + final RecordField noAlias = new RecordField("hello", RecordFieldType.STRING.getDataType()); + assertNull(record.getValue(noAlias)); + + final RecordField withAlias = new RecordField("hello", RecordFieldType.STRING.getDataType(), null, set("baz")); + assertEquals(33, record.getValue(withAlias)); + assertEquals("33", record.getAsString(withAlias, withAlias.getDataType().getFormat())); + } + + + @Test + public void testDefaultValueWithAliasValue() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("foo", RecordFieldType.STRING.getDataType(), "hello", set("bar", "baz"))); + + final RecordSchema schema = new SimpleRecordSchema(fields); + final Map<String, Object> values = new HashMap<>(); + values.put("baz", 1); + values.put("bar", 33); + + final Record record = new MapRecord(schema, values); + assertEquals(33, record.getValue("foo")); + assertEquals(33, record.getValue("bar")); + assertEquals(33, record.getValue("baz")); + } + + @Test + public void testDefaultValueWithAliasesDefined() { + final List<RecordField> fields = new ArrayList<>(); + fields.add(new RecordField("foo", RecordFieldType.STRING.getDataType(), "hello", set("bar", "baz"))); + + final RecordSchema schema = new SimpleRecordSchema(fields); + final Map<String, Object> values = new HashMap<>(); + final Record record = new MapRecord(schema, values); + assertEquals("hello", record.getValue("foo")); + assertEquals("hello", record.getValue("bar")); + assertEquals("hello", record.getValue("baz")); + } +} http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/pom.xml ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/pom.xml b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/pom.xml index d86a8c5..16479f1 100644 --- a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/pom.xml +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/pom.xml @@ -70,7 +70,11 @@ <dependency> <groupId>org.apache.avro</groupId> <artifactId>avro</artifactId> - <version>1.8.1</version> + </dependency> + <dependency> + <groupId>org.apache.nifi</groupId> + <artifactId>nifi-mock</artifactId> + <scope>test</scope> </dependency> </dependencies> @@ -96,6 +100,7 @@ <exclude>src/test/resources/json/json-with-unicode.json</exclude> <exclude>src/test/resources/json/primitive-type-array.json</exclude> <exclude>src/test/resources/json/single-bank-account.json</exclude> + <exclude>src/test/resources/json/single-bank-account-wrong-field-type.json</exclude> <exclude>src/test/resources/json/single-element-nested-array.json</exclude> <exclude>src/test/resources/json/single-element-nested.json</exclude> <exclude>src/test/resources/json/output/dataTypes.json</exclude> http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReader.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReader.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReader.java index f92816f..f5b4373 100644 --- a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReader.java +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReader.java @@ -19,31 +19,47 @@ package org.apache.nifi.avro; import java.io.IOException; import java.io.InputStream; +import java.util.ArrayList; +import java.util.List; import org.apache.nifi.annotation.documentation.CapabilityDescription; import org.apache.nifi.annotation.documentation.Tags; -import org.apache.nifi.controller.AbstractControllerService; +import org.apache.nifi.components.AllowableValue; import org.apache.nifi.flowfile.FlowFile; import org.apache.nifi.logging.ComponentLog; +import org.apache.nifi.schema.access.SchemaNotFoundException; import org.apache.nifi.serialization.MalformedRecordException; import org.apache.nifi.serialization.RecordReader; -import org.apache.nifi.serialization.RowRecordReaderFactory; -import org.apache.nifi.serialization.record.RecordSchema; +import org.apache.nifi.serialization.RecordReaderFactory; +import org.apache.nifi.serialization.SchemaRegistryService; @Tags({"avro", "parse", "record", "row", "reader", "delimited", "comma", "separated", "values"}) -@CapabilityDescription("Parses Avro data and returns each Avro record as an separate Record object. The Avro data must contain " - + "the schema itself.") -public class AvroReader extends AbstractControllerService implements RowRecordReaderFactory { +@CapabilityDescription("Parses Avro data and returns each Avro record as an separate Record object. The Avro data may contain the schema itself, " + + "or the schema can be externalized and accessed by one of the methods offered by the 'Schema Access Strategy' property.") +public class AvroReader extends SchemaRegistryService implements RecordReaderFactory { + private final AllowableValue EMBEDDED_AVRO_SCHEMA = new AllowableValue("embedded-avro-schema", + "Use Embedded Avro Schema", "The FlowFile has the Avro Schema embedded within the content, and this schema will be used."); + @Override - public RecordReader createRecordReader(final FlowFile flowFile, final InputStream in, final ComponentLog logger) throws MalformedRecordException, IOException { - return new AvroRecordReader(in); + protected List<AllowableValue> getSchemaAccessStrategyValues() { + final List<AllowableValue> allowableValues = new ArrayList<>(super.getSchemaAccessStrategyValues()); + allowableValues.add(EMBEDDED_AVRO_SCHEMA); + return allowableValues; } @Override - public RecordSchema getSchema(final FlowFile flowFile) throws MalformedRecordException, IOException { - // TODO: Need to support retrieving schema from registry instead of requiring that it be in the Avro file. - return null; + public RecordReader createRecordReader(final FlowFile flowFile, final InputStream in, final ComponentLog logger) throws MalformedRecordException, IOException, SchemaNotFoundException { + final String schemaAccessStrategy = getConfigurationContext().getProperty(SCHEMA_ACCESS_STRATEGY).getValue(); + if (EMBEDDED_AVRO_SCHEMA.getValue().equals(schemaAccessStrategy)) { + return new AvroReaderWithEmbeddedSchema(in); + } else { + return new AvroReaderWithExplicitSchema(in, getSchema(flowFile, in)); + } } + @Override + protected AllowableValue getDefaultSchemaAccessStrategy() { + return EMBEDDED_AVRO_SCHEMA; + } } http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReaderWithEmbeddedSchema.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReaderWithEmbeddedSchema.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReaderWithEmbeddedSchema.java new file mode 100644 index 0000000..aa61e4c --- /dev/null +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReaderWithEmbeddedSchema.java @@ -0,0 +1,62 @@ +/* + * 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.avro; + +import java.io.IOException; +import java.io.InputStream; + +import org.apache.avro.Schema; +import org.apache.avro.file.DataFileStream; +import org.apache.avro.generic.GenericDatumReader; +import org.apache.avro.generic.GenericRecord; +import org.apache.nifi.serialization.MalformedRecordException; +import org.apache.nifi.serialization.record.RecordSchema; + +public class AvroReaderWithEmbeddedSchema extends AvroRecordReader { + private final DataFileStream<GenericRecord> dataFileStream; + private final InputStream in; + private final Schema avroSchema; + private final RecordSchema recordSchema; + + public AvroReaderWithEmbeddedSchema(final InputStream in) throws IOException { + this.in = in; + dataFileStream = new DataFileStream<>(in, new GenericDatumReader<GenericRecord>()); + this.avroSchema = dataFileStream.getSchema(); + recordSchema = AvroTypeUtil.createSchema(avroSchema); + } + + @Override + public void close() throws IOException { + dataFileStream.close(); + in.close(); + } + + @Override + protected GenericRecord nextAvroRecord() { + if (!dataFileStream.hasNext()) { + return null; + } + + return dataFileStream.next(); + } + + @Override + public RecordSchema getSchema() throws MalformedRecordException { + return recordSchema; + } +} http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReaderWithExplicitSchema.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReaderWithExplicitSchema.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReaderWithExplicitSchema.java new file mode 100644 index 0000000..104214c --- /dev/null +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroReaderWithExplicitSchema.java @@ -0,0 +1,75 @@ +/* + * 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.avro; + +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericDatumReader; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.BinaryDecoder; +import org.apache.avro.io.DatumReader; +import org.apache.avro.io.DecoderFactory; +import org.apache.nifi.schema.access.SchemaNotFoundException; +import org.apache.nifi.serialization.MalformedRecordException; +import org.apache.nifi.serialization.record.RecordSchema; + +public class AvroReaderWithExplicitSchema extends AvroRecordReader { + private final InputStream in; + private final Schema avroSchema; + private final RecordSchema recordSchema; + private final DatumReader<GenericRecord> datumReader; + private final BinaryDecoder decoder; + private GenericRecord genericRecord; + + public AvroReaderWithExplicitSchema(final InputStream in, final RecordSchema recordSchema) throws IOException, SchemaNotFoundException { + this.in = in; + this.recordSchema = recordSchema; + + this.avroSchema = AvroTypeUtil.extractAvroSchema(recordSchema); + datumReader = new GenericDatumReader<GenericRecord>(avroSchema); + decoder = DecoderFactory.get().binaryDecoder(in, null); + } + + @Override + public void close() throws IOException { + in.close(); + } + + @Override + protected GenericRecord nextAvroRecord() throws IOException { + if (decoder.isEnd()) { + return null; + } + + try { + genericRecord = datumReader.read(genericRecord, decoder); + } catch (final EOFException eof) { + return null; + } + + return genericRecord; + } + + @Override + public RecordSchema getSchema() throws MalformedRecordException { + return recordSchema; + } +} http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordReader.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordReader.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordReader.java index d725cbf..621ec74 100644 --- a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordReader.java +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordReader.java @@ -18,7 +18,6 @@ package org.apache.nifi.avro; import java.io.IOException; -import java.io.InputStream; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.HashMap; @@ -30,11 +29,8 @@ import org.apache.avro.LogicalType; import org.apache.avro.LogicalTypes; import org.apache.avro.Schema; import org.apache.avro.Schema.Field; -import org.apache.avro.file.DataFileStream; import org.apache.avro.generic.GenericData; import org.apache.avro.generic.GenericData.Array; -import org.apache.avro.generic.GenericData.StringType; -import org.apache.avro.generic.GenericDatumReader; import org.apache.avro.generic.GenericFixed; import org.apache.avro.generic.GenericRecord; import org.apache.avro.util.Utf8; @@ -48,38 +44,19 @@ import org.apache.nifi.serialization.record.RecordField; import org.apache.nifi.serialization.record.RecordSchema; import org.apache.nifi.serialization.record.util.DataTypeUtils; -public class AvroRecordReader implements RecordReader { - private final InputStream in; - private final Schema avroSchema; - private final DataFileStream<GenericRecord> dataFileStream; - private RecordSchema recordSchema; +public abstract class AvroRecordReader implements RecordReader { - public AvroRecordReader(final InputStream in) throws IOException, MalformedRecordException { - this.in = in; + protected abstract GenericRecord nextAvroRecord() throws IOException; - dataFileStream = new DataFileStream<>(in, new GenericDatumReader<GenericRecord>()); - this.avroSchema = dataFileStream.getSchema(); - GenericData.setStringType(this.avroSchema, StringType.String); - } - - @Override - public void close() throws IOException { - dataFileStream.close(); - in.close(); - } @Override public Record nextRecord() throws IOException, MalformedRecordException { - if (!dataFileStream.hasNext()) { + GenericRecord record = nextAvroRecord(); + if (record == null) { return null; } - GenericRecord record = null; - while (record == null && dataFileStream.hasNext()) { - record = dataFileStream.next(); - } - final RecordSchema schema = getSchema(); final Map<String, Object> values = convertAvroRecordToMap(record, schema); return new MapRecord(schema, values); @@ -89,9 +66,18 @@ public class AvroRecordReader implements RecordReader { private Map<String, Object> convertAvroRecordToMap(final GenericRecord avroRecord, final RecordSchema recordSchema) { final Map<String, Object> values = new HashMap<>(recordSchema.getFieldCount()); - for (final String fieldName : recordSchema.getFieldNames()) { - final Object value = avroRecord.get(fieldName); + for (final RecordField recordField : recordSchema.getFields()) { + Object value = avroRecord.get(recordField.getFieldName()); + if (value == null) { + for (final String alias : recordField.getAliases()) { + value = avroRecord.get(alias); + if (value != null) { + break; + } + } + } + final String fieldName = recordField.getFieldName(); final Field avroField = avroRecord.getSchema().getField(fieldName); if (avroField == null) { values.put(fieldName, null); @@ -101,8 +87,8 @@ public class AvroRecordReader implements RecordReader { final Schema fieldSchema = avroField.schema(); final Object rawValue = normalizeValue(value, fieldSchema); - final DataType desiredType = recordSchema.getDataType(fieldName).get(); - final Object coercedValue = DataTypeUtils.convertType(rawValue, desiredType); + final DataType desiredType = recordField.getDataType(); + final Object coercedValue = DataTypeUtils.convertType(rawValue, desiredType, fieldName); values.put(fieldName, coercedValue); } @@ -215,13 +201,5 @@ public class AvroRecordReader implements RecordReader { } - @Override - public RecordSchema getSchema() throws MalformedRecordException { - if (recordSchema != null) { - return recordSchema; - } - recordSchema = AvroTypeUtil.createSchema(avroSchema); - return recordSchema; - } } http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordSetWriter.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordSetWriter.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordSetWriter.java index 03d766c..381e978 100644 --- a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordSetWriter.java +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroRecordSetWriter.java @@ -17,48 +17,93 @@ package org.apache.nifi.avro; +import java.io.IOException; +import java.io.InputStream; import java.util.ArrayList; +import java.util.EnumSet; import java.util.List; +import java.util.Set; import org.apache.avro.Schema; 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.AllowableValue; import org.apache.nifi.components.PropertyDescriptor; -import org.apache.nifi.controller.AbstractControllerService; -import org.apache.nifi.controller.ConfigurationContext; +import org.apache.nifi.components.ValidationContext; +import org.apache.nifi.flowfile.FlowFile; import org.apache.nifi.logging.ComponentLog; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.schema.access.SchemaField; +import org.apache.nifi.schema.access.SchemaNotFoundException; +import org.apache.nifi.schemaregistry.services.SchemaRegistry; import org.apache.nifi.serialization.RecordSetWriter; import org.apache.nifi.serialization.RecordSetWriterFactory; +import org.apache.nifi.serialization.SchemaRegistryRecordSetWriter; +import org.apache.nifi.serialization.record.RecordSchema; @Tags({"avro", "result", "set", "writer", "serializer", "record", "recordset", "row"}) @CapabilityDescription("Writes the contents of a RecordSet in Binary Avro format.") -public class AvroRecordSetWriter extends AbstractControllerService implements RecordSetWriterFactory { - static final PropertyDescriptor SCHEMA = new PropertyDescriptor.Builder() - .name("Avro Schema") - .description("The Avro Schema to use when writing out the Result Set") - .addValidator(new AvroSchemaValidator()) - .expressionLanguageSupported(false) - .required(true) +public class AvroRecordSetWriter extends SchemaRegistryRecordSetWriter implements RecordSetWriterFactory { + private static final Set<SchemaField> requiredSchemaFields = EnumSet.of(SchemaField.SCHEMA_TEXT, SchemaField.SCHEMA_TEXT_FORMAT); + + static final AllowableValue AVRO_EMBEDDED = new AllowableValue("avro-embedded", "Embed Avro Schema", + "The FlowFile will have the Avro schema embedded into the content, as is typical with Avro"); + + protected static final PropertyDescriptor SCHEMA_REGISTRY = new PropertyDescriptor.Builder() + .name("Schema Registry") + .description("Specifies the Controller Service to use for the Schema Registry") + .identifiesControllerService(SchemaRegistry.class) + .required(false) .build(); - private volatile Schema schema; @Override protected List<PropertyDescriptor> getSupportedPropertyDescriptors() { - final List<PropertyDescriptor> properties = new ArrayList<>(); - properties.add(SCHEMA); + final List<PropertyDescriptor> properties = new ArrayList<>(super.getSupportedPropertyDescriptors()); + properties.add(SCHEMA_ACCESS_STRATEGY); + properties.add(SCHEMA_REGISTRY); return properties; } - @OnEnabled - public void storePropertyValues(final ConfigurationContext context) { - schema = new Schema.Parser().parse(context.getProperty(SCHEMA).getValue()); + + @Override + public RecordSetWriter createWriter(final ComponentLog logger, final FlowFile flowFile, final InputStream in) throws IOException { + final String strategyValue = getConfigurationContext().getProperty(SCHEMA_WRITE_STRATEGY).getValue(); + + try { + final RecordSchema recordSchema = getSchema(flowFile, in); + final Schema avroSchema = AvroTypeUtil.extractAvroSchema(recordSchema); + + if (AVRO_EMBEDDED.getValue().equals(strategyValue)) { + return new WriteAvroResultWithSchema(avroSchema); + } else { + return new WriteAvroResultWithExternalSchema(avroSchema, recordSchema, getSchemaAccessWriter(recordSchema)); + } + } catch (final SchemaNotFoundException e) { + throw new ProcessException("Could not determine the Avro Schema to use for writing the content", e); + } + } + + @Override + protected List<AllowableValue> getSchemaWriteStrategyValues() { + final List<AllowableValue> allowableValues = new ArrayList<>(); + allowableValues.add(AVRO_EMBEDDED); + allowableValues.addAll(super.getSchemaWriteStrategyValues()); + return allowableValues; } @Override - public RecordSetWriter createWriter(final ComponentLog logger) { - return new WriteAvroResult(schema); + protected AllowableValue getDefaultSchemaWriteStrategy() { + return AVRO_EMBEDDED; } + @Override + protected Set<SchemaField> getRequiredSchemaFields(final ValidationContext validationContext) { + final String writeStrategyValue = validationContext.getProperty(SCHEMA_WRITE_STRATEGY).getValue(); + if (writeStrategyValue.equalsIgnoreCase(AVRO_EMBEDDED.getValue())) { + return requiredSchemaFields; + } + + return super.getRequiredSchemaFields(validationContext); + } } http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroSchemaValidator.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroSchemaValidator.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroSchemaValidator.java index 7151348..4449afc 100644 --- a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroSchemaValidator.java +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroSchemaValidator.java @@ -26,6 +26,15 @@ public class AvroSchemaValidator implements Validator { @Override public ValidationResult validate(final String subject, final String input, final ValidationContext context) { + if (context.isExpressionLanguageSupported(subject) && context.isExpressionLanguagePresent(input)) { + return new ValidationResult.Builder() + .input(input) + .subject(subject) + .valid(true) + .explanation("Expression Language is present") + .build(); + } + try { new Schema.Parser().parse(input); http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroTypeUtil.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroTypeUtil.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroTypeUtil.java index 1810c83..b65026a 100644 --- a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroTypeUtil.java +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/AvroTypeUtil.java @@ -20,6 +20,7 @@ package org.apache.nifi.avro; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.stream.Collectors; import org.apache.avro.LogicalType; @@ -27,14 +28,37 @@ import org.apache.avro.LogicalTypes; import org.apache.avro.Schema; import org.apache.avro.Schema.Field; import org.apache.avro.Schema.Type; +import org.apache.nifi.schema.access.SchemaNotFoundException; import org.apache.nifi.serialization.SimpleRecordSchema; import org.apache.nifi.serialization.record.DataType; 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.SchemaIdentifier; import org.apache.nifi.serialization.record.util.IllegalTypeConversionException; public class AvroTypeUtil { + public static final String AVRO_SCHEMA_FORMAT = "avro"; + + public static Schema extractAvroSchema(final RecordSchema recordSchema) throws SchemaNotFoundException { + final Optional<String> schemaFormatOption = recordSchema.getSchemaFormat(); + if (!schemaFormatOption.isPresent()) { + throw new SchemaNotFoundException("No Schema Format was present in the RecordSchema"); + } + + final String schemaFormat = schemaFormatOption.get(); + if (!schemaFormat.equals(AVRO_SCHEMA_FORMAT)) { + throw new SchemaNotFoundException("Schema provided is not in Avro format"); + } + + final Optional<String> textOption = recordSchema.getSchemaText(); + if (!textOption.isPresent()) { + throw new SchemaNotFoundException("No Schema text was present in the RecordSchema"); + } + + final String text = textOption.get(); + return new Schema.Parser().parse(text); + } public static DataType determineDataType(final Schema avroSchema) { final Type avroType = avroSchema.getType(); @@ -93,15 +117,18 @@ public class AvroTypeUtil { final String fieldName = field.name(); final Schema fieldSchema = field.schema(); final DataType fieldType = determineDataType(fieldSchema); - recordFields.add(new RecordField(fieldName, fieldType)); + recordFields.add(new RecordField(fieldName, fieldType, field.defaultVal(), field.aliases())); } - final RecordSchema recordSchema = new SimpleRecordSchema(recordFields); + final RecordSchema recordSchema = new SimpleRecordSchema(recordFields, avroSchema.toString(), AVRO_SCHEMA_FORMAT, SchemaIdentifier.EMPTY); return RecordFieldType.RECORD.getRecordDataType(recordSchema); } case NULL: + return RecordFieldType.STRING.getDataType(); case MAP: - return RecordFieldType.RECORD.getDataType(); + final Schema valueSchema = avroSchema.getValueType(); + final DataType valueType = determineDataType(valueSchema); + return RecordFieldType.MAP.getMapDataType(valueType); case UNION: { final List<Schema> nonNullSubSchemas = avroSchema.getTypes().stream() .filter(s -> s.getType() != Type.NULL) @@ -129,10 +156,11 @@ public class AvroTypeUtil { for (final Field field : avroSchema.getFields()) { final String fieldName = field.name(); final DataType dataType = AvroTypeUtil.determineDataType(field.schema()); - recordFields.add(new RecordField(fieldName, dataType)); + + recordFields.add(new RecordField(fieldName, dataType, field.defaultVal(), field.aliases())); } - final RecordSchema recordSchema = new SimpleRecordSchema(recordFields); + final RecordSchema recordSchema = new SimpleRecordSchema(recordFields, avroSchema.toString(), AVRO_SCHEMA_FORMAT, SchemaIdentifier.EMPTY); return recordSchema; } http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResult.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResult.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResult.java index b512b82..55f796a 100644 --- a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResult.java +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResult.java @@ -42,61 +42,44 @@ import org.apache.avro.io.DatumWriter; import org.apache.nifi.serialization.RecordSetWriter; import org.apache.nifi.serialization.WriteResult; 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.serialization.record.util.DataTypeUtils; import org.apache.nifi.serialization.record.util.IllegalTypeConversionException; -public class WriteAvroResult implements RecordSetWriter { +public abstract class WriteAvroResult implements RecordSetWriter { private final Schema schema; public WriteAvroResult(final Schema schema) { this.schema = schema; } - @Override - public WriteResult write(final RecordSet rs, final OutputStream outStream) throws IOException { - Record record = rs.next(); - if (record == null) { - return WriteResult.of(0, Collections.emptyMap()); - } - - int nrOfRows = 0; - final DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema); - try (final DataFileWriter<GenericRecord> dataFileWriter = new DataFileWriter<>(datumWriter)) { - dataFileWriter.create(schema, outStream); - - do { - final GenericRecord rec = createAvroRecord(record, schema); - dataFileWriter.append(rec); - nrOfRows++; - } while ((record = rs.next()) != null); - } - - return WriteResult.of(nrOfRows, Collections.emptyMap()); + protected Schema getSchema() { + return schema; } - private GenericRecord createAvroRecord(final Record record, final Schema avroSchema) throws IOException { + protected GenericRecord createAvroRecord(final Record record, final Schema avroSchema) throws IOException { final GenericRecord rec = new GenericData.Record(avroSchema); final RecordSchema recordSchema = record.getSchema(); - for (final String fieldName : recordSchema.getFieldNames()) { - final Object rawValue = record.getValue(fieldName); + for (final RecordField recordField : recordSchema.getFields()) { + final Object rawValue = record.getValue(recordField); + final String fieldName = recordField.getFieldName(); final Field field = avroSchema.getField(fieldName); if (field == null) { continue; } - final Object converted = convertToAvroObject(rawValue, field.schema()); + final Object converted = convertToAvroObject(rawValue, field.schema(), fieldName); rec.put(fieldName, converted); } return rec; } - private Object convertToAvroObject(final Object rawValue, final Schema fieldSchema) throws IOException { + protected Object convertToAvroObject(final Object rawValue, final Schema fieldSchema, final String fieldName) throws IOException { if (rawValue == null) { return null; } @@ -105,43 +88,43 @@ public class WriteAvroResult implements RecordSetWriter { case INT: { final LogicalType logicalType = fieldSchema.getLogicalType(); if (logicalType == null) { - return DataTypeUtils.toInteger(rawValue); + return DataTypeUtils.toInteger(rawValue, fieldName); } if (LogicalTypes.date().getName().equals(logicalType.getName())) { - final long longValue = DataTypeUtils.toLong(rawValue); + final long longValue = DataTypeUtils.toLong(rawValue, fieldName); final Date date = new Date(longValue); final Duration duration = Duration.between(new Date(0L).toInstant(), date.toInstant()); final long days = duration.toDays(); return (int) days; } else if (LogicalTypes.timeMillis().getName().equals(logicalType.getName())) { - final long longValue = DataTypeUtils.toLong(rawValue); + final long longValue = DataTypeUtils.toLong(rawValue, fieldName); final Date date = new Date(longValue); final Duration duration = Duration.between(date.toInstant().truncatedTo(ChronoUnit.DAYS), date.toInstant()); final long millisSinceMidnight = duration.toMillis(); return (int) millisSinceMidnight; } - return DataTypeUtils.toInteger(rawValue); + return DataTypeUtils.toInteger(rawValue, fieldName); } case LONG: { final LogicalType logicalType = fieldSchema.getLogicalType(); if (logicalType == null) { - return DataTypeUtils.toLong(rawValue); + return DataTypeUtils.toLong(rawValue, fieldName); } if (LogicalTypes.timeMicros().getName().equals(logicalType.getName())) { - final long longValue = DataTypeUtils.toLong(rawValue); + final long longValue = DataTypeUtils.toLong(rawValue, fieldName); final Date date = new Date(longValue); final Duration duration = Duration.between(date.toInstant().truncatedTo(ChronoUnit.DAYS), date.toInstant()); return duration.toMillis() * 1000L; } else if (LogicalTypes.timestampMillis().getName().equals(logicalType.getName())) { - return DataTypeUtils.toLong(rawValue); + return DataTypeUtils.toLong(rawValue, fieldName); } else if (LogicalTypes.timestampMicros().getName().equals(logicalType.getName())) { - return DataTypeUtils.toLong(rawValue) * 1000L; + return DataTypeUtils.toLong(rawValue, fieldName) * 1000L; } - return DataTypeUtils.toLong(rawValue); + return DataTypeUtils.toLong(rawValue, fieldName); } case BYTES: case FIXED: @@ -157,10 +140,10 @@ public class WriteAvroResult implements RecordSetWriter { if (rawValue instanceof Record) { final Record recordValue = (Record) rawValue; final Map<String, Object> map = new HashMap<>(); - for (final String recordFieldName : recordValue.getSchema().getFieldNames()) { - final Object v = recordValue.getValue(recordFieldName); + for (final RecordField recordField : recordValue.getSchema().getFields()) { + final Object v = recordValue.getValue(recordField); if (v != null) { - map.put(recordFieldName, v); + map.put(recordField.getFieldName(), v); } } @@ -172,15 +155,16 @@ public class WriteAvroResult implements RecordSetWriter { final GenericData.Record avroRecord = new GenericData.Record(fieldSchema); final Record record = (Record) rawValue; - for (final String recordFieldName : record.getSchema().getFieldNames()) { - final Object recordFieldValue = record.getValue(recordFieldName); + for (final RecordField recordField : record.getSchema().getFields()) { + final Object recordFieldValue = record.getValue(recordField); + final String recordFieldName = recordField.getFieldName(); final Field field = fieldSchema.getField(recordFieldName); if (field == null) { continue; } - final Object converted = convertToAvroObject(recordFieldValue, field.schema()); + final Object converted = convertToAvroObject(recordFieldValue, field.schema(), fieldName); avroRecord.put(recordFieldName, converted); } return avroRecord; @@ -188,16 +172,16 @@ public class WriteAvroResult implements RecordSetWriter { final Object[] objectArray = (Object[]) rawValue; final List<Object> list = new ArrayList<>(objectArray.length); for (final Object o : objectArray) { - final Object converted = convertToAvroObject(o, fieldSchema.getElementType()); + final Object converted = convertToAvroObject(o, fieldSchema.getElementType(), fieldName); list.add(converted); } return list; case BOOLEAN: - return DataTypeUtils.toBoolean(rawValue); + return DataTypeUtils.toBoolean(rawValue, fieldName); case DOUBLE: - return DataTypeUtils.toDouble(rawValue); + return DataTypeUtils.toDouble(rawValue, fieldName); case FLOAT: - return DataTypeUtils.toFloat(rawValue); + return DataTypeUtils.toFloat(rawValue, fieldName); case NULL: return null; case ENUM: http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResultWithExternalSchema.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResultWithExternalSchema.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResultWithExternalSchema.java new file mode 100644 index 0000000..74306e4 --- /dev/null +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResultWithExternalSchema.java @@ -0,0 +1,75 @@ +/* + * 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.avro; + +import java.io.BufferedOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.util.Collections; + +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.BinaryEncoder; +import org.apache.avro.io.DatumWriter; +import org.apache.avro.io.EncoderFactory; +import org.apache.nifi.schema.access.SchemaAccessWriter; +import org.apache.nifi.serialization.WriteResult; +import org.apache.nifi.serialization.record.Record; +import org.apache.nifi.serialization.record.RecordSchema; +import org.apache.nifi.serialization.record.RecordSet; + +public class WriteAvroResultWithExternalSchema extends WriteAvroResult { + private final SchemaAccessWriter schemaAccessWriter; + private final RecordSchema recordSchema; + + public WriteAvroResultWithExternalSchema(final Schema avroSchema, final RecordSchema recordSchema, final SchemaAccessWriter schemaAccessWriter) { + super(avroSchema); + this.recordSchema = recordSchema; + this.schemaAccessWriter = schemaAccessWriter; + } + + @Override + public WriteResult write(final RecordSet rs, final OutputStream outStream) throws IOException { + Record record = rs.next(); + if (record == null) { + return WriteResult.of(0, Collections.emptyMap()); + } + + int nrOfRows = 0; + final Schema schema = getSchema(); + final DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema); + + final BufferedOutputStream bufferedOut = new BufferedOutputStream(outStream); + schemaAccessWriter.writeHeader(recordSchema, bufferedOut); + + final BinaryEncoder encoder = EncoderFactory.get().blockingBinaryEncoder(bufferedOut, null); + + do { + final GenericRecord rec = createAvroRecord(record, schema); + + datumWriter.write(rec, encoder); + encoder.flush(); + nrOfRows++; + } while ((record = rs.next()) != null); + + bufferedOut.flush(); + + return WriteResult.of(nrOfRows, schemaAccessWriter.getAttributes(recordSchema)); + } +} http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResultWithSchema.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResultWithSchema.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResultWithSchema.java new file mode 100644 index 0000000..dca8aac --- /dev/null +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/avro/WriteAvroResultWithSchema.java @@ -0,0 +1,62 @@ +/* + * 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.avro; + +import java.io.IOException; +import java.io.OutputStream; +import java.util.Collections; + +import org.apache.avro.Schema; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.DatumWriter; +import org.apache.nifi.serialization.WriteResult; +import org.apache.nifi.serialization.record.Record; +import org.apache.nifi.serialization.record.RecordSet; + +public class WriteAvroResultWithSchema extends WriteAvroResult { + + public WriteAvroResultWithSchema(final Schema schema) { + super(schema); + } + + @Override + public WriteResult write(final RecordSet rs, final OutputStream outStream) throws IOException { + Record record = rs.next(); + if (record == null) { + return WriteResult.of(0, Collections.emptyMap()); + } + + int nrOfRows = 0; + final Schema schema = getSchema(); + final DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema); + + try (final DataFileWriter<GenericRecord> dataFileWriter = new DataFileWriter<>(datumWriter)) { + dataFileWriter.create(schema, outStream); + + do { + final GenericRecord rec = createAvroRecord(record, schema); + dataFileWriter.append(rec); + nrOfRows++; + } while ((record = rs.next()) != null); + } + + return WriteResult.of(nrOfRows, Collections.emptyMap()); + } +} http://git-wip-us.apache.org/repos/asf/nifi/blob/50ea1083/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/csv/CSVHeaderSchemaStrategy.java ---------------------------------------------------------------------- diff --git a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/csv/CSVHeaderSchemaStrategy.java b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/csv/CSVHeaderSchemaStrategy.java new file mode 100644 index 0000000..71093de --- /dev/null +++ b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/csv/CSVHeaderSchemaStrategy.java @@ -0,0 +1,67 @@ +/* + * 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.csv; + +import java.io.InputStream; +import java.io.InputStreamReader; +import java.io.Reader; +import java.util.ArrayList; +import java.util.EnumSet; +import java.util.List; +import java.util.Set; + +import org.apache.commons.csv.CSVFormat; +import org.apache.commons.csv.CSVParser; +import org.apache.commons.io.input.BOMInputStream; +import org.apache.nifi.controller.ConfigurationContext; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.schema.access.SchemaAccessStrategy; +import org.apache.nifi.schema.access.SchemaField; +import org.apache.nifi.schema.access.SchemaNotFoundException; +import org.apache.nifi.serialization.SimpleRecordSchema; +import org.apache.nifi.serialization.record.RecordField; +import org.apache.nifi.serialization.record.RecordFieldType; +import org.apache.nifi.serialization.record.RecordSchema; + +public class CSVHeaderSchemaStrategy implements SchemaAccessStrategy { + private static final Set<SchemaField> schemaFields = EnumSet.noneOf(SchemaField.class); + + @Override + public RecordSchema getSchema(final FlowFile flowFile, final InputStream contentStream, final ConfigurationContext context) throws SchemaNotFoundException { + try { + final CSVFormat csvFormat = CSVUtils.createCSVFormat(context).withFirstRecordAsHeader(); + try (final Reader reader = new InputStreamReader(new BOMInputStream(contentStream)); + final CSVParser csvParser = new CSVParser(reader, csvFormat)) { + + final List<RecordField> fields = new ArrayList<>(); + for (final String columnName : csvParser.getHeaderMap().keySet()) { + fields.add(new RecordField(columnName, RecordFieldType.STRING.getDataType())); + } + + return new SimpleRecordSchema(fields); + } + } catch (final Exception e) { + throw new SchemaNotFoundException("Failed to read Header line from CSV", e); + } + } + + @Override + public Set<SchemaField> getSuppliedSchemaFields() { + return schemaFields; + } +}
