diff --git a/docs/openapi.json b/docs/openapi.json index 75cbfafad..5f81bda4b 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -944,6 +944,16 @@ } } } + }, + "503": { + "description": "Service is not alive", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/LivenessResponse" + } + } + } } } } @@ -1708,6 +1718,17 @@ "alive": { "type": "boolean", "title": "Alive" + }, + "reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Reason" } }, "type": "object", @@ -1715,7 +1736,7 @@ "alive" ], "title": "LivenessResponse", - "description": "Model representing a response to a liveness request.\n\nAttributes:\n alive: If app is alive.\n\nExample:\n ```python\n liveness_response = LivenessResponse(alive=True)\n ```", + "description": "Model representing a response to a liveness request.\n\nAttributes:\n alive: If app is alive.\n reason: Optional reason when not alive.\n\nExample:\n ```python\n liveness_response = LivenessResponse(alive=True)\n ```", "examples": [ { "alive": true diff --git a/ols/app/endpoints/health.py b/ols/app/endpoints/health.py index 10bb72c0a..671a337b1 100644 --- a/ols/app/endpoints/health.py +++ b/ols/app/endpoints/health.py @@ -9,7 +9,7 @@ import time from typing import Any -from fastapi import APIRouter, HTTPException, status +from fastapi import APIRouter, HTTPException, Response, status from langchain_core.messages.ai import AIMessage from ols import config @@ -18,6 +18,7 @@ NotAvailableResponse, ReadinessResponse, ) +from ols.src.cache.postgres_cache import PostgresCache from ols.src.llms.llm_loader import load_llm router = APIRouter(tags=["health"]) @@ -122,10 +123,20 @@ def readiness_probe_get_method() -> ReadinessResponse: "description": "Service is alive", "model": LivenessResponse, }, + 503: { + "description": "Service is not alive", + "model": LivenessResponse, + }, } @router.get("/liveness", responses=get_liveness_responses) -def liveness_probe_get_method() -> LivenessResponse: +def liveness_probe_get_method(response: Response) -> LivenessResponse: """Live status of service.""" + cache = config._conversation_cache + if isinstance(cache, PostgresCache): + threshold = config.ols_config.liveness_db_failure_threshold + if cache.consecutive_failures >= threshold: + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE + return LivenessResponse(alive=False, reason="database unreachable") return LivenessResponse(alive=True) diff --git a/ols/app/models/config.py b/ols/app/models/config.py index 83009a372..1a8e7e9c0 100644 --- a/ols/app/models/config.py +++ b/ols/app/models/config.py @@ -11,6 +11,7 @@ BaseModel, Field, FilePath, + NonNegativeInt, PositiveInt, PrivateAttr, field_validator, @@ -52,6 +53,21 @@ def validate_tool_round_cap_fraction_config(v: float) -> float: return v +def validate_liveness_db_failure_threshold(raw_value: Any) -> int: + """Validate ``liveness_db_failure_threshold`` for ``OLSConfig``.""" + try: + threshold = int(raw_value) + except (TypeError, ValueError) as e: + raise checks.InvalidConfigurationError( + f"liveness_db_failure_threshold must be a positive integer, got {raw_value!r}" + ) from e + if threshold < 1: + raise checks.InvalidConfigurationError( + f"liveness_db_failure_threshold must be at least 1, got {threshold}" + ) + return threshold + + class ModelParameters(BaseModel): """Model parameters.""" @@ -798,6 +814,12 @@ class PostgresConfig(BaseModel): gss_encmode: str = constants.POSTGRES_CACHE_GSSENCMODE ca_cert_path: Optional[FilePath] = None max_entries: PositiveInt = constants.POSTGRES_CACHE_MAX_ENTRIES + statement_timeout: NonNegativeInt = constants.POSTGRES_STATEMENT_TIMEOUT + lock_timeout: PositiveInt = constants.POSTGRES_LOCK_TIMEOUT + health_check_interval: PositiveInt = constants.CACHE_HEALTH_CHECK_INTERVAL + health_check_connect_timeout: PositiveInt = ( + constants.CACHE_HEALTH_CHECK_CONNECT_TIMEOUT + ) tls_security_profile: Optional["TLSSecurityProfile"] = None def __init__(self, **data: Any) -> None: @@ -1160,6 +1182,8 @@ class OLSConfig(BaseModel): offload_storage_path: str = constants.DEFAULT_OFFLOAD_STORAGE_PATH + liveness_db_failure_threshold: int = constants.LIVENESS_DB_FAILURE_THRESHOLD + def __init__( self, data: Optional[dict] = None, ignore_missing_certs: bool = False ) -> None: @@ -1234,6 +1258,12 @@ def __init__( "offload_storage_path", constants.DEFAULT_OFFLOAD_STORAGE_PATH ) + self.liveness_db_failure_threshold = validate_liveness_db_failure_threshold( + data.get( + "liveness_db_failure_threshold", constants.LIVENESS_DB_FAILURE_THRESHOLD + ) + ) + def _propagate_tls_profile(self) -> None: """Set the TLS security profile on all PostgresConfig instances.""" if ( diff --git a/ols/app/models/models.py b/ols/app/models/models.py index 29883e115..f1bba8dfc 100644 --- a/ols/app/models/models.py +++ b/ols/app/models/models.py @@ -475,6 +475,7 @@ class LivenessResponse(BaseModel): Attributes: alive: If app is alive. + reason: Optional reason when not alive. Example: ```python @@ -483,6 +484,7 @@ class LivenessResponse(BaseModel): """ alive: bool + reason: Optional[str] = None # provides examples for /docs endpoint model_config = { diff --git a/ols/constants.py b/ols/constants.py index 4245a3761..5898f67a4 100644 --- a/ols/constants.py +++ b/ols/constants.py @@ -128,6 +128,8 @@ class GenericLLMParameters: POSTGRES_CACHE_DBNAME = "cache" POSTGRES_CACHE_USER = "postgres" POSTGRES_CACHE_MAX_ENTRIES = 1000 +POSTGRES_STATEMENT_TIMEOUT = 5000 +POSTGRES_LOCK_TIMEOUT = 10 # look at https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNECT-SSLMODE # for all possible options @@ -137,6 +139,10 @@ class GenericLLMParameters: # for all possible options POSTGRES_CACHE_GSSENCMODE = "prefer" +CACHE_HEALTH_CHECK_INTERVAL = 30 +CACHE_HEALTH_CHECK_CONNECT_TIMEOUT = 10 +LIVENESS_DB_FAILURE_THRESHOLD = 3 + # default indentity for local testing and deployment # "nil" UUID is used on purpose, because it will be easier to diff --git a/ols/src/cache/postgres_cache.py b/ols/src/cache/postgres_cache.py index a936ba906..a948ab8ed 100644 --- a/ols/src/cache/postgres_cache.py +++ b/ols/src/cache/postgres_cache.py @@ -17,6 +17,7 @@ from ols.src.cache.cache import Cache from ols.src.cache.cache_error import CacheError from ols.utils.postgres import PostgresBase, connection +from ols.utils.ssl import libpq_tls_params logger = logging.getLogger(__name__) @@ -154,9 +155,24 @@ class PostgresCache(Cache, PostgresBase): def __init__(self, config: PostgresConfig) -> None: """Create a new instance of Postgres cache.""" self._tx_lock = threading.Lock() + self._lock_timeout = config.lock_timeout self.capacity = config.max_entries + + self._health_status = True + self._consecutive_failures = 0 + self._health_lock = threading.Lock() + self._health_check_interval = config.health_check_interval + self._health_check_connect_timeout = config.health_check_connect_timeout + self._health_connection: Any = None + self._shutdown_event = threading.Event() + super().__init__(config) + self._health_thread = threading.Thread( + target=self._health_check_loop, daemon=True + ) + self._health_thread.start() + @property def _ddl_statements(self) -> list[str]: """Return DDL statements for cache tables and indexes.""" @@ -166,6 +182,107 @@ def _ddl_statements(self) -> list[str]: self.CREATE_INDEX, ] + def _mark_unhealthy(self) -> None: + """Mark health status as unhealthy (dual-feed from cache operations).""" + with self._health_lock: + self._consecutive_failures += 1 + self._health_status = False + + def _connect_health(self) -> None: + """Establish a dedicated lightweight connection for health checks.""" + if self._health_connection is not None: + try: + self._health_connection.close() + except Exception as close_err: + logger.debug("Failed to close old health connection: %s", close_err) + self._health_connection = None + config = self.connection_config + connect_kwargs: dict[str, Any] = { + "host": config.host, + "port": config.port, + "user": config.user, + "password": config.password, + "dbname": config.dbname, + "sslmode": config.ssl_mode, + "sslrootcert": config.ca_cert_path, + "gssencmode": config.gss_encmode, + "connect_timeout": self._health_check_connect_timeout, + **libpq_tls_params(config.tls_security_profile), + } + self._health_connection = psycopg2.connect(**connect_kwargs) + self._health_connection.autocommit = True + with self._health_connection.cursor() as cursor: + cursor.execute( + "SET statement_timeout = %s", (str(config.statement_timeout),) + ) + + def shutdown(self) -> None: + """Stop the health-check thread and close its connection.""" + self._shutdown_event.set() + self._health_thread.join(timeout=self._health_check_connect_timeout) + if self._health_connection is not None: + try: + self._health_connection.close() + except Exception as close_err: + logger.debug( + "Failed to close health connection on shutdown: %s", close_err + ) + self._health_connection = None + + def _health_check_loop(self) -> None: + """Background loop that periodically checks DB connectivity.""" + first_check = True + while not self._shutdown_event.is_set(): + try: + if self._health_connection is None or self._health_connection.closed: + self._connect_health() + with self._health_connection.cursor() as cursor: + cursor.execute("SELECT 1") + with self._health_lock: + if not self._health_status: + logger.info("Database connection recovered") + self._health_status = True + self._consecutive_failures = 0 + except Exception as e: + if first_check: + logger.info("Initial health check failed, will retry: %s", e) + else: + with self._health_lock: + self._consecutive_failures += 1 + self._health_status = False + failures = self._consecutive_failures + logger.warning( + "Health check failed (consecutive=%d): %s", + failures, + e, + ) + try: + self._connect_health() + except Exception as reconnect_err: + logger.debug("Health reconnect attempt failed: %s", reconnect_err) + first_check = False + self._shutdown_event.wait(self._health_check_interval) + + def _safe_rollback(self) -> None: + """Rollback the current transaction, ignoring errors on dead connections.""" + try: + self.connection.rollback() + except Exception as e: + logger.debug("Rollback failed on dead connection: %s", e) + + def _safe_set_autocommit(self) -> None: + """Restore autocommit, ignoring errors on dead connections.""" + try: + self.connection.autocommit = True + except Exception as e: + logger.debug("Failed to restore autocommit: %s", e) + + def _acquire_lock(self) -> None: + """Acquire _tx_lock with bounded timeout.""" + # pylint: disable-next=consider-using-with + if not self._tx_lock.acquire(timeout=self._lock_timeout): + raise CacheError("lock acquisition timeout") + @connection def get( self, user_id: str, conversation_id: str, skip_user_id_check: bool = False @@ -183,7 +300,8 @@ def get( # just check if user_id and conversation_id are UUIDs super().construct_key(user_id, conversation_id, skip_user_id_check) - with self._tx_lock: + self._acquire_lock() + try: with self.connection.cursor() as cursor: try: value = PostgresCache._select(cursor, user_id, conversation_id) @@ -191,9 +309,13 @@ def get( return [] history = [CacheEntry.from_dict(ce) for ce in value] return history + except (psycopg2.OperationalError, psycopg2.InterfaceError): + raise except psycopg2.DatabaseError as e: logger.error("PostgresCache.get %s", e) raise CacheError("PostgresCache.get", e) from e + finally: + self._tx_lock.release() @connection def insert_or_append( @@ -217,7 +339,8 @@ def insert_or_append( # pg_advisory_xact_lock would be released immediately after the first # execute. Disable autocommit for a real multi-statement transaction. # The lock serialises concurrent callers that share this connection. - with self._tx_lock: + self._acquire_lock() + try: self.connection.autocommit = False with self.connection.cursor() as cursor: try: @@ -247,12 +370,17 @@ def insert_or_append( ) PostgresCache._cleanup(cursor, self.capacity) self.connection.commit() + except (psycopg2.OperationalError, psycopg2.InterfaceError): + self._safe_rollback() + raise except psycopg2.DatabaseError as e: - self.connection.rollback() + self._safe_rollback() logger.error("PostgresCache.insert_or_append: %s", e) raise CacheError("PostgresCache.insert_or_append", e) from e finally: - self.connection.autocommit = True + self._safe_set_autocommit() + finally: + self._tx_lock.release() @connection def delete( @@ -269,7 +397,8 @@ def delete( bool: True if the conversation was deleted, False if not found. """ - with self._tx_lock: + self._acquire_lock() + try: self.connection.autocommit = False with self.connection.cursor() as cursor: try: @@ -280,12 +409,17 @@ def delete( ) self.connection.commit() return deleted + except (psycopg2.OperationalError, psycopg2.InterfaceError): + self._safe_rollback() + raise except psycopg2.DatabaseError as e: - self.connection.rollback() + self._safe_rollback() logger.error("PostgresCache.delete: %s", e) raise CacheError("PostgresCache.delete", e) from e finally: - self.connection.autocommit = True + self._safe_set_autocommit() + finally: + self._tx_lock.release() @connection def list( @@ -302,7 +436,8 @@ def list( topic_summary, last_message_timestamp, and message_count. """ - with self._tx_lock: + self._acquire_lock() + try: with self.connection.cursor() as cursor: try: cursor.execute( @@ -318,9 +453,13 @@ def list( ) for row in rows ] + except (psycopg2.OperationalError, psycopg2.InterfaceError): + raise except psycopg2.DatabaseError as e: logger.error("PostgresCache.list: %s", e) raise CacheError("PostgresCache.list", e) from e + finally: + self._tx_lock.release() @connection def set_topic_summary( @@ -338,35 +477,39 @@ def set_topic_summary( topic_summary: The topic summary to store. skip_user_id_check: Skip user_id suid check. """ - with self._tx_lock: + self._acquire_lock() + try: with self.connection.cursor() as cursor: try: cursor.execute( PostgresCache.INSERT_OR_UPDATE_TOPIC_SUMMARY_STATEMENT, (user_id, conversation_id, topic_summary), ) + except (psycopg2.OperationalError, psycopg2.InterfaceError): + raise except psycopg2.DatabaseError as e: logger.error("PostgresCache.set_topic_summary: %s", e) raise CacheError("PostgresCache.set_topic_summary", e) from e + finally: + self._tx_lock.release() + + @property + def consecutive_failures(self) -> int: + """Return the number of consecutive health-check failures (thread-safe).""" + with self._health_lock: + return self._consecutive_failures def ready(self) -> bool: """Check if the cache is ready. - Postgres cache checks if the connection is alive. + Returns the background health loop's last-known status. + Non-blocking and deadlock-immune. Returns: True if the cache is ready, False otherwise. """ - # TODO: when the connection is closed and the database is back online, - # we need to reestablish the connection => implement this - if not self.connection or self.connection.closed == 1: - return False - try: - return self.connection.poll() == psycopg2.extensions.POLL_OK - except (psycopg2.OperationalError, psycopg2.InterfaceError): - # OperationalError - the once alive connection is closed - # InterfaceError - cannot reach the database server - return False + with self._health_lock: + return self._health_status @staticmethod def _select( diff --git a/ols/utils/postgres.py b/ols/utils/postgres.py index c69fca824..e900cb2a1 100644 --- a/ols/utils/postgres.py +++ b/ols/utils/postgres.py @@ -11,6 +11,7 @@ import psycopg2 from ols.app.models.config import PostgresConfig +from ols.src.cache.cache_error import CacheError from ols.utils.ssl import libpq_tls_params logger = logging.getLogger(__name__) @@ -19,13 +20,31 @@ def connection(f: Callable) -> Callable: """Ensure the object is connected before calling the wrapped method. - If the connection is lost, reconnect transparently. + On connection-level errors (OperationalError, InterfaceError), attempt + a single reconnect and retry. On SQL/data errors (DatabaseError), wrap + in CacheError and propagate immediately. """ def wrapper(connectable: Any, *args: Any, **kwargs: Any) -> Callable: - if not connectable.connected(): - connectable.connect() - return f(connectable, *args, **kwargs) + try: + if not connectable.connected(): + connectable.connect() + return f(connectable, *args, **kwargs) + except (psycopg2.OperationalError, psycopg2.InterfaceError) as e: + logger.warning( + "Connection error in %s, attempting reconnect: %s", f.__name__, e + ) + if hasattr(connectable, "_mark_unhealthy"): + connectable._mark_unhealthy() + try: + connectable.connect() + except Exception as reconnect_err: + raise CacheError( + f"reconnect failed in {f.__name__}", reconnect_err + ) from reconnect_err + return f(connectable, *args, **kwargs) + except psycopg2.DatabaseError as e: + raise CacheError(f"{f.__name__}", e) from e return wrapper @@ -84,6 +103,13 @@ def connect(self) -> None: logger.exception("Error initializing Postgres schema:\n%s", e) raise self.connection.autocommit = True + cursor = self.connection.cursor() + try: + cursor.execute( + "SET statement_timeout = %s", (str(config.statement_timeout),) + ) + finally: + cursor.close() def connected(self) -> bool: """Check if the connection to Postgres is alive.""" diff --git a/tests/integration/test_liveness_readiness.py b/tests/integration/test_liveness_readiness.py index d3ba16a30..535a3a8b1 100644 --- a/tests/integration/test_liveness_readiness.py +++ b/tests/integration/test_liveness_readiness.py @@ -38,7 +38,7 @@ def test_liveness(): """Test handler for /liveness REST API endpoint.""" response = pytest.client.get("/liveness") assert response.status_code == requests.codes.ok - assert response.json() == {"alive": True} + assert response.json() == {"alive": True, "reason": None} def test_readiness(): diff --git a/tests/unit/app/endpoints/test_health.py b/tests/unit/app/endpoints/test_health.py index f5278c659..17648afdc 100644 --- a/tests/unit/app/endpoints/test_health.py +++ b/tests/unit/app/endpoints/test_health.py @@ -17,6 +17,14 @@ from ols.app.models.config import InMemoryCacheConfig from ols.app.models.models import LivenessResponse, ReadinessResponse from ols.src.cache.in_memory_cache import InMemoryCache +from ols.src.cache.postgres_cache import PostgresCache + + +@pytest.fixture(autouse=True) +def _suppress_health_loop(): + """Prevent the background health-check thread from making real DB calls.""" + with patch.object(PostgresCache, "_health_check_loop"): + yield def mock_cache(): @@ -213,8 +221,63 @@ def test_readiness_probe_get_method_cache_not_ready(): def test_liveness_probe_get_method(): - """Test the liveness_probe function.""" - # the tested function returns constant right now - # i.e. it does not depend on application state - response = liveness_probe_get_method() - assert response == LivenessResponse(alive=True) + """Test the liveness_probe function when no postgres cache is configured.""" + mock_response = Mock() + result = liveness_probe_get_method(mock_response) + assert result == LivenessResponse(alive=True) + + +def test_liveness_probe_returns_alive_when_postgres_healthy(): + """Test liveness probe returns alive when postgres failures below threshold.""" + with patch("psycopg2.connect"): + from ols.app.models.config import PostgresConfig + + pg_config = PostgresConfig() + cache = PostgresCache(pg_config) + + with ( + patch.object(config, "_conversation_cache", cache), + patch.object(config.ols_config, "liveness_db_failure_threshold", 3), + ): + mock_response = Mock() + result = liveness_probe_get_method(mock_response) + assert result == LivenessResponse(alive=True) + + +def test_liveness_probe_returns_503_when_postgres_unhealthy(): + """Test liveness probe returns 503 when postgres failures reach threshold.""" + with patch("psycopg2.connect"): + from ols.app.models.config import PostgresConfig + + pg_config = PostgresConfig() + cache = PostgresCache(pg_config) + for _ in range(3): + cache._mark_unhealthy() + + with ( + patch.object(config, "_conversation_cache", cache), + patch.object(config.ols_config, "liveness_db_failure_threshold", 3), + ): + mock_response = Mock() + result = liveness_probe_get_method(mock_response) + assert mock_response.status_code == 503 + assert result == LivenessResponse(alive=False, reason="database unreachable") + + +def test_liveness_probe_returns_alive_when_below_threshold(): + """Test liveness probe returns alive when failures below threshold.""" + with patch("psycopg2.connect"): + from ols.app.models.config import PostgresConfig + + pg_config = PostgresConfig() + cache = PostgresCache(pg_config) + for _ in range(2): + cache._mark_unhealthy() + + with ( + patch.object(config, "_conversation_cache", cache), + patch.object(config.ols_config, "liveness_db_failure_threshold", 3), + ): + mock_response = Mock() + result = liveness_probe_get_method(mock_response) + assert result == LivenessResponse(alive=True) diff --git a/tests/unit/app/models/test_config.py b/tests/unit/app/models/test_config.py index 46416290f..5bb8ab163 100644 --- a/tests/unit/app/models/test_config.py +++ b/tests/unit/app/models/test_config.py @@ -2008,6 +2008,21 @@ def test_postgres_config_default_values(): assert postgres_config.dbname == constants.POSTGRES_CACHE_DBNAME assert postgres_config.user == constants.POSTGRES_CACHE_USER assert postgres_config.max_entries == constants.POSTGRES_CACHE_MAX_ENTRIES + assert postgres_config.statement_timeout == constants.POSTGRES_STATEMENT_TIMEOUT + assert postgres_config.lock_timeout == constants.POSTGRES_LOCK_TIMEOUT + assert ( + postgres_config.health_check_interval == constants.CACHE_HEALTH_CHECK_INTERVAL + ) + + +def test_postgres_config_custom_timeout_values(): + """Test the PostgresConfig model with custom timeout values.""" + postgres_config = PostgresConfig( + statement_timeout=10000, lock_timeout=30, health_check_interval=60 + ) + assert postgres_config.statement_timeout == 10000 + assert postgres_config.lock_timeout == 30 + assert postgres_config.health_check_interval == 60 def test_postgres_config_correct_values(): @@ -2242,6 +2257,23 @@ def test_ols_config(tmpdir): ols_config.tool_round_cap_fraction == constants.DEFAULT_TOOL_ROUND_CAP_FRACTION ) assert ols_config.offload_storage_path == constants.DEFAULT_OFFLOAD_STORAGE_PATH + assert ( + ols_config.liveness_db_failure_threshold + == constants.LIVENESS_DB_FAILURE_THRESHOLD + ) + + +def test_ols_config_with_custom_liveness_threshold(): + """Test OLSConfig with custom liveness threshold.""" + ols_config = OLSConfig( + { + "default_provider": "test_default_provider", + "default_model": "test_default_model", + "conversation_cache": {"type": "memory", "memory": {"max_entries": 100}}, + "liveness_db_failure_threshold": 5, + } + ) + assert ols_config.liveness_db_failure_threshold == 5 def test_ols_config_with_custom_offload_storage_path(): diff --git a/tests/unit/cache/test_postgres_cache.py b/tests/unit/cache/test_postgres_cache.py index 3a88dcebd..cc4bc9a29 100644 --- a/tests/unit/cache/test_postgres_cache.py +++ b/tests/unit/cache/test_postgres_cache.py @@ -1,6 +1,7 @@ """Unit tests for PostgresCache class.""" import json +from typing import Generator from unittest.mock import MagicMock, call, patch import psycopg2 @@ -13,6 +14,16 @@ from ols.src.cache.postgres_cache import PostgresCache from ols.utils import suid +pytestmark = pytest.mark.usefixtures("_suppress_health_loop") + + +@pytest.fixture(autouse=True) +def _suppress_health_loop() -> Generator[None, None, None]: + """Prevent the background health-check thread from making real DB calls.""" + with patch.object(PostgresCache, "_health_check_loop"): + yield + + user_id = suid.get_suid() conversation_id = suid.get_suid() cache_entry_1 = CacheEntry( @@ -634,7 +645,7 @@ def test_delete_operation_not_found(): def test_delete_operation_on_exception(): - """Test the Cache.delete operation when an exception is raised.""" + """Test the Cache.delete operation when exception is thrown.""" # Mock the database cursor behavior to raise an exception mock_cursor = MagicMock() mock_cursor.execute.side_effect = psycopg2.DatabaseError("PLSQL error") @@ -649,8 +660,9 @@ def test_delete_operation_on_exception(): config = PostgresConfig() cache = PostgresCache(config) - # Verify that the exception is raised - with pytest.raises(psycopg2.DatabaseError, match="PLSQL error"): + # DatabaseError is now caught by the @connection decorator and + # wrapped in CacheError (consistent with other operations) + with pytest.raises(CacheError, match="PLSQL error"): cache.delete(user_id, conversation_id) @@ -779,32 +791,76 @@ def test_cleanup_method_when_clean_performed(): mock_cursor.execute.assert_has_calls(calls, any_order=False) -def test_ready(): - """Test the Cache.ready operation.""" - # do not use real PostgreSQL instance +def test_ready_returns_health_status() -> None: + """Test that ready() returns the background health loop status.""" + with patch("psycopg2.connect"): + config = PostgresConfig() + cache = PostgresCache(config) + + # initially healthy + assert cache.ready() is True + + # mark unhealthy + cache._mark_unhealthy() + assert cache.ready() is False + + # restore healthy + with cache._health_lock: + cache._health_status = True + assert cache.ready() is True + + +def test_mark_unhealthy() -> None: + """Test that _mark_unhealthy sets health status to False.""" with patch("psycopg2.connect"): - # initialize Postgres cache config = PostgresConfig() cache = PostgresCache(config) - # mock the connection state 0 - open - cache.connection.closed = 0 - # patch the poll function to return POLL_OK - cache.connection.poll = MagicMock(return_value=psycopg2.extensions.POLL_OK) - # cache is ready - assert cache.ready() - - # mock the connection state 1 - closed - cache.connection.closed = 1 - # cache is not ready - assert not cache.ready() - - for error_type in (psycopg2.OperationalError, psycopg2.InterfaceError): - # mock the connection state 0 - open - cache.connection.closed = 0 - # patch the poll function to raise OperationalError - cache.connection.poll = MagicMock( - side_effect=error_type("Connection closed") - ) - # cache is not ready - assert not cache.ready() + assert cache.ready() is True + cache._mark_unhealthy() + assert cache.ready() is False + + +def test_lock_timeout() -> None: + """Test that lock acquisition timeout raises CacheError.""" + with patch("psycopg2.connect"): + config = PostgresConfig(lock_timeout=1) + cache = PostgresCache(config) + + # Manually acquire the lock so the next acquire times out + cache._tx_lock.acquire() + try: + with pytest.raises(CacheError, match="lock acquisition timeout"): + cache._acquire_lock() + finally: + cache._tx_lock.release() + + +def test_health_check_loop_recovers() -> None: + """Test that the health check loop sets health status on success.""" + with patch("psycopg2.connect"): + config = PostgresConfig() + cache = PostgresCache(config) + + # Mark unhealthy, then simulate a health check success + cache._mark_unhealthy() + assert cache.ready() is False + + # Simulate the health check by directly calling the relevant logic + with cache._health_lock: + cache._health_status = True + cache._consecutive_failures = 0 + assert cache.ready() is True + + +def test_consecutive_failures_tracked() -> None: + """Test that consecutive failures are tracked.""" + with patch("psycopg2.connect"): + config = PostgresConfig() + cache = PostgresCache(config) + + for _ in range(5): + cache._mark_unhealthy() + + assert cache.ready() is False + assert cache.consecutive_failures == 5 diff --git a/tests/unit/utils/test_postgres.py b/tests/unit/utils/test_postgres.py index 452fe6456..8728cb9c6 100644 --- a/tests/unit/utils/test_postgres.py +++ b/tests/unit/utils/test_postgres.py @@ -6,6 +6,7 @@ import pytest from ols.app.models.config import TLSSecurityProfile +from ols.src.cache.cache_error import CacheError from ols.utils.postgres import PostgresBase, connection @@ -30,6 +31,8 @@ def __init__(self, raise_on_call: bool = False): """Initialize connectable.""" self._connected = False self._raise_on_call = raise_on_call + self._call_count = 0 + self._unhealthy_marked = False def connected(self) -> bool: """Check connection status.""" @@ -43,6 +46,10 @@ def disconnect(self) -> None: """Drop connection.""" self._connected = False + def _mark_unhealthy(self) -> None: + """Mark as unhealthy.""" + self._unhealthy_marked = True + @connection def do_work(self) -> str: """Perform work requiring a connection.""" @@ -50,7 +57,28 @@ def do_work(self) -> str: raise RuntimeError("work failed") return "done" - def test_auto_reconnects_when_disconnected(self): + @connection + def do_work_operational_error(self) -> str: + """Raise OperationalError on first call, succeed on retry.""" + self._call_count += 1 + if self._call_count == 1: + raise psycopg2.OperationalError("connection lost") + return "recovered" + + @connection + def do_work_interface_error(self) -> str: + """Raise InterfaceError on first call, succeed on retry.""" + self._call_count += 1 + if self._call_count == 1: + raise psycopg2.InterfaceError("cannot reach server") + return "recovered" + + @connection + def do_work_database_error(self) -> str: + """Raise DatabaseError (SQL error, no retry).""" + raise psycopg2.DatabaseError("SQL syntax error") + + def test_auto_reconnects_when_disconnected(self) -> None: """Decorator calls connect() when not connected.""" c = self.Connectable() c.disconnect() @@ -60,7 +88,7 @@ def test_auto_reconnects_when_disconnected(self): assert c.connected() is True assert result == "done" - def test_does_not_reconnect_when_connected(self): + def test_does_not_reconnect_when_connected(self) -> None: """Decorator skips connect() when already connected.""" c = self.Connectable() c.connect() @@ -68,7 +96,7 @@ def test_does_not_reconnect_when_connected(self): c.do_work() mock_connect.assert_not_called() - def test_propagates_exception_after_reconnect(self): + def test_propagates_exception_after_reconnect(self) -> None: """Decorator reconnects then lets the wrapped exception propagate.""" c = self.Connectable(raise_on_call=True) c.disconnect() @@ -77,6 +105,49 @@ def test_propagates_exception_after_reconnect(self): c.do_work() assert c.connected() is True + def test_operational_error_triggers_reconnect_and_retry(self) -> None: + """OperationalError causes reconnect + retry, succeeding on second attempt.""" + c = self.Connectable() + c.connect() + result = c.do_work_operational_error() + assert result == "recovered" + assert c._unhealthy_marked is True + + def test_interface_error_triggers_reconnect_and_retry(self) -> None: + """InterfaceError causes reconnect + retry, succeeding on second attempt.""" + c = self.Connectable() + c.connect() + result = c.do_work_interface_error() + assert result == "recovered" + assert c._unhealthy_marked is True + + def test_database_error_wraps_in_cache_error_no_retry(self) -> None: + """DatabaseError is wrapped in CacheError without reconnect attempt.""" + c = self.Connectable() + c.connect() + with patch.object(c, "connect") as mock_connect: + with pytest.raises(CacheError, match="SQL syntax error"): + c.do_work_database_error() + mock_connect.assert_not_called() + assert c._unhealthy_marked is False + + def test_connection_error_reconnect_failure_raises_cache_error(self) -> None: + """When reconnect fails after OperationalError, CacheError is raised.""" + c = self.Connectable() + c.connect() + original_connect = c.connect + call_count = [0] + + def failing_connect() -> None: + call_count[0] += 1 + if call_count[0] > 0: + raise psycopg2.OperationalError("cannot connect") + original_connect() + + c.connect = failing_connect + with pytest.raises(CacheError, match="reconnect failed"): + c.do_work_operational_error() + class TestPostgresBaseConnect: """Tests for PostgresBase.connect().""" @@ -93,7 +164,8 @@ def test_connect_executes_ddl_in_order_and_commits(self): call("CREATE INDEX IF NOT EXISTS i1 ON t1 (id)"), ] ) - cursor.close.assert_called_once() + # close is called twice: once for DDL cursor, once for statement_timeout cursor + assert cursor.close.call_count == 2 mock_connect.return_value.commit.assert_called_once() def test_connect_sets_autocommit_after_init(self): @@ -103,6 +175,16 @@ def test_connect_sets_autocommit_after_init(self): assert mock_connect.return_value.autocommit is True + def test_connect_sets_statement_timeout(self): + """Statement timeout is set after autocommit is enabled.""" + mock_config = MagicMock() + mock_config.statement_timeout = 5000 + with patch("psycopg2.connect") as mock_connect: + cursor = mock_connect.return_value.cursor.return_value + FakeComponent(config=mock_config) + + cursor.execute.assert_any_call("SET statement_timeout = %s", ("5000",)) + def test_connect_closes_connection_on_ddl_failure(self): """Connection is closed and exception propagates when DDL fails.""" with patch("psycopg2.connect") as mock_connect: