feat: allow modifying LogEntry data using extra argument (#129)

daniel-sanche · web-flow · commit 92b287f42441 · 2021-01-08T16:19:44.000-08:00
diff --git a/google/cloud/logging_v2/handlers/app_engine.py b/google/cloud/logging_v2/handlers/app_engine.py
@@ -113,15 +113,25 @@ def emit(self, record):
             record (logging.LogRecord): The record to be logged.
         """
         message = super(AppEngineHandler, self).format(record)
+        inferred_http, inferred_trace = get_request_data()
+        if inferred_trace is not None:
+            inferred_trace = f"projects/{self.project_id}/traces/{inferred_trace}"
+        # allow user overrides
+        trace = getattr(record, "trace", inferred_trace)
+        span_id = getattr(record, "span_id", None)
+        http_request = getattr(record, "http_request", inferred_http)
+        resource = getattr(record, "resource", self.resource)
+        user_labels = getattr(record, "labels", {})
+        # merge labels
         gae_labels = self.get_gae_labels()
-        http_request, trace_id = get_request_data()
-        if trace_id is not None:
-            trace_id = f"projects/{self.project_id}/traces/{trace_id}"
+        gae_labels.update(user_labels)
+        # send off request
         self.transport.send(
             record,
             message,
-            resource=self.resource,
+            resource=resource,
             labels=gae_labels,
-            trace=trace_id,
+            trace=trace,
+            span_id=span_id,
             http_request=http_request,
         )
diff --git a/google/cloud/logging_v2/handlers/handlers.py b/google/cloud/logging_v2/handlers/handlers.py
@@ -87,6 +87,7 @@ def __init__(
         self.name = name
         self.client = client
         self.transport = transport(client, name)
+        self.project_id = client.project
         self.resource = resource
         self.labels = labels
 
@@ -101,7 +102,26 @@ def emit(self, record):
             record (logging.LogRecord): The record to be logged.
         """
         message = super(CloudLoggingHandler, self).format(record)
-        self.transport.send(record, message, resource=self.resource, labels=self.labels)
+        trace_id = getattr(record, "trace", None)
+        span_id = getattr(record, "span_id", None)
+        http_request = getattr(record, "http_request", None)
+        resource = getattr(record, "resource", self.resource)
+        user_labels = getattr(record, "labels", {})
+        # merge labels
+        total_labels = self.labels if self.labels is not None else {}
+        total_labels.update(user_labels)
+        if len(total_labels) == 0:
+            total_labels = None
+        # send off request
+        self.transport.send(
+            record,
+            message,
+            resource=resource,
+            labels=(total_labels if total_labels else None),
+            trace=trace_id,
+            span_id=span_id,
+            http_request=http_request,
+        )
 
 
 def setup_logging(
diff --git a/tests/system/test_system.py b/tests/system/test_system.py
@@ -27,7 +27,8 @@
 from google.api_core.exceptions import ServiceUnavailable
 import google.cloud.logging
 from google.cloud._helpers import UTC
-from google.cloud.logging_v2.handlers.handlers import CloudLoggingHandler
+from google.cloud.logging_v2.handlers import AppEngineHandler
+from google.cloud.logging_v2.handlers import CloudLoggingHandler
 from google.cloud.logging_v2.handlers.transports import SyncTransport
 from google.cloud.logging_v2 import client
 from google.cloud.logging_v2.resource import Resource
@@ -308,6 +309,39 @@ def test_log_handler_sync(self):
         self.assertEqual(len(entries), 1)
         self.assertEqual(entries[0].payload, expected_payload)
 
+    def test_handlers_w_extras(self):
+        LOG_MESSAGE = "Testing with injected extras."
+
+        for cls in [CloudLoggingHandler, AppEngineHandler]:
+            LOGGER_NAME = f"{cls.__name__}-handler_extras"
+            handler_name = self._logger_name(LOGGER_NAME)
+
+            handler = cls(Config.CLIENT, name=handler_name, transport=SyncTransport)
+
+            # only create the logger to delete, hidden otherwise
+            logger = Config.CLIENT.logger(handler.name)
+            self.to_delete.append(logger)
+
+            cloud_logger = logging.getLogger(LOGGER_NAME)
+            cloud_logger.addHandler(handler)
+            expected_request = {"requestUrl": "localhost"}
+            extra = {
+                "trace": "123",
+                "span_id": "456",
+                "http_request": expected_request,
+                "resource": Resource(type="cloudiot_device", labels={}),
+                "labels": {"test-label": "manual"},
+            }
+            cloud_logger.warn(LOG_MESSAGE, extra=extra)
+
+            entries = _list_entries(logger)
+            self.assertEqual(len(entries), 1)
+            self.assertEqual(entries[0].trace, extra["trace"])
+            self.assertEqual(entries[0].span_id, extra["span_id"])
+            self.assertEqual(entries[0].http_request, expected_request)
+            self.assertEqual(entries[0].labels, extra["labels"])
+            self.assertEqual(entries[0].resource.type, extra["resource"].type)
+
     def test_log_root_handler(self):
         LOG_MESSAGE = "It was the best of times."
 
diff --git a/tests/unit/handlers/test_app_engine.py b/tests/unit/handlers/test_app_engine.py
@@ -118,10 +118,60 @@ def test_emit(self):
                     gae_resource,
                     gae_labels,
                     expected_trace_id,
+                    None,
                     expected_http_request,
                 ),
             )
 
+    def test_emit_manual_field_override(self):
+        from google.cloud.logging_v2.resource import Resource
+
+        inferred_http_request = {"request_url": "test"}
+        inferred_trace_id = "trace-test"
+        get_request_patch = mock.patch(
+            "google.cloud.logging_v2.handlers.app_engine.get_request_data",
+            return_value=(inferred_http_request, inferred_trace_id),
+        )
+        with get_request_patch:
+            # library integrations mocked to return test data
+            client = mock.Mock(project=self.PROJECT, spec=["project"])
+            handler = self._make_one(client, transport=_Transport)
+            gae_labels = handler.get_gae_labels()
+            logname = "app"
+            message = "hello world"
+            record = logging.LogRecord(
+                logname, logging, None, None, message, None, None
+            )
+            handler.project_id = self.PROJECT
+            # set attributes manually
+            expected_trace = "123"
+            setattr(record, "trace", expected_trace)
+            expected_span = "456"
+            setattr(record, "span_id", expected_span)
+            expected_http = {"reuqest_url": "manual"}
+            setattr(record, "http_request", expected_http)
+            expected_resource = Resource(type="test", labels={})
+            setattr(record, "resource", expected_resource)
+            additional_labels = {"test-label": "manual"}
+            expected_labels = dict(gae_labels)
+            expected_labels.update(additional_labels)
+            setattr(record, "labels", additional_labels)
+            handler.emit(record)
+            self.assertIs(handler.transport.client, client)
+            self.assertEqual(handler.transport.name, logname)
+            self.assertEqual(
+                handler.transport.send_called_with,
+                (
+                    record,
+                    message,
+                    expected_resource,
+                    expected_labels,
+                    expected_trace,
+                    expected_span,
+                    expected_http,
+                ),
+            )
+
     def _get_gae_labels_helper(self, trace_id):
         get_request_patch = mock.patch(
             "google.cloud.logging_v2.handlers.app_engine.get_request_data",
@@ -156,5 +206,13 @@ def __init__(self, client, name):
         self.client = client
         self.name = name
 
-    def send(self, record, message, resource, labels, trace, http_request):
-        self.send_called_with = (record, message, resource, labels, trace, http_request)
+    def send(self, record, message, resource, labels, trace, span_id, http_request):
+        self.send_called_with = (
+            record,
+            message,
+            resource,
+            labels,
+            trace,
+            span_id,
+            http_request,
+        )
diff --git a/tests/unit/handlers/test_handlers.py b/tests/unit/handlers/test_handlers.py
@@ -85,7 +85,44 @@ def test_emit(self):
 
         self.assertEqual(
             handler.transport.send_called_with,
-            (record, message, _GLOBAL_RESOURCE, None),
+            (record, message, _GLOBAL_RESOURCE, None, None, None, None),
+        )
+
+    def test_emit_manual_field_override(self):
+        from google.cloud.logging_v2.logger import _GLOBAL_RESOURCE
+        from google.cloud.logging_v2.resource import Resource
+
+        client = _Client(self.PROJECT)
+        handler = self._make_one(
+            client, transport=_Transport, resource=_GLOBAL_RESOURCE
+        )
+        logname = "loggername"
+        message = "hello world"
+        record = logging.LogRecord(logname, logging, None, None, message, None, None)
+        # set attributes manually
+        expected_trace = "123"
+        setattr(record, "trace", expected_trace)
+        expected_span = "456"
+        setattr(record, "span_id", expected_span)
+        expected_http = {"reuqest_url": "manual"}
+        setattr(record, "http_request", expected_http)
+        expected_resource = Resource(type="test", labels={})
+        setattr(record, "resource", expected_resource)
+        expected_labels = {"test-label": "manual"}
+        setattr(record, "labels", expected_labels)
+        handler.emit(record)
+
+        self.assertEqual(
+            handler.transport.send_called_with,
+            (
+                record,
+                message,
+                expected_resource,
+                expected_labels,
+                expected_trace,
+                expected_span,
+                expected_http,
+            ),
         )
 
 
@@ -148,5 +185,22 @@ def __init__(self, client, name):
         self.client = client
         self.name = name
 
-    def send(self, record, message, resource, labels=None):
-        self.send_called_with = (record, message, resource, labels)
+    def send(
+        self,
+        record,
+        message,
+        resource,
+        labels=None,
+        trace=None,
+        span_id=None,
+        http_request=None,
+    ):
+        self.send_called_with = (
+            record,
+            message,
+            resource,
+            labels,
+            trace,
+            span_id,
+            http_request,
+        )
