NIFI-4319 - Fixed ArrayIndexOutOfBoundsException in QueryCassandra

Signed-off-by: Matthew Burgess <mattyb149@apache.org>

This closes #2112

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-processors/src/main/java/org/apache/nifi/processors/cassandra/QueryCassandra.java

@@ -40,6 +40,7 @@ import org.apache.nifi.annotation.lifecycle.OnShutdown;
 import org.apache.nifi.annotation.lifecycle.OnUnscheduled;
 import org.apache.nifi.components.PropertyDescriptor;
 import org.apache.nifi.flowfile.FlowFile;
+import org.apache.nifi.flowfile.attributes.CoreAttributes;
 import org.apache.nifi.annotation.behavior.WritesAttribute;
 import org.apache.nifi.annotation.behavior.WritesAttributes;
 import org.apache.nifi.annotation.lifecycle.OnScheduled;
@@ -263,6 +264,10 @@ public class QueryCassandra extends AbstractCassandraProcessor {
             // set attribute how many rows were selected
             fileToProcess = session.putAttribute(fileToProcess, RESULT_ROW_COUNT, String.valueOf(nrOfRows.get()));
 
+            // set mime.type based on output format
+            fileToProcess = session.putAttribute(fileToProcess, CoreAttributes.MIME_TYPE.key(),
+                    JSON_FORMAT.equals(outputFormat) ? "application/json" : "application/avro-binary");
+
             logger.info("{} contains {} Avro records; transferring to 'success'",
                     new Object[]{fileToProcess, nrOfRows.get()});
             session.getProvenanceReporter().modifyContent(fileToProcess, "Retrieved " + nrOfRows.get() + " rows",
@@ -510,7 +515,7 @@ public class QueryCassandra extends AbstractCassandraProcessor {
         final int nrOfColumns = (columnDefinitions == null ? 0 : columnDefinitions.size());
         String tableName = "NiFi_Cassandra_Query_Record";
         if (nrOfColumns > 0) {
-            String tableNameFromMeta = columnDefinitions.getTable(1);
+            String tableNameFromMeta = columnDefinitions.getTable(0);
             if (!StringUtils.isBlank(tableNameFromMeta)) {
                 tableName = tableNameFromMeta;
             }

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-processors/src/test/java/org/apache/nifi/processors/cassandra/CassandraQueryTestUtil.java

@@ -61,7 +61,7 @@ public class CassandraQueryTestUtil {
             }
         });
 
-        when(columnDefinitions.getTable(1)).thenReturn("users");
+        when(columnDefinitions.getTable(0)).thenReturn("users");
 
         when(columnDefinitions.getType(anyInt())).thenAnswer(new Answer<DataType>() {
 
@@ -103,6 +103,43 @@ public class CassandraQueryTestUtil {
         return resultSet;
     }
 
+    public static ResultSet createMockResultSetOneColumn() throws Exception {
+        ResultSet resultSet = mock(ResultSet.class);
+        ColumnDefinitions columnDefinitions = mock(ColumnDefinitions.class);
+        when(columnDefinitions.size()).thenReturn(1);
+        when(columnDefinitions.getName(anyInt())).thenAnswer(new Answer<String>() {
+            List<String> colNames = Arrays.asList("user_id");
+            @Override
+            public String answer(InvocationOnMock invocationOnMock) throws Throwable {
+                return colNames.get((Integer) invocationOnMock.getArguments()[0]);
+
+            }
+        });
+
+        when(columnDefinitions.getTable(0)).thenReturn("users");
+
+        when(columnDefinitions.getType(anyInt())).thenAnswer(new Answer<DataType>() {
+            List<DataType> dataTypes = Arrays.asList(DataType.text());
+            @Override
+            public DataType answer(InvocationOnMock invocationOnMock) throws Throwable {
+                return dataTypes.get((Integer) invocationOnMock.getArguments()[0]);
+
+            }
+        });
+
+        List<Row> rows = Arrays.asList(
+                createRow("user1"),
+                createRow("user2")
+        );
+
+        when(resultSet.iterator()).thenReturn(rows.iterator());
+        when(resultSet.all()).thenReturn(rows);
+        when(resultSet.getAvailableWithoutFetching()).thenReturn(rows.size());
+        when(resultSet.isFullyFetched()).thenReturn(false).thenReturn(true);
+        when(resultSet.getColumnDefinitions()).thenReturn(columnDefinitions);
+        return resultSet;
+    }
+
     public static Row createRow(String user_id, String first_name, String last_name, Set<String> emails,
                                 List<String> top_places, Map<Date, String> todo, boolean registered,
                                 float scale, double metric) {
@@ -119,4 +156,10 @@ public class CassandraQueryTestUtil {
 
         return row;
     }
+
+    public static Row createRow(String user_id) {
+        Row row = mock(Row.class);
+        when(row.getString(0)).thenReturn(user_id);
+        return row;
+    }
 }

nifi-nar-bundles/nifi-cassandra-bundle/nifi-cassandra-processors/src/test/java/org/apache/nifi/processors/cassandra/QueryCassandraTest.java

@@ -240,6 +240,14 @@ public class QueryCassandraTest {
         testRunner.assertAllFlowFilesTransferred(QueryCassandra.REL_FAILURE, 1);
     }
 
+    @Test
+    public void testCreateSchemaOneColumn() throws Exception {
+        ResultSet rs = CassandraQueryTestUtil.createMockResultSetOneColumn();
+        Schema schema = QueryCassandra.createSchema(rs);
+        assertNotNull(schema);
+        assertEquals(schema.getName(), "users");
+    }
+
     @Test
     public void testCreateSchema() throws Exception {
         ResultSet rs = CassandraQueryTestUtil.createMockResultSet();