[AIRFLOW-5065] Add colors to console log (#5681)

This commit adds custom formatter for Airflow console log to display colours
when connected to a TTY

airflow/config_templates/airflow_local_settings.py

@@ -36,6 +36,12 @@ FAB_LOG_LEVEL = conf.get('core', 'FAB_LOGGING_LEVEL').upper()
 
 LOG_FORMAT = conf.get('core', 'LOG_FORMAT')
 
+COLORED_LOG_FORMAT = conf.get('core', 'COLORED_LOG_FORMAT')
+
+COLORED_LOG = conf.getboolean('core', 'COLORED_CONSOLE_LOG')
+
+COLORED_FORMATTER_CLASS = conf.get('core', 'COLORED_FORMATTER_CLASS')
+
 BASE_LOG_FOLDER = conf.get('core', 'BASE_LOG_FOLDER')
 
 PROCESSOR_LOG_FOLDER = conf.get('scheduler', 'CHILD_PROCESS_LOG_DIRECTORY')
@@ -72,13 +78,17 @@ DEFAULT_LOGGING_CONFIG = {
     'disable_existing_loggers': False,
     'formatters': {
         'airflow': {
-            'format': LOG_FORMAT,
+            'format': LOG_FORMAT
+        },
+        'airflow_coloured': {
+            'format': COLORED_LOG_FORMAT if COLORED_LOG else LOG_FORMAT,
+            'class': COLORED_FORMATTER_CLASS if COLORED_LOG else 'logging.Formatter'
         },
     },
     'handlers': {
         'console': {
             'class': 'airflow.utils.log.logging_mixin.RedirectStdHandler',
-            'formatter': 'airflow',
+            'formatter': 'airflow_coloured',
             'stream': 'sys.stdout'
         },
         'task': {

airflow/config_templates/default_airflow.cfg

@@ -59,6 +59,11 @@ fab_logging_level = WARN
 logging_config_class =
 
 # Log format
+# Colour the logs when the controlling terminal is a TTY.
+colored_console_log = True
+colored_log_format = [%%(blue)s%%(asctime)s%%(reset)s] {{%%(blue)s%%(filename)s:%%(reset)s%%(lineno)d}} %%(log_color)s%%(levelname)s%%(reset)s - %%(log_color)s%%(message)s%%(reset)s
+colored_formatter_class = airflow.utils.log.colored_log.CustomTTYColoredFormatter
+
 log_format = [%%(asctime)s] {{%%(filename)s:%%(lineno)d}} %%(levelname)s - %%(message)s
 simple_log_format = %%(asctime)s %%(levelname)s - %%(message)s
 

airflow/utils/dag_processing.py

@@ -581,6 +581,7 @@ class DagFileProcessorAgent(LoggingMixin):
         # e.g. RotatingFileHandler. And it can cause connection corruption if we
         # do not recreate the SQLA connection pool.
         os.environ['CONFIG_PROCESSOR_MANAGER_LOGGER'] = 'True'
+        os.environ['AIRFLOW__CORE__COLORED_CONSOLE_LOG'] = 'False'
         # Replicating the behavior of how logging module was loaded
         # in logging_config.py
         reload_module(import_module(airflow.settings.LOGGING_CLASS_PATH.rsplit('.', 1)[0]))

airflow/utils/log/colored_log.py

@@ -0,0 +1,84 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Class responsible for colouring logs based on log level.
+"""
+import sys
+from typing import Any, Union
+
+from logging import LogRecord
+from colorlog import TTYColoredFormatter
+from termcolor import colored
+
+ARGS = {"attrs": ["bold"]}
+
+DEFAULT_COLORS = {
+    "DEBUG": "red",
+    "INFO": "",
+    "WARNING": "yellow",
+    "ERROR": "red",
+    "CRITICAL": "red",
+}
+
+
+class CustomTTYColoredFormatter(TTYColoredFormatter):
+    """
+    Custom log formatter which extends `colored.TTYColoredFormatter`
+    by adding attributes to message arguments and coloring error
+    traceback.
+    """
+    def __init__(self, *args, **kwargs):
+        kwargs["stream"] = sys.stdout or kwargs.get("stream")
+        kwargs["log_colors"] = DEFAULT_COLORS
+        super().__init__(*args, **kwargs)
+
+    @staticmethod
+    def _color_arg(arg: Any) -> Union[str, float, int]:
+        if isinstance(arg, (int, float)):
+            # In case of %d or %f formatting
+            return arg
+        return colored(str(arg), **ARGS)  # type: ignore
+
+    def _color_record_args(self, record: LogRecord) -> LogRecord:
+        if isinstance(record.args, (tuple, list)):
+            record.args = tuple(self._color_arg(arg) for arg in record.args)
+        elif isinstance(record.args, dict):
+            # Case of logging.debug("a %(a)d b %(b)s", {'a':1, 'b':2})
+            record.args = {
+                key: self._color_arg(value) for key, value in record.args.items()
+            }
+        elif isinstance(record.args, str):
+            record.args = self._color_arg(record.args)
+        return record
+
+    def _color_record_traceback(self, record: LogRecord) -> LogRecord:
+        if record.exc_info:
+            # Cache the traceback text to avoid converting it multiple times
+            # (it's constant anyway)
+            if not record.exc_text:
+                record.exc_text = self.formatException(record.exc_info)
+
+            if record.exc_text:
+                record.exc_text = colored(record.exc_text, DEFAULT_COLORS["ERROR"])
+        return record
+
+    def format(self, record: LogRecord) -> str:
+        record = self._color_record_args(record)
+        record = self._color_record_traceback(record)
+        return super().format(record)

setup.py

@@ -323,6 +323,7 @@ def do_setup():
         install_requires=[
             'alembic>=1.0, <2.0',
             'cached_property~=1.5',
+            'colorlog==4.0.2',
             'configparser>=3.5.0, <3.6.0',
             'croniter>=0.3.17, <0.4',
             'dill>=0.2.2, <0.3',
@@ -351,6 +352,7 @@ def do_setup():
             'sqlalchemy~=1.3',
             'tabulate>=0.7.5, <0.9',
             'tenacity==4.12.0',
+            'termcolor==1.1.0',
             'text-unidecode==1.2',
             'typing;python_version<"3.5"',
             'thrift>=0.9.2',