Refactor auth cache with type hints, RateLimiterService, token cleanup, Lua preload, and expanded test coverage.

Author: yihong1120

examples/auth/cache.py

@@ -7,6 +7,7 @@
 from fastapi import FastAPI
 from fastapi_limiter import FastAPILimiter
 
+from examples.auth.cache import _DEFAULT_SERVICE
 from examples.auth.database import engine
 from examples.auth.jwt_scheduler import start_jwt_scheduler
 from examples.auth.models import Base
@@ -26,10 +27,10 @@ async def global_lifespan(app: FastAPI) -> AsyncGenerator[None]:
         None: Control is yielded back to the application
             after performing startup tasks.
     """
-    # Step 1: Start the scheduler (e.g., for rotating JWT secret keys).
+    # Start the scheduler (e.g., for rotating JWT secret keys).
     scheduler = start_jwt_scheduler(app)
 
-    # Step 2: Initialise Redis connection and rate limiter.
+    # Initialise Redis connection and rate limiter.
     redis_host: str = os.getenv('REDIS_HOST', '127.0.0.1')
     redis_port: str = os.getenv('REDIS_PORT', '6379')
     redis_password: str = os.getenv('REDIS_PASSWORD', '')
@@ -39,7 +40,14 @@ async def global_lifespan(app: FastAPI) -> AsyncGenerator[None]:
     redis_conn = await app.state.redis_client.connect()
     await FastAPILimiter.init(redis_conn)
 
-    # Step 3: Optionally create database tables on startup
+    # Preload Lua scripts into Redis (if any).
+    try:
+        await _DEFAULT_SERVICE.preload_script(redis_conn)
+    except Exception:
+        # Log the error or handle it as needed
+        pass
+
+    # Optionally create database tables on startup
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.create_all)
 

examples/auth/lifespan.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import time
+
+import jwt
+from redis.asyncio import Redis
+
+from examples.auth.cache import get_user_data
+from examples.auth.cache import set_user_data
+from examples.auth.config import Settings
+
+
+settings = Settings()
+
+
+async def prune_user_cache(
+    redis_pool: Redis,
+    username: str,
+) -> dict[str, object] | None:
+    """
+    Prune a user's cached authentication data in Redis.
+
+    Args:
+        redis_pool: Asynchronous Redis client/connection.
+        username: Username whose cache entry should be pruned.
+
+    Returns:
+        The updated cache dictionary if present, otherwise ``None`` when no
+        cache entry exists.
+    """
+    cache: dict[str, object] | None = await get_user_data(redis_pool, username)
+    if not cache:
+        return None
+
+    changed: bool = False
+    now: int = int(time.time())
+
+    # Refresh tokens pruning
+    # The cache may contain a list of refresh tokens; keep those that are
+    # valid according to ``jwt.decode`` and drop expired/invalid ones.
+    refresh_raw = cache.get('refresh_tokens', [])
+    refresh_tokens: list[str]
+    if isinstance(refresh_raw, list):
+        # Ensure a typed list of strings only
+        refresh_tokens = [t for t in refresh_raw if isinstance(t, str)]
+    else:
+        refresh_tokens = []
+    new_refresh_tokens: list[str] = []
+    for tok in refresh_tokens:
+        try:
+            # Decode to verify validity and expiry
+            jwt.decode(
+                tok,
+                settings.authjwt_secret_key,
+                algorithms=[settings.ALGORITHM],
+            )
+            # Keep token only if decode succeeds
+            new_refresh_tokens.append(tok)
+        except (jwt.ExpiredSignatureError, jwt.InvalidTokenError):
+            changed = True
+            # Drop expired/invalid tokens silently
+            continue
+    if new_refresh_tokens != refresh_tokens:
+        cache['refresh_tokens'] = new_refresh_tokens
+
+    # JTI metadata pruning
+    # jti_list holds active JWT IDs; jti_meta maps JTI -> expiry timestamp.
+    jti_list_raw = cache.get('jti_list', [])
+    jti_list: list[str]
+    if isinstance(jti_list_raw, list):
+        jti_list = [j for j in jti_list_raw if isinstance(j, str)]
+    else:
+        jti_list = []
+
+    jti_meta_raw = cache.get('jti_meta', {})
+    jti_meta: dict[str, int]
+    if isinstance(jti_meta_raw, dict):
+        # Build a strictly typed mapping of str -> int
+        jti_meta = {
+            k: int(v)
+            for k, v in jti_meta_raw.items()
+            if isinstance(k, str) and isinstance(v, int)
+        }
+    else:
+        jti_meta = {}
+    if jti_meta:
+        new_jti_list: list[str] = []
+        for j in jti_list:
+            exp_ts: int = int(jti_meta.get(j, 0))
+            # Keep if no expiry is tracked (0) or it is still in the future
+            if exp_ts == 0 or exp_ts > now:
+                new_jti_list.append(j)
+            else:
+                changed = True
+
+        # Remove stale jti_meta entries not in list or already expired
+        new_jti_meta: dict[str, int] = {}
+        for j, exp in jti_meta.items():
+            if j in new_jti_list and exp > now:
+                new_jti_meta[j] = int(exp)
+            else:
+                changed = True
+
+        if new_jti_list != jti_list:
+            cache['jti_list'] = new_jti_list
+        cache['jti_meta'] = new_jti_meta
+
+    # Persist only if an actual change occurred
+    if changed:
+        await set_user_data(redis_pool, username, cache)
+
+    return cache

examples/auth/token_cleanup.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import time
+from typing import TypeAlias
 
 from fastapi import HTTPException
 from sqlalchemy import select
@@ -9,23 +10,40 @@
 
 from examples.auth.models import User
 
-# Cache mechanism for storing user site information
-_user_sites_cache: dict[str, tuple[list[str], float]] = {}
-_cache_ttl: int = 300  # Cache time-to-live in seconds (5 minutes)
+# A cache entry stores: (list of site names, cached-at epoch seconds).
+CacheEntry: TypeAlias = tuple[list[str], float]
+
+# Process-local cache for storing user site information.
+_user_sites_cache: dict[str, CacheEntry] = {}
+
+# Cache time-to-live in seconds (5 minutes).
+_cache_ttl: int = 300
 
 
 async def get_user_sites_cached(username: str, db: AsyncSession) -> list[str]:
-    """Return site names the user can access, with simple in-memory caching.
+    """
+    Return site names the user may access, with simple in-memory caching.
+
+    Args:
+        username: The unique username to resolve.
+        db: An asynchronous SQLAlchemy session used for the lookup.
+
+    Returns:
+        A list of site names that the user may access. The list order follows
+        the ORM relationship ordering as returned by the database.
 
-    Raises HTTPException(404) if the user is not found.
+    Raises:
+        HTTPException: With status code 404 if the user is not found.
     """
     current_time: float = time.time()
 
     if username in _user_sites_cache:
+        # Fast path: honour TTL and return cached site names when still fresh.
         cached_names, cached_time = _user_sites_cache[username]
         if current_time - cached_time < _cache_ttl:
             return cached_names
 
+    # Query the user and their sites in one round-trip.
     stmt_user = (
         select(User)
         .where(User.username == username)
@@ -35,6 +53,40 @@ async def get_user_sites_cached(username: str, db: AsyncSession) -> list[str]:
     if not user_obj:
         raise HTTPException(status_code=404, detail='User not found')
 
+    # Extract and cache the site names with the current timestamp.
     site_names: list[str] = [site.name for site in user_obj.sites]
     _user_sites_cache[username] = (site_names, current_time)
     return site_names
+
+
+async def get_user_and_sites(
+    db: AsyncSession, username: str,
+) -> tuple[User, list[str], str]:
+    """
+    Fetch the user, their site names, and role from the database.
+
+    Args:
+        db: An asynchronous SQLAlchemy session.
+        username: The username to query.
+
+    Returns:
+        A 3-tuple of ``(user, site_names, role)`` where:
+        - ``user`` is the fully loaded ``User`` ORM instance,
+        - ``site_names`` is a list of the user's site names, and
+        - ``role`` is the user's role as a string.
+
+    Raises:
+        HTTPException: With status code 401 if the user cannot be found.
+    """
+    stmt_user = (
+        select(User)
+        .where(User.username == username)
+        .options(selectinload(User.sites))
+    )
+    result = await db.execute(stmt_user)
+    user: User | None = result.scalars().first()
+    if not user:
+        raise HTTPException(status_code=401, detail='Invalid user')
+    user_role: str = user.role
+    user_site_names: list[str] = [site.name for site in user.sites]
+    return user, user_site_names, user_role