NIFI-16177 - JSON Record Writer - Add epoch-second timestamp representation

nifi-extension-bundles/nifi-extension-utils/nifi-record-utils/nifi-json-record-utils/src/main/java/org/apache/nifi/json/TimestampRepresentation.java

@@ -0,0 +1,50 @@
+/*
+ * 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.json;
+
+import org.apache.nifi.components.DescribedValue;
+
+public enum TimestampRepresentation implements DescribedValue {
+    AUTO("Automatic", "Uses the configured Timestamp Format when present, otherwise writes epoch milliseconds"),
+    FORMATTED_STRING("Formatted String", "Writes timestamps as strings using the configured Timestamp Format"),
+    EPOCH_MILLISECONDS("Epoch Milliseconds", "Writes timestamps as integer JSON numbers containing milliseconds since the Unix epoch"),
+    EPOCH_SECONDS("Epoch Seconds", "Writes timestamps as decimal JSON numbers containing seconds since the Unix epoch with millisecond precision");
+
+    private final String displayName;
+    private final String description;
+
+    TimestampRepresentation(final String displayName, final String description) {
+        this.displayName = displayName;
+        this.description = description;
+    }
+
+    @Override
+    public String getDisplayName() {
+        return displayName;
+    }
+
+    @Override
+    public String getDescription() {
+        return description;
+    }
+
+    @Override
+    public String getValue() {
+        return name();
+    }
+}

nifi-extension-bundles/nifi-extension-utils/nifi-record-utils/nifi-json-record-utils/src/main/java/org/apache/nifi/json/WriteJsonResult.java

@@ -45,6 +45,7 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.sql.Time;
 import java.util.Map;
@@ -69,23 +70,34 @@ public class WriteJsonResult extends AbstractRecordSetWriter implements RecordSe
     private final boolean prettyPrint;
     private final boolean allowScientificNotation;
     private final boolean serializedInputHandlingEnabled;
+    private final TimestampRepresentation timestampRepresentation;
 
     private static final ObjectMapper objectMapper = new ObjectMapper();
 
     public WriteJsonResult(final ComponentLog logger, final RecordSchema recordSchema, final SchemaAccessWriter schemaAccess, final OutputStream out, final boolean prettyPrint,
             final NullSuppression nullSuppression, final OutputGrouping outputGrouping, final String dateFormat, final String timeFormat, final String timestampFormat) throws IOException {
-        this(logger, recordSchema, schemaAccess, out, prettyPrint, nullSuppression, outputGrouping, dateFormat, timeFormat, timestampFormat, "application/json", false, true);
+        this(logger, recordSchema, schemaAccess, out, prettyPrint, nullSuppression, outputGrouping, dateFormat, timeFormat, timestampFormat, "application/json", false, true,
+                TimestampRepresentation.AUTO);
     }
 
     public WriteJsonResult(final ComponentLog logger, final RecordSchema recordSchema, final SchemaAccessWriter schemaAccess, final OutputStream out, final boolean prettyPrint,
         final NullSuppression nullSuppression, final OutputGrouping outputGrouping, final String dateFormat, final String timeFormat, final String timestampFormat,
         final String mimeType, final boolean allowScientificNotation) throws IOException {
-        this(logger, recordSchema, schemaAccess, out, prettyPrint, nullSuppression, outputGrouping, dateFormat, timeFormat, timestampFormat, mimeType, allowScientificNotation, true);
+        this(logger, recordSchema, schemaAccess, out, prettyPrint, nullSuppression, outputGrouping, dateFormat, timeFormat, timestampFormat, mimeType, allowScientificNotation, true,
+                TimestampRepresentation.AUTO);
     }
 
     public WriteJsonResult(final ComponentLog logger, final RecordSchema recordSchema, final SchemaAccessWriter schemaAccess, final OutputStream out, final boolean prettyPrint,
         final NullSuppression nullSuppression, final OutputGrouping outputGrouping, final String dateFormat, final String timeFormat, final String timestampFormat,
         final String mimeType, final boolean allowScientificNotation, final boolean serializedInputHandlingEnabled) throws IOException {
+        this(logger, recordSchema, schemaAccess, out, prettyPrint, nullSuppression, outputGrouping, dateFormat, timeFormat, timestampFormat, mimeType, allowScientificNotation,
+                serializedInputHandlingEnabled, TimestampRepresentation.AUTO);
+    }
+
+    public WriteJsonResult(final ComponentLog logger, final RecordSchema recordSchema, final SchemaAccessWriter schemaAccess, final OutputStream out, final boolean prettyPrint,
+        final NullSuppression nullSuppression, final OutputGrouping outputGrouping, final String dateFormat, final String timeFormat, final String timestampFormat,
+        final String mimeType, final boolean allowScientificNotation, final boolean serializedInputHandlingEnabled,
+        final TimestampRepresentation timestampRepresentation) throws IOException {
 
         super(out);
         this.logger = logger;
@@ -96,6 +108,7 @@ public WriteJsonResult(final ComponentLog logger, final RecordSchema recordSchem
         this.mimeType = mimeType;
         this.allowScientificNotation = allowScientificNotation;
         this.serializedInputHandlingEnabled = serializedInputHandlingEnabled;
+        this.timestampRepresentation = timestampRepresentation;
 
         this.dateFormat = dateFormat;
         this.timeFormat = timeFormat;
@@ -199,7 +212,7 @@ public WriteResult writeRawRecord(final Record record) throws IOException {
      * this writer with {@code serializedInputHandlingEnabled = false}.
      */
     private boolean isUseSerializeForm(final Record record, final RecordSchema writeSchema) {
-        if (!serializedInputHandlingEnabled) {
+        if (!serializedInputHandlingEnabled || timestampRepresentation != TimestampRepresentation.AUTO) {
             return false;
         }
 
@@ -348,8 +361,7 @@ private void writeRawValue(final JsonGenerator generator, final Object value, fi
             return;
         }
         if (value instanceof java.util.Date) {
-            final Object formatted = STRING_FIELD_CONVERTER.convertField(value, Optional.ofNullable(timestampFormat), fieldName);
-            generator.writeObject(formatted);
+            writeTimestamp(generator, value, fieldName);
             return;
         }
         if (!allowScientificNotation) {
@@ -404,12 +416,7 @@ private void writeValue(final JsonGenerator generator, final Object value, final
                 break;
             }
             case TIMESTAMP: {
-                final String stringValue = STRING_FIELD_CONVERTER.convertField(coercedValue, Optional.ofNullable(timestampFormat), fieldName);
-                if (DataTypeUtils.isLongTypeCompatible(stringValue)) {
-                    generator.writeNumber(DataTypeUtils.toLong(coercedValue, fieldName));
-                } else {
-                    generator.writeString(stringValue);
-                }
+                writeTimestamp(generator, coercedValue, fieldName);
                 break;
             }
             case DOUBLE:
@@ -502,6 +509,29 @@ private void writeArray(final Object[] values, final String fieldName, final Jso
         generator.writeEndArray();
     }
 
+    private void writeTimestamp(final JsonGenerator generator, final Object value, final String fieldName) throws IOException {
+        switch (timestampRepresentation) {
+            case FORMATTED_STRING:
+                generator.writeString(STRING_FIELD_CONVERTER.convertField(value, Optional.ofNullable(timestampFormat), fieldName));
+                break;
+            case EPOCH_MILLISECONDS:
+                generator.writeNumber(DataTypeUtils.toLong(value, fieldName));
+                break;
+            case EPOCH_SECONDS:
+                generator.writeNumber(BigDecimal.valueOf(DataTypeUtils.toLong(value, fieldName), 3));
+                break;
+            case AUTO:
+            default:
+                final String stringValue = STRING_FIELD_CONVERTER.convertField(value, Optional.ofNullable(timestampFormat), fieldName);
+                if (DataTypeUtils.isLongTypeCompatible(stringValue)) {
+                    generator.writeNumber(DataTypeUtils.toLong(value, fieldName));
+                } else {
+                    generator.writeString(stringValue);
+                }
+                break;
+        }
+    }
+
     @Override
     public String getMimeType() {
         return this.mimeType;

nifi-extension-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/main/java/org/apache/nifi/json/JsonRecordSetWriter.java

@@ -33,6 +33,7 @@
 import org.apache.nifi.migration.PropertyConfiguration;
 import org.apache.nifi.schema.access.SchemaNotFoundException;
 import org.apache.nifi.serialization.DateTimeTextRecordSetWriter;
+import org.apache.nifi.serialization.DateTimeUtils;
 import org.apache.nifi.serialization.RecordSetWriter;
 import org.apache.nifi.serialization.RecordSetWriterFactory;
 import org.apache.nifi.serialization.record.RecordSchema;
@@ -143,6 +144,14 @@ public class JsonRecordSetWriter extends DateTimeTextRecordSetWriter implements
             .defaultValue(HANDLING_ENABLED.getValue())
             .required(true)
             .build();
+    public static final PropertyDescriptor TIMESTAMP_REPRESENTATION = new PropertyDescriptor.Builder()
+            .name("Timestamp Representation")
+            .description("Specifies how Timestamp logical fields are represented in JSON. Date and Time logical fields are not affected.")
+            .expressionLanguageSupported(ExpressionLanguageScope.NONE)
+            .allowableValues(TimestampRepresentation.class)
+            .defaultValue(TimestampRepresentation.AUTO.getValue())
+            .required(true)
+            .build();
 
     private volatile boolean prettyPrint;
     private volatile boolean allowScientificNotation;
@@ -151,6 +160,7 @@ public class JsonRecordSetWriter extends DateTimeTextRecordSetWriter implements
     private volatile String compressionFormat;
     private volatile int compressionLevel;
     private volatile boolean serializedInputHandlingEnabled;
+    private volatile TimestampRepresentation timestampRepresentation;
 
     @Override
     protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
@@ -162,6 +172,7 @@ protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
         properties.add(COMPRESSION_FORMAT);
         properties.add(COMPRESSION_LEVEL);
         properties.add(SERIALIZED_JSON_INPUT_HANDLING);
+        properties.add(TIMESTAMP_REPRESENTATION);
         return properties;
     }
 
@@ -189,6 +200,11 @@ protected Collection<ValidationResult> customValidate(ValidationContext context)
             problems.add(new ValidationResult.Builder().input("Pretty Print").valid(false)
                     .explanation("Pretty Print JSON must be false when 'Output Grouping' is set to 'One Line Per Object'").build());
         }
+        if (context.getProperty(TIMESTAMP_REPRESENTATION).asAllowableValue(TimestampRepresentation.class) == TimestampRepresentation.FORMATTED_STRING
+                && !context.getProperty(DateTimeUtils.TIMESTAMP_FORMAT).isSet()) {
+            problems.add(new ValidationResult.Builder().subject(TIMESTAMP_REPRESENTATION.getDisplayName()).input(TimestampRepresentation.FORMATTED_STRING.getDisplayName()).valid(false)
+                    .explanation("Timestamp Format must be configured when Timestamp Representation is set to Formatted String").build());
+        }
         return problems;
     }
 
@@ -220,6 +236,7 @@ public void onEnabled(final ConfigurationContext context) {
         this.compressionFormat = context.getProperty(COMPRESSION_FORMAT).getValue();
         this.compressionLevel = context.getProperty(COMPRESSION_LEVEL).asInteger();
         this.serializedInputHandlingEnabled = HANDLING_ENABLED.getValue().equals(context.getProperty(SERIALIZED_JSON_INPUT_HANDLING).getValue());
+        this.timestampRepresentation = context.getProperty(TIMESTAMP_REPRESENTATION).asAllowableValue(TimestampRepresentation.class);
     }
 
     @Override
@@ -264,7 +281,8 @@ public RecordSetWriter createWriter(final ComponentLog logger, final RecordSchem
         }
 
         return new WriteJsonResult(logger, schema, getSchemaAccessWriter(schema, variables), compressionOut, prettyPrint, nullSuppression, outputGrouping,
-                getDateFormat().orElse(null), getTimeFormat().orElse(null), getTimestampFormat().orElse(null), mimeType, allowScientificNotation, serializedInputHandlingEnabled);
+                getDateFormat().orElse(null), getTimeFormat().orElse(null), getTimestampFormat().orElse(null), mimeType, allowScientificNotation, serializedInputHandlingEnabled,
+                timestampRepresentation);
     }
 
 }

nifi-extension-bundles/nifi-standard-services/nifi-record-serialization-services-bundle/nifi-record-serialization-services/src/test/java/org/apache/nifi/json/TestJsonRecordSetWriter.java

@@ -16,10 +16,15 @@
  */
 package org.apache.nifi.json;
 
+import org.apache.nifi.reporting.InitializationException;
 import org.apache.nifi.schema.access.SchemaAccessUtils;
+import org.apache.nifi.serialization.DateTimeUtils;
 import org.apache.nifi.serialization.SchemaRegistryRecordSetWriter;
 import org.apache.nifi.util.MockPropertyConfiguration;
+import org.apache.nifi.util.NoOpProcessor;
 import org.apache.nifi.util.PropertyMigrationResult;
+import org.apache.nifi.util.TestRunner;
+import org.apache.nifi.util.TestRunners;
 import org.junit.jupiter.api.Test;
 
 import java.util.Map;
@@ -33,9 +38,23 @@
 import static org.apache.nifi.schema.access.SchemaAccessUtils.SCHEMA_TEXT;
 import static org.apache.nifi.schema.access.SchemaAccessUtils.SCHEMA_VERSION;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class TestJsonRecordSetWriter {
 
+    @Test
+    void testFormattedStringRequiresTimestampFormat() throws InitializationException {
+        final JsonRecordSetWriter service = new JsonRecordSetWriter();
+        final TestRunner runner = TestRunners.newTestRunner(NoOpProcessor.class);
+        runner.addControllerService("writer", service);
+        runner.setProperty(service, JsonRecordSetWriter.TIMESTAMP_REPRESENTATION, TimestampRepresentation.FORMATTED_STRING.name());
+
+        assertThrows(IllegalStateException.class, () -> runner.enableControllerService(service));
+
+        runner.setProperty(service, DateTimeUtils.TIMESTAMP_FORMAT, "yyyy-MM-dd'T'HH:mm:ss.SSSX");
+        runner.enableControllerService(service);
+    }
+
     @Test
     void testMigrateProperties() {
         final JsonRecordSetWriter service = new JsonRecordSetWriter();