[nifi] branch main updated: NIFI-10996: Write out CSV header even if there are no records

[mthomsen]

This is an automated email from the ASF dual-hosted git repository.

mthomsen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git


The following commit(s) were added to refs/heads/main by this push:
     new a9baa21f87 NIFI-10996: Write out CSV header even if there are no 
records
a9baa21f87 is described below

commit a9baa21f87eac5f2992c42ea4a45af00a3f19af6
Author: Matthew Burgess <mattyb...@apache.org>
AuthorDate: Tue Jan 3 09:12:14 2023 -0500

    NIFI-10996: Write out CSV header even if there are no records
    
    This closes #6818
    
    Signed-off-by: Mike Thomsen <mthom...@apache.org>
---
 .../java/org/apache/nifi/csv/WriteCSVResult.java   |  2 ++
 .../org/apache/nifi/csv/TestWriteCSVResult.java    | 22 ++++++++++++++++++++++
 2 files changed, 24 insertions(+)

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/WriteCSVResult.java
 
b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/csv/WriteCSVResult.java
index 733d3a57b1..3523ab39ef 100644
--- 
a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/csv/WriteCSVResult.java
+++ 
b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/csv/WriteCSVResult.java
@@ -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);
     }
 
diff --git 
a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/test/java/org/apache/nifi/csv/TestWriteCSVResult.java
 
b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/test/java/org/apache/nifi/csv/TestWriteCSVResult.java
index 96efe4a155..d03d8385f5 100644
--- 
a/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/test/java/org/apache/nifi/csv/TestWriteCSVResult.java
+++ 
b/nifi-nar-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/test/java/org/apache/nifi/csv/TestWriteCSVResult.java
@@ -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);