Repository: PrefectHQ/prefect

Issue:
Fix JSON logging formatter
<!-- 
Thanks for opening a pull request to Prefect! We've got a few requests to help us review contributions:

- Make sure that your title neatly summarizes the proposed changes.
- Provide a short overview of the change and the value it adds.
- Share an example to help us understand the change in user experience.
- Confirm that you've done common tasks so we can give a timely review.

Happy engineering!
-->

<!-- Include an overview here -->

### Example

Currently the "json" logger outputs a str() representation of a Python object. This is not strict JSON and as such may not be parsed correctly by downstream log collectors (i.e. Datadog).

This PR uses `json.dumps` to convert the dictionary into valid JSON.

**Before:**
```
{'name': 'prefect.flow_runs', 'msg': "Finished in state Completed('All states completed.')", 'args': [], 'levelname': 'INFO', 'levelno': 20, 'pathname': '/Users/adam/Library/Caches/pypoetry/virtualenvs/sett-prefect-vThVQ5x2-py3.9/lib/python3.9/site-packages/prefect/engine.py', 'filename': 'engine.py', 'module': 'engine', 'exc_info': None, 'exc_text': None, 'stack_info': None, 'lineno': 366, 'funcName': 'begin_flow_run', 'created': 1667049056.934597, 'msecs': 934.5970153808594, 'relativeCreated': 4141.36004447937, 'thread': 4377200000, 'threadName': 'MainThread', 'processName': 'MainProcess', 'process': 99513, 'flow_run_name': 'magenta-porpoise', 'flow_run_id': 'ee770a88-9a37-4ce3-b97b-5a3240cc6fab', 'flow_name': 'healthcheck', 'send_to_orion': False}
```

**After:**
```
{"name": "prefect.flow_runs", "msg": "Finished in state Completed('All states completed.')", "args": [], "levelname": "INFO", "levelno": 20, "pathname": "/Users/adam/Library/Caches/pypoetry/virtualenvs/sett-prefect-vThVQ5x2-py3.9/lib/python3.9/site-packages/prefect/engine.py", "filename": "engine.py", "module": "engine", "exc_info": null, "exc_text": null, "stack_info": null, "lineno": 366, "funcName": "begin_flow_run", "created": 1667048352.241838, "msecs": 241.8379783630371, "relativeCreated": 3872.573137283325, "thread": 4338468224, "threadName": "MainThread", "processName": "MainProcess", "process": 99058, "flow_run_name": "impressive-hyrax", "flow_run_id": "2cd95e62-b09e-4339-928d-ba55b968aeb5", "flow_name": "healthcheck", "send_to_orion": false}
```

Note the double quotes instead of single quotes and the boolean types (`False` vs `false`)



### Checklist
<!-- These boxes may be checked after opening the pull request. -->

- [ ] This pull request references any related issue by including "closes `<link to issue>`"
	- If no issue exists and your change is not a small fix, please [create an issue](https://github.com/PrefectHQ/prefect/issues/new/choose) first.
- [ ] This pull request includes tests or only affects documentation.
- [ ] This pull request includes a label categorizing the change e.g. `fix`, `feature`, `enhancement`
  <!--  If you do not have permission to add a label, a maintainer will add one for you and check this box. -->



Discussion:
### <span aria-hidden="true">✅</span> Deploy Preview for *prefect-orion* ready!
Built [without sensitive environment variables](https://docs.netlify.com/configure-builds/environment-variables/#sensitive-variable-policy)

|  Name | Link |
|---------------------------------|------------------------|
|<span aria-hidden="true">🔨</span> Latest commit | 863b35a782bf5b6ca1dbd9268d369b257c5d0a92 |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/sites/prefect-orion/deploys/636a4e689f60020008cc1ca3 |
|<span aria-hidden="true">😎</span> Deploy Preview | [https://deploy-preview-7377--prefect-orion.netlify.app/api-ref/prefect/logging](https://deploy-preview-7377--prefect-orion.netlify.app/api-ref/prefect/logging) |
|<span aria-hidden="true">📱</span> Preview on mobile | <details><summary> Toggle QR Code... </summary><br /><br />![QR Code](https://app.netlify.com/qr-code/eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJ1cmwiOiJodHRwczovL2RlcGxveS1wcmV2aWV3LTczNzctLXByZWZlY3Qtb3Jpb24ubmV0bGlmeS5hcHAifQ.d5t2ZgpjNYavoyIwuHjkzdF-q70YCl4VEEnQfUvVa7c)<br /><br />_Use your smartphone camera to open QR code link._</details> |
---

_To edit notification comments on pull requests, go to your [Netlify site settings](https://app.netlify.com/sites/prefect-orion/settings/deploys#deploy-notifications)._
@mousetree I'm not sure this is what we want. We are calling `str` on the dictionary after it has gone through a `jsonable_encoder`. I don't remember why we did this instead of dumping directly to JSON, but I think this has different implications.

@chrisguidry are we using this JSON formatter internally w/ DataDog?
@madkinsz We are not using this formatter for Prefect Cloud.
Hey @madkinsz - any ideas to move this forward? We're trying to get nicely formatted JSON so that Datadog will understand our logs.
@mousetree let's use `prefect.serializer.JSONSerializer` and drop use of the `jsonable_encoder`. We can initialize this in the formatter's `__init__`. For `"pretty"` formatting let's print the content after a JSON roundtrip `dumps/loads`.

Note you can also just configure logging to use a formatter defined in your library if you need to.

Resolving patch:
diff --git a/src/prefect/logging/formatters.py b/src/prefect/logging/formatters.py
--- a/src/prefect/logging/formatters.py
+++ b/src/prefect/logging/formatters.py
@@ -1,36 +1,64 @@
-import logging
-import logging.config
 import logging.handlers
-from pprint import pformat
+import traceback
+from types import TracebackType
+from typing import Optional, Tuple, Type, Union
 
-from fastapi.encoders import jsonable_encoder
+import orjson
 
+from prefect.serializers import JSONSerializer
 
-def safe_jsonable(obj):
-    try:
-        return jsonable_encoder(obj)
-    except:
-        return "<unserializable>"
+ExceptionInfoType = Union[
+    Tuple[Type[BaseException], BaseException, Optional[TracebackType]],
+    Tuple[None, None, None],
+]
+
+
+def format_exception_info(exc_info: ExceptionInfoType) -> dict:
+    # if sys.exc_info() returned a (None, None, None) tuple,
+    # then there's nothing to format
+    if exc_info[0] is None:
+        return {}
+
+    (exception_type, exception_obj, exception_traceback) = exc_info
+    return {
+        "type": exception_type.__name__,
+        "message": str(exception_obj),
+        "traceback": "".join(traceback.format_tb(exception_traceback))
+        if exception_traceback
+        else None,
+    }
 
 
 class JsonFormatter(logging.Formatter):
     """
     Formats log records as a JSON string.
 
-    If a log record attribute is not JSON serialiazable, it will be dropped from the
-    output.
-
     The format may be specified as "pretty" to format the JSON with indents and
     newlines.
     """
 
-    def __init__(self, fmt, dmft, style) -> None:
+    def __init__(self, fmt, dmft, style) -> None:  # noqa
+        super().__init__()
+
         if fmt not in ["pretty", "default"]:
             raise ValueError("Format must be either 'pretty' or 'default'.")
-        self.fmt = fmt
 
-    def format(self, record: logging.LogRecord):
-        format_fn = pformat if self.fmt == "pretty" else str
-        return format_fn(
-            {key: safe_jsonable(val) for key, val in record.__dict__.items()}
+        self.serializer = JSONSerializer(
+            jsonlib="orjson",
+            object_encoder="pydantic.json.pydantic_encoder",
+            dumps_kwargs={"option": orjson.OPT_INDENT_2} if fmt == "pretty" else {},
         )
+
+    def format(self, record: logging.LogRecord) -> str:
+        record_dict = record.__dict__.copy()
+
+        # replace any exception tuples returned by `sys.exc_info()`
+        # with a JSON-serializable `dict`.
+        if record.exc_info:
+            record_dict["exc_info"] = format_exception_info(record.exc_info)
+
+        log_json_bytes = self.serializer.dumps(record_dict)
+
+        # JSONSerializer returns bytes; decode to string to conform to
+        # the `logging.Formatter.format` interface
+        return log_json_bytes.decode()


Test patch:
diff --git a/tests/test_logging.py b/tests/test_logging.py
--- a/tests/test_logging.py
+++ b/tests/test_logging.py
@@ -1,4 +1,5 @@
 import asyncio
+import json
 import logging
 import queue
 import sys
@@ -30,6 +31,7 @@
     load_logging_config,
     setup_logging,
 )
+from prefect.logging.formatters import JsonFormatter
 from prefect.logging.handlers import OrionHandler, OrionLogWorker, PrefectConsoleHandler
 from prefect.logging.highlighters import PrefectConsoleHighlighter
 from prefect.logging.loggers import (
@@ -1136,6 +1138,68 @@ def test_does_not_word_wrap_or_crop_messages(self, capsys):
         assert "x" * 1000 in stderr
 
 
+class TestJsonFormatter:
+    def test_json_log_formatter(self):
+        formatter = JsonFormatter("default", None, "%")
+        record = logging.LogRecord(
+            name="Test Log",
+            level=1,
+            pathname="/path/file.py",
+            lineno=1,
+            msg="log message",
+            args=None,
+            exc_info=None,
+        )
+
+        formatted = formatter.format(record)
+
+        # we should be able to load the formatted JSON successfully
+        deserialized = json.loads(formatted)
+
+        # we can't check for an exact JSON string because some attributes vary at
+        # runtime, so check some known attributes instead
+        assert deserialized["name"] == "Test Log"
+        assert deserialized["levelname"] == "Level 1"
+        assert deserialized["filename"] == "file.py"
+        assert deserialized["lineno"] == 1
+
+    def test_json_log_formatter_with_exception(self):
+
+        exc_info = None
+        try:
+            raise Exception("test exception")  # noqa
+        except Exception as exc:  # noqa
+            exc_info = sys.exc_info()
+
+        formatter = JsonFormatter("default", None, "%")
+        record = logging.LogRecord(
+            name="Test Log",
+            level=1,
+            pathname="/path/file.py",
+            lineno=1,
+            msg="log message",
+            args=None,
+            exc_info=exc_info,
+        )
+
+        formatted = formatter.format(record)
+
+        # we should be able to load the formatted JSON successfully
+        deserialized = json.loads(formatted)
+
+        # we can't check for an exact JSON string because some attributes vary at
+        # runtime, so check some known attributes instead
+        assert deserialized["name"] == "Test Log"
+        assert deserialized["levelname"] == "Level 1"
+        assert deserialized["filename"] == "file.py"
+        assert deserialized["lineno"] == 1
+        assert deserialized["exc_info"] is not None
+        assert deserialized["exc_info"]["type"] == "Exception"
+        assert deserialized["exc_info"]["message"] == "test exception"
+        assert deserialized["exc_info"]["traceback"] is not None
+        assert len(deserialized["exc_info"]["traceback"]) > 0
+
+
 def test_log_in_flow(caplog):
     msg = "Hello world!"