Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
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
5 changes: 5 additions & 0 deletions apps/api/app/api/dependencies/current_user.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,21 @@
)
from app.services.rate_limit.job_admission_service import JobAdmissionService
from fastapi import Depends
from sqlalchemy.ext.asyncio import AsyncSession

from shared.core.database import get_db

_job_admission_service = JobAdmissionService()


async def with_current_user(
route_context: RouteAdmissionContext = Depends(get_route_admission_context),
user_id: str = Depends(get_current_user_id),
db: AsyncSession = Depends(get_db),
) -> AsyncGenerator[CurrentUser, None]:
current_user = await _job_admission_service.resolve_current_user(
route_context=route_context,
user_id=user_id,
db=db,
)
yield current_user
55 changes: 42 additions & 13 deletions apps/api/app/services/auth/api_key_authentication_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

from __future__ import annotations

import asyncio
import json
from datetime import datetime, timezone

Expand All @@ -11,11 +10,11 @@
from sqlalchemy.ext.asyncio import AsyncSession

from shared.core.config import redis_pool_manager
from shared.core.database import get_db_context
from shared.services.redis.redis_service import RedisService
from shared.utils.api_keys import hash_api_key

_API_KEY_USER_CACHE_TTL_SECONDS: int = 3600
_LAST_USED_DEBOUNCE_SECONDS: int = 300


class APIKeyAuthenticationService:
Expand Down Expand Up @@ -44,8 +43,12 @@
if not api_key_record or not api_key_record.is_valid():
return None

self._schedule_last_used_update(str(api_key_record.id))
user_id = str(api_key_record.user_id)
await self._update_last_used_best_effort(
session,
redis_service,
str(api_key_record.id),
)
await self._set_cached_user_id(
redis_service,
key_hash,
Expand Down Expand Up @@ -75,6 +78,10 @@
def _get_user_api_keys_key(user_id: str) -> str:
return f"api-key:user-hashes:{user_id}"

@staticmethod
def _get_last_used_debounce_key(api_key_id: str) -> str:
return f"api-key:last-used-debounce:{api_key_id}"

async def _get_cached_user_id(
self,
redis_service: RedisService,
Expand Down Expand Up @@ -159,20 +166,42 @@
remaining_seconds = int((expires_at_utc - now).total_seconds())
return max(1, min(_API_KEY_USER_CACHE_TTL_SECONDS, remaining_seconds))

def _schedule_last_used_update(self, api_key_id: str) -> None:
async def _update_last_used_best_effort(
self,
session: AsyncSession,
redis_service: RedisService,
api_key_id: str,
) -> None:
"""Update last_used_at on the request session without a nested checkout.

Redis SET NX debounce skips redundant writes within the debounce window
so job polls do not compete for QueuePool via create_task+get_db_context.
"""
debounce_key = self._get_last_used_debounce_key(api_key_id)
try:
asyncio.create_task(
self._update_last_used_best_effort(api_key_id),
name=f"api_key_last_used:{api_key_id}",
acquired = await redis_service.set_nx(
debounce_key,
"1",
ex=_LAST_USED_DEBOUNCE_SECONDS,
)
except Exception as exc:
if not acquired:
return
except Exception:
logger.warning(
f"Failed to schedule API key last-used update (ignored): {exc}"
"api_key_authentication: last-used debounce failed for api_key_id={}; "
"updating anyway",
api_key_id,
)

async def _update_last_used_best_effort(self, api_key_id: str) -> None:
try:
async with get_db_context() as db:
await self._repository.update_last_used(db, api_key_id)
await self._repository.update_last_used(session, api_key_id)
# Request-scoped sessions do not auto-commit.
await session.commit()
except Exception as exc:
logger.warning(f"Failed to update API key last-used time (ignored): {exc}")
logger.warning(
f"Failed to update API key last-used time (ignored): {exc}"
)
try:
await session.rollback()
except Exception:
Comment thread
github-advanced-security[bot] marked this conversation as resolved.
Fixed
pass
3 changes: 2 additions & 1 deletion apps/api/app/services/rate_limit/job_admission_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,9 @@ async def resolve_current_user(
*,
route_context: RouteAdmissionContext,
user_id: str,
db: AsyncSession,
) -> CurrentUser:
user_tier = await TierService.get_tier(user_id)
user_tier = await TierService.get_tier(user_id, session=db)
self._route_policy_service.enforce_guest_api_key_scope(
route_context=route_context,
user_tier=user_tier,
Expand Down
45 changes: 38 additions & 7 deletions apps/api/app/services/rate_limit/tier_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,15 @@ class TierService:
"""Manages user tier lookup, caching, and refresh."""

@staticmethod
async def get_tier(user_id: str) -> str:
async def get_tier(
user_id: str,
session: AsyncSession | None = None,
) -> str:
"""Return a user's tier from cache or database.

When ``session`` is provided, reuse it instead of opening a nested
``get_db_context`` checkout (critical on the job-poll auth path).

Missing user tier state is treated as invalid data and raises directly;
this method never falls back to a default tier for user lookup.
"""
Expand All @@ -42,18 +48,43 @@ async def get_tier(user_id: str) -> str:
if cached_tier is not None:
return cached_tier

async with get_db_context() as session:
try:
user_tier: str = await TierService._get_tier_from_db(session, user_id)
except NotFoundException:
user_tier = await TierService._initialize_missing_user_tier(
session,
if session is not None:
user_tier = await TierService._resolve_tier_from_db(
session,
user_id,
commit_on_initialize=True,
)
else:
async with get_db_context() as owned_session:
user_tier = await TierService._resolve_tier_from_db(
owned_session,
user_id,
commit_on_initialize=False,
)

await TierService._set_cached_tier(redis_service, user_id, user_tier)
return user_tier

@staticmethod
async def _resolve_tier_from_db(
session: AsyncSession,
user_id: str,
*,
commit_on_initialize: bool,
) -> str:
"""Load tier from DB, initializing missing first-use billing state."""
try:
return await TierService._get_tier_from_db(session, user_id)
except NotFoundException:
user_tier = await TierService._initialize_missing_user_tier(
session,
user_id,
)
# Request-scoped sessions do not auto-commit; get_db_context does.
if commit_on_initialize:
await session.commit()
return user_tier

@staticmethod
async def refresh_tier(user_id: str, session: AsyncSession) -> str:
"""Called on payment success.
Expand Down
Empty file added apps/api/tests/unit/__init__.py
Empty file.
Loading
Loading