Add rid to logs with asgi-correlation-id (#782)

* Add `asgi_correlation_id` package

* Replace our RID middleware with `asgi_correlation_id`'s middleware

* Provide logging outside uvicorn

We now call `logging.dictConfig` in app.py, rather that passing the config 
dict to the `log_config` param of uvicorn. This is so that the configuration
is picked up both in tests and when running the app.

* Avoid using TestClient as context manager

I don't know _why_ this was happening, but trying to use a breakpoint in
test_request_summary_is_logged caused the test to hang. I found that using
the fixture directly (i.e. not as a context manager) made the test work as
expected

Author: grahamalama

asgi.py

@@ -13,7 +13,7 @@
             host=settings.host,
             port=settings.port,
             reload=settings.app_reload,
-            log_config=LOGGING_CONFIG,
+            log_config=None,
         )
     )
     server.run()

jbi/app.py

@@ -8,21 +8,24 @@
 from typing import Any, Awaitable, Callable
 
 import sentry_sdk
+from asgi_correlation_id import CorrelationIdMiddleware
 from fastapi import FastAPI, Request, Response, status
 from fastapi.encoders import jsonable_encoder
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
 from fastapi.staticfiles import StaticFiles
 
 from jbi.environment import get_settings, get_version
-from jbi.log import format_request_summary_fields
+from jbi.log import CONFIG, format_request_summary_fields
 from jbi.router import router
 
 SRC_DIR = Path(__file__).parent
 
 settings = get_settings()
 version_info = get_version()
 
+logging.config.dictConfig(CONFIG)
+
 logger = logging.getLogger(__name__)
 
 
@@ -54,16 +57,6 @@ def traces_sampler(sampling_context: dict[str, Any]) -> float:
 app.mount("/static", StaticFiles(directory=SRC_DIR / "static"), name="static")
 
 
-@app.middleware("http")
-async def request_id(
-    request: Request, call_next: Callable[[Request], Awaitable[Response]]
-) -> Response:
-    """Read the request id from headers. This is set by NGinx."""
-    request.state.rid = request.headers.get("X-Request-Id", token_hex(16))
-    response = await call_next(request)
-    return response
-
-
 @app.middleware("http")
 async def request_summary(
     request: Request, call_next: Callable[[Request], Awaitable[Response]]
@@ -106,3 +99,11 @@ async def validation_exception_handler(
         status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
         content={"detail": jsonable_encoder(exc.errors())},
     )
+
+
+app.add_middleware(
+    CorrelationIdMiddleware,
+    header_name="X-Request-Id",
+    generator=lambda: token_hex(16),
+    validator=None,
+)

jbi/log.py

@@ -8,30 +8,57 @@
 from datetime import datetime
 from typing import Optional
 
+from asgi_correlation_id import CorrelationIdFilter
 from fastapi import Request
 from pydantic import BaseModel
 
 from jbi.environment import get_settings
 
 settings = get_settings()
 
+
+class RequestIdFilter(CorrelationIdFilter):
+    """Renames `correlation_id` log field to `rid`"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def filter(self, record) -> bool:
+        result = super().filter(record)  # Apply the existing filter
+
+        # Rename the field from 'correlation_id' to 'rid'
+        if hasattr(record, "correlation_id"):
+            setattr(record, "rid", getattr(record, "correlation_id"))
+            delattr(record, "correlation_id")
+
+        return result
+
+
 CONFIG = {
     "version": 1,
     "disable_existing_loggers": False,
+    "filters": {
+        "request_id": {
+            "()": RequestIdFilter,
+            "uuid_length": 32,
+            "default_value": "-",
+        },
+    },
     "formatters": {
         "mozlog_json": {
             "()": "dockerflow.logging.JsonLogFormatter",
             "logger_name": "jbi",
         },
         "text": {
-            "format": "%(asctime)s %(levelname)-8s %(name)-15s %(message)s",
+            "format": "%(asctime)s %(levelname)-8s [%(rid)s] %(name)-15s %(message)s",
             "datefmt": "%Y-%m-%d %H:%M:%S",
         },
     },
     "handlers": {
         "console": {
             "level": settings.log_level.upper(),
             "class": "logging.StreamHandler",
+            "filters": ["request_id"],
             "formatter": "text"
             if settings.log_format.lower() == "text"
             else "mozlog_json",
@@ -63,7 +90,6 @@ class RequestSummary(BaseModel):
     t: int
     time: str
     status_code: int
-    rid: str
 
 
 def format_request_summary_fields(
@@ -82,5 +108,4 @@ def format_request_summary_fields(
         t=int((current_time - request_time) * 1000.0),
         time=datetime.fromtimestamp(current_time).isoformat(),
         status_code=status_code,
-        rid=request.state.rid,
     ).model_dump()

jbi/models.py

@@ -286,7 +286,6 @@ def lookup_action(self, actions: Actions) -> Action:
 class BugzillaWebhookRequest(BaseModel):
     """Bugzilla Webhook Request Object"""
 
-    rid: str = ""  # This field has a default since it's not parsed from body.
     webhook_id: int
     webhook_name: str
     event: BugzillaWebhookEvent
@@ -362,7 +361,6 @@ class JiraContext(Context):
 class RunnerContext(Context, extra="forbid"):
     """Logging context from runner"""
 
-    rid: str
     operation: Operation
     event: BugzillaWebhookEvent
     action: Optional[Action] = None
@@ -373,7 +371,6 @@ class ActionContext(Context, extra="forbid"):
     """Logging context from actions"""
 
     action: Action
-    rid: str
     operation: Operation
     current_step: Optional[str] = None
     event: BugzillaWebhookEvent

jbi/router.py

@@ -79,7 +79,6 @@ def bugzilla_webhook(
     webhook_request: BugzillaWebhookRequest = Body(..., embed=False),
 ):
     """API endpoint that Bugzilla Webhook Events request"""
-    webhook_request.rid = request.state.rid
     try:
         result = execute_action(webhook_request, actions)
         return result

jbi/runner.py

@@ -152,7 +152,6 @@ def execute_action(
     """
     bug, event = request.bug, request.event
     runner_context = RunnerContext(
-        rid=request.rid,
         bug=bug,
         event=event,
         operation=Operation.HANDLE,
@@ -190,7 +189,6 @@ def execute_action(
 
         action_context = ActionContext(
             action=action,
-            rid=request.rid,
             bug=bug,
             event=event,
             operation=Operation.IGNORE,

poetry.lock

@@ -19,6 +19,7 @@ backoff = "^2.2.1"
 statsd = "^4.0.1"
 requests = "^2.31.0"
 pydantic-settings = "^2.1.0"
+asgi-correlation-id = "^4.2.0"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.5.0"

pyproject.toml

@@ -139,7 +139,6 @@ class Meta:
         model = models.ActionContext
 
     action = factory.SubFactory(ActionFactory)
-    rid = factory.LazyFunction(lambda: token_hex(16))
     operation = Operation.IGNORE
     bug = factory.SubFactory(BugFactory)
     event = factory.SubFactory(WebhookEventFactory)

tests/fixtures/factories.py

@@ -161,7 +161,6 @@ def test_get_issue(mocked_responses, action_context_factory, capturelogs):
 
     assert response == mock_response_data
     for record in capturelogs.records:
-        assert record.rid == context.rid
         assert record.action["whiteboard_tag"] == context.action.whiteboard_tag
 
     before, after = capturelogs.messages
@@ -179,9 +178,6 @@ def test_get_issue_handles_404(mocked_responses, action_context_factory, capture
 
     assert return_val is None
 
-    for record in capturelogs.records:
-        assert record.rid == context.rid
-
     before, after = capturelogs.records
     assert before.levelno == logging.DEBUG
     assert before.message == "Getting issue JBI-234"
@@ -202,7 +198,6 @@ def test_get_issue_reraises_other_erroring_status_codes(
             jira.get_service().get_issue(context=context, issue_key="JBI-234")
 
     [record] = capturelogs.records
-    assert record.rid == context.rid
     assert record.levelno == logging.DEBUG
     assert record.message == "Getting issue JBI-234"
 
@@ -223,10 +218,6 @@ def test_update_issue_resolution(mocked_responses, action_context_factory, captu
             context=context, jira_resolution="DONEZO"
         )
 
-    for record in capturelogs.records:
-        assert record.rid == context.rid
-        assert record.levelno == logging.DEBUG
-
     before, after = capturelogs.messages
     assert before == "Updating resolution of Jira issue JBI-234 to DONEZO"
     assert after == "Updated resolution of Jira issue JBI-234 to DONEZO"
@@ -252,10 +243,8 @@ def test_update_issue_resolution_raises(
                 context=context, jira_resolution="DONEZO"
             )
 
-    [record] = capturelogs.records
-    assert record.rid == context.rid
-    assert record.levelno == logging.DEBUG
-    assert record.message == "Updating resolution of Jira issue JBI-234 to DONEZO"
+    [message] = capturelogs.messages
+    assert message == "Updating resolution of Jira issue JBI-234 to DONEZO"
 
 
 def test_create_jira_issue(mocked_responses, action_context_factory, capturelogs):
@@ -287,10 +276,6 @@ def test_create_jira_issue(mocked_responses, action_context_factory, capturelogs
 
     assert response == mocked_response_data
 
-    for record in capturelogs.records:
-        assert record.rid == context.rid
-        assert record.levelno == logging.DEBUG
-
     before, after = capturelogs.records
     assert before.message == f"Creating new Jira issue for Bug {context.bug.id}"
     assert before.fields == issue_fields
@@ -333,12 +318,10 @@ def test_create_jira_issue_when_list_is_returned(
     before, after = capturelogs.records
     assert before.message == f"Creating new Jira issue for Bug {context.bug.id}"
     assert before.levelno == logging.DEBUG
-    assert before.rid == context.rid
     assert before.fields == issue_fields
 
     assert after.message == f"Jira issue JBI-234 created for Bug {context.bug.id}"
     assert after.levelno == logging.DEBUG
-    assert after.rid == context.rid
     assert after.response == [mocked_issue_data]
 
 
@@ -369,11 +352,9 @@ def test_create_jira_issue_returns_errors(
 
     before, after = capturelogs.records
     assert before.message == f"Creating new Jira issue for Bug {context.bug.id}"
-    assert before.rid == context.rid
     assert before.levelno == logging.DEBUG
     assert before.fields == issue_fields
 
     assert after.message == f"Failed to create issue for Bug {context.bug.id}"
-    assert after.rid == context.rid
     assert after.levelno == logging.ERROR
     assert after.response == fake_error_data

tests/unit/services/test_jira.py

@@ -9,23 +9,32 @@
 from jbi.models import BugzillaWebhookRequest
 
 
-def test_request_summary_is_logged(caplog):
+def test_request_summary_is_logged(caplog, anon_client):
     with caplog.at_level(logging.INFO):
-        with TestClient(app) as anon_client:
-            # https://fastapi.tiangolo.com/advanced/testing-events/
-            anon_client.get(
-                "/",
-                headers={
-                    "X-Request-Id": "foo-bar",
-                },
-            )
+        # https://fastapi.tiangolo.com/advanced/testing-events/
+        anon_client.get(
+            "/",
+            headers={
+                "X-Request-Id": "foo-bar",
+            },
+        )
+
+    summary = [r for r in caplog.records if r.name == "request.summary"][0]
+
+    assert summary.rid == "foo-bar"
+    assert summary.method == "GET"
+    assert summary.path == "/"
+    assert summary.querystring == "{}"
+
 
-            summary = [r for r in caplog.records if r.name == "request.summary"][0]
+def test_request_summary_defaults_user_agent_to_empty_string(caplog, anon_client):
+    with caplog.at_level(logging.INFO):
+        del anon_client.headers["User-Agent"]
+        anon_client.get("/")
+
+        summary = [r for r in caplog.records if r.name == "request.summary"][0]
 
-            assert summary.rid == "foo-bar"
-            assert summary.method == "GET"
-            assert summary.path == "/"
-            assert summary.querystring == "{}"
+        assert summary.agent == ""
 
 
 def test_422_errors_are_logged(webhook_create_example, caplog):
@@ -43,17 +52,6 @@ def test_422_errors_are_logged(webhook_create_example, caplog):
     )
 
 
-def test_request_summary_defaults_user_agent_to_empty_string(caplog):
-    with caplog.at_level(logging.INFO):
-        with TestClient(app) as anon_client:
-            del anon_client.headers["User-Agent"]
-            anon_client.get("/")
-
-            summary = [r for r in caplog.records if r.name == "request.summary"][0]
-
-            assert summary.agent == ""
-
-
 @pytest.mark.parametrize(
     "sampling_context,expected",
     [