From 6e37719dd06d84fd08cfbde3a1687b8053812df5 Mon Sep 17 00:00:00 2001 From: Sri Roopa Ramesh Babu Date: Wed, 24 Jun 2026 15:59:32 -0400 Subject: [PATCH 1/4] OLS-3221 Add PostgreSQL auto-recovery after DB restart Implement background health-check loop, smarter error classification in the @connection decorator, operation timeouts, and enhanced liveness/readiness probes so OLS can recover automatically when the backing PostgreSQL database is restarted. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/openapi.json | 10 ++ ols/app/endpoints/health.py | 15 +++ ols/app/models/config.py | 9 ++ ols/constants.py | 5 + ols/src/cache/postgres_cache.py | 147 +++++++++++++++++++++--- ols/utils/postgres.py | 26 ++++- tests/unit/app/endpoints/test_health.py | 66 ++++++++++- tests/unit/app/models/test_config.py | 32 ++++++ tests/unit/cache/test_postgres_cache.py | 106 +++++++++++++---- tests/unit/utils/test_postgres.py | 82 ++++++++++++- 10 files changed, 449 insertions(+), 49 deletions(-) diff --git a/docs/openapi.json b/docs/openapi.json index 75cbfafad..dce54a6c1 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" + } + } + } } } } diff --git a/ols/app/endpoints/health.py b/ols/app/endpoints/health.py index 10bb72c0a..9511110fa 100644 --- a/ols/app/endpoints/health.py +++ b/ols/app/endpoints/health.py @@ -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,24 @@ 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: """Live status of service.""" + cache = config._conversation_cache + if isinstance(cache, PostgresCache): + threshold = config.ols_config.liveness_db_failure_threshold + with cache._health_lock: + failures = cache._consecutive_failures + if failures >= threshold: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={"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..75cb13ec6 100644 --- a/ols/app/models/config.py +++ b/ols/app/models/config.py @@ -798,6 +798,9 @@ 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: int = constants.POSTGRES_STATEMENT_TIMEOUT + lock_timeout: int = constants.POSTGRES_LOCK_TIMEOUT + health_check_interval: int = constants.CACHE_HEALTH_CHECK_INTERVAL tls_security_profile: Optional["TLSSecurityProfile"] = None def __init__(self, **data: Any) -> None: @@ -1160,6 +1163,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 +1239,10 @@ def __init__( "offload_storage_path", constants.DEFAULT_OFFLOAD_STORAGE_PATH ) + self.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/constants.py b/ols/constants.py index 4245a3761..d21aac602 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,9 @@ class GenericLLMParameters: # for all possible options POSTGRES_CACHE_GSSENCMODE = "prefer" +CACHE_HEALTH_CHECK_INTERVAL = 30 +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..f19f6e537 100644 --- a/ols/src/cache/postgres_cache.py +++ b/ols/src/cache/postgres_cache.py @@ -3,6 +3,7 @@ import json import logging import threading +import time from typing import Any import psycopg2 @@ -17,6 +18,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 +156,22 @@ 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_connection: Any = None + 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 +181,82 @@ 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._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) + 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, + **libpq_tls_params(config.tls_security_profile), + } + self._health_connection = psycopg2.connect(**connect_kwargs) + self._health_connection.autocommit = True + + def _health_check_loop(self) -> None: + """Background loop that periodically checks DB connectivity.""" + while True: + time.sleep(self._health_check_interval) + 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: + 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) + + 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 +274,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 +283,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 +313,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 +344,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() 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 +371,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 +383,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() 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 +410,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 +427,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 +451,33 @@ 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() 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..fca9bf188 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: + 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,9 @@ def connect(self) -> None: logger.exception("Error initializing Postgres schema:\n%s", e) raise self.connection.autocommit = True + cursor = self.connection.cursor() + cursor.execute("SET statement_timeout = %s", (str(config.statement_timeout),)) + cursor.close() def connected(self) -> bool: """Check if the connection to Postgres is alive.""" diff --git a/tests/unit/app/endpoints/test_health.py b/tests/unit/app/endpoints/test_health.py index f5278c659..85b23ddfd 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,60 @@ 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 + """Test the liveness_probe function when no postgres cache is configured.""" response = liveness_probe_get_method() assert response == 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) + cache._consecutive_failures = 0 + + with ( + patch.object(config, "_conversation_cache", cache), + patch.object(config.ols_config, "liveness_db_failure_threshold", 3), + ): + response = liveness_probe_get_method() + assert response == 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) + cache._consecutive_failures = 3 + cache._health_status = False + + with ( + patch.object(config, "_conversation_cache", cache), + patch.object(config.ols_config, "liveness_db_failure_threshold", 3), + ): + with pytest.raises(HTTPException) as exc_info: + liveness_probe_get_method() + assert exc_info.value.status_code == 503 + assert exc_info.value.detail["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) + cache._consecutive_failures = 2 + + with ( + patch.object(config, "_conversation_cache", cache), + patch.object(config.ols_config, "liveness_db_failure_threshold", 3), + ): + response = liveness_probe_get_method() + assert response == 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..c1755295a 100644 --- a/tests/unit/cache/test_postgres_cache.py +++ b/tests/unit/cache/test_postgres_cache.py @@ -13,6 +13,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(): + """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( @@ -779,32 +789,78 @@ 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(): + """Test that ready() returns the background health loop status.""" 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() + # 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(): + """Test that _mark_unhealthy sets health status to False.""" + with patch("psycopg2.connect"): + config = PostgresConfig() + cache = PostgresCache(config) + + assert cache.ready() is True + cache._mark_unhealthy() + assert cache.ready() is False + + +def test_lock_timeout(): + """Test that lock acquisition timeout raises CacheError.""" + with patch("psycopg2.connect"): + config = PostgresConfig(lock_timeout=0) + 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(): + """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(): + """Test that consecutive failures are tracked.""" + with patch("psycopg2.connect"): + config = PostgresConfig() + cache = PostgresCache(config) + + with cache._health_lock: + cache._consecutive_failures = 5 + cache._health_status = False + + assert cache.ready() is False + with cache._health_lock: + assert cache._consecutive_failures == 5 diff --git a/tests/unit/utils/test_postgres.py b/tests/unit/utils/test_postgres.py index 452fe6456..f8e4c508a 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,6 +57,27 @@ def do_work(self) -> str: raise RuntimeError("work failed") return "done" + @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): """Decorator calls connect() when not connected.""" c = self.Connectable() @@ -77,6 +105,47 @@ def test_propagates_exception_after_reconnect(self): c.do_work() assert c.connected() is True + def test_operational_error_triggers_reconnect_and_retry(self): + """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): + """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): + """DatabaseError is wrapped in CacheError without reconnect attempt.""" + c = self.Connectable() + c.connect() + with pytest.raises(CacheError, match="SQL syntax error"): + c.do_work_database_error() + + def test_connection_error_reconnect_failure_raises_cache_error(self): + """When reconnect fails after OperationalError, CacheError is raised.""" + c = self.Connectable() + c.connect() + # Make connect() fail on reconnect attempt + original_connect = c.connect + call_count = [0] + + def failing_connect(): + 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 +162,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 +173,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: From c81358c649f81242d7054a3a708e629e5c4e5fa5 Mon Sep 17 00:00:00 2001 From: Sri Roopa Ramesh Babu Date: Thu, 2 Jul 2026 09:00:02 -0400 Subject: [PATCH 2/4] OLS-3221 Address CodeRabbit review comments - Align 503 liveness response with LivenessResponse schema by adding optional reason field to LivenessResponse model and using Response parameter to set status code - Add NonNegativeInt/PositiveInt validation for PostgresConfig timeout fields (statement_timeout, lock_timeout, health_check_interval) - Validate liveness_db_failure_threshold in OLSConfig with helper function to ensure positive integer value - Add connect_timeout and statement_timeout to health check connection - Increment _consecutive_failures in _mark_unhealthy for proper liveness threshold tracking - Use _safe_rollback in DatabaseError exception paths - Close cursor with try/finally in postgres.py for proper cleanup - Add type annotations to test fixtures and functions - Tighten no-retry assertion in test_postgres.py to verify reconnect is not called on DatabaseError Co-authored-by: Cursor --- ols/app/endpoints/health.py | 10 ++++----- ols/app/models/config.py | 28 ++++++++++++++++++++----- ols/app/models/models.py | 2 ++ ols/src/cache/postgres_cache.py | 10 +++++++-- ols/utils/postgres.py | 8 +++++-- tests/unit/app/endpoints/test_health.py | 23 +++++++++++--------- tests/unit/cache/test_postgres_cache.py | 13 ++++++------ tests/unit/utils/test_postgres.py | 24 +++++++++++---------- 8 files changed, 76 insertions(+), 42 deletions(-) diff --git a/ols/app/endpoints/health.py b/ols/app/endpoints/health.py index 9511110fa..4f5130075 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 @@ -131,7 +131,7 @@ def readiness_probe_get_method() -> ReadinessResponse: @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): @@ -139,8 +139,6 @@ def liveness_probe_get_method() -> LivenessResponse: with cache._health_lock: failures = cache._consecutive_failures if failures >= threshold: - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail={"alive": False, "reason": "database unreachable"}, - ) + 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 75cb13ec6..3278af39e 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,9 +814,9 @@ 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: int = constants.POSTGRES_STATEMENT_TIMEOUT - lock_timeout: int = constants.POSTGRES_LOCK_TIMEOUT - health_check_interval: int = constants.CACHE_HEALTH_CHECK_INTERVAL + statement_timeout: NonNegativeInt = constants.POSTGRES_STATEMENT_TIMEOUT + lock_timeout: NonNegativeInt = constants.POSTGRES_LOCK_TIMEOUT + health_check_interval: PositiveInt = constants.CACHE_HEALTH_CHECK_INTERVAL tls_security_profile: Optional["TLSSecurityProfile"] = None def __init__(self, **data: Any) -> None: @@ -1239,8 +1255,10 @@ def __init__( "offload_storage_path", constants.DEFAULT_OFFLOAD_STORAGE_PATH ) - self.liveness_db_failure_threshold = data.get( - "liveness_db_failure_threshold", constants.LIVENESS_DB_FAILURE_THRESHOLD + 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: 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/src/cache/postgres_cache.py b/ols/src/cache/postgres_cache.py index f19f6e537..0b90f7fcc 100644 --- a/ols/src/cache/postgres_cache.py +++ b/ols/src/cache/postgres_cache.py @@ -184,6 +184,7 @@ def _ddl_statements(self) -> list[str]: 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: @@ -203,10 +204,15 @@ def _connect_health(self) -> None: "sslmode": config.ssl_mode, "sslrootcert": config.ca_cert_path, "gssencmode": config.gss_encmode, + "connect_timeout": 10, **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 _health_check_loop(self) -> None: """Background loop that periodically checks DB connectivity.""" @@ -348,7 +354,7 @@ def insert_or_append( 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: @@ -387,7 +393,7 @@ def delete( 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: diff --git a/ols/utils/postgres.py b/ols/utils/postgres.py index fca9bf188..0d5f1c1b3 100644 --- a/ols/utils/postgres.py +++ b/ols/utils/postgres.py @@ -104,8 +104,12 @@ def connect(self) -> None: raise self.connection.autocommit = True cursor = self.connection.cursor() - cursor.execute("SET statement_timeout = %s", (str(config.statement_timeout),)) - cursor.close() + 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/unit/app/endpoints/test_health.py b/tests/unit/app/endpoints/test_health.py index 85b23ddfd..f4d479177 100644 --- a/tests/unit/app/endpoints/test_health.py +++ b/tests/unit/app/endpoints/test_health.py @@ -222,8 +222,9 @@ def test_readiness_probe_get_method_cache_not_ready(): def test_liveness_probe_get_method(): """Test the liveness_probe function when no postgres cache is configured.""" - response = liveness_probe_get_method() - assert response == LivenessResponse(alive=True) + mock_response = Mock() + result = liveness_probe_get_method(mock_response) + assert result == LivenessResponse(alive=True) def test_liveness_probe_returns_alive_when_postgres_healthy(): @@ -239,8 +240,9 @@ def test_liveness_probe_returns_alive_when_postgres_healthy(): patch.object(config, "_conversation_cache", cache), patch.object(config.ols_config, "liveness_db_failure_threshold", 3), ): - response = liveness_probe_get_method() - assert response == LivenessResponse(alive=True) + mock_response = Mock() + result = liveness_probe_get_method(mock_response) + assert result == LivenessResponse(alive=True) def test_liveness_probe_returns_503_when_postgres_unhealthy(): @@ -257,10 +259,10 @@ def test_liveness_probe_returns_503_when_postgres_unhealthy(): patch.object(config, "_conversation_cache", cache), patch.object(config.ols_config, "liveness_db_failure_threshold", 3), ): - with pytest.raises(HTTPException) as exc_info: - liveness_probe_get_method() - assert exc_info.value.status_code == 503 - assert exc_info.value.detail["reason"] == "database unreachable" + 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(): @@ -276,5 +278,6 @@ def test_liveness_probe_returns_alive_when_below_threshold(): patch.object(config, "_conversation_cache", cache), patch.object(config.ols_config, "liveness_db_failure_threshold", 3), ): - response = liveness_probe_get_method() - assert response == LivenessResponse(alive=True) + mock_response = Mock() + result = liveness_probe_get_method(mock_response) + assert result == LivenessResponse(alive=True) diff --git a/tests/unit/cache/test_postgres_cache.py b/tests/unit/cache/test_postgres_cache.py index c1755295a..d4126a498 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 @@ -17,7 +18,7 @@ @pytest.fixture(autouse=True) -def _suppress_health_loop(): +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 @@ -789,7 +790,7 @@ def test_cleanup_method_when_clean_performed(): mock_cursor.execute.assert_has_calls(calls, any_order=False) -def test_ready_returns_health_status(): +def test_ready_returns_health_status() -> None: """Test that ready() returns the background health loop status.""" with patch("psycopg2.connect"): config = PostgresConfig() @@ -808,7 +809,7 @@ def test_ready_returns_health_status(): assert cache.ready() is True -def test_mark_unhealthy(): +def test_mark_unhealthy() -> None: """Test that _mark_unhealthy sets health status to False.""" with patch("psycopg2.connect"): config = PostgresConfig() @@ -819,7 +820,7 @@ def test_mark_unhealthy(): assert cache.ready() is False -def test_lock_timeout(): +def test_lock_timeout() -> None: """Test that lock acquisition timeout raises CacheError.""" with patch("psycopg2.connect"): config = PostgresConfig(lock_timeout=0) @@ -834,7 +835,7 @@ def test_lock_timeout(): cache._tx_lock.release() -def test_health_check_loop_recovers(): +def test_health_check_loop_recovers() -> None: """Test that the health check loop sets health status on success.""" with patch("psycopg2.connect"): config = PostgresConfig() @@ -851,7 +852,7 @@ def test_health_check_loop_recovers(): assert cache.ready() is True -def test_consecutive_failures_tracked(): +def test_consecutive_failures_tracked() -> None: """Test that consecutive failures are tracked.""" with patch("psycopg2.connect"): config = PostgresConfig() diff --git a/tests/unit/utils/test_postgres.py b/tests/unit/utils/test_postgres.py index f8e4c508a..8728cb9c6 100644 --- a/tests/unit/utils/test_postgres.py +++ b/tests/unit/utils/test_postgres.py @@ -78,7 +78,7 @@ 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): + def test_auto_reconnects_when_disconnected(self) -> None: """Decorator calls connect() when not connected.""" c = self.Connectable() c.disconnect() @@ -88,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() @@ -96,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() @@ -105,7 +105,7 @@ def test_propagates_exception_after_reconnect(self): c.do_work() assert c.connected() is True - def test_operational_error_triggers_reconnect_and_retry(self): + def test_operational_error_triggers_reconnect_and_retry(self) -> None: """OperationalError causes reconnect + retry, succeeding on second attempt.""" c = self.Connectable() c.connect() @@ -113,7 +113,7 @@ def test_operational_error_triggers_reconnect_and_retry(self): assert result == "recovered" assert c._unhealthy_marked is True - def test_interface_error_triggers_reconnect_and_retry(self): + def test_interface_error_triggers_reconnect_and_retry(self) -> None: """InterfaceError causes reconnect + retry, succeeding on second attempt.""" c = self.Connectable() c.connect() @@ -121,22 +121,24 @@ def test_interface_error_triggers_reconnect_and_retry(self): assert result == "recovered" assert c._unhealthy_marked is True - def test_database_error_wraps_in_cache_error_no_retry(self): + 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 pytest.raises(CacheError, match="SQL syntax error"): - c.do_work_database_error() + 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): + def test_connection_error_reconnect_failure_raises_cache_error(self) -> None: """When reconnect fails after OperationalError, CacheError is raised.""" c = self.Connectable() c.connect() - # Make connect() fail on reconnect attempt original_connect = c.connect call_count = [0] - def failing_connect(): + def failing_connect() -> None: call_count[0] += 1 if call_count[0] > 0: raise psycopg2.OperationalError("cannot connect") From f7d770073cd5bdc19dc499c1c85523e2b5e9d5a3 Mon Sep 17 00:00:00 2001 From: Sri Roopa Ramesh Babu Date: Thu, 2 Jul 2026 13:54:48 -0400 Subject: [PATCH 3/4] OLS-3221 Address adversarial review feedback - Fix encapsulation: add public consecutive_failures property, remove direct access to private _health_lock/_consecutive_failures from liveness endpoint - Fix initial connection bypass: move connect() inside try/except in @connection decorator so failures get _mark_unhealthy() and CacheError - Make health check connect timeout configurable via PostgresConfig - Add graceful shutdown: shutdown() joins thread and closes connection - Fix delayed initial check: run health probe before sleep, use Event.wait() instead of time.sleep() - Fix startup race: skip failure counter on first health check to avoid premature liveness failures - Fix stale connection reference: set _health_connection = None after close before reconnect - Reject zero lock_timeout: change from NonNegativeInt to PositiveInt Co-Authored-By: Claude Opus 4.6 (1M context) --- ols/app/endpoints/health.py | 4 +- ols/app/models/config.py | 5 ++- ols/constants.py | 1 + ols/src/cache/postgres_cache.py | 52 ++++++++++++++++++------- ols/utils/postgres.py | 4 +- tests/unit/app/endpoints/test_health.py | 8 ++-- tests/unit/cache/test_postgres_cache.py | 17 ++++---- 7 files changed, 59 insertions(+), 32 deletions(-) diff --git a/ols/app/endpoints/health.py b/ols/app/endpoints/health.py index 4f5130075..671a337b1 100644 --- a/ols/app/endpoints/health.py +++ b/ols/app/endpoints/health.py @@ -136,9 +136,7 @@ def liveness_probe_get_method(response: Response) -> LivenessResponse: cache = config._conversation_cache if isinstance(cache, PostgresCache): threshold = config.ols_config.liveness_db_failure_threshold - with cache._health_lock: - failures = cache._consecutive_failures - if failures >= 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 3278af39e..1a8e7e9c0 100644 --- a/ols/app/models/config.py +++ b/ols/app/models/config.py @@ -815,8 +815,11 @@ class PostgresConfig(BaseModel): ca_cert_path: Optional[FilePath] = None max_entries: PositiveInt = constants.POSTGRES_CACHE_MAX_ENTRIES statement_timeout: NonNegativeInt = constants.POSTGRES_STATEMENT_TIMEOUT - lock_timeout: NonNegativeInt = constants.POSTGRES_LOCK_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: diff --git a/ols/constants.py b/ols/constants.py index d21aac602..5898f67a4 100644 --- a/ols/constants.py +++ b/ols/constants.py @@ -140,6 +140,7 @@ class GenericLLMParameters: POSTGRES_CACHE_GSSENCMODE = "prefer" CACHE_HEALTH_CHECK_INTERVAL = 30 +CACHE_HEALTH_CHECK_CONNECT_TIMEOUT = 10 LIVENESS_DB_FAILURE_THRESHOLD = 3 diff --git a/ols/src/cache/postgres_cache.py b/ols/src/cache/postgres_cache.py index 0b90f7fcc..a948ab8ed 100644 --- a/ols/src/cache/postgres_cache.py +++ b/ols/src/cache/postgres_cache.py @@ -3,7 +3,6 @@ import json import logging import threading -import time from typing import Any import psycopg2 @@ -163,7 +162,9 @@ def __init__(self, config: PostgresConfig) -> None: 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) @@ -194,6 +195,7 @@ def _connect_health(self) -> None: 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, @@ -204,7 +206,7 @@ def _connect_health(self) -> None: "sslmode": config.ssl_mode, "sslrootcert": config.ca_cert_path, "gssencmode": config.gss_encmode, - "connect_timeout": 10, + "connect_timeout": self._health_check_connect_timeout, **libpq_tls_params(config.tls_security_profile), } self._health_connection = psycopg2.connect(**connect_kwargs) @@ -214,10 +216,23 @@ def _connect_health(self) -> None: "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.""" - while True: - time.sleep(self._health_check_interval) + 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() @@ -229,19 +244,24 @@ def _health_check_loop(self) -> None: self._health_status = True self._consecutive_failures = 0 except Exception as e: - 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, - ) + 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.""" @@ -473,6 +493,12 @@ def set_topic_summary( 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. diff --git a/ols/utils/postgres.py b/ols/utils/postgres.py index 0d5f1c1b3..e900cb2a1 100644 --- a/ols/utils/postgres.py +++ b/ols/utils/postgres.py @@ -26,9 +26,9 @@ def connection(f: Callable) -> Callable: """ def wrapper(connectable: Any, *args: Any, **kwargs: Any) -> Callable: - if not connectable.connected(): - connectable.connect() try: + if not connectable.connected(): + connectable.connect() return f(connectable, *args, **kwargs) except (psycopg2.OperationalError, psycopg2.InterfaceError) as e: logger.warning( diff --git a/tests/unit/app/endpoints/test_health.py b/tests/unit/app/endpoints/test_health.py index f4d479177..17648afdc 100644 --- a/tests/unit/app/endpoints/test_health.py +++ b/tests/unit/app/endpoints/test_health.py @@ -234,7 +234,6 @@ def test_liveness_probe_returns_alive_when_postgres_healthy(): pg_config = PostgresConfig() cache = PostgresCache(pg_config) - cache._consecutive_failures = 0 with ( patch.object(config, "_conversation_cache", cache), @@ -252,8 +251,8 @@ def test_liveness_probe_returns_503_when_postgres_unhealthy(): pg_config = PostgresConfig() cache = PostgresCache(pg_config) - cache._consecutive_failures = 3 - cache._health_status = False + for _ in range(3): + cache._mark_unhealthy() with ( patch.object(config, "_conversation_cache", cache), @@ -272,7 +271,8 @@ def test_liveness_probe_returns_alive_when_below_threshold(): pg_config = PostgresConfig() cache = PostgresCache(pg_config) - cache._consecutive_failures = 2 + for _ in range(2): + cache._mark_unhealthy() with ( patch.object(config, "_conversation_cache", cache), diff --git a/tests/unit/cache/test_postgres_cache.py b/tests/unit/cache/test_postgres_cache.py index d4126a498..cc4bc9a29 100644 --- a/tests/unit/cache/test_postgres_cache.py +++ b/tests/unit/cache/test_postgres_cache.py @@ -645,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") @@ -660,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) @@ -823,7 +824,7 @@ def test_mark_unhealthy() -> None: def test_lock_timeout() -> None: """Test that lock acquisition timeout raises CacheError.""" with patch("psycopg2.connect"): - config = PostgresConfig(lock_timeout=0) + config = PostgresConfig(lock_timeout=1) cache = PostgresCache(config) # Manually acquire the lock so the next acquire times out @@ -858,10 +859,8 @@ def test_consecutive_failures_tracked() -> None: config = PostgresConfig() cache = PostgresCache(config) - with cache._health_lock: - cache._consecutive_failures = 5 - cache._health_status = False + for _ in range(5): + cache._mark_unhealthy() assert cache.ready() is False - with cache._health_lock: - assert cache._consecutive_failures == 5 + assert cache.consecutive_failures == 5 From cdb26ca59982afe9dafdda7b186aae07f923b0da Mon Sep 17 00:00:00 2001 From: Sri Roopa Ramesh Babu Date: Thu, 2 Jul 2026 15:04:47 -0400 Subject: [PATCH 4/4] OLS-3221 Fix integration tests for LivenessResponse schema - Update test_liveness assertion to include reason: null field added to LivenessResponse model - Regenerate docs/openapi.json to reflect config and model changes Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/openapi.json | 13 ++++++++++++- tests/integration/test_liveness_readiness.py | 2 +- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/docs/openapi.json b/docs/openapi.json index dce54a6c1..5f81bda4b 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -1718,6 +1718,17 @@ "alive": { "type": "boolean", "title": "Alive" + }, + "reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Reason" } }, "type": "object", @@ -1725,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/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():