Add AppendLabelLoggingAdapter to have an ability to adapt logger with predefined labels

Author: Alexey Nikitin

google/cloud/logging_v2/handlers/structured_log.py

@@ -19,10 +19,11 @@
 import logging
 import logging.handlers
 
-from google.cloud.logging_v2.handlers.handlers import CloudLoggingFilter
-from google.cloud.logging_v2.handlers.handlers import _format_and_parse_message
 import google.cloud.logging_v2
 from google.cloud.logging_v2._instrumentation import _create_diagnostic_entry
+from google.cloud.logging_v2.handlers.handlers import (
+    CloudLoggingFilter, _format_and_parse_message
+)
 
 GCP_FORMAT = (
     "{%(_payload_str)s"
@@ -136,3 +137,62 @@ def emit_instrumentation_info(self):
         struct_logger.setLevel(logging.INFO)
         struct_logger.info(diagnostic_object.payload)
         struct_logger.handlers.clear()
+
+
+class AppendLabelLoggingAdapter(logging.LoggerAdapter):
+    """
+    Logging adapter that allows to add required
+    constant key/value to the labels part of log record.
+    Example:
+
+    .. code-block:: python
+
+        import logging
+        from google.cloud.logging_v2.handlers.structured_log import AppendLabelLoggingAdapter
+        from google.cloud.logging_v2.handlers.structured_log import StructuredLogHandler
+        logging.root.setLevel(logging.INFO)
+        logging.root.handlers = [StructuredLogHandler()]
+        first_adapter = AppendLabelLoggingAdapter(logging.root, {'a': 5, 'b': 6})
+        first_adapter.info('first info')
+        {
+            "message": "first info",
+            "severity": "INFO",
+            "logging.googleapis.com/labels": {"python_logger": "root", "a": 5, "b": 6}
+            [...]
+        }
+        # Could be stacked
+        second_adapter=AppendLabelLoggingAdapter(first_adapter, {'hello': 'world'})
+        second_adapter.info('second info')
+        {
+            "message": "second info",
+            "severity": "INFO",
+            "logging.googleapis.com/labels": {"python_logger": "root", "hello": "world", "a": 5, "b": 6}
+            [...]
+        }
+    """
+
+    def __init__(self, logger, append_labels):
+        """
+        Args:
+            logger (~logging.Logger):
+                The Logger for this adapter to use.
+            append_labels (~typing.Dict[str, str]): the required data to be added to logger "labels" field.
+        """
+        self.append_labels = append_labels
+        super().__init__(logger)
+
+    def process(self, msg, kwargs):
+        """
+        Args:
+            msg (str):
+                Log message
+            kwargs (dict):
+                logging kwargs
+        """
+        extra = kwargs.get("extra", {})
+        labels = extra.get("labels", {})
+        for label_key, label_value in self.append_labels.items():
+            labels.setdefault(label_key, label_value)
+        extra["labels"] = labels
+        kwargs["extra"] = extra
+        return msg, kwargs

tests/unit/handlers/test_structured_log.py

@@ -667,3 +667,93 @@ def test_valid_instrumentation_info(self):
                 inst_source_dict,
                 "instrumentation payload not logged properly",
             )
+
+    def test_append_labels_adapter(self):
+        import logging
+
+        import mock
+
+        from google.cloud.logging_v2.handlers.structured_log import \
+            AppendLabelLoggingAdapter
+
+        logger = logging.getLogger("google.cloud.logging_v2.handlers.structured_log")
+        handler = self._make_one()
+        with mock.patch.object(handler, "emit_instrumentation_info"):
+            with mock.patch.object(logger, "_log") as mock_log:
+                logger.addHandler(handler)
+                logger.setLevel(logging.INFO)
+                adapted_logger = AppendLabelLoggingAdapter(
+                    logger, append_labels={"service_id": 1, "another_value": "foo"}
+                )
+                adapted_logger.info("test message")
+                mock_log.assert_called_once()
+                self.assertEqual(
+                    mock_log.call_args_list[0].kwargs,
+                    {"extra": {"labels": {"service_id": 1, "another_value": "foo"}}},
+                )
+
+    def test_append_labels_adapter_override_defaults(self):
+        import logging
+
+        import mock
+
+        from google.cloud.logging_v2.handlers.structured_log import \
+            AppendLabelLoggingAdapter
+
+        logger = logging.getLogger("google.cloud.logging_v2.handlers.structured_log")
+        handler = self._make_one()
+        with mock.patch.object(handler, "emit_instrumentation_info"):
+            with mock.patch.object(logger, "_log") as mock_log:
+                logger.addHandler(handler)
+                logger.setLevel(logging.INFO)
+                adapted_logger = AppendLabelLoggingAdapter(
+                    logger, append_labels={"service_id": 1, "another_value": "foo"}
+                )
+                adapted_logger.info(
+                    "test message", extra={"labels": {"another_value": "baz"}}
+                )
+                mock_log.assert_called_once()
+                # the default value was overridden
+                self.assertEqual(
+                    mock_log.call_args_list[0].kwargs,
+                    {"extra": {"labels": {"service_id": 1, "another_value": "baz"}}},
+                )
+
+    def test_append_labels_adapter_stacked(self):
+        import logging
+
+        import mock
+
+        from google.cloud.logging_v2.handlers.structured_log import \
+            AppendLabelLoggingAdapter
+
+        logger = logging.getLogger("google.cloud.logging_v2.handlers.structured_log")
+        handler = self._make_one()
+        with mock.patch.object(handler, "emit_instrumentation_info"):
+            with mock.patch.object(logger, "_log") as mock_log:
+                logger.addHandler(handler)
+                logger.setLevel(logging.INFO)
+                adapted_logger = AppendLabelLoggingAdapter(
+                    logger, append_labels={"service_id": 1, "another_value": "foo"}
+                )
+                twice_adapted_logger = AppendLabelLoggingAdapter(
+                    adapted_logger,
+                    # one fields is new, another was adapted already
+                    append_labels={"new_field": "new_value", "another_value": "baz"},
+                )
+                twice_adapted_logger.info(
+                    "test message", extra={"labels": {"another_value": "baz"}}
+                )
+                mock_log.assert_called_once()
+                self.assertEqual(
+                    mock_log.call_args_list[0].kwargs,
+                    {
+                        "extra": {
+                            "labels": {
+                                "another_value": "baz",  # value is changed by the second adapter
+                                "new_field": "new_value",  # introduced by the second adapter
+                                "service_id": 1,  # left as is from the first adapter configuration
+                            }
+                        }
+                    },
+                )