From 287e41d8a69f9d201e77840a3cced6fd576387ed Mon Sep 17 00:00:00 2001
From: Joe Gresock <jgresock@gmail.com>
Date: Thu, 6 Jan 2022 15:51:13 -0500
Subject: [PATCH] NIFI-9547: Allowing more datatypes in LookupRecord value
 substitution

Signed-off-by: Pierre Villard <pierre.villard.fr@gmail.com>

This closes #5640.
---
 .../processors/standard/LookupRecord.java     |  6 ++--
 .../processors/standard/TestLookupRecord.java | 33 ++++++++++++++++---
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/LookupRecord.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/LookupRecord.java
index e6910f4275..0132cd8652 100644
--- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/LookupRecord.java
+++ b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/LookupRecord.java
@@ -325,8 +325,7 @@ public class LookupRecord extends AbstractRouteRecord<Tuple<Map<String, RecordPa
             }
 
             for (FieldValue fieldValue : lookupFieldValues) {
-                final Object coordinateValue = (fieldValue.getValue() instanceof Number || fieldValue.getValue() instanceof Boolean)
-                        ? fieldValue.getValue() : DataTypeUtils.toString(fieldValue.getValue(), (String) null);
+                final Object coordinateValue = DataTypeUtils.convertType(fieldValue.getValue(), fieldValue.getField().getDataType(), null, null, null, fieldValue.getField().getFieldName());
 
                 lookupCoordinates.clear();
                 lookupCoordinates.put(coordinateKey, coordinateValue);
@@ -382,8 +381,7 @@ public class LookupRecord extends AbstractRouteRecord<Tuple<Map<String, RecordPa
             }
 
             final FieldValue fieldValue = lookupFieldValues.get(0);
-            final Object coordinateValue = (fieldValue.getValue() instanceof Number || fieldValue.getValue() instanceof Boolean)
-                    ? fieldValue.getValue() : DataTypeUtils.toString(fieldValue.getValue(), (String) null);
+            final Object coordinateValue = DataTypeUtils.convertType(fieldValue.getValue(), fieldValue.getField().getDataType(), null, null, null, fieldValue.getField().getFieldName());
             lookupCoordinates.put(coordinateKey, coordinateValue);
         }
 
diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestLookupRecord.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestLookupRecord.java
index 85aa0b6bba..48200ea8a6 100644
--- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestLookupRecord.java
+++ b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestLookupRecord.java
@@ -43,6 +43,7 @@ import org.junit.Test;
 
 import java.io.File;
 import java.io.IOException;
+import java.sql.Timestamp;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -85,9 +86,9 @@ public class TestLookupRecord {
         recordReader.addSchemaField("age", RecordFieldType.INT);
         recordReader.addSchemaField("sport", RecordFieldType.STRING);
 
-        recordReader.addRecord("John Doe", 48, null);
-        recordReader.addRecord("Jane Doe", 47, null);
-        recordReader.addRecord("Jimmy Doe", 14, null);
+        recordReader.addRecord("John Doe", 48, null, null);
+        recordReader.addRecord("Jane Doe", 47, null, null);
+        recordReader.addRecord("Jimmy Doe", 14, null, null);
     }
 
     @Test
@@ -133,6 +134,30 @@ public class TestLookupRecord {
         out.assertContentEquals("John Doe,48,Soccer\nJane Doe,47,Basketball\nJimmy Doe,14,Football\n");
     }
 
+    @Test
+    public void testLookupWithTimestamp() {
+        recordReader.addSchemaField("record_timestamp", RecordFieldType.TIMESTAMP);
+        runner.setProperty("lookup", "/record_timestamp");
+
+        final Timestamp timestamp = new Timestamp(0L);
+        final String timestampKey = timestamp.toString();
+        recordReader.addRecord("Jason Doe", 15, null, timestamp);
+
+        lookupService.addValue(timestampKey, "Bowling");
+
+        runner.enqueue("");
+        runner.run();
+
+        runner.assertTransferCount(LookupRecord.REL_MATCHED, 1);
+        runner.assertTransferCount(LookupRecord.REL_UNMATCHED, 1);
+        final MockFlowFile out = runner.getFlowFilesForRelationship(LookupRecord.REL_MATCHED).get(0);
+
+        out.assertAttributeEquals("record.count", "1");
+        out.assertAttributeEquals("mime.type", "text/plain");
+        String contents = out.getContent();
+        assertTrue(contents.matches("Jason Doe,15,Bowling,19[0-9]{2}-[0-9]{2}-[0-9]{2} [0-9]{2}:[0-9]{2}:[0-9]{2}\n"));
+    }
+
     @Test
     public void testAllUnmatched() {
         runner.enqueue("");
@@ -549,7 +574,7 @@ public class TestLookupRecord {
                 return Optional.empty();
             }
 
-            final String key = (String)coordinates.get("lookup");
+            final String key = coordinates.containsKey("lookup") ? coordinates.get("lookup").toString() : null;
             if (key == null) {
                 return Optional.empty();
             }