Repository: snok/asgi-correlation-id (language: python)

Issue:
Add generated request ID to the request headers
It might be a good idea to add the newly generated request ID also to the request headers. There may be cases where we can't or don't want to use the value from context variable and rely solely on the value in the request headers (e.g. other third party middlewares). This is especially true when there is a header in the incoming request but it doesn't pass the validator and thus a new value is generated. In that case, the header value is inconsistent with the context variable value and the value present in logs.

If this sounds reasonable, I could try to prepare a PR.


Discussion / hints:
I think writing your own middleware that sets a request header is probably a better solution, unless you have a more general motivation for this feature that is clearly useful for other users. I'm not against it, but if you feel this is useful, could you elaborate on how?
I actually initially did just that, wrote my own middleware that adds the request ID to request headers. It is an ASGI version of the WSGI middleware I use on another project for years. Then I found this project which seems for cover all use cases one might possibly need, with this little exception.

As an additional motivation, consider various tracing middlewares which use request headers to set the span context. For example these two:
- https://github.com/identixone/fastapi_contrib/blob/master/fastapi_contrib/tracing/middlewares.py
- https://github.com/wesdu/fastapi-opentracing/blob/master/fastapi_opentracing/middleware.py
So you're saying, when an incoming request has no (for example) `'x-request-id'` header present, and we generate one OR for the code path where it is present and we replace it, it would be useful to overwrite this request header with our set value?

I can see how it would be hard to automate tracing without this, so I'd consider myself convinced :slightly_smiling_face: 

Before moving forward with this, do you see any drawbacks? It feels a bit weird to modify the request object when there wasn't a header there to begin with, but I can't think of a specific reason why that would be true.
Took me a few reads to really understand the use cases behind this, but I am convinced as well.   
Altering incoming headers may not be what people want tho, so I suggest having it as an optional. 
Yeah, that's what I'm trying to say :-) To reiterate, I find this especially useful in the case when you replace the header value with a new one should the original value not pass the validator. In that case, tracing middlewares would contain the original value in their context, while response headers (and presumably logs) would contain the new one. That would be misleading and would make it quite hard to debug issues in any observability platform correlating metrics, logs and traces (see e.g. https://grafana.com/blog/2020/03/31/how-to-successfully-correlate-metrics-logs-and-traces-in-grafana/).

Modifying request headers is IMHO quite common. I've been using it in a whole stack which consists of several components where each of them either reuses the request ID from the previous one or creates a new one and passes it in the request ID header to the next one. For example, if the top-level component is a HAProxy acting as a load-balancer, I'd use something along these lines:
```
...
unique-id-format %{+X}o\ %Ts%pid%fi%fp%ci%cp%rt
acl request_id_exists req.hdr(X-Request-ID) -m found
http-request set-header X-Request-ID %[unique-id,lower] unless request_id_exists
http-request capture hdr(X-Request-ID) len 64
http-response set-header X-Request-ID %[capture.req.hdr(0)]
...
```
Next, there might be an nginx server in front of the WSGI/ASGI server. I'd use:
```
map $http_x_request_id $req_id {
    ~* $http_x_request_id;
    default $request_id;
}

server {
    location / {
        ...
        proxy_hide_header X-Request-ID;
        add_header X-Request-ID $req_id always;
        ...
    }
}
```
So I don't really see any problems with this approach. The only thing I think should be taken care of is the order of middlewares so the request ID header is guaranteed to be present before any other middleware relying on it tries to access it.
Sounds good :+1: 
> Took me a few reads to really understand the use cases behind this, but I am convinced as well. Altering incoming headers may not be what people want tho, so I suggest having it as an optional.

Missed this, yeah I think an opt-in might be best for a first iteration, but would be happy to see your suggested implementation first @tlinhart, then we can discuss implementation details in the PR 

Fix patch:
diff --git a/README.md b/README.md
index 92e8c16..98af75a 100644
--- a/README.md
+++ b/README.md
@@ -148,6 +148,7 @@ The middleware can be configured in a few ways, but there are no required argume
 app.add_middleware(
     CorrelationIdMiddleware,
     header_name='X-Request-ID',
+    update_request_header=False,
     generator=lambda: uuid4().hex,
     validator=is_valid_uuid4,
     transformer=lambda a: a,
@@ -163,6 +164,13 @@ Configurable middleware arguments include:
 - Description: The header name decides which HTTP header value to read correlation IDs from. `X-Request-ID` and
   `X-Correlation-ID` are common choices.
 
+**update_request_header**
+
+- Type: `bool`
+- Default: `False`
+- Description: Whether to update incoming request's header value with the generated correlation ID. This is to support
+  use cases where it's relied on the presence of the request header (like various tracing middlewares).
+
 **generator**
 
 - Type: `Callable[[], str]`
diff --git a/asgi_correlation_id/middleware.py b/asgi_correlation_id/middleware.py
index c8ba187..a083387 100644
--- a/asgi_correlation_id/middleware.py
+++ b/asgi_correlation_id/middleware.py
@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING, Callable, Optional
 from uuid import UUID, uuid4
 
-from starlette.datastructures import Headers, MutableHeaders
+from starlette.datastructures import MutableHeaders
 
 from asgi_correlation_id.context import correlation_id
 from asgi_correlation_id.extensions.sentry import get_sentry_extension
@@ -31,6 +31,7 @@ def is_valid_uuid4(uuid_: str) -> bool:
 class CorrelationIdMiddleware:
     app: 'ASGIApp'
     header_name: str = 'X-Request-ID'
+    update_request_header: bool = False
 
     # ID-generating callable
     generator: Callable[[], str] = field(default=lambda: uuid4().hex)
@@ -50,7 +51,8 @@ async def __call__(self, scope: 'Scope', receive: 'Receive', send: 'Send') -> No
             return
 
         # Try to load request ID from the request headers
-        header_value = Headers(scope=scope).get(self.header_name.lower())
+        headers = MutableHeaders(scope=scope)
+        header_value = headers.get(self.header_name.lower())
 
         if not header_value:
             # Generate request ID if none was found
@@ -67,6 +69,10 @@ async def __call__(self, scope: 'Scope', receive: 'Receive', send: 'Send') -> No
         if self.transformer:
             id_value = self.transformer(id_value)
 
+        # Update the request headers if needed
+        if id_value != header_value and self.update_request_header is True:
+            headers[self.header_name] = id_value
+
         correlation_id.set(id_value)
         self.sentry_extension(id_value)
 


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 281dbeb..6f81b1f 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -48,6 +48,7 @@ def _configure_logging():
 TRANSFORMER_VALUE = 'some-id'
 
 default_app = FastAPI(middleware=[Middleware(CorrelationIdMiddleware)])
+update_request_header_app = FastAPI(middleware=[Middleware(CorrelationIdMiddleware, update_request_header=True)])
 no_validator_or_transformer_app = FastAPI(
     middleware=[Middleware(CorrelationIdMiddleware, validator=None, transformer=None)]
 )
diff --git a/tests/test_middleware.py b/tests/test_middleware.py
index 99e979e..09df78a 100644
--- a/tests/test_middleware.py
+++ b/tests/test_middleware.py
@@ -3,17 +3,18 @@
 from uuid import uuid4
 
 import pytest
-from fastapi import Response
+from fastapi import Request, Response
 from httpx import AsyncClient
 from starlette.testclient import TestClient
 
-from asgi_correlation_id.middleware import FAILED_VALIDATION_MESSAGE
+from asgi_correlation_id.middleware import FAILED_VALIDATION_MESSAGE, is_valid_uuid4
 from tests.conftest import (
     TRANSFORMER_VALUE,
     default_app,
     generator_app,
     no_validator_or_transformer_app,
     transformer_app,
+    update_request_header_app,
 )
 
 if TYPE_CHECKING:
@@ -21,7 +22,7 @@
 
 logger = logging.getLogger('asgi_correlation_id')
 
-apps = [default_app, no_validator_or_transformer_app, transformer_app, generator_app]
+apps = [default_app, update_request_header_app, no_validator_or_transformer_app, transformer_app, generator_app]
 
 pytestmark = pytest.mark.asyncio
 
@@ -58,6 +59,35 @@ async def test_view() -> dict:
         assert third_response.headers['X-Request-ID'] not in [correlation_id, second_correlation_id]
 
 
+@pytest.mark.parametrize('app', [update_request_header_app])
+async def test_update_request_header(app):
+    """
+    We expect the middleware to update the request header with the request ID
+    value.
+    """
+
+    @app.get('/test', status_code=200)
+    async def test_view(request: Request) -> dict:
+        logger.debug('Test view')
+        return {'correlation_id': request.headers.get('X-Request-ID')}
+
+    async with AsyncClient(app=app, base_url='http://test') as client:
+        # Check for newly generated request ID in the request header if none
+        # was initially provided.
+        response = await client.get('test')
+        assert is_valid_uuid4(response.json()['correlation_id'])
+
+        # Check for newly generated request ID in the request header if it
+        # initially contains an invalid value.
+        response = await client.get('test', headers={'X-Request-ID': 'invalid'})
+        assert is_valid_uuid4(response.json()['correlation_id'])
+
+        # Check for our request ID value in the request header.
+        correlation_id = uuid4().hex
+        response = await client.get('test', headers={'X-Request-ID': correlation_id})
+        assert response.json()['correlation_id'] == correlation_id
+
+
 bad_uuids = [
     'test',
     'bad-uuid',