feat: add support to custom JSON encoders (#657)

googleapis · Oct 26, 2022 · 77e621c · 77e621c
1 parent 061a381
commit 77e621c
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 3 deletions.
diff --git a/google/cloud/logging_v2/handlers/structured_log.py b/google/cloud/logging_v2/handlers/structured_log.py
@@ -62,12 +62,15 @@ class StructuredLogHandler(logging.StreamHandler):
     and write them to standard output
     """
 
-    def __init__(self, *, labels=None, stream=None, project_id=None):
+    def __init__(
+        self, *, labels=None, stream=None, project_id=None, json_encoder_cls=None
+    ):
         """
         Args:
             labels (Optional[dict]): Additional labels to attach to logs.
             stream (Optional[IO]): Stream to be used by the handler.
             project (Optional[str]): Project Id associated with the logs.
+            json_encoder_cls (Optional[Type[JSONEncoder]]): Custom JSON encoder. Defaults to json.JSONEncoder
         """
         super(StructuredLogHandler, self).__init__(stream=stream)
         self.project_id = project_id
@@ -79,6 +82,8 @@ def __init__(self, *, labels=None, stream=None, project_id=None):
         # make logs appear in GCP structured logging format
         self._gcp_formatter = logging.Formatter(GCP_FORMAT)
 
+        self._json_encoder_cls = json_encoder_cls or json.JSONEncoder
+
     def format(self, record):
         """Format the message into structured log JSON.
         Args:
@@ -95,14 +100,18 @@ def format(self, record):
                 if key in GCP_STRUCTURED_LOGGING_FIELDS:
                     del message[key]
             # if input is a dictionary, encode it as a json string
-            encoded_msg = json.dumps(message, ensure_ascii=False)
+            encoded_msg = json.dumps(
+                message, ensure_ascii=False, cls=self._json_encoder_cls
+            )
             # all json.dumps strings should start and end with parentheses
             # strip them out to embed these fields in the larger JSON payload
             if len(encoded_msg) > 2:
                 payload = encoded_msg[1:-1] + ","
         elif message:
             # properly break any formatting in string to make it json safe
-            encoded_message = json.dumps(message, ensure_ascii=False)
+            encoded_message = json.dumps(
+                message, ensure_ascii=False, cls=self._json_encoder_cls
+            )
             payload = '"message": {},'.format(encoded_message)
 
         record._payload_str = payload or ""

diff --git a/tests/unit/handlers/test_structured_log.py b/tests/unit/handlers/test_structured_log.py
@@ -46,6 +46,15 @@ def test_ctor_w_project(self):
         handler = self._make_one(project_id="foo")
         self.assertEqual(handler.project_id, "foo")
 
+    def test_ctor_w_encoder(self):
+        import json
+
+        class CustomJSONEncoder(json.JSONEncoder):
+            pass
+
+        handler = self._make_one(json_encoder_cls=CustomJSONEncoder)
+        self.assertEqual(handler._json_encoder_cls, CustomJSONEncoder)
+
     def test_format(self):
         import logging
         import json
@@ -207,6 +216,51 @@ def test_format_with_custom_formatter(self):
         self.assertIn(expected_result, result)
         self.assertIn("message", result)
 
+    def test_format_with_custom_json_encoder(self):
+        import json
+        import logging
+
+        from pathlib import Path
+        from typing import Any
+
+        class CustomJSONEncoder(json.JSONEncoder):
+            def default(self, obj: Any) -> Any:
+                if isinstance(obj, Path):
+                    return str(obj)
+                return json.JSONEncoder.default(self, obj)
+
+        handler = self._make_one(json_encoder_cls=CustomJSONEncoder)
+
+        message = "hello world"
+        json_fields = {"path": Path("/path")}
+        record = logging.LogRecord(
+            None,
+            logging.INFO,
+            None,
+            None,
+            message,
+            None,
+            None,
+        )
+        setattr(record, "json_fields", json_fields)
+        expected_payload = {
+            "message": message,
+            "severity": "INFO",
+            "logging.googleapis.com/trace": "",
+            "logging.googleapis.com/spanId": "",
+            "logging.googleapis.com/trace_sampled": False,
+            "logging.googleapis.com/sourceLocation": {},
+            "httpRequest": {},
+            "logging.googleapis.com/labels": {},
+            "path": "/path",
+        }
+        handler.filter(record)
+
+        result = json.loads(handler.format(record))
+
+        self.assertEqual(set(expected_payload.keys()), set(result.keys()))
+        self.assertEqual(result["path"], "/path")
+
     def test_format_with_reserved_json_field(self):
         # drop json_field data with reserved names
         # related issue: https://github.com/googleapis/python-logging/issues/543