OLS-3221 Add PostgreSQL auto-recovery after DB restart - #2964
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughChangesPostgres cache health and liveness
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant LivenessEndpoint
participant PostgresCache
participant PostgreSQL
Client->>LivenessEndpoint: GET /liveness
LivenessEndpoint->>PostgresCache: read consecutive_failures
PostgresCache-->>LivenessEndpoint: failure count
PostgresCache->>PostgreSQL: background health check
PostgreSQL-->>PostgresCache: connection health
LivenessEndpoint-->>Client: 200 alive or 503 database unreachable
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (1)
tests/unit/cache/test_postgres_cache.py (1)
836-865: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftExercise the health-check behavior instead of assigning its result.
These tests directly set
_health_statusand_consecutive_failures, so they would still pass if_health_check_loop()stopped updating either field. Extract a single-iteration health-check helper or inject a controllable sleep/event so the tests drive the success and failure paths. As per coding guidelines,tests/**/*.py: Assert specific values and behaviors in tests, not just that code runs without error.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/cache/test_postgres_cache.py` around lines 836 - 865, The health-check tests are bypassing the real behavior by mutating PostgresCache internals directly, so they won’t catch regressions in _health_check_loop(). Update the tests to drive the loop logic through a single-iteration helper or a controllable sleep/event on PostgresCache, then assert the resulting _health_status and _consecutive_failures values after simulated success/failure. Use the existing PostgresCache, _health_check_loop, _mark_unhealthy, and ready() symbols to locate the behavior and keep the assertions focused on observable state changes rather than direct field assignment.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ols/app/endpoints/health.py`:
- Around line 126-129: The 503 liveness response in the health endpoint is
returning an HTTPException shape instead of the declared LivenessResponse
contract. Update the liveness failure path in the health endpoint so the 503
body matches LivenessResponse (for example, by returning that model directly),
or change the documented 503 response model to the actual error payload; keep
the response schema and the OpenAPI declaration in sync around the liveness
handler.
In `@ols/app/models/config.py`:
- Around line 801-803: The config model currently accepts negative values for
statement_timeout, lock_timeout, and health_check_interval, so tighten
validation in the Pydantic config class that defines these fields. Add explicit
field constraints or validators on statement_timeout and lock_timeout to reject
negative integers, and make health_check_interval strictly positive so the
health-check loop and PostgresCache lock acquisition keep their timing
guarantees. Use the existing config model/field definitions in Config to enforce
these contracts at load time rather than downstream.
- Around line 1166-1167: `OLSConfig.__init__` is currently taking
`liveness_db_failure_threshold` directly from `data.get(...)`, so it can remain
a string or accept invalid values. Parse this field as an int when loading
config, and validate that the value is at least 1; if it is missing,
non-numeric, zero, or negative, raise `checks.InvalidConfigurationError` before
the value is used by `/liveness`.
In `@ols/src/cache/postgres_cache.py`:
- Around line 159-165: Validate the Postgres cache timing settings before they
reach the daemon: `health_check_interval` must not allow negative values or
zero-hot-loop behavior, and `lock_timeout` should reject negative values unless
indefinite blocking is explicitly intended. Add Pydantic validation in
`PostgresConfig` (using appropriate constrained types such as
non-negative/positive ints) and keep the runtime consumers in `PostgresCache`
safe by relying on those validated values in the initialization and health-check
logic. Update any related uses in `PostgresCache` and its config wiring so the
daemon never calls `time.sleep()` or lock acquisition with invalid values.
- Around line 189-219: The health-check path in _connect_health and
_health_check_loop can block indefinitely because neither the psycopg2.connect
call nor the SELECT 1 probe has a timeout. Update the PostgresCache health
connection setup to pass an explicit connect timeout in connect_kwargs and add a
statement timeout before executing the cursor probe, so ready() does not wait on
a hung network or server and _health_status can recover promptly.
- Around line 194-195: The broad exception handling in PostgresCache needs to be
narrowed to satisfy Ruff BLE001. Update the catch blocks in _connect_health,
_health_check_loop, _safe_rollback, and _safe_set_autocommit to handle only the
specific psycopg2 or attribute-related exceptions they can actually raise, and
keep the existing logger/debug behavior in those paths. If any catch-all is
intentionally required to keep the daemon alive, add an explicit # noqa: BLE001
on that except clause rather than leaving a bare except Exception.
In `@ols/utils/postgres.py`:
- Around line 106-108: The new cursor in the statement-timeout setup can leak if
`cursor.execute` fails during schema initialization. Update the `Postgres`
initialization flow around the `cursor = self.connection.cursor()` and `SET
statement_timeout` call to use a `try/finally` so `cursor.close()` always runs,
even on exceptions; keep the fix localized to this cursor-handling block.
- Around line 29-45: The wrapper in postgres connection handling is calling
connectable.connect() before the try/except, so a failed initial connect skips
the same recovery path as query failures. Move the initial connect() call into
the connection-error handling flow inside wrapper() so both the first connect
and the retry go through the same psycopg2.OperationalError/InterfaceError
handling, including _mark_unhealthy() and CacheError wrapping in the reconnect
path.
- Around line 33-45: The `@connection` retry logic in the postgres decorator is
causing duplicate writes for methods like PostgresCache.insert_or_append because
the operation may have already committed before reconnect and retry. Update the
retry behavior in the connection wrapper to avoid automatically retrying
write/mutating methods after OperationalError/InterfaceError, and keep the retry
path only for safe read operations or explicitly idempotent calls identified by
the wrapped function name.
In `@tests/unit/cache/test_postgres_cache.py`:
- Around line 20-23: The new pytest fixture and several test functions in
Postgres cache tests are missing type annotations. Update _suppress_health_loop
to declare Generator[None, None, None] as its return type, and add -> None to
the affected test functions in test_postgres_cache so the new autouse fixture
and tests are fully typed and consistent with the rest of the module.
In `@tests/unit/utils/test_postgres.py`:
- Around line 108-147: Add explicit -> None return annotations to the new test
methods in test_postgres.py, including the helper failing_connect nested in
test_connection_error_reconnect_failure_raises_cache_error, to keep the test
signatures consistent. Also tighten
test_database_error_wraps_in_cache_error_no_retry by asserting the Connectable
instance does not attempt reconnect/retry when do_work_database_error is called;
use the existing Connectable, connect, and do_work_database_error symbols to
place the assertion near the current CacheError check.
- Around line 124-129: The test
`test_database_error_wraps_in_cache_error_no_retry` only verifies that
`do_work_database_error()` raises `CacheError`, but it does not prove the
reconnect path was skipped. After the initial `c.connect()` call, patch or mock
`Connectable.connect` and assert it is not called when
`do_work_database_error()` raises, so this test checks the “no retry” behavior
as well as the `CacheError` wrapping.
---
Nitpick comments:
In `@tests/unit/cache/test_postgres_cache.py`:
- Around line 836-865: The health-check tests are bypassing the real behavior by
mutating PostgresCache internals directly, so they won’t catch regressions in
_health_check_loop(). Update the tests to drive the loop logic through a
single-iteration helper or a controllable sleep/event on PostgresCache, then
assert the resulting _health_status and _consecutive_failures values after
simulated success/failure. Use the existing PostgresCache, _health_check_loop,
_mark_unhealthy, and ready() symbols to locate the behavior and keep the
assertions focused on observable state changes rather than direct field
assignment.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3804766b-967a-49bb-b307-520f135d2133
📒 Files selected for processing (10)
ols/app/endpoints/health.pyols/app/models/config.pyols/app/models/models.pyols/constants.pyols/src/cache/postgres_cache.pyols/utils/postgres.pytests/unit/app/endpoints/test_health.pytests/unit/app/models/test_config.pytests/unit/cache/test_postgres_cache.pytests/unit/utils/test_postgres.py
| 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 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Validate timeout and interval values before the daemon uses them.
health_check_interval reaches time.sleep() before the try, so a negative value kills the health thread and 0 creates a hot polling loop. lock_timeout should also reject negative values unless indefinite blocking is intentional. The related PostgresConfig contract currently exposes these as plain int; validate them there with PositiveInt/non-negative types. As per coding guidelines, ols/**/*.py: Use Pydantic models for configuration and data validation.
Also applies to: 211-215, 254-257
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ols/src/cache/postgres_cache.py` around lines 159 - 165, Validate the
Postgres cache timing settings before they reach the daemon:
`health_check_interval` must not allow negative values or zero-hot-loop
behavior, and `lock_timeout` should reject negative values unless indefinite
blocking is explicitly intended. Add Pydantic validation in `PostgresConfig`
(using appropriate constrained types such as non-negative/positive ints) and
keep the runtime consumers in `PostgresCache` safe by relying on those validated
values in the initialization and health-check logic. Update any related uses in
`PostgresCache` and its config wiring so the daemon never calls `time.sleep()`
or lock acquisition with invalid values.
Source: Coding guidelines
15f8455 to
6e37719
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ols/src/cache/postgres_cache.py`:
- Around line 350-353: The DatabaseError handling in PostgresCache should use
the same rollback protection as the operational-error path, because calling
connection.rollback() directly can hide the original psycopg2.DatabaseError and
prevent the CacheError wrapping from happening. Update the exception blocks in
PostgresCache.insert_or_append and the other matching DatabaseError path to call
_safe_rollback() instead of self.connection.rollback(), while keeping the
existing logger.error and CacheError("PostgresCache.insert_or_append", e)
behavior intact.
- Around line 184-188: The `_mark_unhealthy()` path only updates
`_health_status`, so repeated request-path failures from the `@connection` retry
flow do not advance `_consecutive_failures` toward the liveness threshold.
Update `_mark_unhealthy()` in `PostgresCache` to also increment the failure
counter under `_health_lock`, keeping the unhealthy flag and failure tracking in
sync so `is_healthy`/liveness checks can trip promptly after repeated
cache-operation DB failures.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c84be3c2-a66c-475a-91d3-5d332a00fb83
📒 Files selected for processing (10)
docs/openapi.jsonols/app/endpoints/health.pyols/app/models/config.pyols/constants.pyols/src/cache/postgres_cache.pyols/utils/postgres.pytests/unit/app/endpoints/test_health.pytests/unit/app/models/test_config.pytests/unit/cache/test_postgres_cache.pytests/unit/utils/test_postgres.py
🚧 Files skipped from review as they are similar to previous changes (8)
- ols/app/endpoints/health.py
- ols/constants.py
- tests/unit/cache/test_postgres_cache.py
- tests/unit/app/endpoints/test_health.py
- ols/app/models/config.py
- tests/unit/utils/test_postgres.py
- ols/utils/postgres.py
- tests/unit/app/models/test_config.py
There was a problem hiding this comment.
🧹 Nitpick comments (2)
ols/app/endpoints/health.py (1)
134-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winEncapsulate health-state access instead of reaching into PostgresCache internals.
The endpoint directly touches
cache._health_lockandcache._consecutive_failures, private attributes ofPostgresCache. This crosses the app/cache layer boundary and couples the endpoint to the cache's internal locking scheme.Consider exposing a public method (e.g.
PostgresCache.is_healthy(threshold) -> bool) that encapsulates the lock + comparison, sohealth.pyonly calls a stable public API.As per coding guidelines, "Respect architectural boundaries: Do not cross module or layer boundaries, even when it is the shorter path."
♻️ Proposed refactor
- 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: - response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE - return LivenessResponse(alive=False, reason="database unreachable") + cache = config._conversation_cache + if isinstance(cache, PostgresCache): + threshold = config.ols_config.liveness_db_failure_threshold + if not cache.is_healthy(threshold): + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE + return LivenessResponse(alive=False, reason="database unreachable")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ols/app/endpoints/health.py` around lines 134 - 144, The liveness probe in liveness_probe_get_method is reaching into PostgresCache private state (_health_lock and _consecutive_failures), so move that lock-and-compare logic behind a public API on PostgresCache such as is_healthy(threshold) -> bool. Update health.py to call only the new public method and keep the endpoint from depending on cache internals.Source: Coding guidelines
tests/unit/app/endpoints/test_health.py (1)
230-283: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHoist repeated inline
PostgresConfigimport to module level.
from ols.app.models.config import PostgresConfigis repeated inline in three test functions (lines 233, 251, 271). Unless this is needed to avoid a circular import, move it to the top of the file.As per coding guidelines, "module-level imports by default, inline imports only when needed for circular dependencies or deferred optional heavy dependencies."
♻️ Proposed refactor
+from ols.app.models.config import PostgresConfig from ols.src.cache.postgres_cache import PostgresCacheAnd remove the three inline
from ols.app.models.config import PostgresConfigoccurrences inside the test bodies.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/app/endpoints/test_health.py` around lines 230 - 283, The three liveness probe tests repeat an inline PostgresConfig import inside each test body; hoist that import to the module level in the test file and remove the repeated local imports from test_liveness_probe_returns_alive_when_postgres_healthy, test_liveness_probe_returns_503_when_postgres_unhealthy, and test_liveness_probe_returns_alive_when_below_threshold. Keep the tests using PostgresConfig exactly as before, just reference the shared module-level import to follow the import guideline unless a circular dependency requires otherwise.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@ols/app/endpoints/health.py`:
- Around line 134-144: The liveness probe in liveness_probe_get_method is
reaching into PostgresCache private state (_health_lock and
_consecutive_failures), so move that lock-and-compare logic behind a public API
on PostgresCache such as is_healthy(threshold) -> bool. Update health.py to call
only the new public method and keep the endpoint from depending on cache
internals.
In `@tests/unit/app/endpoints/test_health.py`:
- Around line 230-283: The three liveness probe tests repeat an inline
PostgresConfig import inside each test body; hoist that import to the module
level in the test file and remove the repeated local imports from
test_liveness_probe_returns_alive_when_postgres_healthy,
test_liveness_probe_returns_503_when_postgres_unhealthy, and
test_liveness_probe_returns_alive_when_below_threshold. Keep the tests using
PostgresConfig exactly as before, just reference the shared module-level import
to follow the import guideline unless a circular dependency requires otherwise.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository: openshift/coderabbit/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 51571fc0-22a7-4977-978e-38af5cbf15fd
📒 Files selected for processing (8)
ols/app/endpoints/health.pyols/app/models/config.pyols/app/models/models.pyols/src/cache/postgres_cache.pyols/utils/postgres.pytests/unit/app/endpoints/test_health.pytests/unit/cache/test_postgres_cache.pytests/unit/utils/test_postgres.py
🚧 Files skipped from review as they are similar to previous changes (5)
- ols/app/models/config.py
- ols/src/cache/postgres_cache.py
- tests/unit/cache/test_postgres_cache.py
- ols/utils/postgres.py
- tests/unit/utils/test_postgres.py
7e46a50 to
c81358c
Compare
Adversarial Code Review - PostgreSQL Auto-Recovery (OLS-3221)Critical Issues1. Race Condition: Failure Counter Reset vs Increment (
|
Adversarial Review - Updated Code (Follow-up)Issues Addressed ✅Great progress! The following issues from my previous review have been addressed:
Remaining Issues1. Health Connection State Inconsistency (
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ols/src/cache/postgres_cache.py`:
- Around line 219-230: The Postgres cache shutdown path is implemented in
PostgresCache.shutdown() but is never invoked by the app lifecycle. Update the
app owner in main application startup/shutdown flow to call
config.conversation_cache.shutdown() during process teardown so the health
thread is joined and the connection is closed. Use the existing
CacheFactory-created cache instance and wire the call into the app
shutdown/lifespan handling in ols/app/main.py.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b907267e-0034-4eef-b88d-2b21f0a353dc
📒 Files selected for processing (9)
ols/app/endpoints/health.pyols/app/models/config.pyols/app/models/models.pyols/constants.pyols/src/cache/postgres_cache.pyols/utils/postgres.pytests/unit/app/endpoints/test_health.pytests/unit/cache/test_postgres_cache.pytests/unit/utils/test_postgres.py
✅ Files skipped from review due to trivial changes (1)
- ols/constants.py
🚧 Files skipped from review as they are similar to previous changes (7)
- ols/app/models/models.py
- tests/unit/app/endpoints/test_health.py
- ols/app/endpoints/health.py
- ols/utils/postgres.py
- tests/unit/cache/test_postgres_cache.py
- ols/app/models/config.py
- tests/unit/utils/test_postgres.py
cdb26ca to
c1bdf14
Compare
openshift-ci
Bot
added
the
needs-rebase
c1bdf14 to
ff59440
Compare
openshift-ci
Bot
removed
the
needs-rebase
|
/retest |
e909bf9 to
302a4f0
Compare
6d3dcb6 to
de68693
Compare
de68693 to
148d383
Compare
|
/retest |
148d383 to
58fa0a1
Compare
|
/test e2e-ols-cluster |
c0dda09 to
ae8a382
Compare
|
@CodeRabbit review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/unit/cache/test_postgres_cache_transaction_fix.py (1)
32-32: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winAdd return annotations to all test functions.
Add
-> Noneto each test function. Strict MyPy requires complete function signatures.Proposed fix
-def test_insert_or_append_transaction_status_check_on_success(): +def test_insert_or_append_transaction_status_check_on_success() -> None:As per coding guidelines, “Use type hints for all function signatures.”
Also applies to: 60-60, 93-93, 117-117, 145-145
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/cache/test_postgres_cache_transaction_fix.py` at line 32, Add -> None return annotations to every test function in this file, including test_insert_or_append_transaction_status_check_on_success and the other functions identified in the review, so all test signatures satisfy strict MyPy.Source: Coding guidelines
ols/app/endpoints/health.py (1)
55-60: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftStop timed-out LLM invocations from accumulating worker threads.
future.result(timeout=...)stops waiting, butexecutor.shutdown(wait=False, cancel_futures=True)cannot cancel an already-runningbare_llm.invoke(). A blocked provider call can keep its worker alive afterllm_is_ready()returnsFalse, and repeated probes can accumulate live workers. Enforce a deadline on the underlying provider I/O and use native async cancellation when supported. Add a regression test for repeated timed-out probes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ols/app/endpoints/health.py` around lines 55 - 60, Update llm_is_ready and the underlying bare_llm provider invocation to enforce a real I/O deadline, rather than relying only on future.result(timeout=...). Use native async cancellation when the provider supports it, and ensure timed-out probes do not leave running worker threads behind. Add a regression test that performs repeated timed-out readiness probes and verifies workers do not accumulate.Source: Coding guidelines
🧹 Nitpick comments (1)
tests/unit/cache/test_postgres_cache_transaction_fix.py (1)
38-43: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove comments that restate the test code.
The comments repeat mock setup, assertions, and control flow already expressed by the test names and statements. Keep the test docstrings and remove these redundant comments.
As per coding guidelines, “Avoid comments unless explicitly requested; make code self-documenting.”
Also applies to: 51-56, 63-67, 83-89, 101-103, 121-123, 128-130, 140-141, 153-168
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/cache/test_postgres_cache_transaction_fix.py` around lines 38 - 43, Remove the redundant explanatory comments throughout the affected test blocks, including the setup and transaction-status comments around mock_connection and mock_cursor. Keep the test docstrings and all executable statements unchanged so the tests remain self-documenting.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ols/src/cache/postgres_cache.py`:
- Around line 227-238: Move health-connection closing into a finally block owned
by _health_check_loop, ensuring every connection created by the thread is closed
when the loop exits. Synchronize all reads, writes, and closes of
_health_connection so shutdown cannot race with probe creation or cleanup;
update shutdown to signal and join the thread without independently closing a
potentially active connection.
- Around line 391-393: Update the mutation method containing the psycopg2
OperationalError/InterfaceError handler so ambiguous commit failures are not
retried by the `@connection` decorator: either add an idempotency key that makes
repeated appends and message_count updates safe, or convert the ambiguous write
outcome into a CacheError. Preserve rollback behavior for clearly failed writes
and ensure a committed-before-disconnect retry cannot duplicate cache_entry or
message_count.
- Around line 161-162: Update the first failed health-probe path in the
surrounding health-check method to immediately set _health_status to False and
increment _consecutive_failures, including the corresponding path around the
second referenced location. Preserve the existing success-state reset behavior
and failure logging.
In `@ols/utils/postgres.py`:
- Around line 33-61: Update the connection wrapper around
connectable.connected() so an initially disconnected client acquires _tx_lock
when available, reconnects only after rechecking connectivity, and calls
_mark_healthy after successful recovery. Also wrap OperationalError and
InterfaceError raised by the retry invocation of f(connectable, *args, **kwargs)
in CacheError, preserving the original exception as the cause.
---
Outside diff comments:
In `@ols/app/endpoints/health.py`:
- Around line 55-60: Update llm_is_ready and the underlying bare_llm provider
invocation to enforce a real I/O deadline, rather than relying only on
future.result(timeout=...). Use native async cancellation when the provider
supports it, and ensure timed-out probes do not leave running worker threads
behind. Add a regression test that performs repeated timed-out readiness probes
and verifies workers do not accumulate.
In `@tests/unit/cache/test_postgres_cache_transaction_fix.py`:
- Line 32: Add -> None return annotations to every test function in this file,
including test_insert_or_append_transaction_status_check_on_success and the
other functions identified in the review, so all test signatures satisfy strict
MyPy.
---
Nitpick comments:
In `@tests/unit/cache/test_postgres_cache_transaction_fix.py`:
- Around line 38-43: Remove the redundant explanatory comments throughout the
affected test blocks, including the setup and transaction-status comments around
mock_connection and mock_cursor. Keep the test docstrings and all executable
statements unchanged so the tests remain self-documenting.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 2a1fd023-159d-49c9-8797-03cd76d24cfa
📒 Files selected for processing (11)
docs/openapi.jsonols/app/endpoints/health.pyols/app/models/config.pyols/app/models/models.pyols/constants.pyols/src/cache/postgres_cache.pyols/utils/postgres.pytests/unit/app/endpoints/test_health.pytests/unit/app/models/test_config.pytests/unit/cache/test_postgres_cache.pytests/unit/cache/test_postgres_cache_transaction_fix.py
🚧 Files skipped from review as they are similar to previous changes (7)
- ols/app/models/models.py
- docs/openapi.json
- ols/constants.py
- tests/unit/app/models/test_config.py
- tests/unit/app/endpoints/test_health.py
- tests/unit/cache/test_postgres_cache.py
- ols/app/models/config.py
| 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() | ||
| tx_lock = getattr(connectable, "_tx_lock", None) | ||
| if tx_lock is not None: | ||
| with tx_lock: | ||
| try: | ||
| connectable.connect() | ||
| except Exception as reconnect_err: | ||
| raise CacheError( | ||
| f"reconnect failed in {f.__name__}", reconnect_err | ||
| ) from reconnect_err | ||
| else: | ||
| try: | ||
| connectable.connect() | ||
| except Exception as reconnect_err: | ||
| raise CacheError( | ||
| f"reconnect failed in {f.__name__}", reconnect_err | ||
| ) from reconnect_err | ||
| if hasattr(connectable, "_mark_healthy"): | ||
| connectable._mark_healthy() | ||
| return f(connectable, *args, **kwargs) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- file map ---'
fd -i 'postgres.py|postgres|cache' . | head -80
printf '%s\n' '--- target file ---'
wc -l ols/utils/postgres.py
cat -n ols/utils/postgres.py
printf '%s\n' '--- related symbols and decorator usages ---'
rg -n --glob '*.py' '_tx_lock|_mark_unhealthy|_mark_healthy|connected\(\)|CacheError|postgres_retry|retry' ols tests 2>/dev/null | head -300Repository: openshift/lightspeed-service
Length of output: 26131
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- PostgresCache structure and lifecycle ---'
ast-grep outline ols/src/cache/postgres_cache.py --match 'class PostgresCache' --view expanded
sed -n '130,220p' ols/src/cache/postgres_cache.py
sed -n '270,345p' ols/src/cache/postgres_cache.py
sed -n '360,455p' ols/src/cache/postgres_cache.py
sed -n '455,535p' ols/src/cache/postgres_cache.py
printf '%s\n' '--- decorator tests ---'
cat -n tests/unit/utils/test_postgres.py | sed -n '1,180p'
printf '%s\n' '--- transaction-fix tests ---'
cat -n tests/unit/cache/test_postgres_cache_transaction_fix.py | sed -n '1,180p'
printf '%s\n' '--- health endpoint and health-state consumers ---'
rg -n -C 5 --glob '*.py' 'health|_healthy|_mark_unhealthy|_mark_healthy|503' ols/src/cache/postgres_cache.py ols/src/app tests/unit/cache tests/unit/app/endpoints/test_health.pyRepository: openshift/lightspeed-service
Length of output: 50384
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- readiness implementation ---'
fd -i 'health.py' ols
health_file="$(fd -i -t f 'health.py' ols | head -1)"
wc -l "$health_file"
rg -n -C 8 'ready\(|consecutive_failures|503|HTTPException|cache' "$health_file"
printf '%s\n' '--- exact decorated cache method boundaries ---'
sed -n '292,410p' ols/src/cache/postgres_cache.py
sed -n '430,525p' ols/src/cache/postgres_cache.py
printf '%s\n' '--- all decorator definitions and uses in application code ---'
rg -n -C 2 --glob '*.py' '`@connection`|def connection\(' ols tests/unit
printf '%s\n' '--- focused health tests involving cache readiness ---'
rg -n -C 10 'ready\(|_health_status|503|consecutive_failures|_mark_unhealthy|_mark_healthy' tests/unit/app/endpoints/test_health.py tests/unit/cache/test_postgres_cache.pyRepository: openshift/lightspeed-service
Length of output: 45421
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
import ast
import threading
from dataclasses import dataclass
source = open("ols/utils/postgres.py", encoding="utf-8").read()
tree = ast.parse(source)
connection = next(
node for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name == "connection"
)
wrapper = next(
node for node in ast.walk(connection)
if isinstance(node, ast.FunctionDef) and node.name == "wrapper"
)
outer_try = next(node for node in wrapper.body if isinstance(node, ast.Try))
print("outer_try_handlers:", [
ast.unparse(handler.type) for handler in outer_try.handlers
])
print("outer_try_has_retry_return:", any(
isinstance(node, ast.Return)
and isinstance(node.value, ast.Call)
and isinstance(node.value.func, ast.Name)
and node.value.func.id == "f"
for node in ast.walk(outer_try)
))
print("retry_return_is_inside_operational_handler:", any(
isinstance(node, ast.Return)
and isinstance(node.value, ast.Call)
and isinstance(node.value.func, ast.Name)
and node.value.func.id == "f"
for node in ast.walk(outer_try.handlers[0])
))
class OperationalError(Exception):
pass
class InterfaceError(Exception):
pass
class DatabaseError(Exception):
pass
class CacheError(Exception):
pass
def exact_wrapper(f):
def wrapper(connectable, *args, **kwargs):
try:
if not connectable.connected():
connectable.connect()
return f(connectable, *args, **kwargs)
except (OperationalError, InterfaceError) as e:
connectable.mark_unhealthy()
connectable.connect()
connectable.mark_healthy()
return f(connectable, *args, **kwargs)
except DatabaseError as e:
raise CacheError(str(e)) from e
return wrapper
`@dataclass`
class RetryProbe:
calls: int = 0
healthy: bool = True
reconnects: int = 0
def connected(self):
return True
def connect(self):
self.reconnects += 1
def mark_unhealthy(self):
self.healthy = False
def mark_healthy(self):
self.healthy = True
probe = RetryProbe()
`@exact_wrapper`
def fails_twice(connectable):
connectable.calls += 1
raise OperationalError("retry failure")
try:
fails_twice(probe)
except Exception as exc:
print("second_retry_exception_type:", type(exc).__name__)
print("second_retry_calls:", probe.calls)
print("second_retry_is_cache_error:", isinstance(exc, CacheError))
class RaceProbe:
def __init__(self):
self.connection = None
self.connect_calls = 0
self.started = threading.Barrier(2)
self.release = threading.Barrier(2)
self.lock = threading.Lock()
def connected(self):
return self.connection is not None
def connect(self):
with self.lock:
self.connect_calls += 1
connection_id = self.connect_calls
self.started.wait()
self.release.wait()
self.connection = connection_id
def work(self):
return self.connection
race = RaceProbe()
def preflight_only():
if not race.connected():
race.connect()
race.work()
threads = [threading.Thread(target=preflight_only) for _ in range(2)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
print("concurrent_preflight_connect_calls:", race.connect_calls)
print("concurrent_preflight_final_connection:", race.connection)
PYRepository: openshift/lightspeed-service
Length of output: 530
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
import ast
import threading
source = open("ols/utils/postgres.py", encoding="utf-8").read()
tree = ast.parse(source)
connection = next(
node for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name == "connection"
)
wrapper = next(
node for node in ast.walk(connection)
if isinstance(node, ast.FunctionDef) and node.name == "wrapper"
)
outer_try = next(node for node in wrapper.body if isinstance(node, ast.Try))
operational_handler = outer_try.handlers[0]
retry_return = next(
node for node in ast.walk(operational_handler)
if isinstance(node, ast.Return)
and isinstance(node.value, ast.Call)
and isinstance(node.value.func, ast.Name)
and node.value.func.id == "f"
)
print("retry_return_line:", retry_return.lineno)
print("retry_handler_catches_retry_exception:", any(
isinstance(node, ast.ExceptHandler)
for node in ast.walk(operational_handler)
))
class OperationalError(Exception):
pass
class InterfaceError(Exception):
pass
class DatabaseError(Exception):
pass
class CacheError(Exception):
pass
def exact_wrapper(f):
def wrapper(connectable, *args, **kwargs):
try:
if not connectable.connected():
connectable.connect()
return f(connectable, *args, **kwargs)
except (OperationalError, InterfaceError):
connectable.mark_unhealthy()
connectable.connect()
connectable.mark_healthy()
return f(connectable, *args, **kwargs)
except DatabaseError as error:
raise CacheError(str(error)) from error
return wrapper
class RetryProbe:
def __init__(self):
self.calls = 0
self.reconnects = 0
def connected(self):
return True
def connect(self):
self.reconnects += 1
def mark_unhealthy(self):
pass
def mark_healthy(self):
pass
probe = RetryProbe()
`@exact_wrapper`
def fails_twice(connectable):
connectable.calls += 1
raise OperationalError("retry failure")
retry_exception = None
try:
fails_twice(probe)
except Exception as error:
retry_exception = error
print("second_retry_exception_type:", type(retry_exception).__name__)
print("second_retry_calls:", probe.calls)
print("second_retry_is_cache_error:", isinstance(retry_exception, CacheError))
class RaceProbe:
def __init__(self):
self.connection = None
self.connect_calls = 0
self.started = threading.Barrier(2)
self.release = threading.Barrier(2)
self.lock = threading.Lock()
def connected(self):
return self.connection is not None
def connect(self):
with self.lock:
self.connect_calls += 1
connection_id = self.connect_calls
self.started.wait()
self.release.wait()
self.connection = connection_id
race = RaceProbe()
def preflight_only():
if not race.connected():
race.connect()
threads = [threading.Thread(target=preflight_only) for _ in range(2)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
print("concurrent_preflight_connect_calls:", race.connect_calls)
print("concurrent_preflight_final_connection:", race.connection)
PYRepository: openshift/lightspeed-service
Length of output: 414
Synchronize reconnects and wrap retry failures.
If connected() returns False, acquire _tx_lock before calling connect(), then recheck connectivity. Mark the cache healthy after recovery so failure counters reset.
The retry call can raise another OperationalError or InterfaceError, which escapes the sibling handlers as a raw exception. Catch and wrap retry failures in CacheError.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ols/utils/postgres.py` around lines 33 - 61, Update the connection wrapper
around connectable.connected() so an initially disconnected client acquires
_tx_lock when available, reconnects only after rechecking connectivity, and
calls _mark_healthy after successful recovery. Also wrap OperationalError and
InterfaceError raised by the retry invocation of f(connectable, *args, **kwargs)
in CacheError, preserving the original exception as the cause.
3e2971d to
c554b58
Compare
|
/retest |
55faccd to
0580826
Compare
Adversarial Review — OLS-3221 PostgreSQL Auto-Recovery🟡 Thread-Safety: Double-Counting Inflates
|
Follow-up: Fixes for Medium IssuesI've addressed the three medium issues from the adversarial review: 1. ✅ Fixed: Double-counting of
|
48ca650 to
114d40f
Compare
|
/retest |
|
/test ols-evaluation |
db94976 to
9509747
Compare
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: Cursor <cursoragent@cursor.com>
9509747 to
e748e2c
Compare
|
/retest |
2 similar comments
|
/retest |
|
/retest |
|
@sriroopar: The following tests failed, say
Full PR test history. Your PR dashboard. DetailsInstructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. I understand the commands that are listed here. |
2 participants
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.
Description
Type of change
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit
New Features
/livenessendpoint to detect repeated PostgreSQL connectivity failures.503with a clear “database unreachable” reason when the failure threshold is reached.Bug Fixes