NIFI-15024 - Add executed query as a FlowFile attribute for ExecuteSQL and ExecuteSQLRecord

Author: pvillard31

nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/AbstractExecuteSQL.java

@@ -65,6 +65,7 @@ public abstract class AbstractExecuteSQL extends AbstractProcessor {
     public static final String RESULT_QUERY_DURATION = "executesql.query.duration";
     public static final String RESULT_QUERY_EXECUTION_TIME = "executesql.query.executiontime";
     public static final String RESULT_QUERY_FETCH_TIME = "executesql.query.fetchtime";
+    public static final String RESULT_QUERY_STATEMENT = "executesql.query";
     public static final String RESULTSET_INDEX = "executesql.resultset.index";
     public static final String RESULT_ERROR_MESSAGE = "executesql.error.message";
     public static final String INPUT_FLOWFILE_UUID = "input.flowfile.uuid";
@@ -272,6 +273,10 @@ public void onTrigger(final ProcessContext context, final ProcessSession session
             selectQuery = queryContents.toString();
         }
 
+        if (fileToProcess != null) {
+            fileToProcess = session.putAttribute(fileToProcess, RESULT_QUERY_STATEMENT, selectQuery);
+        }
+
         int resultCount = 0;
         try (final Connection con = dbcpService.getConnection(fileToProcess == null ? Collections.emptyMap() : fileToProcess.getAttributes())) {
             final boolean isAutoCommit = con.getAutoCommit();
@@ -378,6 +383,7 @@ public void onTrigger(final ProcessContext context, final ProcessSession session
                                     attributesToAdd.put(RESULT_QUERY_DURATION, String.valueOf(executionTimeElapsed + fetchTimeElapsed));
                                     attributesToAdd.put(RESULT_QUERY_EXECUTION_TIME, String.valueOf(executionTimeElapsed));
                                     attributesToAdd.put(RESULT_QUERY_FETCH_TIME, String.valueOf(fetchTimeElapsed));
+                                    attributesToAdd.put(RESULT_QUERY_STATEMENT, selectQuery);
                                     attributesToAdd.put(RESULTSET_INDEX, String.valueOf(resultCount));
                                     if (inputFileUUID != null) {
                                         attributesToAdd.put(INPUT_FLOWFILE_UUID, inputFileUUID);
@@ -479,14 +485,14 @@ public void onTrigger(final ProcessContext context, final ProcessSession session
                             session.transfer(fileToProcess, REL_SUCCESS);
                         } else {
                             // Set Empty Results as the default behavior based on strategy or null property
-                            session.transfer(setFlowFileEmptyResults(session, fileToProcess, sqlWriter), REL_SUCCESS);
+                            session.transfer(setFlowFileEmptyResults(session, fileToProcess, sqlWriter, selectQuery), REL_SUCCESS);
                         }
                     }
                 } else if (resultCount == 0) {
                     // If we had no inbound FlowFile, no exceptions, and the SQL generated no result sets (Insert/Update/Delete statements only)
                     // Then generate an empty Output FlowFile
                     FlowFile resultSetFF = session.create();
-                    session.transfer(setFlowFileEmptyResults(session, resultSetFF, sqlWriter), REL_SUCCESS);
+                    session.transfer(setFlowFileEmptyResults(session, resultSetFF, sqlWriter, selectQuery), REL_SUCCESS);
                 }
             }
         } catch (final ProcessException | SQLException e) {
@@ -505,16 +511,20 @@ public void onTrigger(final ProcessContext context, final ProcessSession session
                     context.yield();
                 }
                 session.putAttribute(fileToProcess, RESULT_ERROR_MESSAGE, e.getMessage());
+                session.putAttribute(fileToProcess, RESULT_QUERY_STATEMENT, selectQuery);
                 session.transfer(fileToProcess, REL_FAILURE);
             }
         }
     }
 
-    protected FlowFile setFlowFileEmptyResults(final ProcessSession session, FlowFile flowFile, SqlWriter sqlWriter) {
+    protected FlowFile setFlowFileEmptyResults(final ProcessSession session, FlowFile flowFile, SqlWriter sqlWriter, final String selectQuery) {
         flowFile = session.write(flowFile, out -> sqlWriter.writeEmptyResultSet(out, getLogger()));
         final Map<String, String> attributesToAdd = new HashMap<>();
         attributesToAdd.put(RESULT_ROW_COUNT, "0");
         attributesToAdd.put(CoreAttributes.MIME_TYPE.key(), sqlWriter.getMimeType());
+        if (selectQuery != null) {
+            attributesToAdd.put(RESULT_QUERY_STATEMENT, selectQuery);
+        }
         return session.putAllAttributes(flowFile, attributesToAdd);
     }
 

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

@@ -83,6 +83,7 @@
 @WritesAttributes({
         @WritesAttribute(attribute = "executesql.row.count", description = "Contains the number of rows returned by the query. "
                 + "If 'Max Rows Per Flow File' is set, then this number will reflect the number of rows in the Flow File instead of the entire result set."),
+        @WritesAttribute(attribute = "executesql.query", description = "SQL query executed for the FlowFile."),
         @WritesAttribute(attribute = "executesql.query.duration", description = "Combined duration of the query execution time and fetch time in milliseconds. "
                 + "If 'Max Rows Per Flow File' is set, then this number will reflect only the fetch time for the rows in the Flow File instead of the entire result set."),
         @WritesAttribute(attribute = "executesql.query.executiontime", description = "Duration of the query execution time in milliseconds. "

nifi-extension-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/ExecuteSQLRecord.java

@@ -82,6 +82,7 @@
 })
 @WritesAttributes({
         @WritesAttribute(attribute = "executesql.row.count", description = "Contains the number of rows returned in the select query"),
+        @WritesAttribute(attribute = "executesql.query", description = "SQL query executed for the FlowFile."),
         @WritesAttribute(attribute = "executesql.query.duration", description = "Combined duration of the query execution time and fetch time in milliseconds"),
         @WritesAttribute(attribute = "executesql.query.executiontime", description = "Duration of the query execution time in milliseconds"),
         @WritesAttribute(attribute = "executesql.query.fetchtime", description = "Duration of the result set fetch time in milliseconds"),