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

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

View File

@ -89,6 +89,8 @@ public class WriteCSVResult extends AbstractRecordSetWriter implements RecordSet
@Override
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);
}

View File

@ -387,6 +387,28 @@ public class TestWriteCSVResult {
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) {
final DateFormat df = new SimpleDateFormat(format);