diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py index 69859af3f5610..e6837f13fb15c 100644 --- a/airflow-core/src/airflow/cli/commands/task_command.py +++ b/airflow-core/src/airflow/cli/commands/task_command.py @@ -59,7 +59,6 @@ from sqlalchemy.orm.session import Session from airflow.models.operator import Operator - from airflow.typing_compat import Self CreateIfNecessary = Literal[False, "db", "memory"] @@ -448,53 +447,3 @@ def task_clear(args) -> None: only_running=args.only_running, confirm_prompt=not args.yes, ) - - -class LoggerMutationHelper: - """ - Helper for moving and resetting handlers and other logger attrs. - - :meta private: - """ - - def __init__(self, logger: logging.Logger) -> None: - self.handlers = logger.handlers[:] - self.level = logger.level - self.propagate = logger.propagate - self.source_logger = logger - - def apply(self, logger: logging.Logger, replace: bool = True) -> None: - """ - Set ``logger`` with attrs stored on instance. - - If ``logger`` is root logger, don't change propagate. - """ - if replace: - logger.handlers[:] = self.handlers - else: - for h in self.handlers: - if h not in logger.handlers: - logger.addHandler(h) - logger.level = self.level - if logger is not logging.getLogger(): - logger.propagate = self.propagate - - def move(self, logger: logging.Logger, replace: bool = True) -> None: - """ - Replace ``logger`` attrs with those from source. - - :param logger: target logger - :param replace: if True, remove all handlers from target first; otherwise add if not present. - """ - self.apply(logger, replace=replace) - self.source_logger.propagate = True - self.source_logger.handlers[:] = [] - - def reset(self) -> None: - self.apply(self.source_logger) - - def __enter__(self) -> Self: - return self - - def __exit__(self, exc_type, exc_val, exc_tb) -> None: - self.reset() diff --git a/airflow-core/tests/unit/cli/commands/test_task_command.py b/airflow-core/tests/unit/cli/commands/test_task_command.py index a71cf02930117..faad02503e407 100644 --- a/airflow-core/tests/unit/cli/commands/test_task_command.py +++ b/airflow-core/tests/unit/cli/commands/test_task_command.py @@ -29,13 +29,11 @@ from pathlib import Path from typing import TYPE_CHECKING from unittest import mock -from unittest.mock import sentinel import pytest from airflow.cli import cli_parser from airflow.cli.commands import task_command -from airflow.cli.commands.task_command import LoggerMutationHelper from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.configuration import conf from airflow.exceptions import DagRunNotFound @@ -480,70 +478,3 @@ def assert_log_line(self, text, logs_list, expect_from_logging_mixin=False): # Example: [2020-06-24 17:07:00,482] {logging_mixin.py:91} INFO - Log from Print statement assert "logging_mixin.py" not in log_line return log_line - - -class TestLoggerMutationHelper: - @pytest.mark.parametrize("target_name", ["test_apply_target", None]) - def test_apply(self, target_name): - """ - Handlers, level and propagate should be applied on target. - """ - src = logging.getLogger(f"test_apply_source_{target_name}") - src.propagate = False - src.addHandler(sentinel.handler) - src.setLevel(-1) - obj = LoggerMutationHelper(src) - tgt = logging.getLogger("test_apply_target") - obj.apply(tgt) - assert tgt.handlers == [sentinel.handler] - assert tgt.propagate is False if target_name else True # root propagate unchanged - assert tgt.level == -1 - - def test_apply_no_replace(self, clear_all_logger_handlers): - """ - Handlers, level and propagate should be applied on target. - """ - src = logging.getLogger("test_apply_source_no_repl") - tgt = logging.getLogger("test_apply_target_no_repl") - h1 = logging.Handler() - h1.name = "h1" - h2 = logging.Handler() - h2.name = "h2" - h3 = logging.Handler() - h3.name = "h3" - src.handlers[:] = [h1, h2] - tgt.handlers[:] = [h2, h3] - LoggerMutationHelper(src).apply(tgt, replace=False) - assert tgt.handlers == [h2, h3, h1] - - def test_move(self): - """Move should apply plus remove source handler, set propagate to True""" - src = logging.getLogger("test_move_source") - src.propagate = False - src.addHandler(sentinel.handler) - src.setLevel(-1) - obj = LoggerMutationHelper(src) - tgt = logging.getLogger("test_apply_target") - obj.move(tgt) - assert tgt.handlers == [sentinel.handler] - assert tgt.propagate is False - assert tgt.level == -1 - assert src.propagate is True - assert obj.propagate is False - assert src.level == obj.level - assert src.handlers == [] - assert obj.handlers == tgt.handlers - - def test_reset(self): - src = logging.getLogger("test_move_reset") - src.propagate = True - src.addHandler(sentinel.h1) - src.setLevel(-1) - obj = LoggerMutationHelper(src) - src.propagate = False - src.addHandler(sentinel.h2) - src.setLevel(-2) - obj.reset() - assert src.propagate is True - assert src.handlers == [sentinel.h1] - assert src.level == -1