use dataclass

Author: RayBB

openlibrary/core/fulltext.py

@@ -8,7 +8,7 @@
 from openlibrary.core.lending import get_availability
 from openlibrary.plugins.openlibrary.home import format_book_data
 from openlibrary.utils import async_utils
-from openlibrary.utils.async_utils import async_bridge, x_forwarded_for
+from openlibrary.utils.async_utils import async_bridge, req_context
 
 logger = logging.getLogger("openlibrary.inside")
 
@@ -24,7 +24,7 @@ async def fulltext_search_api(params):
     search_endpoint = config.plugin_inside['search_endpoint']
     search_select = search_endpoint + '?' + urlencode(params, 'utf-8')
     headers = {
-        "x-preferred-client-id": x_forwarded_for.get(),
+        "x-preferred-client-id": req_context.get().x_forwarded_for or "ol-internal",
         "x-application-id": "openlibrary",
     }
     if config_fts_context is not None:

openlibrary/core/lending.py

@@ -21,7 +21,7 @@
 from openlibrary.core import cache, stats
 from openlibrary.plugins.upstream.utils import urlencode
 from openlibrary.utils import dateutil, uniq
-from openlibrary.utils.async_utils import user_agent, x_forwarded_for
+from openlibrary.utils.async_utils import req_context
 
 from . import helpers as h
 from . import ia
@@ -298,7 +298,7 @@ def get_available(
         # Internet Archive Elastic Search (which powers some of our
         # carousel queries) needs Open Library to forward user IPs so
         # we can attribute requests to end-users
-        client_ip = x_forwarded_for.get()
+        client_ip = req_context.get().x_forwarded_for or "ol-internal"
         headers = {
             "x-client-id": client_ip,
             "x-preferred-client-id": client_ip,
@@ -440,8 +440,8 @@ def key_func(_id: str) -> str:
 
     try:
         headers = {
-            "x-preferred-client-id": x_forwarded_for.get(),
-            "x-preferred-client-useragent": user_agent.get(),
+            "x-preferred-client-id": req_context.get().x_forwarded_for or "ol-internal",
+            "x-preferred-client-useragent": req_context.get().user_agent or "",
             "x-application-id": "openlibrary",
             "user-agent": "Open Library Site",
         }

openlibrary/tests/core/test_lending.py

@@ -1,7 +1,7 @@
 from unittest.mock import Mock, patch
 
 from openlibrary.core import lending
-from openlibrary.utils.async_utils import user_agent, x_forwarded_for
+from openlibrary.utils.async_utils import RequestContextVars, req_context
 
 
 class TestAddAvailability:
@@ -44,8 +44,12 @@ def mock_get_availability(id_type, ocaids):
 class TestGetAvailability:
     def test_cache(self):
         # Set needed context variables
-        x_forwarded_for.set("ol-internal")
-        user_agent.set("test-user-agent")
+        req_context.set(
+            RequestContextVars(
+                x_forwarded_for="ol-internal",
+                user_agent="test-user-agent",
+            )
+        )
 
         with patch("openlibrary.core.ia.session.get") as mock_get:
             mock_get.return_value = Mock()

openlibrary/utils/async_utils.py

@@ -2,6 +2,7 @@
 import threading
 from collections.abc import Callable, Coroutine
 from contextvars import ContextVar
+from dataclasses import dataclass
 from typing import Any, ParamSpec, TypeVar
 
 import web
@@ -48,8 +49,20 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
 # These are scoped to a specific request
 # Defining them here for now to keep things close together
 
-x_forwarded_for: ContextVar[str] = ContextVar("x_forwarded_for")
-user_agent: ContextVar[str] = ContextVar("user_agent")
+
+@dataclass(frozen=True)
+class RequestContextVars:
+    """
+    These are scoped to a specific request and should be derived directly from the request.
+    Use sparingly.
+    """
+
+    x_forwarded_for: str | None
+    user_agent: str | None
+
+
+req_context: ContextVar[RequestContextVars] = ContextVar("req_context")
+
 # TODO: Create an async and stateless version of site so we don't have to do this
 site: ContextVar[Site] = ContextVar("site")
 
@@ -62,8 +75,12 @@ def set_context_for_sync_request():
     That version should be usable without instantiating a site
     """
     site.set(create_site())
-    user_agent.set(web.ctx.env.get("HTTP_USER_AGENT", ""))
-    x_forwarded_for.set(web.ctx.env.get("HTTP_X_FORWARDED_FOR", "ol-internal"))
+    req_context.set(
+        RequestContextVars(
+            x_forwarded_for=web.ctx.env.get("HTTP_X_FORWARDED_FOR"),
+            user_agent=web.ctx.env.get("HTTP_USER_AGENT"),
+        )
+    )
 
 
 def set_context_for_async_request(request: Request):
@@ -76,5 +93,9 @@ def set_context_for_async_request(request: Request):
     This is preferable for testable and maintainable code.
     """
     site.set(create_site())
-    user_agent.set(request.headers.get("User-Agent", ""))
-    x_forwarded_for.set(request.headers.get("X-Forwarded-For", "ol-internal"))
+    req_context.set(
+        RequestContextVars(
+            x_forwarded_for=request.headers.get("X-Forwarded-For"),
+            user_agent=request.headers.get("User-Agent"),
+        )
+    )