Apache/nifi
1
0
Fork 0
mirror of https://github.com/apache/nifi.git synced 2025-03-01 06:59:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

NIFI-10996: Write out CSV header even if there are no records

Browse Source 
This closes #6818

Signed-off-by: Mike Thomsen <mthomsen@apache.org>
This commit is contained in:
Matthew Burgess 2023-01-03 09:12:14 -05:00 committed by Mike Thomsen
parent 40ccb71f85
commit a9baa21f87
2 changed files with 24 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/csv/WriteCSVResult.java
View File

@ -89,6 +89,8 @@ public class WriteCSVResult extends AbstractRecordSetWriter implements RecordSet
@Override @Override
protected Map<String, String> onFinishRecordSet() throws IOException { protected Map<String, String> onFinishRecordSet() throws IOException {
// If the header has not yet been written (but should be), write it out now
includeHeaderIfNecessary(null, true);
return schemaWriter.getAttributes(recordSchema); return schemaWriter.getAttributes(recordSchema);
} }

22
nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/test/java/org/apache/nifi/csv/TestWriteCSVResult.java
View File

@ -387,6 +387,28 @@ public class TestWriteCSVResult {
assertEquals("id,name\n1\\,John Doe\n", output); assertEquals("id,name\n1\\,John Doe\n", output);
} }
@Test
public void testWriteHeaderWithNoRecords() throws IOException {
final CSVFormat csvFormat = CSVFormat.DEFAULT.builder().setEscape('\\').setQuoteMode(QuoteMode.NONE).setRecordSeparator(",").build();
final List<RecordField> fields = new ArrayList<>();
fields.add(new RecordField("id", RecordFieldType.STRING.getDataType()));
fields.add(new RecordField("name", RecordFieldType.STRING.getDataType()));
final RecordSchema schema = new SimpleRecordSchema(fields);
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
final String output;
try (final WriteCSVResult writer = new WriteCSVResult(csvFormat, schema, new SchemaNameAsAttribute(), baos,
RecordFieldType.DATE.getDefaultFormat(), RecordFieldType.TIME.getDefaultFormat(), RecordFieldType.TIMESTAMP.getDefaultFormat(), true, "ASCII")) {
writer.beginRecordSet();
writer.finishRecordSet();
writer.flush();
output = baos.toString();
}
assertEquals("id,name,", output);
}
private DateFormat getDateFormat(final String format) { private DateFormat getDateFormat(final String format) {
final DateFormat df = new SimpleDateFormat(format); final DateFormat df = new SimpleDateFormat(format);