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 <cursoragent@cursor.com>

Author: sriroopar

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,16 +131,14 @@ 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):
         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"},
-            )
+            response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
+            return LivenessResponse(alive=False, reason="database unreachable")
     return LivenessResponse(alive=True)

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:

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 = {

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:

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."""

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)

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()

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,15 +88,15 @@ 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()
         with patch.object(c, "connect") as mock_connect:
             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,38 +105,40 @@ 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()
         result = c.do_work_operational_error()
         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()
         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):
+    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")