From 8caa11893b33028d6ae9f36adb6b644345e53149 Mon Sep 17 00:00:00 2001 From: charles Date: Wed, 16 Jul 2025 18:44:35 -0700 Subject: [PATCH] Update the cache system * Refactor API key handling to use asynchronous cache functions and update caching logic for Forge API keys * This includes the addition of new async cache methods and updates to existing API key routes to support async operations. --- app/api/dependencies.py | 10 +++--- app/api/routes/api_keys.py | 68 +++++++++++++++++++------------------- app/core/async_cache.py | 61 +++++++++++++++++++++++++++++++--- app/core/cache.py | 34 +++++++++++++------ 4 files changed, 118 insertions(+), 55 deletions(-) diff --git a/app/api/dependencies.py b/app/api/dependencies.py index 6519b3a..dcd6084 100644 --- a/app/api/dependencies.py +++ b/app/api/dependencies.py @@ -17,10 +17,11 @@ from app.api.schemas.user import TokenData from app.core.async_cache import ( - async_provider_service_cache, cache_user_async, get_cached_user_async, invalidate_user_cache_async, + forge_scope_cache_async, + get_forge_scope_cache_async, ) from app.core.database import get_db from app.core.logger import get_logger @@ -180,8 +181,7 @@ async def get_user_by_api_key( # Try scope cache first – this doesn't remove the need to verify the key, but it # avoids an extra query later in /models. - scope_cache_key = f"forge_scope:{api_key}" - cached_scope = await async_provider_service_cache.get(scope_cache_key) + cached_scope = await get_forge_scope_cache_async(api_key) api_key_record = ( db.query(ForgeApiKey) @@ -219,9 +219,7 @@ async def get_user_by_api_key( pk.provider_name for pk in api_key_record.allowed_provider_keys ] # Cache it (short TTL – scope changes are rare) - await async_provider_service_cache.set( - scope_cache_key, allowed_provider_names, ttl=300 - ) + await forge_scope_cache_async(api_key, allowed_provider_names, ttl=300) else: allowed_provider_names = cached_scope diff --git a/app/api/routes/api_keys.py b/app/api/routes/api_keys.py index b6b8c9b..723a34e 100644 --- a/app/api/routes/api_keys.py +++ b/app/api/routes/api_keys.py @@ -13,7 +13,7 @@ ForgeApiKeyResponse, ForgeApiKeyUpdate, ) -from app.core.cache import invalidate_provider_service_cache, invalidate_user_cache, invalidate_forge_scope_cache +from app.core.async_cache import invalidate_forge_scope_cache_async, invalidate_user_cache_async, invalidate_provider_service_cache_async from app.core.database import get_db from app.core.security import generate_forge_api_key from app.models.forge_api_key import ForgeApiKey @@ -25,7 +25,7 @@ # --- Internal Service Functions --- -def _get_api_keys_internal( +async def _get_api_keys_internal( db: Session, current_user: UserModel ) -> list[ForgeApiKeyMasked]: """ @@ -47,7 +47,7 @@ def _get_api_keys_internal( return masked_keys -def _create_api_key_internal( +async def _create_api_key_internal( api_key_create: ForgeApiKeyCreate, db: Session, current_user: UserModel ) -> ForgeApiKeyResponse: """ @@ -91,7 +91,7 @@ def _create_api_key_internal( return ForgeApiKeyResponse(**response_data) -def _update_api_key_internal( +async def _update_api_key_internal( key_id: int, api_key_update: ForgeApiKeyUpdate, db: Session, current_user: UserModel ) -> ForgeApiKeyResponse: """ @@ -112,7 +112,7 @@ def _update_api_key_internal( old_active_state = db_api_key.is_active db_api_key.is_active = update_data["is_active"] if old_active_state and not db_api_key.is_active: - invalidate_user_cache(db_api_key.key) + await invalidate_user_cache_async(db_api_key.key) if api_key_update.allowed_provider_key_ids is not None: db_api_key.allowed_provider_keys.clear() @@ -139,7 +139,7 @@ def _update_api_key_internal( # Invalidate forge scope cache if the scope was updated if api_key_update.allowed_provider_key_ids is not None: - invalidate_forge_scope_cache(db_api_key.key) + await invalidate_forge_scope_cache_async(db_api_key.key) response_data = db_api_key.__dict__.copy() response_data["allowed_provider_key_ids"] = [ @@ -148,7 +148,7 @@ def _update_api_key_internal( return ForgeApiKeyResponse(**response_data) -def _delete_api_key_internal( +async def _delete_api_key_internal( key_id: int, db: Session, current_user: UserModel ) -> ForgeApiKeyResponse: """ @@ -177,13 +177,13 @@ def _delete_api_key_internal( db.delete(db_api_key) db.commit() - invalidate_user_cache(key_to_invalidate) - invalidate_forge_scope_cache(key_to_invalidate) - invalidate_provider_service_cache(current_user.id) + await invalidate_user_cache_async(key_to_invalidate) + await invalidate_forge_scope_cache_async(key_to_invalidate) + await invalidate_provider_service_cache_async(current_user.id) return ForgeApiKeyResponse(**response_data) -def _regenerate_api_key_internal( +async def _regenerate_api_key_internal( key_id: int, db: Session, current_user: UserModel ) -> ForgeApiKeyResponse: """ @@ -199,9 +199,9 @@ def _regenerate_api_key_internal( # Invalidate caches for the old key old_key = db_api_key.key - invalidate_user_cache(old_key) - invalidate_forge_scope_cache(old_key) - invalidate_provider_service_cache(current_user.id) + await invalidate_user_cache_async(old_key) + await invalidate_forge_scope_cache_async(old_key) + await invalidate_provider_service_cache_async(current_user.id) # Generate and set new key new_key_value = generate_forge_api_key() @@ -221,91 +221,91 @@ def _regenerate_api_key_internal( @router.get("/", response_model=list[ForgeApiKeyMasked]) -def get_api_keys( +async def get_api_keys( db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user), ) -> Any: - return _get_api_keys_internal(db, current_user) + return await _get_api_keys_internal(db, current_user) @router.post("/", response_model=ForgeApiKeyResponse) -def create_api_key( +async def create_api_key( api_key_create: ForgeApiKeyCreate, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user), ) -> Any: - return _create_api_key_internal(api_key_create, db, current_user) + return await _create_api_key_internal(api_key_create, db, current_user) @router.put("/{key_id}", response_model=ForgeApiKeyResponse) -def update_api_key( +async def update_api_key( key_id: int, api_key_update: ForgeApiKeyUpdate, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user), ) -> Any: - return _update_api_key_internal(key_id, api_key_update, db, current_user) + return await _update_api_key_internal(key_id, api_key_update, db, current_user) @router.delete("/{key_id}", response_model=ForgeApiKeyResponse) -def delete_api_key( +async def delete_api_key( key_id: int, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user), ) -> Any: - return _delete_api_key_internal(key_id, db, current_user) + return await _delete_api_key_internal(key_id, db, current_user) @router.post("/{key_id}/regenerate", response_model=ForgeApiKeyResponse) -def regenerate_api_key( +async def regenerate_api_key( key_id: int, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user), ) -> Any: - return _regenerate_api_key_internal(key_id, db, current_user) + return await _regenerate_api_key_internal(key_id, db, current_user) # Clerk versions of the routes @router.get("/clerk", response_model=list[ForgeApiKeyMasked]) -def get_api_keys_clerk( +async def get_api_keys_clerk( db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user_from_clerk), ) -> Any: - return _get_api_keys_internal(db, current_user) + return await _get_api_keys_internal(db, current_user) @router.post("/clerk", response_model=ForgeApiKeyResponse) -def create_api_key_clerk( +async def create_api_key_clerk( api_key_create: ForgeApiKeyCreate, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user_from_clerk), ) -> Any: - return _create_api_key_internal(api_key_create, db, current_user) + return await _create_api_key_internal(api_key_create, db, current_user) @router.put("/clerk/{key_id}", response_model=ForgeApiKeyResponse) -def update_api_key_clerk( +async def update_api_key_clerk( key_id: int, api_key_update: ForgeApiKeyUpdate, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user_from_clerk), ) -> Any: - return _update_api_key_internal(key_id, api_key_update, db, current_user) + return await _update_api_key_internal(key_id, api_key_update, db, current_user) @router.delete("/clerk/{key_id}", response_model=ForgeApiKeyResponse) -def delete_api_key_clerk( +async def delete_api_key_clerk( key_id: int, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user_from_clerk), ) -> Any: - return _delete_api_key_internal(key_id, db, current_user) + return await _delete_api_key_internal(key_id, db, current_user) @router.post("/clerk/{key_id}/regenerate", response_model=ForgeApiKeyResponse) -def regenerate_api_key_clerk( +async def regenerate_api_key_clerk( key_id: int, db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_active_user_from_clerk), ) -> Any: - return _regenerate_api_key_internal(key_id, db, current_user) + return await _regenerate_api_key_internal(key_id, db, current_user) diff --git a/app/core/async_cache.py b/app/core/async_cache.py index c08e222..a3e34e6 100644 --- a/app/core/async_cache.py +++ b/app/core/async_cache.py @@ -159,9 +159,12 @@ async def wrapper(*args, **kwargs): # User-specific functions async def get_cached_user_async(api_key: str) -> CachedUser | None: - """Get a user from cache by API key asynchronously""" + """Get a user from cache by Forge API key asynchronously""" if not api_key: return None + # Remove the forge- prefix for caching from the API key + if api_key.startswith("forge-"): + api_key = api_key[6:] cached_data = await async_user_cache.get(f"user:{api_key}") if cached_data: return CachedUser.model_validate(cached_data) @@ -169,17 +172,23 @@ async def get_cached_user_async(api_key: str) -> CachedUser | None: async def cache_user_async(api_key: str, user: User) -> None: - """Cache a user by API key asynchronously""" + """Cache a user by Forge API key asynchronously""" if not api_key or user is None: return cached_user = CachedUser.model_validate(user) + # Remove the forge- prefix for caching from the API key + if api_key.startswith("forge-"): + api_key = api_key[6:] await async_user_cache.set(f"user:{api_key}", cached_user.model_dump()) async def invalidate_user_cache_async(api_key: str) -> None: - """Invalidate user cache for a specific API key asynchronously""" + """Invalidate user cache for a specific Forge API key asynchronously""" if not api_key: return + # Remove the forge- prefix for caching from the API key + if api_key.startswith("forge-"): + api_key = api_key[6:] await async_user_cache.delete(f"user:{api_key}") @@ -246,6 +255,51 @@ async def invalidate_user_cache_by_id_async(user_id: int) -> None: if DEBUG_CACHE: logger.debug(f"Cache: Invalidated user cache for key: {key[:8]}...") +async def get_forge_scope_cache_async(api_key: str) -> list[str] | None: + """Get the forge scope cache for a specific Forge API key asynchronously""" + if not api_key: + return None + # Remove the forge- prefix for caching from the API key + cache_key = api_key + if cache_key.startswith("forge-"): + cache_key = cache_key[6:] + return await async_provider_service_cache.get(f"forge_scope:{cache_key}") + + +async def forge_scope_cache_async(api_key: str, allowed_provider_names: list[str], ttl: int = 300) -> None: + """Cache the forge scope cache for a specific Forge API key asynchronously""" + if not api_key: + return None + # Remove the forge- prefix for caching from the API key + cache_key = api_key + if cache_key.startswith("forge-"): + cache_key = cache_key[6:] + await async_provider_service_cache.set(f"forge_scope:{cache_key}", allowed_provider_names, ttl=ttl) + if DEBUG_CACHE: + # Mask the API key for logging + masked_key = cache_key[:8] + "..." if len(cache_key) > 8 else cache_key + logger.debug(f"Cache: set forge scope cache for Forge API key: {masked_key} (async)") + + +async def invalidate_forge_scope_cache_async(api_key: str) -> None: + """Invalidate forge scope cache for a specific API key asynchronously. + + Args: + api_key (str): The API key to invalidate cache for. Can include or exclude 'forge-' prefix. + """ + if not api_key: + return + + cache_key = api_key + if cache_key.startswith("forge-"): + cache_key = cache_key[6:] # Remove "forge-" prefix to match cache setting format + + await async_provider_service_cache.delete(f"forge_scope:{cache_key}") + + if DEBUG_CACHE: + # Mask the API key for logging + masked_key = cache_key[:8] + "..." if len(cache_key) > 8 else cache_key + logger.debug(f"Cache: Invalidated forge scope cache for Forge API key: {masked_key} (async)") # Provider service functions async def get_cached_provider_service_async(user_id: int) -> Any: @@ -355,7 +409,6 @@ async def warm_cache_async(db: Session) -> None: .all() ) for key in forge_api_keys: - # Cache user with their Forge API key await cache_user_async(key.key, user) # Cache provider services for active users diff --git a/app/core/cache.py b/app/core/cache.py index 01e4d15..26f5693 100644 --- a/app/core/cache.py +++ b/app/core/cache.py @@ -146,9 +146,12 @@ def wrapper(*args, **kwargs): # User-specific functions def get_cached_user(api_key: str) -> CachedUser | None: - """Get a user from cache by API key""" + """Get a user from cache by Forge API key""" if not api_key: return None + # Remove the forge- prefix for caching from the API key + if api_key.startswith("forge-"): + api_key = api_key[6:] cached_data = user_cache.get(f"user:{api_key}") if cached_data: return CachedUser.model_validate(cached_data) @@ -156,17 +159,23 @@ def get_cached_user(api_key: str) -> CachedUser | None: def cache_user(api_key: str, user: User) -> None: - """Cache a user by API key""" + """Cache a user by Forge API key""" if not api_key or user is None: return cached_user = CachedUser.model_validate(user) + # Remove the forge- prefix for caching from the API key + if api_key.startswith("forge-"): + api_key = api_key[6:] user_cache.set(f"user:{api_key}", cached_user.model_dump()) def invalidate_user_cache(api_key: str) -> None: - """Invalidate user cache for a specific API key""" + """Invalidate user cache for a specific Forge API key""" if not api_key: return + # Remove the forge- prefix for caching from the API key + if api_key.startswith("forge-"): + api_key = api_key[6:] user_cache.delete(f"user:{api_key}") @@ -190,7 +199,7 @@ def invalidate_forge_scope_cache(api_key: str) -> None: if DEBUG_CACHE: # Mask the API key for logging masked_key = cache_key[:8] + "..." if len(cache_key) > 8 else cache_key - logger.debug(f"Cache: Invalidated forge scope cache for API key: {masked_key}") + logger.debug(f"Cache: Invalidated forge scope cache for Forge API key: {masked_key}") # Provider service functions @@ -323,7 +332,7 @@ def invalidate_all_caches() -> None: async def warm_cache(db: Session) -> None: """Pre-cache frequently accessed data""" from app.core.security import decrypt_api_key - from app.models.provider_key import ProviderKey + from app.models.forge_api_key import ForgeApiKey from app.models.user import User from app.services.provider_service import ProviderService @@ -333,12 +342,15 @@ async def warm_cache(db: Session) -> None: # Cache active users active_users = db.query(User).filter(User.is_active).all() for user in active_users: - # Get user's API keys - api_keys = db.query(ProviderKey).filter(ProviderKey.user_id == user.id).all() - for key in api_keys: - # Decrypt the API key before caching - decrypted_key = decrypt_api_key(key.encrypted_api_key) - cache_user(decrypted_key, user) + # Get user's Forge API keys + forge_api_keys = ( + db.query(ForgeApiKey) + .filter(ForgeApiKey.user_id == user.id, ForgeApiKey.is_active) + .all() + ) + for key in forge_api_keys: + # Cache user with their Forge API key + cache_user(key.key, user) # Cache provider services for active users for user in active_users: