Repository: apache/airflow

Issue:
task_command.task_run log handling is either broken or redundant
**Apache Airflow version**: 2.0.0a2

As of #9363, the `tasks run` command attempts to restore the root logger handlers to the previous state after running the task with the handlers all replaced by those from the `airflow.tasks` logger.

However, because no actual *copy* is created of the root handlers list, what you effectively end up with is either an empty list of handlers and nothing is restored. To make matters worse, the block ends with a `logging.shutdown()` call, which completely closes and releases all handlers.

So, either the code to 'restore' can just be removed, or a proper copy needs to be made and the `logging.shutdown()` call needs to be removed from the command. `logging.shutdown` is going to be called when the interpreter exits, anyway.

In detail, this is the code section:

https://github.com/apache/airflow/blob/bec9f3b29fd42ecd1beae3db75784b9a726caf15/airflow/cli/commands/task_command.py#L194-L222

The first problem here is that `root_logger_handlers = root_logger.handlers` merely creates a reference to the logger `handlers` list. The `root_logger.removeHandler(handler)` in a loop further on remove the handlers from that same list as you iterate, and this causes it to [only remove *every second handler*](https://sopython.com/canon/95/removing-items-from-a-list-while-iterating-over-the-list/). Luckily, this was never a problem because in the standard configuration there is only a single handler on the root. If there are more, there is a bigger problem further down.

Continuing on, the next loop, adding handlers from `airflow_logger_handlers`, causes those same handlers to show up in the `root_logger_handlers` reference, but those very same handlers are removed again in a second loop over the same `airflow_logger_handlers` list.

So when you hit the final `for handler in root_logger_handlers: root_logger.addHandler(handler)` loop, either `root_logger_handlers` is an empty list (no handlers added), or, if you started with more than one root handler and every second handler was left in place, you now are adding handlers to the list that are already in the list.  The `Logger.addHandler()` method uses `self.handler.append()`, appending to a list you are iterating over, and so you have an infinite loop on your hands.

Either just remove all handlers from the root logger (taking care to loop over the list in reverse, e.g. `for h in reversed(root_logger.handlers): root_logger.removeHandler(h)`), or create proper copy of the list with `root_logger_handlers = root_logger.handlers[:]`. The `logging.shutdown()` call has to be removed then too.






Discussion:
The `shutdown()` call was added [in this pull request](https://github.com/apache/airflow/pull/3259). I don't know if this still applies in Airflow 2.0.0, but I'm of the opinion that any logging shutdown should only be done when the process exits. It should not be up to `task_command.task_run()` to assume it is only ever used as a top-level call.

In other words, if a `logging.shutdown()` call is necessary when, say, a forked child is about to exit, then the code triggering the exit should call `logging.shutdown()` just before it calls `os._exit()`.
I found [this review remark](https://github.com/apache/airflow/pull/3259#discussion_r184536876):

> I removed this line in the Elasticsearch commit (because I saw this method registered in atexit()) and it causes NoneType error. My theory is that logging.shutdown() called in atexit() need settings.Session() to be defined but at the point it was called according to the order exit functions were registered, the orm was disposed. Explicitly calling it here is safe as all the task handler have logic prevent duplicate close() calls.

The correct fix would have been to fix the  logger that requires a live session object, rather than shut down all logging. Barring that, only shut down the handler that causes the problems, not all logging.

What exactly was causing a `NoneType` error at exit there, do you remember, @KevinYang21?
Btw. if there will be more handlers for `airflow.task` I think we will end up with duplicated logs in webserver view, won't we?
> Btw. if there will be more handlers for `airflow.task` I think we will end up with duplicated logs in webserver view, won't we?

Only if those handlers write to the same log file. The webviewer displays for specific log files (via a Flask file-serving route), not to anything a handler might do. There are all sorts of handlers, not all write to log files.

Or were you thinking of something else?
> Or were you thinking of something else?

I was confused by custom code I was fixing on fork 😄 You are right, if we have separate handlers (different streams/files) everything works as expected. 


Resolving patch:
diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py
--- a/airflow/cli/commands/task_command.py
+++ b/airflow/cli/commands/task_command.py
@@ -21,7 +21,7 @@
 import logging
 import os
 import textwrap
-from contextlib import redirect_stderr, redirect_stdout
+from contextlib import contextmanager, redirect_stderr, redirect_stdout
 from typing import List
 
 from tabulate import tabulate
@@ -142,6 +142,41 @@ def _run_raw_task(args, ti):
     )
 
 
+@contextmanager
+def _capture_task_logs(ti):
+    """Manage logging context for a task run
+
+    - Replace the root logger configuration with the airflow.task configuration
+      so we can capture logs from any custom loggers used in the task.
+
+    - Redirect stdout and stderr to the task instance log, as INFO and WARNING
+      level messages, respectively.
+
+    """
+    modify = not settings.DONOT_MODIFY_HANDLERS
+
+    if modify:
+        root_logger, task_logger = logging.getLogger(), logging.getLogger('airflow.task')
+
+        orig_level = root_logger.level
+        root_logger.setLevel(task_logger.level)
+        orig_handlers = root_logger.handlers.copy()
+        root_logger.handlers[:] = task_logger.handlers
+
+    try:
+        info_writer = StreamLogWriter(ti.log, logging.INFO)
+        warning_writer = StreamLogWriter(ti.log, logging.WARNING)
+
+        with redirect_stdout(info_writer), redirect_stderr(warning_writer):
+            yield
+
+    finally:
+        if modify:
+            # Restore the root logger to its original state.
+            root_logger.setLevel(orig_level)
+            root_logger.handlers[:] = orig_handlers
+
+
 @cli_utils.action_logging
 def task_run(args, dag=None):
     """Runs a single task instance"""
@@ -185,41 +220,8 @@ def task_run(args, dag=None):
     if args.interactive:
         _run_task_by_selected_method(args, dag, ti)
     else:
-        if settings.DONOT_MODIFY_HANDLERS:
-            with redirect_stdout(StreamLogWriter(ti.log, logging.INFO)), redirect_stderr(
-                StreamLogWriter(ti.log, logging.WARN)
-            ):
-                _run_task_by_selected_method(args, dag, ti)
-        else:
-            # Get all the Handlers from 'airflow.task' logger
-            # Add these handlers to the root logger so that we can get logs from
-            # any custom loggers defined in the DAG
-            airflow_logger_handlers = logging.getLogger('airflow.task').handlers
-            root_logger = logging.getLogger()
-            root_logger_handlers = root_logger.handlers
-
-            # Remove all handlers from Root Logger to avoid duplicate logs
-            for handler in root_logger_handlers:
-                root_logger.removeHandler(handler)
-
-            for handler in airflow_logger_handlers:
-                root_logger.addHandler(handler)
-            root_logger.setLevel(logging.getLogger('airflow.task').level)
-
-            with redirect_stdout(StreamLogWriter(ti.log, logging.INFO)), redirect_stderr(
-                StreamLogWriter(ti.log, logging.WARN)
-            ):
-                _run_task_by_selected_method(args, dag, ti)
-
-            # We need to restore the handlers to the loggers as celery worker process
-            # can call this command multiple times,
-            # so if we don't reset this then logs from next task would go to the wrong place
-            for handler in airflow_logger_handlers:
-                root_logger.removeHandler(handler)
-            for handler in root_logger_handlers:
-                root_logger.addHandler(handler)
-
-    logging.shutdown()
+        with _capture_task_logs(ti):
+            _run_task_by_selected_method(args, dag, ti)
 
 
 @cli_utils.action_logging


Test patch:
diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py
--- a/tests/cli/commands/test_task_command.py
+++ b/tests/cli/commands/test_task_command.py
@@ -17,6 +17,7 @@
 # under the License.
 #
 import io
+import logging
 import os
 import unittest
 from contextlib import redirect_stdout
@@ -316,16 +317,28 @@ def setUp(self) -> None:
         reset(self.dag_id)
         self.execution_date = timezone.make_aware(datetime(2017, 1, 1))
         self.execution_date_str = self.execution_date.isoformat()
+        self.task_args = ['tasks', 'run', self.dag_id, self.task_id, '--local', self.execution_date_str]
         self.log_dir = conf.get('logging', 'base_log_folder')
         self.log_filename = f"{self.dag_id}/{self.task_id}/{self.execution_date_str}/1.log"
         self.ti_log_file_path = os.path.join(self.log_dir, self.log_filename)
         self.parser = cli_parser.get_parser()
+
+        root = self.root_logger = logging.getLogger()
+        self.root_handlers = root.handlers.copy()
+        self.root_filters = root.filters.copy()
+        self.root_level = root.level
+
         try:
             os.remove(self.ti_log_file_path)
         except OSError:
             pass
 
     def tearDown(self) -> None:
+        root = self.root_logger
+        root.setLevel(self.root_level)
+        root.handlers[:] = self.root_handlers
+        root.filters[:] = self.root_filters
+
         reset(self.dag_id)
         try:
             os.remove(self.ti_log_file_path)
@@ -354,11 +367,7 @@ def test_logging_with_run_task(self):
         # as that is what gets displayed
 
         with conf_vars({('core', 'dags_folder'): self.dag_path}):
-            task_command.task_run(
-                self.parser.parse_args(
-                    ['tasks', 'run', self.dag_id, self.task_id, '--local', self.execution_date_str]
-                )
-            )
+            task_command.task_run(self.parser.parse_args(self.task_args))
 
         with open(self.ti_log_file_path) as l_file:
             logs = l_file.read()
@@ -390,11 +399,7 @@ def test_logging_with_run_task_subprocess(self):
         # We are not using self.assertLogs as we want to verify what actually is stored in the Log file
         # as that is what gets displayed
         with conf_vars({('core', 'dags_folder'): self.dag_path}):
-            task_command.task_run(
-                self.parser.parse_args(
-                    ['tasks', 'run', self.dag_id, self.task_id, '--local', self.execution_date_str]
-                )
-            )
+            task_command.task_run(self.parser.parse_args(self.task_args))
 
         with open(self.ti_log_file_path) as l_file:
             logs = l_file.read()
@@ -443,11 +448,7 @@ def test_log_file_template_with_run_task(self):
                 log_file_path = os.path.join(os.path.dirname(self.ti_log_file_path), "2.log")
 
                 try:
-                    task_command.task_run(
-                        self.parser.parse_args(
-                            ['tasks', 'run', self.dag_id, self.task_id, '--local', self.execution_date_str]
-                        )
-                    )
+                    task_command.task_run(self.parser.parse_args(self.task_args))
 
                     assert os.path.exists(log_file_path)
                 finally:
@@ -456,6 +457,59 @@ def test_log_file_template_with_run_task(self):
                     except OSError:
                         pass
 
+    @mock.patch.object(task_command, "_run_task_by_selected_method")
+    def test_root_logger_restored(self, run_task_mock):
+        """Verify that the root logging context is restored"""
+
+        logger = logging.getLogger("foo.bar")
+
+        def task_inner(*args, **kwargs):
+            logger.warning("redirected log message")
+
+        run_task_mock.side_effect = task_inner
+
+        config = {
+            ('core', 'dags_folder'): self.dag_path,
+            ('logging', 'logging_level'): "INFO",
+        }
+
+        with conf_vars(config):
+            with self.assertLogs(level=logging.WARNING) as captured:
+                logger.warning("not redirected")
+                task_command.task_run(self.parser.parse_args(self.task_args))
+
+                assert captured.output == ["WARNING:foo.bar:not redirected"]
+                assert self.root_logger.level == logging.WARNING
+
+        assert self.root_logger.handlers == self.root_handlers
+
+    @mock.patch.object(task_command, "_run_task_by_selected_method")
+    def test_disable_handler_modifying(self, run_task_mock):
+        """If [core] donot_modify_handlers is set to True, the root logger is untouched"""
+        from airflow import settings
+
+        logger = logging.getLogger("foo.bar")
+
+        def task_inner(*args, **kwargs):
+            logger.warning("not redirected")
+
+        run_task_mock.side_effect = task_inner
+
+        config = {
+            ('core', 'dags_folder'): self.dag_path,
+            ('logging', 'logging_level'): "INFO",
+        }
+        old_value = settings.DONOT_MODIFY_HANDLERS
+        settings.DONOT_MODIFY_HANDLERS = True
+
+        with conf_vars(config):
+            with self.assertLogs(level=logging.WARNING) as captured:
+                task_command.task_run(self.parser.parse_args(self.task_args))
+
+                assert captured.output == ["WARNING:foo.bar:not redirected"]
+
+        settings.DONOT_MODIFY_HANDLERS = old_value
+
 
 class TestCliTaskBackfill(unittest.TestCase):
     @classmethod