Repository: PrefectHQ/prefect

Issue:
allow ShellTask to return stdout when it fails
**The context**
Currently, if the shell process returns non-zero exit code, It will raise a FAIL signal, without any result attached. 

https://github.com/PrefectHQ/prefect/blob/b939c7d043f5de910a54ad95e07873a835f48b31/src/prefect/tasks/shell.py#L113-L122

**The probem**
Sometime, the downstream task need to interpret the output of a failed shell task (might be from both stdout and stderr) to handle the the failure properly. 

In my use case, I use DBTShellTask (which is a subclass of ShellTask) for execute a  `dbt run` as shell sub-process. I then use another task, called `parse_dbt_result`, to take the output of the dbt task and handle both success and failure results, such as knowing which models in the dbt-run process failed.

When some models fails in a `dbt run`, those model errors are shown in the stdout, and  the shell subprocess returns non-zero exit code (1). The ShellTask raises a FAIL signal without any stdout result.  As a result, the `parse_dbt_result` task is not able to interpret the log of failed run and not know which model fails, hereby not be able to handle the failure properly. 

**Solution**
please provide a flag that allows ShellTask to return stdout(or/and stderror) result when the shell process fails and FAIL signal is raised.


Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2f1a4572cf..02c95f551d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,7 +23,7 @@ Released on November 11, 2020.
 
 - Fix issue with retrying mapped pipelines on dask - [#3519](https://github.com/PrefectHQ/prefect/issues/3519)
 - Task arguments take precedence when generating `task_run_name` - [#3605](https://github.com/PrefectHQ/prefect/issues/3605)
-- Fix breaking change in flow registration with old server versions - [#3642](https://github.com/PrefectHQ/prefect/pull/3642)
+- Fix breaking change in flow registration with old server versions - [#3624](https://github.com/PrefectHQ/prefect/pull/3624)
 - Task arguments take precedence when generating templated targets and locations - [#3627](https://github.com/PrefectHQ/prefect/pull/3627)
 
 ### Breaking Changes
diff --git a/changes/pr3649.yaml b/changes/pr3649.yaml
new file mode 100644
index 0000000000..d1791f458c
--- /dev/null
+++ b/changes/pr3649.yaml
@@ -0,0 +1,3 @@
+enhancement:
+  - "ShellTask returns output on failure - [#3649](https://github.com/PrefectHQ/prefect/pull/3649)"
+  - "ShellTask allows streaming of output independently of the number of lines returned - [#3649](https://github.com/PrefectHQ/prefect/pull/3649)"
diff --git a/src/prefect/tasks/shell.py b/src/prefect/tasks/shell.py
index ff4f9420a7..307f2264c9 100644
--- a/src/prefect/tasks/shell.py
+++ b/src/prefect/tasks/shell.py
@@ -1,7 +1,7 @@
 import os
 import tempfile
 from subprocess import PIPE, STDOUT, Popen
-from typing import Any
+from typing import Any, List, Union, Optional
 
 import prefect
 from prefect.utilities.tasks import defaults_from_attrs
@@ -11,6 +11,9 @@ class ShellTask(prefect.Task):
     """
     Task for running arbitrary shell commands.
 
+    NOTE: This task combines stderr and stdout because reading from both
+          streams without blocking is tricky.
+
     Args:
         - command (string, optional): shell command to be executed; can also be
             provided post-initialization by calling this task instance
@@ -24,9 +27,13 @@ class ShellTask(prefect.Task):
         - return_all (bool, optional): boolean specifying whether this task
             should return all lines of stdout as a list, or just the last line
             as a string; defaults to `False`
-        - log_stderr (bool, optional): boolean specifying whether this task
-            should log the output from stderr in the case of a non-zero exit code;
-            defaults to `False`
+        - log_stderr (bool, optional): boolean specifying whether this task should log
+            the output in the case of a non-zero exit code; defaults to `False`. This
+             actually logs both stderr and stdout and will only log the last line of
+             output unless `return_all` is `True`
+        - stream_output (bool, optional): specifies whether this task should log
+            the output as it occurs. If enabled, `log_stderr` will be ignored as the
+            output will have already been logged. defaults to `False`
         - **kwargs: additional keyword arguments to pass to the Task constructor
 
     Example:
@@ -52,6 +59,7 @@ class ShellTask(prefect.Task):
         shell: str = "bash",
         return_all: bool = False,
         log_stderr: bool = False,
+        stream_output: bool = False,
         **kwargs: Any
     ):
         self.command = command
@@ -60,10 +68,13 @@ class ShellTask(prefect.Task):
         self.shell = shell
         self.return_all = return_all
         self.log_stderr = log_stderr
+        self.stream_output = stream_output
         super().__init__(**kwargs)
 
     @defaults_from_attrs("command", "env")
-    def run(self, command: str = None, env: dict = None) -> str:
+    def run(
+        self, command: str = None, env: dict = None
+    ) -> Optional[Union[str, List[str]]]:
         """
         Run the shell command.
 
@@ -76,8 +87,8 @@ class ShellTask(prefect.Task):
                 the subprocess
 
         Returns:
-            - stdout (string): if `return_all` is `False` (the default), only
-                the last line of stdout is returned, otherwise all lines are
+            - output (string): if `return_all` is `False` (the default), only
+                the last line of output is returned, otherwise all lines are
                 returned, which is useful for passing result of shell command
                 to other downstream tasks. If there is no output, `None` is
                 returned.
@@ -100,27 +111,31 @@ class ShellTask(prefect.Task):
             with Popen(
                 [self.shell, tmp.name], stdout=PIPE, stderr=STDOUT, env=current_env
             ) as sub_process:
-                lines = []
                 line = None
+                lines = []
                 for raw_line in iter(sub_process.stdout.readline, b""):
                     line = raw_line.decode("utf-8").rstrip()
+
                     if self.return_all:
                         lines.append(line)
-                    else:
-                        # if we're returning all, we don't log every line
+
+                    if self.stream_output:
                         self.logger.debug(line)
+
                 sub_process.wait()
                 if sub_process.returncode:
                     msg = "Command failed with exit code {}".format(
                         sub_process.returncode,
                     )
+
                     self.logger.error(msg)
 
-                    if self.log_stderr:
-                        self.logger.error("\n".join(lines))
+                    if self.log_stderr and not self.stream_output:
+                        self.logger.error("\n".join(lines) if self.return_all else line)
+
+                    raise prefect.engine.signals.FAIL(
+                        msg,
+                        result=lines if self.return_all else line,
+                    )
 
-                    raise prefect.engine.signals.FAIL(msg) from None  # type: ignore
-        if self.return_all:
-            return lines
-        else:
-            return line
+        return lines if self.return_all else line


Test patch:
diff --git a/tests/tasks/test_shell.py b/tests/tasks/test_shell.py
index 94c7c70ec0..b287e6994d 100644
--- a/tests/tasks/test_shell.py
+++ b/tests/tasks/test_shell.py
@@ -108,6 +108,34 @@ def test_shell_logs_non_zero_exit(caplog):
     assert "Command failed" in error_log[0].message
 
 
+def test_shell_attaches_result_to_failure(caplog):
+    def assert_fail_result(task, state):
+        assert (
+            state.result == "foo"
+            "ls: surely_a_dir_that_doesnt_exist: No such file or directory"
+        )
+
+    with Flow(name="test") as f:
+        task = ShellTask(on_failure=assert_fail_result)(
+            command="echo foo && ls surely_a_dir_that_doesnt_exist"
+        )
+    out = f.run()
+    assert out.is_failed()
+
+
+@pytest.mark.parametrize("stream_output", [True, False])
+def test_shell_respects_stream_output(caplog, stream_output):
+
+    with Flow(name="test") as f:
+        ShellTask(stream_output=stream_output)(
+            command="echo foo && echo bar",
+        )
+    f.run()
+
+    stdout_in_log = "foo" in caplog.messages and "bar" in caplog.messages
+    assert stdout_in_log == stream_output
+
+
 def test_shell_logs_stderr_on_non_zero_exit(caplog):
     with Flow(name="test") as f:
         task = ShellTask(log_stderr=True, return_all=True)(