[v3-0-test] Separate configurations for colorized and json logs in Task SDK / Celery Executor (#51082) (#51344)

* Enhance logging configuration by adding 'enable_colors' parameter to 'logging_processors' and 'configure_logging' functions. This allows for customizable console output with or without colors based on user preference or configuration settings.
(cherry picked from commit e25351d)

Co-authored-by: Peter Bryant <[email protected]>

Author: 2 people

task-sdk/src/airflow/sdk/log.py

@@ -140,7 +140,7 @@ def emit(self, record: logging.LogRecord):
 
 
 @cache
-def logging_processors(enable_pretty_log: bool, mask_secrets: bool = True):
+def logging_processors(enable_pretty_log: bool, mask_secrets: bool = True, colored_console_log: bool = True):
     if enable_pretty_log:
         timestamper = structlog.processors.MaybeTimeStamper(fmt="%Y-%m-%d %H:%M:%S.%f")
     else:
@@ -176,20 +176,34 @@ def logging_processors(enable_pretty_log: bool, mask_secrets: bool = True):
     )
 
     if enable_pretty_log:
-        rich_exc_formatter = structlog.dev.RichTracebackFormatter(
-            # These values are picked somewhat arbitrarily to produce useful-but-compact tracebacks. If
-            # we ever need to change these then they should be configurable.
-            extra_lines=0,
-            max_frames=30,
-            indent_guides=False,
-            suppress=suppress,
-        )
-        my_styles = structlog.dev.ConsoleRenderer.get_default_level_styles()
-        my_styles["debug"] = structlog.dev.CYAN
+        if colored_console_log:
+            rich_exc_formatter = structlog.dev.RichTracebackFormatter(
+                # These values are picked somewhat arbitrarily to produce useful-but-compact tracebacks. If
+                # we ever need to change these then they should be configurable.
+                extra_lines=0,
+                max_frames=30,
+                indent_guides=False,
+                suppress=suppress,
+            )
+            my_styles = structlog.dev.ConsoleRenderer.get_default_level_styles()
+            my_styles["debug"] = structlog.dev.CYAN
 
-        console = structlog.dev.ConsoleRenderer(
-            exception_formatter=rich_exc_formatter, level_styles=my_styles
-        )
+            console = structlog.dev.ConsoleRenderer(
+                exception_formatter=rich_exc_formatter, level_styles=my_styles
+            )
+        else:
+            # Create a console renderer without colors - use the same RichTracebackFormatter
+            # but rely on ConsoleRenderer(colors=False) to disable colors
+            rich_exc_formatter = structlog.dev.RichTracebackFormatter(
+                extra_lines=0,
+                max_frames=30,
+                indent_guides=False,
+                suppress=suppress,
+            )
+            console = structlog.dev.ConsoleRenderer(
+                colors=False,
+                exception_formatter=rich_exc_formatter,
+            )
         processors.append(console)
         return processors, {
             "timestamper": timestamper,
@@ -252,22 +266,30 @@ def configure_logging(
     output: BinaryIO | TextIO | None = None,
     cache_logger_on_first_use: bool = True,
     sending_to_supervisor: bool = False,
+    colored_console_log: bool | None = None,
 ):
     """Set up struct logging and stdlib logging config."""
     if log_level == "DEFAULT":
         log_level = "INFO"
-        if "airflow.configuration" in sys.modules:
-            from airflow.configuration import conf
+        from airflow.configuration import conf
+
+        log_level = conf.get("logging", "logging_level", fallback="INFO")
 
-            log_level = conf.get("logging", "logging_level", fallback="INFO")
+    # If colored_console_log is not explicitly set, read from configuration
+    if colored_console_log is None:
+        from airflow.configuration import conf
+
+        colored_console_log = conf.getboolean("logging", "colored_console_log", fallback=True)
 
     lvl = structlog.stdlib.NAME_TO_LEVEL[log_level.lower()]
 
     if enable_pretty_log:
         formatter = "colored"
     else:
         formatter = "plain"
-    processors, named = logging_processors(enable_pretty_log, mask_secrets=not sending_to_supervisor)
+    processors, named = logging_processors(
+        enable_pretty_log, mask_secrets=not sending_to_supervisor, colored_console_log=colored_console_log
+    )
     timestamper = named["timestamper"]
 
     pre_chain: list[structlog.typing.Processor] = [

task-sdk/tests/task_sdk/log/test_log.py

@@ -147,3 +147,80 @@ def test_logs_are_masked(captured_logs):
         "try_number=1, map_index=-1, hostname=None, context_carrier=None)",
         "timestamp": "2025-03-25T05:13:27.073918Z",
     }
+
+
+def test_logging_processors_with_colors():
+    """Test that logging_processors creates colored console renderer when colored_console_log=True."""
+    from airflow.sdk.log import logging_processors
+
+    _, named = logging_processors(enable_pretty_log=True, colored_console_log=True)
+    assert "console" in named
+    console_renderer = named["console"]
+    assert hasattr(console_renderer, "_styles")
+
+
+def test_logging_processors_without_colors():
+    """Test that logging_processors creates non-colored console renderer when colored_console_log=False."""
+    from airflow.sdk.log import logging_processors
+
+    _, named = logging_processors(enable_pretty_log=True, colored_console_log=False)
+    assert "console" in named
+    console_renderer = named["console"]
+    assert hasattr(console_renderer, "_styles")
+    assert console_renderer._styles.__name__ == "_PlainStyles"
+
+
+def test_logging_processors_json_format():
+    """Test that logging_processors creates JSON renderer when enable_pretty_log=False."""
+    from airflow.sdk.log import logging_processors
+
+    _, named = logging_processors(enable_pretty_log=False, colored_console_log=True)
+    assert "console" not in named
+    assert "json" in named
+
+
+def test_configure_logging_respects_colored_console_log_config():
+    """Test that configure_logging respects the colored_console_log configuration."""
+    from airflow.sdk.log import configure_logging, reset_logging
+
+    mock_conf = mock.MagicMock()
+    mock_conf.get.return_value = "INFO"
+    mock_conf.getboolean.return_value = False  # colored_console_log = False
+
+    with mock.patch("airflow.configuration.conf", mock_conf):
+        reset_logging()
+        configure_logging(enable_pretty_log=True)
+        # Check that getboolean was called with colored_console_log
+        calls = [call for call in mock_conf.getboolean.call_args_list if call[0][1] == "colored_console_log"]
+        assert len(calls) == 1
+        assert calls[0] == mock.call("logging", "colored_console_log", fallback=True)
+
+
+def test_configure_logging_explicit_colored_console_log():
+    """Test that configure_logging respects explicit colored_console_log parameter."""
+    from airflow.sdk.log import configure_logging, reset_logging
+
+    mock_conf = mock.MagicMock()
+    mock_conf.get.return_value = "INFO"
+    mock_conf.getboolean.return_value = True  # colored_console_log = True
+
+    with mock.patch("airflow.configuration.conf", mock_conf):
+        reset_logging()
+        # Explicitly disable colors despite config saying True
+        configure_logging(enable_pretty_log=True, colored_console_log=False)
+        mock_conf.getboolean.assert_not_called()
+
+
+def test_configure_logging_no_airflow_config():
+    """Test that configure_logging defaults work correctly."""
+    from airflow.sdk.log import configure_logging, reset_logging
+
+    # This test can be removed or repurposed since we now always import airflow.configuration
+    mock_conf = mock.MagicMock()
+    mock_conf.get.return_value = "INFO"
+    mock_conf.getboolean.return_value = True  # colored_console_log = True by default
+
+    with mock.patch("airflow.configuration.conf", mock_conf):
+        reset_logging()
+        configure_logging(enable_pretty_log=True)
+        mock_conf.getboolean.assert_called_with("logging", "colored_console_log", fallback=True)