NIFI-1409: Fix missing transfer on error in ExecuteSql

Signed-off-by: jpercivall <joepercivall@yahoo.com>

nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/ExecuteSQL.java

@@ -53,11 +53,10 @@ import org.apache.nifi.util.StopWatch;
 @InputRequirement(Requirement.INPUT_ALLOWED)
 @Tags({"sql", "select", "jdbc", "query", "database"})
 @CapabilityDescription("Execute provided SQL select query. Query result will be converted to Avro format."
-    + " Streaming is used so arbitrarily large result sets are supported. This processor can be scheduled to run on " +
-    "a timer, or cron expression, using the standard scheduling methods, or it can be triggered by an incoming FlowFile. " +
-    "If it is triggered by an incoming FlowFile, then attributes of that FlowFile will be available when evaluating the " +
-    "select query. " +
-    "FlowFile attribute 'executesql.row.count' indicates how many rows were selected.")
+        + " Streaming is used so arbitrarily large result sets are supported. This processor can be scheduled to run on "
+        + "a timer, or cron expression, using the standard scheduling methods, or it can be triggered by an incoming FlowFile. "
+        + "If it is triggered by an incoming FlowFile, then attributes of that FlowFile will be available when evaluating the "
+        + "select query. FlowFile attribute 'executesql.row.count' indicates how many rows were selected.")
 public class ExecuteSQL extends AbstractProcessor {
 
     public static final String RESULT_ROW_COUNT = "executesql.row.count";
@@ -125,32 +124,32 @@ public class ExecuteSQL extends AbstractProcessor {
 
     @Override
     public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
-        FlowFile incoming = null;
+        FlowFile fileToProcess = null;
         if (context.hasIncomingConnection()) {
-            incoming = session.get();
+            fileToProcess = session.get();
 
             // If we have no FlowFile, and all incoming connections are self-loops then we can continue on.
             // However, if we have no FlowFile and we have connections coming from other Processors, then
             // we know that we should run only if we have a FlowFile.
-            if (incoming == null && context.hasNonLoopConnection()) {
+            if (fileToProcess == null && context.hasNonLoopConnection()) {
                 return;
             }
         }
 
         final ProcessorLog logger = getLogger();
-
         final DBCPService dbcpService = context.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class);
-        final String selectQuery = context.getProperty(SQL_SELECT_QUERY).evaluateAttributeExpressions(incoming).getValue();
+        final String selectQuery = context.getProperty(SQL_SELECT_QUERY).evaluateAttributeExpressions(fileToProcess).getValue();
         final Integer queryTimeout = context.getProperty(QUERY_TIMEOUT).asTimePeriod(TimeUnit.SECONDS).intValue();
-
         final StopWatch stopWatch = new StopWatch(true);
 
         try (final Connection con = dbcpService.getConnection();
             final Statement st = con.createStatement()) {
             st.setQueryTimeout(queryTimeout); // timeout in seconds
             final LongHolder nrOfRows = new LongHolder(0L);
-            FlowFile outgoing = (incoming == null ? session.create() : incoming);
-            outgoing = session.write(outgoing, new OutputStreamCallback() {
+            if (fileToProcess == null) {
+                fileToProcess = session.create();
+            }
+            fileToProcess = session.write(fileToProcess, new OutputStreamCallback() {
                 @Override
                 public void process(final OutputStream out) throws IOException {
                     try {
@@ -164,17 +163,30 @@ public class ExecuteSQL extends AbstractProcessor {
             });
 
             // set attribute how many rows were selected
-            outgoing = session.putAttribute(outgoing, RESULT_ROW_COUNT, nrOfRows.get().toString());
+            fileToProcess = session.putAttribute(fileToProcess, RESULT_ROW_COUNT, nrOfRows.get().toString());
 
-            logger.info("{} contains {} Avro records; transferring to 'success'", new Object[] {outgoing, nrOfRows.get()});
-            session.getProvenanceReporter().modifyContent(outgoing, "Retrieved " + nrOfRows.get() + " rows", stopWatch.getElapsed(TimeUnit.MILLISECONDS));
-            session.transfer(outgoing, REL_SUCCESS);
+            logger.info("{} contains {} Avro records; transferring to 'success'",
+                    new Object[]{fileToProcess, nrOfRows.get()});
+            session.getProvenanceReporter().modifyContent(fileToProcess, "Retrieved " + nrOfRows.get() + " rows",
+                    stopWatch.getElapsed(TimeUnit.MILLISECONDS));
+            session.transfer(fileToProcess, REL_SUCCESS);
         } catch (final ProcessException | SQLException e) {
-            if (incoming == null) {
-                logger.error("Unable to execute SQL select query {} due to {}. No incoming flow file to route to failure", new Object[] {selectQuery, e});
+            if (fileToProcess == null) {
+                // This can happen if any exceptions occur while setting up the connection, statement, etc.
+                logger.error("Unable to execute SQL select query {} due to {}. No FlowFile to route to failure",
+                        new Object[]{selectQuery, e});
+                context.yield();
             } else {
-                logger.error("Unable to execute SQL select query {} for {} due to {}; routing to failure", new Object[] {selectQuery, incoming, e});
-                session.transfer(incoming, REL_FAILURE);
+                if (context.hasIncomingConnection()) {
+                    logger.error("Unable to execute SQL select query {} for {} due to {}; routing to failure",
+                            new Object[]{selectQuery, fileToProcess, e});
+                    fileToProcess = session.penalize(fileToProcess);
+                } else {
+                    logger.error("Unable to execute SQL select query {} due to {}; routing to failure",
+                            new Object[]{selectQuery, e});
+                    context.yield();
+                }
+                session.transfer(fileToProcess, REL_FAILURE);
             }
         }
     }

nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestExecuteSQL.java

@@ -152,6 +152,31 @@ public class TestExecuteSQL {
         runner.getFlowFilesForRelationship(ExecuteSQL.REL_SUCCESS).get(0).assertAttributeEquals(ExecuteSQL.RESULT_ROW_COUNT, "2");
     }
 
+    @Test
+    public void testWithSqlException() throws SQLException {
+        // remove previous test database, if any
+        final File dbLocation = new File(DB_LOCATION);
+        dbLocation.delete();
+
+        // load test data to database
+        final Connection con = ((DBCPService) runner.getControllerService("dbcp")).getConnection();
+        Statement stmt = con.createStatement();
+
+        try {
+            stmt.execute("drop table TEST_NO_ROWS");
+        } catch (final SQLException sqle) {
+        }
+
+        stmt.execute("create table TEST_NO_ROWS (id integer)");
+
+        runner.setIncomingConnection(false);
+        // Try a valid SQL statment that will generate an error (val1 does not exist, e.g.)
+        runner.setProperty(ExecuteSQL.SQL_SELECT_QUERY, "SELECT val1 FROM TEST_NO_ROWS");
+        runner.run();
+
+        runner.assertAllFlowFilesTransferred(ExecuteSQL.REL_FAILURE, 1);
+    }
+
     public void invokeOnTrigger(final Integer queryTimeout, final String query, final boolean incomingFlowFile)
         throws InitializationException, ClassNotFoundException, SQLException, IOException {