⚡️(api) cache logged-in user database request

Getting request user from database is the most frequent request as it's
accessed by every API endpoint. Caching logged user object for a limited
time can significantly decrease database queries.

.github/workflows/api.yml

@@ -365,6 +365,7 @@ jobs:
           QUALICHARGE_DB_HOST: localhost
           QUALICHARGE_DB_NAME: test-qualicharge-api
           QUALICHARGE_TEST_DB_NAME: test-qualicharge-api
+          QUALICHARGE_API_GET_USER_CACHE_INFO: true
           # Speed up tests
           QUALICHARGE_API_STATIQUE_BULK_CREATE_MAX_SIZE: 10
           QUALICHARGE_API_STATUS_BULK_CREATE_MAX_SIZE: 10

env.d/api

@@ -3,6 +3,7 @@ QUALICHARGE_ALLOWED_HOSTS=["http://localhost:8010"]
 QUALICHARGE_API_ADMIN_PASSWORD=admin
 QUALICHARGE_API_ADMIN_USER=admin
 QUALICHARGE_API_STATIQUE_BULK_CREATE_MAX_SIZE=1000
+QUALICHARGE_API_GET_USER_CACHE_INFO=True
 QUALICHARGE_DB_CONNECTION_MAX_OVERFLOW=200
 QUALICHARGE_DB_CONNECTION_POOL_SIZE=50
 QUALICHARGE_DB_ENGINE=postgresql+psycopg

src/api/CHANGELOG.md

@@ -17,6 +17,8 @@ and this project adheres to
 
 - Prefetch user-related groups and operational units in `get_user` dependency
 - Improve bulk endpoints permissions checking
+- Cache logged user object for `API_GET_USER_CACHE_TTL` seconds to decrease the
+  number of database queries
 
 ## [0.16.0] - 2024-12-12
 

src/api/Pipfile

@@ -6,6 +6,7 @@ name = "pypi"
 [packages]
 alembic = "==1.14.0"
 annotated-types = "==0.7.0"
+cachetools = "==5.5.0"
 email-validator = "==2.2.0"
 fastapi = "==0.115.6"
 geoalchemy2 = {extras = ["shapely"], version = "==0.16.0"}
@@ -42,6 +43,7 @@ pytest-cov = "==6.0.0"
 pytest-httpx = "==0.35.0"
 qualicharge = {path = ".", editable = true}
 ruff = "==0.8.2"
+types-cachetools = "==5.5.0.20240820"
 types-passlib = "==1.7.7.20240819"
 types-python-jose = "==3.3.4.20240106"
 types-requests = "==2.32.0.20241016"

src/api/qualicharge/auth/oidc.py

@@ -2,10 +2,12 @@
 
 import logging
 from functools import lru_cache
+from threading import Lock
 from typing import Annotated, Dict, Union
 
 import httpx
 import jwt
+from cachetools import TTLCache, cached
 from fastapi import Depends
 from fastapi.security import (
     HTTPAuthorizationCredentials,
@@ -148,26 +150,46 @@ def get_token(
     return IDToken(**decoded_token)
 
 
-def get_user(
-    security_scopes: SecurityScopes,
-    token: Annotated[IDToken, Depends(get_token)],
+@cached(
+    TTLCache(
+        maxsize=settings.API_GET_USER_CACHE_MAXSIZE,
+        ttl=settings.API_GET_USER_CACHE_TTL,
+    ),
+    lock=Lock(),
+    key=lambda email, session: email,
+    info=settings.API_GET_USER_CACHE_INFO,
+)
+def get_user_from_db(
+    email: str,
     session: Annotated[
         SMSession,
         Depends(get_session),
     ],
-) -> User:
-    """Get request user."""
-    # Get registered user
-    user = (
+):
+    """Fetch user and related objects from database."""
+    logging.debug(f"Getting user from database: {email}")
+    return (
         session.exec(
             select(User)
             .options(joinedload(User.groups).joinedload(Group.operational_units))  # type: ignore[arg-type]
-            .where(User.email == token.email)
+            .where(User.email == email)
         )
         .unique()
         .one_or_none()
     )
 
+
+def get_user(
+    security_scopes: SecurityScopes,
+    token: Annotated[IDToken, Depends(get_token)],
+    session: Annotated[
+        SMSession,
+        Depends(get_session),
+    ],
+) -> User:
+    """Get request user."""
+    user = get_user_from_db(email=token.email, session=session)
+
     # User does not exist: raise an error
     if user is None:
         logger.error(f"User {token.email} tried to login but is not registered yet")

src/api/qualicharge/conf.py

@@ -130,6 +130,9 @@ def PASSWORD_CONTEXT(self) -> CryptContext:
     API_STATIQUE_PAGE_MAX_SIZE: int = 100
     API_STATIQUE_PAGE_SIZE: int = 10
     API_STATUS_BULK_CREATE_MAX_SIZE: int = 10
+    API_GET_USER_CACHE_MAXSIZE: int = 256
+    API_GET_USER_CACHE_TTL: int = 1800
+    API_GET_USER_CACHE_INFO: bool = False
 
     model_config = SettingsConfigDict(
         case_sensitive=True, env_nested_delimiter="__", env_prefix="QUALICHARGE_"

src/api/tests/api/v1/routers/test_auth.py

@@ -8,9 +8,10 @@
 
 from qualicharge.auth.factories import IDTokenFactory
 from qualicharge.auth.models import IDToken, UserCreate, UserRead
-from qualicharge.auth.oidc import discover_provider, get_public_keys
+from qualicharge.auth.oidc import discover_provider, get_public_keys, get_user_from_db
 from qualicharge.auth.schemas import User
 from qualicharge.conf import settings
+from qualicharge.db import SAQueryCounter
 
 
 def setup_function():
@@ -51,6 +52,42 @@ def test_whoami_auth(client_auth):
     assert user.is_staff is True
 
 
+@pytest.mark.parametrize(
+    "client_auth",
+    ((True, {"email": "[email protected]", "username": "jdoe"}),),
+    indirect=True,
+)
+def test_whoami_auth_get_user_cache(client_auth, db_session):
+    """Test the get_user cache on the whoami endpoint."""
+    cache_info = get_user_from_db.cache_info()
+    assert cache_info.hits == 0
+    assert cache_info.currsize == 0
+
+    with SAQueryCounter(db_session.connection()) as counter:
+        response = client_auth.get("/auth/whoami")
+    expected = 2
+    assert counter.count == expected
+    assert response.status_code == status.HTTP_200_OK
+    cache_info = get_user_from_db.cache_info()
+    assert cache_info.hits == 0
+    assert cache_info.currsize == 1
+
+    user = UserRead(**response.json())
+    assert user.email == "[email protected]"
+
+    # Now we should be using cache 10 times
+    for hit in range(1, 10):
+        with SAQueryCounter(db_session.connection()) as counter:
+            response = client_auth.get("/auth/whoami")
+        cache_info = get_user_from_db.cache_info()
+        assert counter.count == 0
+        assert cache_info.hits == hit
+        assert cache_info.currsize == 1
+        assert response.status_code == status.HTTP_200_OK
+        user = UserRead(**response.json())
+        assert user.email == "[email protected]"
+
+
 def test_whoami_expired_signature(
     client, id_token_factory: IDTokenFactory, httpx_mock, monkeypatch
 ):

src/api/tests/auth/test_oidc.py

@@ -15,6 +15,7 @@
     get_public_keys,
     get_token,
     get_user,
+    get_user_from_db,
 )
 from qualicharge.auth.schemas import GroupOperationalUnit, ScopesEnum
 from qualicharge.conf import settings
@@ -349,3 +350,68 @@ def test_get_user_number_of_queries(id_token_factory: IDTokenFactory, db_session
             ou.id for ou in operational_units
         }
     assert counter.count == 0
+
+
+def test_get_user_cache(id_token_factory: IDTokenFactory, db_session):
+    """Test the OIDC get user utility number of queries for a standard user."""
+    UserFactory.__session__ = db_session
+    GroupFactory.__session__ = db_session
+
+    token = id_token_factory.build()
+
+    # Create groups linked to Operational Units
+    groups = GroupFactory.create_batch_sync(3)
+    operational_units = db_session.exec(select(OperationalUnit).limit(3)).all()
+    for group, operational_unit in zip(groups, operational_units, strict=True):
+        db_session.add(
+            GroupOperationalUnit(
+                group_id=group.id, operational_unit_id=operational_unit.id
+            )
+        )
+
+    # Create user linked to this groups and related operational units
+    user = UserFactory.create_sync(
+        email=token.email,
+        is_superuser=False,
+        is_active=True,
+        groups=groups,
+        scopes=[ScopesEnum.ALL_CREATE],
+    )
+    security_scopes = SecurityScopes(scopes=[ScopesEnum.ALL_CREATE])
+
+    # Test the original number of queries
+    with SAQueryCounter(db_session.connection()) as counter:
+        user = get_user(
+            security_scopes=security_scopes,
+            token=token,
+            session=db_session,
+        )
+    cache_info = get_user_from_db.cache_info()  # type: ignore[attr-defined]
+    assert counter.count == 1
+    assert cache_info.hits == 0
+    assert cache_info.currsize == 1
+
+    # User should be cached, we should not hit the database
+    for hit in range(1, 10):
+        with SAQueryCounter(db_session.connection()) as counter:
+            user = get_user(
+                security_scopes=security_scopes,
+                token=token,
+                session=db_session,
+            )
+        cache_info = get_user_from_db.cache_info()  # type: ignore[attr-defined]
+        assert counter.count == 0
+        assert cache_info.hits == hit
+        assert cache_info.currsize == 1
+
+    # When getting groups...
+    with SAQueryCounter(db_session.connection()) as counter:
+        assert {g.id for g in user.groups} == {g.id for g in groups}
+    assert counter.count == 0
+
+    # ... and related operational units
+    with SAQueryCounter(db_session.connection()) as counter:
+        assert {ou.id for g in user.groups for ou in g.operational_units} == {
+            ou.id for ou in operational_units
+        }
+    assert counter.count == 0

src/api/tests/conftest.py

@@ -4,7 +4,7 @@
 
 from qualicharge.auth.factories import IDTokenFactory
 
-from .fixtures.app import client, client_auth
+from .fixtures.app import clear_lru_cache, client, client_auth
 from .fixtures.cli import runner
 from .fixtures.db import (
     db_engine,

src/api/tests/fixtures/app.py

@@ -6,7 +6,7 @@
 
 from qualicharge.api.v1 import app
 from qualicharge.auth.factories import GroupFactory, IDTokenFactory, UserFactory
-from qualicharge.auth.oidc import get_token
+from qualicharge.auth.oidc import get_token, get_user_from_db
 from qualicharge.auth.schemas import UserGroup
 
 
@@ -54,3 +54,16 @@ def client_auth(request, id_token_factory: IDTokenFactory, db_session: Session):
     )
     yield TestClient(app)
     app.dependency_overrides = {}
+
+
+@pytest.fixture(autouse=True)
+def clear_lru_cache():
+    """Taken from codeinthehole.
+
+    https://til.codeinthehole.com/posts/how-to-inspect-and-clear-pythons-functoolslrucache/
+    """
+    # Execute the test...
+    yield
+
+    # Clear the LRU cache.
+    get_user_from_db.cache_clear()