Repository: PrefectHQ/prefect

Issue:
Custom log_level for ShellTask
## Current behavior
All `ShellTask` logs are logged at `DEBUG` level, preventing showing the command output in the UI logs unless the agent's logging threshold is set to DEBUG. But this is undesired for production workloads, as this applies to all log messages

streaming logging for the `ShellTask` was introduced in #3584 


## Proposed behavior
change existing `stream_logs` parameter to `Union[bool, int]` to accept the desired log_level as suggest by @madkinsz.
see discussion at the bottom of #3584


## Example
see #3584



Discussion:

Resolving patch:
diff --git a/src/prefect/tasks/shell.py b/src/prefect/tasks/shell.py
--- a/src/prefect/tasks/shell.py
+++ b/src/prefect/tasks/shell.py
@@ -1,5 +1,6 @@
 import os
 import tempfile
+import logging
 from subprocess import PIPE, STDOUT, Popen
 from typing import Any, List, Union, Optional
 
@@ -31,9 +32,13 @@ class ShellTask(prefect.Task):
             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`
+        - stream_output (Union[bool, int, str], optional): specifies whether this task should log
+            the output as it occurs, and at what logging level. If `True` is passed,
+            the logging level defaults to `INFO`; otherwise, any integer or string
+            value that's passed will be treated as the log level, provided
+            the `logging` library can successfully interpret it. 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:
@@ -49,6 +54,9 @@ class ShellTask(prefect.Task):
 
         out = f.run()
         ```
+    Raises:
+        - TypeError: if `stream_output` is passed in as a string, but cannot
+          successfully be converted to a numeric value by logging.getLevelName()
     """
 
     def __init__(
@@ -59,8 +67,8 @@ def __init__(
         shell: str = "bash",
         return_all: bool = False,
         log_stderr: bool = False,
-        stream_output: bool = False,
-        **kwargs: Any
+        stream_output: Union[bool, int, str] = False,
+        **kwargs: Any,
     ):
         self.command = command
         self.env = env
@@ -68,7 +76,15 @@ def __init__(
         self.shell = shell
         self.return_all = return_all
         self.log_stderr = log_stderr
-        self.stream_output = stream_output
+
+        if isinstance(stream_output, str):
+            stream_output = logging.getLevelName(stream_output)
+            if not isinstance(stream_output, int):
+                raise TypeError(
+                    f"'stream_output': {stream_output} is not a valid log level"
+                )
+
+        self.stream_output = logging.INFO if stream_output is True else stream_output
         super().__init__(**kwargs)
 
     @defaults_from_attrs("command", "env", "helper_script")
@@ -124,7 +140,7 @@ def run(
                         lines.append(line)
 
                     if self.stream_output:
-                        self.logger.debug(line)
+                        self.logger.log(level=self.stream_output, msg=line)
 
                 sub_process.wait()
                 if sub_process.returncode:


Test patch:
diff --git a/tests/tasks/test_shell.py b/tests/tasks/test_shell.py
--- a/tests/tasks/test_shell.py
+++ b/tests/tasks/test_shell.py
@@ -136,6 +136,53 @@ def test_shell_respects_stream_output(caplog, stream_output):
     assert stdout_in_log == stream_output
 
 
+def test_shell_log_stream_default(caplog):
+    with Flow(name="test") as f:
+        ShellTask(stream_output=True)(command="echo foo && echo bar")
+    f.run()
+
+    log_messages = [(r.levelname, r.message) for r in caplog.records]
+    assert ("INFO", "foo") in log_messages and ("INFO", "bar") in log_messages
+
+
+def test_shell_log_stream_as_info(caplog):
+    with Flow(name="test") as f:
+        ShellTask(stream_output=logging.INFO)(command="echo foo && echo bar")
+    f.run()
+
+    log_messages = [(r.levelname, r.message) for r in caplog.records]
+    assert ("INFO", "foo") in log_messages and ("INFO", "bar") in log_messages
+
+
+def test_shell_logs_stream_as_debug(caplog):
+    with Flow(name="test") as f:
+        ShellTask(stream_output=logging.DEBUG)(command="echo foo && echo bar")
+    f.run()
+
+    log_messages = [(r.levelname, r.message) for r in caplog.records]
+    assert ("DEBUG", "foo") in log_messages and ("DEBUG", "bar") in log_messages
+
+
+def test_shell_log_stream_type_error_on_invalid_log_level_string(caplog):
+    with pytest.raises(TypeError):
+        with raise_on_exception():
+            with Flow(name="test") as f:
+                ShellTask(stream_output="FOO")
+
+
+@pytest.mark.parametrize("stream_output", ["INFO", "DEBUG"])
+def test_shell_log_stream_as_info_string_input(caplog, stream_output):
+    with Flow(name="test") as f:
+        ShellTask(stream_output=stream_output)(command="echo foo && echo bar")
+    f.run()
+
+    log_messages = [(r.levelname, r.message) for r in caplog.records]
+    assert (stream_output, "foo") in log_messages and (
+        stream_output,
+        "bar",
+    ) in log_messages
+
+
 def test_shell_logs_stderr_on_non_zero_exit(caplog):
     caplog.set_level(level=logging.ERROR, logger="prefect.ShellTask")
     with Flow(name="test") as f: