NIFI-15209 JoltTransformRecord should not only take schema from first record

nifi-extension-bundles/nifi-jolt-bundle/nifi-jolt-processors/src/main/java/org/apache/nifi/processors/jolt/JoltTransformWritingStrategy.java

@@ -0,0 +1,46 @@
+/*
+ * 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.processors.jolt;
+
+import org.apache.nifi.components.DescribedValue;
+
+public enum JoltTransformWritingStrategy implements DescribedValue {
+    USE_FIRST_SCHEMA("The first successful transformation determines the schema for the entire FlowFile. All subsequent records must adhere to this schema or they may be invalid."),
+    PARTITION_BY_SCHEMA("Each unique schema generated by the transformation results in a separate output FlowFile. Useful when Jolt produces variable output structures.");
+
+    private final String description;
+
+    JoltTransformWritingStrategy(final String description) {
+        this.description = description;
+    }
+
+    @Override
+    public String getValue() {
+        return name();
+    }
+
+    @Override
+    public String getDisplayName() {
+        return name();
+    }
+
+    @Override
+    public String getDescription() {
+        return description;
+    }
+}

nifi-extension-bundles/nifi-jolt-bundle/nifi-jolt-processors/src/test/java/org/apache/nifi/processors/jolt/TestJoltTransformRecordPartitioned.java

@@ -0,0 +1,83 @@
+/*
+ * 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.processors.jolt;
+
+import org.apache.nifi.jolt.util.JoltTransformStrategy;
+import org.apache.nifi.json.JsonRecordSetWriter;
+import org.apache.nifi.json.JsonTreeReader;
+import org.apache.nifi.reporting.InitializationException;
+import org.apache.nifi.schema.access.SchemaAccessUtils;
+import org.apache.nifi.schema.inference.SchemaInferenceUtil;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledOnOs;
+import org.junit.jupiter.api.condition.OS;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+@DisabledOnOs(OS.WINDOWS) //The pretty printed json comparisons don't work on windows
+public class TestJoltTransformRecordPartitioned extends TestBaseJoltTransformRecord {
+
+    @Override
+    protected String getWritingStrategy() {
+        return JoltTransformWritingStrategy.PARTITION_BY_SCHEMA.getValue();
+    }
+
+    @Test
+    public void testTransformInputWithDifferentSchemas() throws InitializationException, IOException {
+        final JsonTreeReader reader = new JsonTreeReader();
+        runner.addControllerService("reader", reader);
+        runner.setProperty(reader, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaInferenceUtil.INFER_SCHEMA);
+        runner.enableControllerService(reader);
+        runner.setProperty(JoltTransformRecord.RECORD_READER, "reader");
+
+        runner.setProperty(writer, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaAccessUtils.INHERIT_RECORD_SCHEMA);
+        runner.setProperty(writer, JsonRecordSetWriter.PRETTY_PRINT_JSON, "true");
+        runner.enableControllerService(writer);
+
+        final String flattenSpec = Files.readString(Paths.get("src/test/resources/TestJoltTransformRecord/multipleSchemasSpec.json"));
+        runner.setProperty(JoltTransformRecord.JOLT_SPEC, flattenSpec);
+        runner.setProperty(JoltTransformRecord.JOLT_TRANSFORM, JoltTransformStrategy.CHAINR);
+
+        final String inputJson = Files.readString(Paths.get("src/test/resources/TestJoltTransformRecord/multipleSchemasInput.json"));
+        runner.enqueue(inputJson);
+
+        runner.run();
+
+        runner.assertTransferCount(JoltTransformRecord.REL_SUCCESS, 2);
+        runner.assertTransferCount(JoltTransformRecord.REL_ORIGINAL, 1);
+
+        final String expectedOutput1 = Files.readString(Paths.get("src/test/resources/TestJoltTransformRecord/multipleSchemasOutput1.json"));
+        final String expectedOutput2 = Files.readString(Paths.get("src/test/resources/TestJoltTransformRecord/multipleSchemasOutput2.json"));
+
+        final java.util.List<org.apache.nifi.util.MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(JoltTransformRecord.REL_SUCCESS);
+        final String result1 = new String(flowFiles.get(0).toByteArray());
+        final String result2 = new String(flowFiles.get(1).toByteArray());
+
+        // Handles non-deterministic order
+        if (result1.contains("TRASH")) {
+            assertEquals(expectedOutput1, result1);
+            assertEquals(expectedOutput2, result2);
+        } else {
+            assertEquals(expectedOutput2, result1);
+            assertEquals(expectedOutput1, result2);
+        }
+    }
+}

nifi-extension-bundles/nifi-jolt-bundle/nifi-jolt-processors/src/test/java/org/apache/nifi/processors/jolt/TestJoltTransformRecordUniform.java

@@ -0,0 +1,69 @@
+/*
+ * 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.processors.jolt;
+
+import org.apache.nifi.jolt.util.JoltTransformStrategy;
+import org.apache.nifi.json.JsonRecordSetWriter;
+import org.apache.nifi.json.JsonTreeReader;
+import org.apache.nifi.schema.access.SchemaAccessUtils;
+import org.apache.nifi.schema.inference.SchemaInferenceUtil;
+import org.apache.nifi.util.MockFlowFile;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledOnOs;
+import org.junit.jupiter.api.condition.OS;
+
+import java.nio.file.Files;
+import java.nio.file.Paths;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+@DisabledOnOs(OS.WINDOWS) //The pretty printed json comparisons don't work on windows
+public class TestJoltTransformRecordUniform extends TestBaseJoltTransformRecord {
+
+    @Override
+    protected String getWritingStrategy() {
+        return JoltTransformWritingStrategy.USE_FIRST_SCHEMA.getValue();
+    }
+
+    @Test
+    public void testJoltComplexChoiceField() throws Exception {
+        final JsonTreeReader reader = new JsonTreeReader();
+        runner.addControllerService("reader", reader);
+        runner.setProperty(reader, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaInferenceUtil.INFER_SCHEMA);
+        runner.enableControllerService(reader);
+        runner.setProperty(JoltTransformRecord.RECORD_READER, "reader");
+
+        runner.setProperty(writer, SchemaAccessUtils.SCHEMA_ACCESS_STRATEGY, SchemaAccessUtils.INHERIT_RECORD_SCHEMA);
+        runner.setProperty(writer, JsonRecordSetWriter.PRETTY_PRINT_JSON, "true");
+        runner.enableControllerService(writer);
+
+        final String flattenSpec = Files.readString(Paths.get("src/test/resources/TestJoltTransformRecord/flattenSpec.json"));
+        runner.setProperty(JoltTransformRecord.JOLT_SPEC, flattenSpec);
+        runner.setProperty(JoltTransformRecord.JOLT_TRANSFORM, JoltTransformStrategy.CHAINR);
+
+        final String inputJson = Files.readString(Paths.get("src/test/resources/TestJoltTransformRecord/input.json"));
+        runner.enqueue(inputJson);
+
+        runner.run();
+
+        runner.assertTransferCount(JoltTransformRecord.REL_SUCCESS, 1);
+        runner.assertTransferCount(JoltTransformRecord.REL_ORIGINAL, 1);
+
+        final MockFlowFile transformed = runner.getFlowFilesForRelationship(JoltTransformRecord.REL_SUCCESS).getFirst();
+        assertEquals(Files.readString(Paths.get("src/test/resources/TestJoltTransformRecord/flattenedOutput.json")), new String(transformed.toByteArray()));
+    }
+}

nifi-extension-bundles/nifi-jolt-bundle/nifi-jolt-processors/src/test/resources/TestJoltTransformRecord/filterOutAllSpec.json

@@ -0,0 +1,8 @@
+[
+  {
+    "operation": "shift",
+    "spec": {
+      "key": "value"
+    }
+  }
+]

nifi-extension-bundles/nifi-jolt-bundle/nifi-jolt-processors/src/test/resources/TestJoltTransformRecord/multipleSchemasInput.json

@@ -0,0 +1,8 @@
+[
+  {
+    "test_field": ""
+  },
+  {
+    "test_field": "value2"
+  }
+]

nifi-extension-bundles/nifi-jolt-bundle/nifi-jolt-processors/src/test/resources/TestJoltTransformRecord/multipleSchemasOutput1.json

@@ -0,0 +1,3 @@
+[ {
+  "TRASH" : null
+} ]

nifi-extension-bundles/nifi-jolt-bundle/nifi-jolt-processors/src/test/resources/TestJoltTransformRecord/multipleSchemasOutput2.json

@@ -0,0 +1,5 @@
+[ {
+  "value" : {
+    "test_field" : "value2"
+  }
+} ]

nifi-extension-bundles/nifi-jolt-bundle/nifi-jolt-processors/src/test/resources/TestJoltTransformRecord/multipleSchemasSpec.json

@@ -0,0 +1,13 @@
+[
+  {
+    "operation": "shift",
+    "spec": {
+      "test_field": {
+        "": "TRASH",
+        "*": {
+          "$": "value.test_field"
+        }
+      }
+    }
+  }
+]