Skip to content

Update the cache system #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wilsonccccc merged 1 commit into from
Jul 17, 2025
Merged

Update the cache system #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 4 additions & 6 deletions app/api/dependencies.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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

Expand Down
68 changes: 34 additions & 34 deletions app/api/routes/api_keys.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]:
"""
Expand All @@ -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:
"""
Expand Down Expand Up @@ -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:
"""
Expand All @@ -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()
Expand All @@ -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"] = [
Expand All @@ -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:
"""
Expand Down Expand Up @@ -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:
"""
Expand All @@ -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()
Expand All @@ -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)
61 changes: 57 additions & 4 deletions app/core/async_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,27 +159,36 @@ 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)
return 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}")


Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading