From 8164f5150b8601b5e27b11ed707af3e4730c93bb Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 25 Jul 2026 20:14:45 +0800 Subject: [PATCH 1/4] fix: stop nested QueuePool checkouts on job-poll auth Reuse the request AsyncSession for tier lookup and API-key last_used updates (with Redis debounce), and wire the async engine to DB_POOL_SIZE/MAX_OVERFLOW defaults of 50/50 so legal poll bursts no longer exhaust the per-process pool via nested get_db_context holds. Co-authored-by: Cursor --- apps/api/app/api/dependencies/current_user.py | 5 + .../auth/api_key_authentication_service.py | 61 +++- .../rate_limit/job_admission_service.py | 3 +- .../app/services/rate_limit/tier_service.py | 35 ++- apps/api/tests/unit/__init__.py | 0 .../unit/test_job_poll_session_hygiene.py | 273 ++++++++++++++++++ .../shared/core/config/database.py | 6 +- .../shared-python/shared/core/database.py | 42 ++- .../shared/tests/test_database_pool_config.py | 22 ++ 9 files changed, 414 insertions(+), 33 deletions(-) create mode 100644 apps/api/tests/unit/__init__.py create mode 100644 apps/api/tests/unit/test_job_poll_session_hygiene.py create mode 100644 packages/shared-python/shared/tests/test_database_pool_config.py diff --git a/apps/api/app/api/dependencies/current_user.py b/apps/api/app/api/dependencies/current_user.py index 4c5593c1a..65448cd00 100644 --- a/apps/api/app/api/dependencies/current_user.py +++ b/apps/api/app/api/dependencies/current_user.py @@ -10,6 +10,9 @@ ) 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() @@ -17,9 +20,11 @@ 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 diff --git a/apps/api/app/services/auth/api_key_authentication_service.py b/apps/api/app/services/auth/api_key_authentication_service.py index 67e677247..bf851db68 100644 --- a/apps/api/app/services/auth/api_key_authentication_service.py +++ b/apps/api/app/services/auth/api_key_authentication_service.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import json from datetime import datetime, timezone @@ -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: @@ -44,8 +43,12 @@ async def validate_api_key( 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, @@ -75,6 +78,10 @@ def _get_user_id_key(api_key_hash: str) -> str: 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, @@ -159,20 +166,48 @@ def _resolve_api_key_cache_ttl_seconds(self, expires_at: datetime | None) -> int 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 debounce skips redundant writes within the debounce window so job + polls do not compete for QueuePool capacity 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}", + if await redis_service.exists(debounce_key): + return + except Exception: + logger.warning( + "api_key_authentication: failed to read last-used debounce for api_key_id={}", + api_key_id, ) + + try: + await self._repository.update_last_used(session, api_key_id) + await session.commit() except Exception as exc: logger.warning( - f"Failed to schedule API key last-used update (ignored): {exc}" + f"Failed to update API key last-used time (ignored): {exc}" ) + try: + await session.rollback() + except Exception: + pass + return - 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) - except Exception as exc: - logger.warning(f"Failed to update API key last-used time (ignored): {exc}") + await redis_service.set( + debounce_key, + "1", + ttl=_LAST_USED_DEBOUNCE_SECONDS, + ) + except Exception: + logger.warning( + "api_key_authentication: failed to write last-used debounce for api_key_id={}", + api_key_id, + ) diff --git a/apps/api/app/services/rate_limit/job_admission_service.py b/apps/api/app/services/rate_limit/job_admission_service.py index 130a880c3..0c0321eed 100644 --- a/apps/api/app/services/rate_limit/job_admission_service.py +++ b/apps/api/app/services/rate_limit/job_admission_service.py @@ -33,8 +33,9 @@ async def resolve_current_user( *, route_context: RouteAdmissionContext, user_id: str, + db: AsyncSession | None = None, ) -> 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, diff --git a/apps/api/app/services/rate_limit/tier_service.py b/apps/api/app/services/rate_limit/tier_service.py index 91bb73d00..dc1255f16 100644 --- a/apps/api/app/services/rate_limit/tier_service.py +++ b/apps/api/app/services/rate_limit/tier_service.py @@ -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. """ @@ -42,18 +48,33 @@ 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) + else: + async with get_db_context() as owned_session: + user_tier = await TierService._resolve_tier_from_db( + owned_session, user_id, ) 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) -> 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, + ) + # Persist first-use init when reusing the request session + # (get_db_context commits on exit when we own the session). + await session.commit() + return user_tier + @staticmethod async def refresh_tier(user_id: str, session: AsyncSession) -> str: """Called on payment success. diff --git a/apps/api/tests/unit/__init__.py b/apps/api/tests/unit/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/apps/api/tests/unit/test_job_poll_session_hygiene.py b/apps/api/tests/unit/test_job_poll_session_hygiene.py new file mode 100644 index 000000000..8df95befb --- /dev/null +++ b/apps/api/tests/unit/test_job_poll_session_hygiene.py @@ -0,0 +1,273 @@ +"""Regression tests: job-poll auth must not open nested DB checkouts.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from tests.support.import_environment import ( + configure_import_environment, + ensure_import_paths, +) + +configure_import_environment() +ensure_import_paths() + +from app.services.auth.api_key_authentication_service import ( # noqa: E402 + APIKeyAuthenticationService, +) +from app.services.rate_limit.data_structures import ( # noqa: E402 + RouteAdmissionContext, +) +from app.services.rate_limit.job_admission_service import ( # noqa: E402 + JobAdmissionService, +) +from app.services.rate_limit.tier_service import TierService # noqa: E402 + + +@pytest.mark.asyncio +async def test_get_tier_reuses_provided_session_without_get_db_context() -> None: + session = AsyncMock() + redis_service = AsyncMock() + redis_service.get = AsyncMock(return_value=None) + redis_service.set = AsyncMock() + + with ( + patch( + "app.services.rate_limit.tier_service.redis_pool_manager.get_redis_service", + return_value=redis_service, + ), + patch( + "app.services.rate_limit.tier_service.get_db_context", + ) as get_db_context_mock, + patch.object( + TierService, + "_get_tier_from_db", + new=AsyncMock(return_value="pro"), + ) as get_tier_from_db, + ): + tier = await TierService.get_tier("user-1", session=session) + + assert tier == "pro" + get_tier_from_db.assert_awaited_once_with(session, "user-1") + get_db_context_mock.assert_not_called() + + +@pytest.mark.asyncio +async def test_resolve_current_user_passes_request_session_to_get_tier() -> None: + session = AsyncMock() + route_context = RouteAdmissionContext( + method="GET", + path="/v1/jobs/job_abc", + limit_identifier="GET:/v1/jobs/{job_id}", + ) + service = JobAdmissionService( + route_policy_service=MagicMock( + enforce_guest_api_key_scope=MagicMock(), + enforce_user_system_limit=AsyncMock(), + ), + ) + + with ( + patch.object( + TierService, + "get_tier", + new=AsyncMock(return_value="free"), + ) as get_tier, + patch( + "app.services.rate_limit.job_admission_service.RateLimitConfig.get_instance", + return_value=SimpleNamespace(is_enabled=False), + ), + ): + current_user = await service.resolve_current_user( + route_context=route_context, + user_id="user-1", + db=session, + ) + + assert current_user.user_id == "user-1" + assert current_user.user_tier == "free" + get_tier.assert_awaited_once_with("user-1", session=session) + + +@pytest.mark.asyncio +async def test_validate_api_key_updates_last_used_on_same_session() -> None: + session = AsyncMock() + session.commit = AsyncMock() + redis_service = AsyncMock() + redis_service.get = AsyncMock(return_value=None) + redis_service.exists = AsyncMock(return_value=False) + redis_service.set = AsyncMock() + redis_service.sadd = AsyncMock() + redis_service.ttl = AsyncMock(return_value=-2) + redis_service.expire = AsyncMock() + + api_key_record = SimpleNamespace( + id="key-1", + user_id="user-1", + expires_at=None, + is_valid=lambda: True, + ) + repository = MagicMock() + repository.get_by_key_hash = AsyncMock(return_value=api_key_record) + repository.update_last_used = AsyncMock(return_value=True) + + service = APIKeyAuthenticationService(repository=repository) + + with ( + patch( + "app.services.auth.api_key_authentication_service.redis_pool_manager.get_redis_service", + return_value=redis_service, + ), + patch( + "app.services.auth.api_key_authentication_service.hash_api_key", + return_value="hash-1", + ), + patch( + "app.services.auth.api_key_authentication_service.get_db_context", + create=True, + ) as get_db_context_mock, + patch("asyncio.create_task") as create_task_mock, + ): + user_id = await service.validate_api_key(session, "kw_test_key") + + assert user_id == "user-1" + repository.update_last_used.assert_awaited_once_with(session, "key-1") + session.commit.assert_awaited_once() + create_task_mock.assert_not_called() + get_db_context_mock.assert_not_called() + redis_service.set.assert_any_await( + "api-key:last-used-debounce:key-1", + "1", + ttl=300, + ) + + +@pytest.mark.asyncio +async def test_validate_api_key_skips_last_used_when_debounced() -> None: + session = AsyncMock() + session.commit = AsyncMock() + redis_service = AsyncMock() + redis_service.get = AsyncMock(return_value=None) + redis_service.exists = AsyncMock(return_value=True) + redis_service.set = AsyncMock() + redis_service.sadd = AsyncMock() + redis_service.ttl = AsyncMock(return_value=-2) + redis_service.expire = AsyncMock() + + api_key_record = SimpleNamespace( + id="key-1", + user_id="user-1", + expires_at=None, + is_valid=lambda: True, + ) + repository = MagicMock() + repository.get_by_key_hash = AsyncMock(return_value=api_key_record) + repository.update_last_used = AsyncMock(return_value=True) + + service = APIKeyAuthenticationService(repository=repository) + + with ( + patch( + "app.services.auth.api_key_authentication_service.redis_pool_manager.get_redis_service", + return_value=redis_service, + ), + patch( + "app.services.auth.api_key_authentication_service.hash_api_key", + return_value="hash-1", + ), + ): + user_id = await service.validate_api_key(session, "kw_test_key") + + assert user_id == "user-1" + repository.update_last_used.assert_not_called() + session.commit.assert_not_called() + + +@pytest.mark.asyncio +async def test_job_poll_auth_path_uses_single_session_factory_checkout() -> None: + """End-to-end hygiene: tier + last-used must not call get_db_context.""" + session = AsyncMock() + session.commit = AsyncMock() + checkout_count = {"n": 0} + + class _CountingContext: + async def __aenter__(self) -> Any: + checkout_count["n"] += 1 + return AsyncMock() + + async def __aexit__(self, *args: object) -> None: + return None + + redis_service = AsyncMock() + redis_service.get = AsyncMock(side_effect=[None, None]) # api-key miss, tier miss + redis_service.exists = AsyncMock(return_value=False) + redis_service.set = AsyncMock() + redis_service.sadd = AsyncMock() + redis_service.ttl = AsyncMock(return_value=-2) + redis_service.expire = AsyncMock() + + api_key_record = SimpleNamespace( + id="key-1", + user_id="user-1", + expires_at=None, + is_valid=lambda: True, + ) + repository = MagicMock() + repository.get_by_key_hash = AsyncMock(return_value=api_key_record) + repository.update_last_used = AsyncMock(return_value=True) + auth_service = APIKeyAuthenticationService(repository=repository) + + route_context = RouteAdmissionContext( + method="GET", + path="/v1/jobs/job_abc", + limit_identifier="GET:/v1/jobs/{job_id}", + ) + admission = JobAdmissionService( + route_policy_service=MagicMock( + enforce_guest_api_key_scope=MagicMock(), + enforce_user_system_limit=AsyncMock(), + ), + ) + + with ( + patch( + "app.services.auth.api_key_authentication_service.redis_pool_manager.get_redis_service", + return_value=redis_service, + ), + patch( + "app.services.rate_limit.tier_service.redis_pool_manager.get_redis_service", + return_value=redis_service, + ), + patch( + "app.services.auth.api_key_authentication_service.hash_api_key", + return_value="hash-1", + ), + patch( + "app.services.rate_limit.tier_service.get_db_context", + side_effect=_CountingContext, + ), + patch.object( + TierService, + "_get_tier_from_db", + new=AsyncMock(return_value="free"), + ), + patch( + "app.services.rate_limit.job_admission_service.RateLimitConfig.get_instance", + return_value=SimpleNamespace(is_enabled=False), + ), + patch("asyncio.create_task") as create_task_mock, + ): + user_id = await auth_service.validate_api_key(session, "kw_test_key") + current_user = await admission.resolve_current_user( + route_context=route_context, + user_id=user_id or "", + db=session, + ) + + assert current_user.user_id == "user-1" + assert checkout_count["n"] == 0 + create_task_mock.assert_not_called() + repository.update_last_used.assert_awaited_once_with(session, "key-1") diff --git a/packages/shared-python/shared/core/config/database.py b/packages/shared-python/shared/core/config/database.py index d78629284..a9baac5f5 100644 --- a/packages/shared-python/shared/core/config/database.py +++ b/packages/shared-python/shared/core/config/database.py @@ -25,8 +25,10 @@ class DatabaseConfig(BaseModel): ) # Async database pool configuration for the API. - DB_POOL_SIZE: int = Field(default=20, description="Connection-pool size") - DB_MAX_OVERFLOW: int = Field(default=30, description="Maximum overflow connections") + # Defaults sized for legal job-poll bursts after nested-checkout hygiene + # (pool_size + max_overflow = 100 checkouts per API process). + DB_POOL_SIZE: int = Field(default=50, description="Connection-pool size") + DB_MAX_OVERFLOW: int = Field(default=50, description="Maximum overflow connections") DB_POOL_RECYCLE: int = Field( default=1800, description="Connection recycle interval in seconds" ) diff --git a/packages/shared-python/shared/core/database.py b/packages/shared-python/shared/core/database.py index cd8a03d6e..80ee99188 100644 --- a/packages/shared-python/shared/core/database.py +++ b/packages/shared-python/shared/core/database.py @@ -25,7 +25,7 @@ os.getenv("DB_USE_NULL_POOL", "false").lower() == "true" ) engine_options: dict[str, Any] = { - "pool_recycle": ProcessingConstants.DB_POOL_RECYCLE, + "pool_recycle": settings.DB_POOL_RECYCLE, "pool_pre_ping": ProcessingConstants.DB_POOL_PRE_PING, "pool_reset_on_return": ProcessingConstants.DB_POOL_RESET_ON_RETURN, "connect_args": { @@ -46,9 +46,9 @@ else: engine_options.update( { - "pool_size": ProcessingConstants.DB_POOL_SIZE, - "max_overflow": ProcessingConstants.DB_MAX_OVERFLOW, - "pool_timeout": ProcessingConstants.DB_POOL_TIMEOUT, + "pool_size": settings.DB_POOL_SIZE, + "max_overflow": settings.DB_MAX_OVERFLOW, + "pool_timeout": settings.DB_POOL_TIMEOUT, } ) @@ -204,12 +204,33 @@ def on_connect(dbapi_connection, connection_record): """Handle new connection events.""" logger.info("New database connection established") - @event.listens_for(engine.sync_engine, "checkout") + @event.listens_for(engine.sync_engine.pool, "checkout") def on_checkout(dbapi_connection, connection_record, connection_proxy): - """Handle connection checkout events.""" - logger.debug("Connection checked out from pool") - - @event.listens_for(engine.sync_engine, "checkin") + """Handle connection checkout events and surface pool pressure.""" + pool = engine.sync_engine.pool + checked_out = pool.checkedout() + overflow = pool.overflow() + pool_size = pool.size() if hasattr(pool, "size") else settings.DB_POOL_SIZE + logger.debug( + "Connection checked out from pool " + "(checkedout=%s overflow=%s pool_size=%s)", + checked_out, + overflow, + pool_size, + ) + # QueuePool does not expose a first-party wait-started hook; treat + # checkedout >= pool_size (overflow in use) as pressure / likely wait. + if checked_out >= pool_size: + logger.warning( + "Database pool under pressure: checkedout=%s overflow=%s " + "pool_size=%s max_overflow=%s (checkout waits may exceed 1s)", + checked_out, + overflow, + pool_size, + settings.DB_MAX_OVERFLOW, + ) + + @event.listens_for(engine.sync_engine.pool, "checkin") def on_checkin(dbapi_connection, connection_record): """Handle connection check-in events.""" logger.debug("Connection checked in to pool") @@ -219,6 +240,7 @@ def on_invalidate(dbapi_connection, connection_record, exception): """Handle connection invalidation events.""" logger.warning(f"Database connection invalidated: {exception}") + setup_pool_event_listeners() @@ -231,7 +253,7 @@ async def prewarm_connection_pool(): logger.info("Starting connection pool prewarming...") try: # Warm the base connection pool. - connections_to_warm = min(ProcessingConstants.DB_POOL_SIZE, 5) + connections_to_warm = min(settings.DB_POOL_SIZE, 5) tasks = [] for _ in range(connections_to_warm): diff --git a/packages/shared-python/shared/tests/test_database_pool_config.py b/packages/shared-python/shared/tests/test_database_pool_config.py new file mode 100644 index 000000000..83cac6680 --- /dev/null +++ b/packages/shared-python/shared/tests/test_database_pool_config.py @@ -0,0 +1,22 @@ +"""Unit tests for DatabaseConfig pool defaults.""" + +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from shared.core.config.database import DatabaseConfig + + +def test_database_pool_defaults_are_fifty() -> None: + config = DatabaseConfig( + DATABASE_URL="postgresql+asyncpg://test:test@localhost/test", + ) + assert config.DB_POOL_SIZE == 50 + assert config.DB_MAX_OVERFLOW == 50 From 1b93fe5885c69c71e48827eea92633af704cf30c Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 25 Jul 2026 20:18:46 +0800 Subject: [PATCH 2/4] fix: harden QueuePool hygiene with atomic last-used debounce Use Redis SET NX for last-used debounce, require the request session in resolve_current_user, and avoid double-commit when TierService owns the session. Co-authored-by: Cursor --- .../auth/api_key_authentication_service.py | 28 ++++++++----------- .../rate_limit/job_admission_service.py | 2 +- .../app/services/rate_limit/tier_service.py | 20 +++++++++---- .../unit/test_job_poll_session_hygiene.py | 15 ++++------ .../shared/services/redis/redis_service.py | 22 +++++++++++++++ 5 files changed, 54 insertions(+), 33 deletions(-) diff --git a/apps/api/app/services/auth/api_key_authentication_service.py b/apps/api/app/services/auth/api_key_authentication_service.py index bf851db68..462afd6e0 100644 --- a/apps/api/app/services/auth/api_key_authentication_service.py +++ b/apps/api/app/services/auth/api_key_authentication_service.py @@ -174,21 +174,28 @@ async def _update_last_used_best_effort( ) -> None: """Update last_used_at on the request session without a nested checkout. - Redis debounce skips redundant writes within the debounce window so job - polls do not compete for QueuePool capacity via create_task+get_db_context. + 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: - if await redis_service.exists(debounce_key): + acquired = await redis_service.set_nx( + debounce_key, + "1", + ex=_LAST_USED_DEBOUNCE_SECONDS, + ) + if not acquired: return except Exception: logger.warning( - "api_key_authentication: failed to read last-used debounce for api_key_id={}", + "api_key_authentication: last-used debounce failed for api_key_id={}; " + "updating anyway", api_key_id, ) try: 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( @@ -198,16 +205,3 @@ async def _update_last_used_best_effort( await session.rollback() except Exception: pass - return - - try: - await redis_service.set( - debounce_key, - "1", - ttl=_LAST_USED_DEBOUNCE_SECONDS, - ) - except Exception: - logger.warning( - "api_key_authentication: failed to write last-used debounce for api_key_id={}", - api_key_id, - ) diff --git a/apps/api/app/services/rate_limit/job_admission_service.py b/apps/api/app/services/rate_limit/job_admission_service.py index 0c0321eed..1e46211da 100644 --- a/apps/api/app/services/rate_limit/job_admission_service.py +++ b/apps/api/app/services/rate_limit/job_admission_service.py @@ -33,7 +33,7 @@ async def resolve_current_user( *, route_context: RouteAdmissionContext, user_id: str, - db: AsyncSession | None = None, + db: AsyncSession, ) -> CurrentUser: user_tier = await TierService.get_tier(user_id, session=db) self._route_policy_service.enforce_guest_api_key_scope( diff --git a/apps/api/app/services/rate_limit/tier_service.py b/apps/api/app/services/rate_limit/tier_service.py index dc1255f16..e62a82124 100644 --- a/apps/api/app/services/rate_limit/tier_service.py +++ b/apps/api/app/services/rate_limit/tier_service.py @@ -49,19 +49,29 @@ async def get_tier( return cached_tier if session is not None: - user_tier = await TierService._resolve_tier_from_db(session, user_id) + 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) -> str: + 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) @@ -70,9 +80,9 @@ async def _resolve_tier_from_db(session: AsyncSession, user_id: str) -> str: session, user_id, ) - # Persist first-use init when reusing the request session - # (get_db_context commits on exit when we own the session). - await session.commit() + # Request-scoped sessions do not auto-commit; get_db_context does. + if commit_on_initialize: + await session.commit() return user_tier @staticmethod diff --git a/apps/api/tests/unit/test_job_poll_session_hygiene.py b/apps/api/tests/unit/test_job_poll_session_hygiene.py index 8df95befb..059c5c80d 100644 --- a/apps/api/tests/unit/test_job_poll_session_hygiene.py +++ b/apps/api/tests/unit/test_job_poll_session_hygiene.py @@ -98,7 +98,7 @@ async def test_validate_api_key_updates_last_used_on_same_session() -> None: session.commit = AsyncMock() redis_service = AsyncMock() redis_service.get = AsyncMock(return_value=None) - redis_service.exists = AsyncMock(return_value=False) + redis_service.set_nx = AsyncMock(return_value=True) redis_service.set = AsyncMock() redis_service.sadd = AsyncMock() redis_service.ttl = AsyncMock(return_value=-2) @@ -125,10 +125,6 @@ async def test_validate_api_key_updates_last_used_on_same_session() -> None: "app.services.auth.api_key_authentication_service.hash_api_key", return_value="hash-1", ), - patch( - "app.services.auth.api_key_authentication_service.get_db_context", - create=True, - ) as get_db_context_mock, patch("asyncio.create_task") as create_task_mock, ): user_id = await service.validate_api_key(session, "kw_test_key") @@ -137,11 +133,10 @@ async def test_validate_api_key_updates_last_used_on_same_session() -> None: repository.update_last_used.assert_awaited_once_with(session, "key-1") session.commit.assert_awaited_once() create_task_mock.assert_not_called() - get_db_context_mock.assert_not_called() - redis_service.set.assert_any_await( + redis_service.set_nx.assert_awaited_once_with( "api-key:last-used-debounce:key-1", "1", - ttl=300, + ex=300, ) @@ -151,7 +146,7 @@ async def test_validate_api_key_skips_last_used_when_debounced() -> None: session.commit = AsyncMock() redis_service = AsyncMock() redis_service.get = AsyncMock(return_value=None) - redis_service.exists = AsyncMock(return_value=True) + redis_service.set_nx = AsyncMock(return_value=False) redis_service.set = AsyncMock() redis_service.sadd = AsyncMock() redis_service.ttl = AsyncMock(return_value=-2) @@ -203,7 +198,7 @@ async def __aexit__(self, *args: object) -> None: redis_service = AsyncMock() redis_service.get = AsyncMock(side_effect=[None, None]) # api-key miss, tier miss - redis_service.exists = AsyncMock(return_value=False) + redis_service.set_nx = AsyncMock(return_value=True) redis_service.set = AsyncMock() redis_service.sadd = AsyncMock() redis_service.ttl = AsyncMock(return_value=-2) diff --git a/packages/shared-python/shared/services/redis/redis_service.py b/packages/shared-python/shared/services/redis/redis_service.py index fa9d39fa9..e83e3a7bc 100644 --- a/packages/shared-python/shared/services/redis/redis_service.py +++ b/packages/shared-python/shared/services/redis/redis_service.py @@ -111,6 +111,28 @@ async def _operation(): original_exception=e, ) + async def set_nx(self, key: str, value: str, ex: int) -> bool: + """Atomic SET NX EX — set only if the key does not already exist. + + Returns ``True`` if the key was written, ``False`` if it already existed. + Does not JSON-encode the value and does not fall back to a default TTL. + """ + try: + client = await self._get_client() + full_key = self._build_key(key) + + async def _operation(): + return await client.set(full_key, value, nx=True, ex=ex) + + return bool(await self._execute_with_retry(_operation)) + except Exception as e: + logger.error(f"Redis SET NX operation failed: {e}") + raise RedisOperationError( + internal_message=f"SET NX operation failed: {str(e)}", + operation="SET_NX", + original_exception=e, + ) + async def get(self, key: str, default: Any = None) -> Any: """Get a key value.""" try: From 01469e6ff52950ec8bfe6fef7ab3f5eb0a00bf62 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 25 Jul 2026 20:59:40 +0800 Subject: [PATCH 3/4] fix: address QueuePool review CI and CodeQL findings Explain the best-effort rollback swallow for CodeQL, narrow pool pressure logging to QueuePool for pyright, and stop unit tests from shadowing the worker `app` package during shared pytest collection. Co-authored-by: Cursor --- .../auth/api_key_authentication_service.py | 2 + apps/api/tests/unit/__init__.py | 0 .../unit/test_job_poll_session_hygiene.py | 62 +++++++++++++++---- .../shared-python/shared/core/database.py | 7 ++- 4 files changed, 58 insertions(+), 13 deletions(-) delete mode 100644 apps/api/tests/unit/__init__.py diff --git a/apps/api/app/services/auth/api_key_authentication_service.py b/apps/api/app/services/auth/api_key_authentication_service.py index 462afd6e0..2625892d4 100644 --- a/apps/api/app/services/auth/api_key_authentication_service.py +++ b/apps/api/app/services/auth/api_key_authentication_service.py @@ -204,4 +204,6 @@ async def _update_last_used_best_effort( try: await session.rollback() except Exception: + # Rollback itself can fail if the session is already closed; + # ignore so last-used remains best-effort. pass diff --git a/apps/api/tests/unit/__init__.py b/apps/api/tests/unit/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/api/tests/unit/test_job_poll_session_hygiene.py b/apps/api/tests/unit/test_job_poll_session_hygiene.py index 059c5c80d..dd9807364 100644 --- a/apps/api/tests/unit/test_job_poll_session_hygiene.py +++ b/apps/api/tests/unit/test_job_poll_session_hygiene.py @@ -2,7 +2,7 @@ from __future__ import annotations -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -12,23 +12,35 @@ ensure_import_paths, ) +# Defer importing apps/api `app` until test bodies run. Module-level imports +# would cache API's `app` package and break apps/worker contract collection in +# the same pytest process (both packages are named `app`). configure_import_environment() ensure_import_paths() -from app.services.auth.api_key_authentication_service import ( # noqa: E402 - APIKeyAuthenticationService, -) -from app.services.rate_limit.data_structures import ( # noqa: E402 - RouteAdmissionContext, -) -from app.services.rate_limit.job_admission_service import ( # noqa: E402 - JobAdmissionService, -) -from app.services.rate_limit.tier_service import TierService # noqa: E402 + +def _load_api_modules() -> tuple[ModuleType, ModuleType, ModuleType, ModuleType]: + ensure_import_paths() + from app.services.auth import api_key_authentication_service + from app.services.rate_limit import ( + data_structures, + job_admission_service, + tier_service, + ) + + return ( + api_key_authentication_service, + data_structures, + job_admission_service, + tier_service, + ) @pytest.mark.asyncio async def test_get_tier_reuses_provided_session_without_get_db_context() -> None: + _, _, _, tier_service = _load_api_modules() + TierService = tier_service.TierService + session = AsyncMock() redis_service = AsyncMock() redis_service.get = AsyncMock(return_value=None) @@ -57,6 +69,11 @@ async def test_get_tier_reuses_provided_session_without_get_db_context() -> None @pytest.mark.asyncio async def test_resolve_current_user_passes_request_session_to_get_tier() -> None: + _, data_structures, job_admission_service, tier_service = _load_api_modules() + RouteAdmissionContext = data_structures.RouteAdmissionContext + JobAdmissionService = job_admission_service.JobAdmissionService + TierService = tier_service.TierService + session = AsyncMock() route_context = RouteAdmissionContext( method="GET", @@ -94,6 +111,11 @@ async def test_resolve_current_user_passes_request_session_to_get_tier() -> None @pytest.mark.asyncio async def test_validate_api_key_updates_last_used_on_same_session() -> None: + api_key_authentication_service, _, _, _ = _load_api_modules() + APIKeyAuthenticationService = ( + api_key_authentication_service.APIKeyAuthenticationService + ) + session = AsyncMock() session.commit = AsyncMock() redis_service = AsyncMock() @@ -142,6 +164,11 @@ async def test_validate_api_key_updates_last_used_on_same_session() -> None: @pytest.mark.asyncio async def test_validate_api_key_skips_last_used_when_debounced() -> None: + api_key_authentication_service, _, _, _ = _load_api_modules() + APIKeyAuthenticationService = ( + api_key_authentication_service.APIKeyAuthenticationService + ) + session = AsyncMock() session.commit = AsyncMock() redis_service = AsyncMock() @@ -184,6 +211,19 @@ async def test_validate_api_key_skips_last_used_when_debounced() -> None: @pytest.mark.asyncio async def test_job_poll_auth_path_uses_single_session_factory_checkout() -> None: """End-to-end hygiene: tier + last-used must not call get_db_context.""" + ( + api_key_authentication_service, + data_structures, + job_admission_service, + tier_service, + ) = _load_api_modules() + APIKeyAuthenticationService = ( + api_key_authentication_service.APIKeyAuthenticationService + ) + RouteAdmissionContext = data_structures.RouteAdmissionContext + JobAdmissionService = job_admission_service.JobAdmissionService + TierService = tier_service.TierService + session = AsyncMock() session.commit = AsyncMock() checkout_count = {"n": 0} diff --git a/packages/shared-python/shared/core/database.py b/packages/shared-python/shared/core/database.py index 80ee99188..bfea49902 100644 --- a/packages/shared-python/shared/core/database.py +++ b/packages/shared-python/shared/core/database.py @@ -12,7 +12,7 @@ create_async_engine, ) from sqlalchemy.orm import declarative_base -from sqlalchemy.pool import NullPool +from sqlalchemy.pool import NullPool, QueuePool from shared.core.config import settings from shared.core.constants import ProcessingConstants @@ -208,9 +208,12 @@ def on_connect(dbapi_connection, connection_record): def on_checkout(dbapi_connection, connection_record, connection_proxy): """Handle connection checkout events and surface pool pressure.""" pool = engine.sync_engine.pool + if not isinstance(pool, QueuePool): + logger.debug("Connection checked out from pool") + return checked_out = pool.checkedout() overflow = pool.overflow() - pool_size = pool.size() if hasattr(pool, "size") else settings.DB_POOL_SIZE + pool_size = pool.size() logger.debug( "Connection checked out from pool " "(checkedout=%s overflow=%s pool_size=%s)", From 13312fff5ada4c1b690419c28e8a668b8e4ef380 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 25 Jul 2026 21:07:54 +0800 Subject: [PATCH 4/4] fix: isolate API unit app imports from worker contracts Harden worker contract eviction to drop deepest API-shadowed modules first, and make job-poll unit tests reassert the API import root then clear cached API app modules so shared pytest runs stay green. Co-authored-by: Cursor --- .../unit/test_job_poll_session_hygiene.py | 52 ++++++++++++++++++- apps/worker/tests/contract/conftest.py | 24 ++++++--- 2 files changed, 69 insertions(+), 7 deletions(-) diff --git a/apps/api/tests/unit/test_job_poll_session_hygiene.py b/apps/api/tests/unit/test_job_poll_session_hygiene.py index dd9807364..3e8ce9b6a 100644 --- a/apps/api/tests/unit/test_job_poll_session_hygiene.py +++ b/apps/api/tests/unit/test_job_poll_session_hygiene.py @@ -2,6 +2,8 @@ from __future__ import annotations +import sys +from pathlib import Path from types import ModuleType, SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -18,9 +20,49 @@ configure_import_environment() ensure_import_paths() +_API_ROOT = str(Path(__file__).resolve().parents[2]) -def _load_api_modules() -> tuple[ModuleType, ModuleType, ModuleType, ModuleType]: + +def _prioritize_api_import_root() -> None: + """Keep apps/api ahead of apps/worker for the shared `app` package name.""" ensure_import_paths() + if _API_ROOT in sys.path: + sys.path.remove(_API_ROOT) + sys.path.insert(0, _API_ROOT) + + +def _is_api_app_module(module: ModuleType | None) -> bool: + if module is None: + return False + module_file = getattr(module, "__file__", None) + if isinstance(module_file, str) and module_file.startswith(_API_ROOT): + return True + module_paths = getattr(module, "__path__", ()) + try: + return any(str(path).startswith(_API_ROOT) for path in module_paths) + except KeyError: + return False + + +def _drop_non_api_app_modules() -> None: + for module_name in sorted(sys.modules, key=len, reverse=True): + if module_name != "app" and not module_name.startswith("app."): + continue + if not _is_api_app_module(sys.modules.get(module_name)): + sys.modules.pop(module_name, None) + + +def _drop_api_app_modules() -> None: + for module_name in sorted(sys.modules, key=len, reverse=True): + if module_name != "app" and not module_name.startswith("app."): + continue + if _is_api_app_module(sys.modules.get(module_name)): + sys.modules.pop(module_name, None) + + +def _load_api_modules() -> tuple[ModuleType, ModuleType, ModuleType, ModuleType]: + _prioritize_api_import_root() + _drop_non_api_app_modules() from app.services.auth import api_key_authentication_service from app.services.rate_limit import ( data_structures, @@ -36,6 +78,14 @@ def _load_api_modules() -> tuple[ModuleType, ModuleType, ModuleType, ModuleType] ) +@pytest.fixture(autouse=True) +def _clear_api_app_modules_after_unit_test(): + """Avoid leaving API's `app` package cached for later worker contract tests.""" + yield + _drop_api_app_modules() + + + @pytest.mark.asyncio async def test_get_tier_reuses_provided_session_without_get_db_context() -> None: _, _, _, tier_service = _load_api_modules() diff --git a/apps/worker/tests/contract/conftest.py b/apps/worker/tests/contract/conftest.py index f9738370f..1a9ae74a9 100644 --- a/apps/worker/tests/contract/conftest.py +++ b/apps/worker/tests/contract/conftest.py @@ -37,7 +37,14 @@ def _module_loaded_from(module_name: str, root: Path) -> bool: return True module_paths = getattr(module, "__path__", ()) - return any(str(module_path).startswith(root_value) for module_path in module_paths) + try: + return any( + str(module_path).startswith(root_value) for module_path in module_paths + ) + except KeyError: + # Namespace path iteration can raise if a parent package was already + # removed from sys.modules mid-eviction. + return False def _ensure_worker_import_context() -> None: @@ -46,11 +53,16 @@ def _ensure_worker_import_context() -> None: sys.path.remove(worker_root_value) sys.path.insert(0, worker_root_value) - cached_module_names = list(sys.modules) - for module_name in cached_module_names: - if module_name == "app" or module_name.startswith("app."): - if _module_loaded_from(module_name, _API_ROOT): - sys.modules.pop(module_name, None) + cached_module_names = [ + module_name + for module_name in sys.modules + if module_name == "app" or module_name.startswith("app.") + ] + # Evict deepest modules first so namespace __path__ checks never observe a + # missing parent `app` entry while walking API-shadowed packages. + for module_name in sorted(cached_module_names, key=len, reverse=True): + if _module_loaded_from(module_name, _API_ROOT): + sys.modules.pop(module_name, None) @pytest.fixture(autouse=True)