[nifi] branch main updated: NIFI-9781: Fix handling when selecting array element via QueryRecord

[jg...@apache.org]

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

jgresock 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 6a1c7c7  NIFI-9781: Fix handling when selecting array element via QueryRecord
6a1c7c7 is described below

commit 6a1c7c72d5b91b9ce5d5cb5b86e3155d21e2c19b
Author: Mark Payne <ma...@hotmail.com>
AuthorDate: Wed Mar 9 11:04:00 2022 -0500

    NIFI-9781: Fix handling when selecting array element via QueryRecord
    
    Signed-off-by: Joe Gresock <jg...@gmail.com>
    
    This closes #5853.
---
 .../nifi/serialization/record/ResultSetRecordSet.java | 19 +++++++++++++++++--
 .../apache/nifi/queryrecord/FlowFileEnumerator.java   |  2 +-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/nifi-commons/nifi-record/src/main/java/org/apache/nifi/serialization/record/ResultSetRecordSet.java b/nifi-commons/nifi-record/src/main/java/org/apache/nifi/serialization/record/ResultSetRecordSet.java
index cfb1e4b..8ce2255 100644
--- a/nifi-commons/nifi-record/src/main/java/org/apache/nifi/serialization/record/ResultSetRecordSet.java
+++ b/nifi-commons/nifi-record/src/main/java/org/apache/nifi/serialization/record/ResultSetRecordSet.java
@@ -18,6 +18,7 @@
 package org.apache.nifi.serialization.record;
 
 import org.apache.nifi.serialization.SimpleRecordSchema;
+import org.apache.nifi.serialization.record.type.ArrayDataType;
 import org.apache.nifi.serialization.record.util.DataTypeUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -201,7 +202,7 @@ public class ResultSetRecordSet implements RecordSet, Closeable {
             throws SQLException {
         switch (sqlType) {
             case Types.ARRAY:
-                return getArrayDataType(rs, columnIndex, useLogicalTypes);
+                return getArrayDataType(rs, readerSchema, columnIndex, useLogicalTypes);
             case Types.BINARY:
             case Types.LONGVARBINARY:
             case Types.VARBINARY:
@@ -282,7 +283,21 @@ public class ResultSetRecordSet implements RecordSet, Closeable {
         }
     }
 
-    private DataType getArrayDataType(final ResultSet rs, final int columnIndex, final boolean useLogicalTypes) throws SQLException {
+    private DataType getArrayDataType(final ResultSet rs, final RecordSchema readerSchema, final int columnIndex, final boolean useLogicalTypes) throws SQLException {
+        // We first want to check if the Reader Schema can tell us what the type of the array is.
+        final String columnName = rs.getMetaData().getColumnName(columnIndex);
+        final Optional<RecordField> optionalRecordField = readerSchema.getField(columnName);
+        if (optionalRecordField.isPresent()) {
+            final RecordField recordField = optionalRecordField.get();
+            final DataType dataType = recordField.getDataType();
+            if (dataType.getFieldType() == RecordFieldType.ARRAY) {
+                final ArrayDataType arrayDataType = (ArrayDataType) dataType;
+                if (arrayDataType.getElementType() != null) {
+                    return dataType;
+                }
+            }
+        }
+
         // The JDBC API does not allow us to know what the base type of an array is through the metadata.
         // As a result, we have to obtain the actual Array for this record. Once we have this, we can determine
         // the base type. However, if the base type is, itself, an array, we will simply return a base type of
diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/queryrecord/FlowFileEnumerator.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/queryrecord/FlowFileEnumerator.java
index e4814ec..db66c5a 100644
--- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/queryrecord/FlowFileEnumerator.java
+++ b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/queryrecord/FlowFileEnumerator.java
@@ -107,7 +107,7 @@ public class FlowFileEnumerator implements Enumerator<Object> {
         // the actual value, NOT a 1-element array of values.
         if (fields.length == 1) {
             final int desiredCellIndex = fields[0];
-            return row[desiredCellIndex];
+            return cast(row[desiredCellIndex]);
         }
 
         // Create a new Object array that contains only the desired fields.