feat: Ability to log WSGI environ in JSON logs, log_extra utility (#55)

Author: khvn26

src/common/gunicorn/constants.py

@@ -1,4 +1,11 @@
-WSGI_DJANGO_ROUTE_ENVIRON_KEY = "wsgi.django_route"
+import re
+
+WSGI_EXTRA_PREFIX = "flagsmith."
+WSGI_EXTRA_SUFFIX_TO_CATEGORY = {
+    "i": "request_headers",
+    "o": "response_headers",
+    "e": "environ_variables",
+}
 HTTP_SERVER_RESPONSE_SIZE_DEFAULT_BUCKETS = (
     # 1 kB, 10 kB, 100 kB, 500 kB, 1 MB, 5 MB, 10 MB
     1 * 1024,
@@ -10,3 +17,7 @@
     10 * 1024 * 1024,
     float("inf"),
 )
+
+wsgi_extra_key_regex = re.compile(
+    r"^{(?P<key>[^}]+)}(?P<suffix>[%s])$" % "".join(WSGI_EXTRA_SUFFIX_TO_CATEGORY)
+)

src/common/gunicorn/logging.py

@@ -13,10 +13,42 @@
 
 from common.core.logging import JsonFormatter
 from common.gunicorn import metrics
-from common.gunicorn.constants import WSGI_DJANGO_ROUTE_ENVIRON_KEY
+from common.gunicorn.constants import (
+    WSGI_EXTRA_PREFIX,
+    WSGI_EXTRA_SUFFIX_TO_CATEGORY,
+    wsgi_extra_key_regex,
+)
+from common.gunicorn.utils import get_extra
 
 
 class GunicornAccessLogJsonFormatter(JsonFormatter):
+    def _get_extra(self, record_args: dict[str, Any]) -> dict[str, Any]:
+        ret: dict[str, dict[str, Any]] = {}
+
+        extra_items_to_log: list[str] | None
+        if extra_items_to_log := getattr(settings, "ACCESS_LOG_EXTRA_ITEMS", None):
+            # We expect the extra items to be in the form of
+            # Gunicorn's access log format string for
+            # request headers, response headers and environ variables
+            # without the % prefix, e.g. "{origin}i" or "{flagsmith.environment_id}e"
+            # https://docs.gunicorn.org/en/stable/settings.html#access-log-format
+            for extra_key in extra_items_to_log:
+                extra_key_lower = extra_key.lower()
+                if (
+                    (extra_value := record_args.get(extra_key_lower))
+                    and (re_match := wsgi_extra_key_regex.match(extra_key_lower))
+                    and (
+                        extra_category := WSGI_EXTRA_SUFFIX_TO_CATEGORY.get(
+                            re_match.group("suffix")
+                        )
+                    )
+                ):
+                    ret.setdefault(extra_category, {})[re_match.group("key")] = (
+                        extra_value
+                    )
+
+        return ret
+
     def get_json_record(self, record: logging.LogRecord) -> dict[str, Any]:
         args = record.args
 
@@ -32,11 +64,13 @@ def get_json_record(self, record: logging.LogRecord) -> dict[str, Any]:
             "time": datetime.strptime(args["t"], "[%d/%b/%Y:%H:%M:%S %z]").isoformat(),
             "path": url,
             "remote_ip": args["h"],
-            "route": args["R"],
+            "route": args.get(f"{{{WSGI_EXTRA_PREFIX}route}}e") or "",
             "method": args["m"],
             "status": str(args["s"]),
             "user_agent": args["a"],
             "duration_in_ms": args["M"],
+            "response_size_in_bytes": args["B"] or 0,
+            **self._get_extra(args),
         }
 
 
@@ -56,7 +90,7 @@ def access(
             # To avoid cardinality explosion, we use a resolved Django route
             # instead of raw path.
             # The Django route is set by `RouteLoggerMiddleware`.
-            "route": environ.get(WSGI_DJANGO_ROUTE_ENVIRON_KEY) or "",
+            "route": get_extra(environ=environ, key="route") or "",
             "method": environ.get("REQUEST_METHOD") or "",
             "response_status": resp.status_code,
         }
@@ -65,22 +99,11 @@ def access(
         )
         metrics.flagsmith_http_server_requests_total.labels(**labels).inc()
         metrics.flagsmith_http_server_response_size_bytes.labels(**labels).observe(
-            resp.response_length or 0,
+            getattr(resp, "sent", 0),
         )
 
 
 class GunicornJsonCapableLogger(PrometheusGunicornLogger):
-    def atoms(
-        self,
-        resp: Response,
-        req: Request,
-        environ: dict[str, Any],
-        request_time: timedelta,
-    ) -> dict[str, str]:
-        atoms: dict[str, str] = super().atoms(resp, req, environ, request_time)
-        atoms["R"] = environ.get(WSGI_DJANGO_ROUTE_ENVIRON_KEY) or "-"
-        return atoms
-
     def setup(self, cfg: Config) -> None:
         super().setup(cfg)
         if getattr(settings, "LOG_FORMAT", None) == "json":

src/common/gunicorn/middleware.py

@@ -2,8 +2,7 @@
 
 from django.http import HttpRequest, HttpResponse
 
-from common.gunicorn.constants import WSGI_DJANGO_ROUTE_ENVIRON_KEY
-from common.gunicorn.utils import get_route_template
+from common.gunicorn.utils import get_route_template, log_extra
 
 
 class RouteLoggerMiddleware:
@@ -22,10 +21,10 @@ def __call__(self, request: HttpRequest) -> HttpResponse:
         response = self.get_response(request)
 
         if resolver_match := request.resolver_match:
-            # https://peps.python.org/pep-3333/#specification-details
-            # "...the application is allowed to modify the dictionary in any way it desires"
-            request.META[WSGI_DJANGO_ROUTE_ENVIRON_KEY] = get_route_template(
-                resolver_match.route
+            log_extra(
+                request=request,
+                key="route",
+                value=get_route_template(resolver_match.route),
             )
 
         return response

src/common/gunicorn/utils.py

@@ -5,13 +5,16 @@
 
 from django.core.handlers.wsgi import WSGIHandler
 from django.core.wsgi import get_wsgi_application
+from django.http import HttpRequest
 from drf_yasg.generators import EndpointEnumerator  # type: ignore[import-untyped]
 from environs import Env
 from gunicorn.app.wsgiapp import (  # type: ignore[import-untyped]
     WSGIApplication as GunicornWSGIApplication,
 )
 from gunicorn.config import Config  # type: ignore[import-untyped]
 
+from common.gunicorn.constants import WSGI_EXTRA_PREFIX
+
 env = Env()
 
 DEFAULT_ACCESS_LOG_FORMAT = '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s "%(f)s" "%(a)s" %({origin}i)s %({access-control-allow-origin}o)s'
@@ -76,3 +79,26 @@ def get_route_template(route: str) -> str:
     """
     route_template: str = EndpointEnumerator().get_path_from_regex(route)
     return route_template
+
+
+def log_extra(
+    request: HttpRequest,
+    key: str,
+    value: Any,
+) -> None:
+    """
+    Store a value in the WSGI request `environ` using a prefixed key.
+
+    https://peps.python.org/pep-3333/#specification-details
+    "...the application is allowed to modify the dictionary in any way it desires"
+    """
+    meta_key = f"{WSGI_EXTRA_PREFIX}{key}"
+    request.META[meta_key] = value
+
+
+def get_extra(environ: dict[str, Any], key: str) -> Any:
+    """
+    Retrieve a value from the WSGI request `environ` using a prefixed key.
+    """
+    meta_key = f"{WSGI_EXTRA_PREFIX}{key}"
+    return environ.get(meta_key)

tests/unit/common/gunicorn/test_logging.py

@@ -17,8 +17,16 @@
 
 
 @pytest.mark.freeze_time("2023-12-08T06:05:47.320000+00:00")
-def test_gunicorn_access_log_json_formatter__outputs_expected() -> None:
+def test_gunicorn_access_log_json_formatter__outputs_expected(
+    settings: SettingsWrapper,
+) -> None:
     # Given
+    settings.ACCESS_LOG_EXTRA_ITEMS = [
+        "{flagsmith.route}e",
+        "{X-LOG-ME-STATUS}o",
+        "{x-log-me}i",
+    ]
+
     gunicorn_access_log_json_formatter = GunicornAccessLogJsonFormatter()
     log_record = logging.LogRecord(
         name="gunicorn.access",
@@ -27,9 +35,13 @@ def test_gunicorn_access_log_json_formatter__outputs_expected() -> None:
         lineno=1,
         msg=AccessLogFormat.default,
         args={
+            "{flagsmith.route}e": "/test/{test_id}",
+            "{wsgi.version}e": (1, 0),
+            "{x-log-me-status}o": "acked",
+            "{x-log-me}i": "42",
             "a": "requests",
-            "b": "-",
-            "B": None,
+            "b": "42",
+            "B": 42,
             "D": 1000000,
             "f": "-",
             "h": "192.168.0.1",
@@ -40,15 +52,14 @@ def test_gunicorn_access_log_json_formatter__outputs_expected() -> None:
             "M": 1000,
             "p": "<42>",
             "q": "foo=bar",
-            "R": "^test/",
             "r": "GET",
             "s": 200,
             "T": 1,
             "t": datetime.fromisoformat("2023-12-08T06:05:47.320000+00:00").strftime(
                 "[%d/%b/%Y:%H:%M:%S %z]"
             ),
             "u": "-",
-            "U": "/test",
+            "U": "/test/42",
         },
         exc_info=None,
     )
@@ -60,14 +71,24 @@ def test_gunicorn_access_log_json_formatter__outputs_expected() -> None:
     # Then
     assert json_log == {
         "duration_in_ms": 1000,
+        "environ_variables": {
+            "flagsmith.route": "/test/{test_id}",
+        },
         "levelname": "INFO",
         "logger_name": "gunicorn.access",
-        "message": '192.168.0.1 - - [08/Dec/2023:06:05:47 +0000] "GET" 200 - "-" "requests"',
+        "message": '192.168.0.1 - - [08/Dec/2023:06:05:47 +0000] "GET" 200 42 "-" "requests"',
         "method": "GET",
-        "path": "/test?foo=bar",
+        "path": "/test/42?foo=bar",
         "pid": expected_pid,
         "remote_ip": "192.168.0.1",
-        "route": "^test/",
+        "request_headers": {
+            "x-log-me": "42",
+        },
+        "response_headers": {
+            "x-log-me-status": "acked",
+        },
+        "response_size_in_bytes": 42,
+        "route": "/test/{test_id}",
         "status": "200",
         "thread_name": "MainThread",
         "time": "2023-12-08T06:05:47+00:00",
@@ -87,13 +108,13 @@ def test_gunicorn_prometheus_gunicorn_logger__expected_metrics(
     response_mock = mocker.Mock()
     response_mock.status = b"200 OK"
     response_mock.status_code = 200
-    response_mock.response_length = 42
+    response_mock.sent = 42
 
     # When
     logger.access(
         response_mock,
         mocker.Mock(),
-        {"wsgi.django_route": "/health", "REQUEST_METHOD": "GET"},
+        {"flagsmith.route": "/health", "REQUEST_METHOD": "GET"},
         timedelta(milliseconds=101),
     )