1from typing import TYPE_CHECKING
 2from weakref import WeakKeyDictionary
 3
 4from opentelemetry import trace
 5from opentelemetry.semconv._incubating.attributes import user_attributes
 6
 7if TYPE_CHECKING:
 8    from app.users.models import User
 9    from plain.http import Request
10
11_request_users: WeakKeyDictionary[Request, User | None] = WeakKeyDictionary()
12
13
14def _stamp_span(user: User | None) -> None:
15    if user is None:
16        return
17    span = trace.get_current_span()
18    if span.is_recording():
19        span.set_attribute(user_attributes.USER_ID, str(user.id))
20
21
22def set_request_user(request: Request, user: User | None) -> None:
23    """Store the authenticated user for this request."""
24    _request_users[request] = user
25    _stamp_span(user)
26
27
28def get_request_user(request: Request) -> User | None:
29    """
30    Get the authenticated user for this request, if any.
31
32    Lazily loads the user from the session on first access.
33    """
34    if request not in _request_users:
35        from .sessions import get_user
36
37        user = get_user(request)
38
39        # Don't need to store a bunch of None values
40        if not user:
41            return None
42
43        _request_users[request] = user
44        _stamp_span(user)
45
46    return _request_users[request]