From 8707879408190cb81793c1f34de89218ef73904c Mon Sep 17 00:00:00 2001
From: John Mulligan <jmulligan@redhat.com>
Date: Tue, 22 Aug 2023 12:42:14 -0400
Subject: [PATCH] cephadm: move colored output support into logging.py

Rewrite cephadm's colored output support such that it abstracts away
the colorization into extra logging metadata. The new code will not
unconditionally put control characters into the log files. It will
only print the control chars if the stderr is a tty.
In theory this is probably more future proof as well, but it's only
got two callers so it is hard to say how useful it'll be.

Signed-off-by: John Mulligan <jmulligan@redhat.com>
---
 src/cephadm/cephadm.py            | 16 ++++--------
 src/cephadm/cephadmlib/logging.py | 42 ++++++++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/src/cephadm/cephadm.py b/src/cephadm/cephadm.py
index da3d03128a29b..fb40a9d67f537 100755
--- a/src/cephadm/cephadm.py
+++ b/src/cephadm/cephadm.py
@@ -137,7 +137,7 @@ from cephadmlib.net_utils import (
 from cephadmlib.locking import FileLock
 from cephadmlib.daemon_identity import DaemonIdentity, DaemonSubIdentity
 from cephadmlib.packagers import create_packager, Packager
-from cephadmlib.logging import cephadm_init_logging
+from cephadmlib.logging import cephadm_init_logging, Highlight
 
 FuncT = TypeVar('FuncT', bound=Callable)
 
@@ -202,12 +202,6 @@ class DeploymentType(Enum):
     # files and potentially restarts daemon.
     RECONFIG = 'Reconfig'
 
-
-class termcolor:
-    yellow = '\033[93m'
-    red = '\033[31m'
-    end = '\033[0m'
-
 ##################################
 
 
@@ -1615,7 +1609,7 @@ For information regarding the latest stable release:
     https://docs.ceph.com/docs/{}/cephadm/install
 """.format(LATEST_STABLE_RELEASE)
         for line in warn.splitlines():
-            logger.warning('{}{}{}'.format(termcolor.yellow, line, termcolor.end))
+            logger.warning(line, extra=Highlight.WARNING.extra())
     return DEFAULT_IMAGE
 
 
@@ -2519,9 +2513,9 @@ def _get_container_mounts_for_type(
                 mounts[ceph_folder + '/monitoring/ceph-mixin/dashboards_out'] = '/etc/grafana/dashboards/ceph-dashboard'
                 mounts[ceph_folder + '/monitoring/ceph-mixin/prometheus_alerts.yml'] = '/etc/prometheus/ceph/ceph_default_alerts.yml'
             else:
-                logger.error('{}{}{}'.format(termcolor.red,
-                                             'Ceph shared source folder does not exist.',
-                                             termcolor.end))
+                logger.error(
+                    'Ceph shared source folder does not exist.',
+                    extra=Highlight.FAILURE.extra())
     except AttributeError:
         pass
     return mounts
diff --git a/src/cephadm/cephadmlib/logging.py b/src/cephadm/cephadmlib/logging.py
index 7369bd26c1ed7..1f919668316b3 100644
--- a/src/cephadm/cephadmlib/logging.py
+++ b/src/cephadm/cephadmlib/logging.py
@@ -1,11 +1,12 @@
 # logging.py - cephadm specific logging behavior
 
+import enum
 import logging
 import logging.config
 import os
 import sys
 
-from typing import List
+from typing import List, Any, Dict, cast
 
 from .context import CephadmContext
 from .constants import QUIET_LOG_LEVEL, LOG_DIR
@@ -17,10 +18,48 @@ class _ExcludeErrorsFilter(logging.Filter):
         return record.levelno < logging.WARNING
 
 
+class _termcolors(str, enum.Enum):
+    yellow = '\033[93m'
+    red = '\033[31m'
+    end = '\033[0m'
+
+
+class Highlight(enum.Enum):
+    FAILURE = 1
+    WARNING = 2
+
+    def extra(self) -> Dict[str, 'Highlight']:
+        """Return logging extra for the current kind of highlight."""
+        return {'highlight': self}
+
+    def _highlight(self, s: str) -> str:
+        color = {
+            self.FAILURE: _termcolors.red.value,
+            self.WARNING: _termcolors.yellow.value,
+        }[
+            cast(int, self)
+        ]  # cast works around mypy confusion wrt enums
+        return f'{color}{s}{_termcolors.end.value}'
+
+
+class _Colorizer(logging.Formatter):
+    def format(self, record: Any) -> str:
+        res = super().format(record)
+        highlight = getattr(record, 'highlight', None)
+        # checking sys.stderr here is a bit dirty but we know exactly
+        # how _Colorizer will be used, so it works for now
+        if highlight is not None and sys.stderr.isatty():
+            res = highlight._highlight(res)
+        return res
+
+
 _common_formatters = {
     'cephadm': {
         'format': '%(asctime)s %(thread)x %(levelname)s %(message)s',
     },
+    'colorized': {
+        '()': _Colorizer,
+    },
 }
 
 
@@ -78,6 +117,7 @@ _interactive_logging_config = {
             'level': 'WARNING',
             'class': 'logging.StreamHandler',
             'stream': sys.stderr,
+            'formatter': 'colorized',
         },
         'log_file': _log_file_handler,
     },
-- 
2.39.5