MPT-26429 Include extension id in mrok access logs

mrok/frontend/app.py

@@ -6,7 +6,7 @@
 from jinja2 import Environment, FileSystemLoader, select_autoescape
 
 from mrok.conf import get_settings
-from mrok.frontend.utils import get_target_name, parse_accept_header
+from mrok.frontend.utils import get_headers, get_target_name, parse_accept_header
 from mrok.proxy.app import ProxyAppBase
 from mrok.proxy.backend import AIOZitiNetworkBackend
 from mrok.proxy.exceptions import InvalidTargetError
@@ -54,13 +54,11 @@ def setup_connection_pool(
         )
 
     def get_upstream_base_url(self, scope: Scope) -> str:
-        target = get_target_name(
-            {k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {})}
-        )
+        target = get_target_name(get_headers(scope))
         if not target:
             raise InvalidTargetError()
 
-        return f"http://{target.lower()}"
+        return f"http://{target}"
 
     async def send_error_response(
         self,
@@ -70,9 +68,7 @@ async def send_error_response(
         body: str,
         headers: list[tuple[bytes, bytes]] | None = None,
     ):
-        request_headers = {
-            k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {})
-        }
+        request_headers = get_headers(scope)
         accept_header = request_headers.get("accept")
         if not (accept_header and str(http_status) in self._templates_by_error):
             return await super().send_error_response(scope, send, http_status, body)
@@ -113,7 +109,7 @@ async def _render_error_template(
         return await template.render_async(context)
 
     def _extract_request_context(self, scope: Scope) -> dict[str, Any]:
-        headers = {k.decode("latin-1"): v.decode("latin-1") for k, v in scope.get("headers", [])}
+        headers = get_headers(scope)
 
         return {
             "method": scope.get("method"),

mrok/frontend/logging.py

@@ -0,0 +1,17 @@
+import contextvars
+import logging
+
+request_target: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "request_target", default=None
+)
+
+
+def setup_log_record_factory() -> None:
+    default_factory = logging.getLogRecordFactory()
+
+    def record_factory(*args, **kwargs):
+        record = default_factory(*args, **kwargs)
+        record.request_target = request_target.get() or "-"
+        return record
+
+    logging.setLogRecordFactory(record_factory)

mrok/frontend/main.py

@@ -7,7 +7,12 @@
 from mrok.authentication import HTTPAuthManager
 from mrok.conf import get_settings
 from mrok.frontend.app import FrontendProxyApp
-from mrok.frontend.middleware import ASGIAuthenticationMiddleware, HealthCheckMiddleware
+from mrok.frontend.logging import setup_log_record_factory
+from mrok.frontend.middleware import (
+    ASGIAuthenticationMiddleware,
+    HealthCheckMiddleware,
+    TargetContextMiddleware,
+)
 from mrok.logging import get_logging_config
 from mrok.proxy.asgi import ASGIAppWrapper
 
@@ -45,6 +50,7 @@ def load(self):
                 ASGIAuthenticationMiddleware,
                 auth_manager=HTTPAuthManager(settings.frontend.auth),
             )
+        app.add_middleware(TargetContextMiddleware)
         return app
 
 
@@ -58,6 +64,8 @@ def run(
     max_keepalive_connections: int | None,
     keepalive_expiry: float | None,
 ):
+    setup_log_record_factory()
+
     options = {
         "bind": f"{host}:{port}",
         "workers": workers,

mrok/frontend/middleware.py

@@ -1,19 +1,33 @@
 import json
 
 from mrok.authentication import HTTPAuthManager
-from mrok.frontend.utils import get_target_name
+from mrok.frontend.logging import request_target
+from mrok.frontend.utils import get_headers, get_target_name
 from mrok.types.proxy import ASGIApp, ASGIReceive, ASGISend, Scope
 
 
+class TargetContextMiddleware:
+    """Expose the target (extension/instance) of the current request to the log record factory."""
+
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend):
+        target = get_target_name(get_headers(scope))
+        token = request_target.set(target)
+        try:
+            await self.app(scope, receive, send)
+        finally:
+            request_target.reset(token)
+
+
 class HealthCheckMiddleware:
     def __init__(self, app: ASGIApp):
         self.app = app
 
     async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend):
         if scope["type"] == "http" and scope["path"] == "/healthcheck":
-            target = get_target_name(
-                {k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {})}
-            )
+            target = get_target_name(get_headers(scope))
 
             if not target:
                 await send(

mrok/frontend/utils.py

@@ -1,6 +1,7 @@
 import re
 
 from mrok.conf import get_settings
+from mrok.types.proxy import Scope
 
 
 def parse_accept_header(accept: str | None) -> list[str]:
@@ -79,5 +80,9 @@ def get_target_name(headers: dict[str, str]) -> str | None:
         re.fullmatch(settings.identifiers.extension.regex, target)
         or re.fullmatch(settings.identifiers.instance.regex, target)
     ):
-        return target
+        return target.lower()
     return None
+
+
+def get_headers(scope: Scope) -> dict[str, str]:
+    return {k.decode("latin-1"): v.decode("latin-1") for k, v in scope.get("headers", [])}

mrok/logging.py

@@ -3,6 +3,8 @@
 
 from mrok.conf import Settings
 
+ACCESS_LOG_FORMAT = "{asctime} {name} {levelname} (pid: {process}) [{request_target}] {message}"
+
 
 class HealthCheckFilter(logging.Filter):
     def filter(self, record):
@@ -12,6 +14,7 @@
 def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict:
     log_level = "DEBUG" if settings.logging.debug else "INFO"
     handler = "rich" if settings.logging.rich else "console"
+    access_handler = "rich" if settings.logging.rich else "access"
 
     if cli_mode:
         mrok_handler = "cli"
@@ -31,15 +34,20 @@
                 "style": "{",
             },
             "plain": {"format": "%(message)s"},
+            "access": {
+                "format": ACCESS_LOG_FORMAT,
+                "style": "{",
+                "defaults": {"request_target": "-"},
+            },
         },
         "filters": {
             "healthcheck_filter": {
                 "()": HealthCheckFilter,
-            }
+            },
         },
         "handlers": {
             "console": {
                 "class": "logging.StreamHandler",
                 "formatter": "verbose",
                 "stream": "ext://sys.stderr",
             },
@@ -55,14 +63,19 @@
                 "log_time_format": lambda x: x.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
                 "rich_tracebacks": True,
             },
+            "access": {
+                "class": "logging.StreamHandler",
+                "formatter": "access",
+                "stream": "ext://sys.stderr",
+            },
         },
         "root": {
             "handlers": [handler],
             "level": "WARNING",
         },
         "loggers": {
             "gunicorn.access": {
-                "handlers": [handler],
+                "handlers": [access_handler],
                 "level": log_level,
                 "propagate": False,
                 "filters": ["healthcheck_filter"],
@@ -73,7 +86,7 @@
                 "propagate": False,
             },
             "uvicorn.access": {
-                "handlers": [handler],
+                "handlers": [access_handler],
                 "level": log_level,
                 "propagate": False,
                 "filters": ["healthcheck_filter"],

tests/frontend/test_logging.py

@@ -0,0 +1,55 @@
+import logging
+
+import pytest
+from pytest_mock import MockerFixture
+
+from mrok.frontend.logging import request_target, setup_log_record_factory
+from mrok.logging import ACCESS_LOG_FORMAT, get_logging_config
+from tests.types import SettingsFactory
+
+
+def _install_factory(mocker: MockerFixture):
+    m_set_factory = mocker.patch("logging.setLogRecordFactory")
+    setup_log_record_factory()
+    return m_set_factory.call_args.args[0]
+
+
+def test_log_record_factory_sets_request_target(mocker: MockerFixture):
+    factory = _install_factory(mocker)
+
+    token = request_target.set("ext-1234-5678")
+    try:
+        record = factory("uvicorn.access", logging.INFO, __file__, 1, "msg", (), None)
+    finally:
+        request_target.reset(token)
+
+    assert isinstance(record, logging.LogRecord)
+    assert record.request_target == "ext-1234-5678"
+
+
+def test_log_record_factory_without_request_target(mocker: MockerFixture):
+    factory = _install_factory(mocker)
+
+    record = factory("uvicorn.access", logging.INFO, __file__, 1, "msg", (), None)
+
+    assert record.request_target == "-"
+
+
+@pytest.mark.parametrize(
+    ("rich", "expected_handler"),
+    [(False, "access"), (True, "rich")],
+)
+def test_get_logging_config_access_handler(
+    settings_factory: SettingsFactory, rich: bool, expected_handler: str
+):
+    settings = settings_factory(logging={"debug": True, "rich": rich})
+
+    config = get_logging_config(settings)
+
+    assert config["loggers"]["gunicorn.access"]["handlers"] == [expected_handler]
+    assert config["loggers"]["uvicorn.access"]["handlers"] == [expected_handler]
+    assert config["formatters"]["access"] == {
+        "format": ACCESS_LOG_FORMAT,
+        "style": "{",
+        "defaults": {"request_target": "-"},
+    }

tests/frontend/test_main.py

@@ -8,6 +8,7 @@ def test_run(mocker: MockerFixture):
         "mrok.frontend.main.get_logging_config",
         return_value={"logging": "config"},
     )
+    m_setup_factory = mocker.patch("mrok.frontend.main.setup_log_record_factory")
 
     mock_settings = mocker.MagicMock()
     mock_settings.controller.auth = {"backends": []}
@@ -32,3 +33,4 @@ def test_run(mocker: MockerFixture):
     assert m_standalone_app.mock_calls[0].args[0]["mrok"]["keepalive_expiry"] == 99.5
 
     m_app.run.assert_called_once()
+    m_setup_factory.assert_called_once()

tests/frontend/test_middleware.py

@@ -3,7 +3,8 @@
 import pytest
 from pytest_mock import MockerFixture
 
-from mrok.frontend.middleware import HealthCheckMiddleware
+from mrok.frontend.logging import request_target
+from mrok.frontend.middleware import HealthCheckMiddleware, TargetContextMiddleware
 
 
 @pytest.mark.asyncio
@@ -58,3 +59,40 @@ async def test_healthcheck_passtrhough(
     )
 
     m_app.assert_awaited_once_with(scope, m_receive, m_send)
+
+
+@pytest.mark.asyncio
+async def test_target_middleware_sets_target(mocker: MockerFixture):
+    mocker.patch("mrok.frontend.utils.get_frontend_domain", return_value=".extdomain")
+    seen: list[str | None] = []
+    m_app = mocker.AsyncMock(side_effect=lambda *args: seen.append(request_target.get()))
+    m_receive = mocker.AsyncMock()
+    m_send = mocker.AsyncMock()
+    middleware = TargetContextMiddleware(m_app)
+    scope = {
+        "type": "http",
+        "path": "/",
+        "headers": [(b"host", b"EXT-1234-5678.extdomain:8000")],
+    }
+
+    await middleware(scope, m_receive, m_send)
+
+    m_app.assert_awaited_once_with(scope, m_receive, m_send)
+    assert seen == ["ext-1234-5678"]
+    assert request_target.get() is None
+
+
+@pytest.mark.asyncio
+async def test_target_middleware_without_target(mocker: MockerFixture):
+    mocker.patch("mrok.frontend.utils.get_frontend_domain", return_value=".extdomain")
+    seen: list[str | None] = []
+    m_app = mocker.AsyncMock(side_effect=lambda *args: seen.append(request_target.get()))
+    m_receive = mocker.AsyncMock()
+    m_send = mocker.AsyncMock()
+    middleware = TargetContextMiddleware(m_app)
+    scope = {"type": "http", "path": "/", "headers": [(b"host", b"127.0.0.1")]}
+
+    await middleware(scope, m_receive, m_send)
+
+    m_app.assert_awaited_once_with(scope, m_receive, m_send)
+    assert seen == [None]