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..2625892d4 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,44 @@ 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 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: + # Rollback itself can fail if the session is already closed; + # ignore so last-used remains best-effort. + pass 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..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,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, diff --git a/apps/api/app/services/rate_limit/tier_service.py b/apps/api/app/services/rate_limit/tier_service.py index 91bb73d00..e62a82124 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,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. 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..3e8ce9b6a --- /dev/null +++ b/apps/api/tests/unit/test_job_poll_session_hygiene.py @@ -0,0 +1,358 @@ +"""Regression tests: job-poll auth must not open nested DB checkouts.""" + +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 + +import pytest +from tests.support.import_environment import ( + configure_import_environment, + 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() + +_API_ROOT = str(Path(__file__).resolve().parents[2]) + + +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, + job_admission_service, + tier_service, + ) + + return ( + api_key_authentication_service, + data_structures, + job_admission_service, + tier_service, + ) + + +@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() + TierService = tier_service.TierService + + 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: + _, 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", + 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: + api_key_authentication_service, _, _, _ = _load_api_modules() + APIKeyAuthenticationService = ( + api_key_authentication_service.APIKeyAuthenticationService + ) + + session = AsyncMock() + session.commit = AsyncMock() + redis_service = AsyncMock() + redis_service.get = AsyncMock(return_value=None) + redis_service.set_nx = 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", + ), + 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() + redis_service.set_nx.assert_awaited_once_with( + "api-key:last-used-debounce:key-1", + "1", + ex=300, + ) + + +@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() + redis_service.get = AsyncMock(return_value=None) + redis_service.set_nx = 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", + ), + ): + 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.""" + ( + 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} + + 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.set_nx = 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) + 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/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) 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..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 @@ -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,36 @@ 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 + if not isinstance(pool, QueuePool): + logger.debug("Connection checked out from pool") + return + checked_out = pool.checkedout() + overflow = pool.overflow() + pool_size = 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 +243,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 +256,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/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: 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