From a72434c6db93480d1980f821b92f44706a8c8d30 Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 20 Aug 2026 12:51:31 +0100 Subject: [PATCH] fix: small pg related fixes for stability --- application/app.py | 14 ++++++++- application/graphrag/store.py | 6 +++- application/sandbox/daytona.py | 6 ++-- application/storage/db/bootstrap.py | 4 ++- application/vectorstore/pgvector.py | 7 ++++- tests/graphrag/test_store.py | 23 ++++++++++++++ tests/sandbox/test_daytona.py | 22 ++++++++++++++ .../db/test_bootstrap_vector_schema.py | 30 ++++++++++++++++++- .../test_pgvector_schema_and_pool.py | 23 ++++++++++++++ 9 files changed, 128 insertions(+), 7 deletions(-) diff --git a/application/app.py b/application/app.py index 719b67dd..466939c1 100644 --- a/application/app.py +++ b/application/app.py @@ -65,7 +65,19 @@ ensure_database_ready( # Own the vector DB's schema here too, so the retrieval hot path is pure reads # instead of re-running DDL for every source of every request. if settings.AUTO_VECTOR_SCHEMA: - ensure_vector_schema(logger=logging.getLogger("application.app")) + _vector_schema_log = logging.getLogger("application.app") + try: + ensure_vector_schema(logger=_vector_schema_log) + except Exception: + # The vector DB is often a separate cluster. This runs at import time, + # so re-raising would stop gunicorn and every Celery worker from + # booting -- taking auth, chat history and webhooks down over a fault + # that only affects retrieval. PGVectorStore re-checks the schema on + # its write path, so degrading here loses nothing. + _vector_schema_log.exception( + "ensure_vector_schema failed; retrieval is degraded until the " + "vector database is reachable and its width matches EMBEDDINGS_NAME." + ) from application.agents.default_tools import ( # noqa: E402 validate_default_chat_tools, diff --git a/application/graphrag/store.py b/application/graphrag/store.py index 120035e8..1ee642c0 100644 --- a/application/graphrag/store.py +++ b/application/graphrag/store.py @@ -126,7 +126,11 @@ class GraphStore: ``PGVectorStore``, so a retrieval that touches both pays one checkout each instead of two fresh connect handshakes. """ - if self._connection is None or self._connection.closed: + if self._connection is not None and self._connection.closed: + # Hand the dead connection back before replacing it; an unreturned + # checkout is a pool slot lost for the life of the process. + self.close() + if self._connection is None: if self._pool_max_size > 0: self._connection = pgconn.pool_for( self._connection_string, self._pool_max_size diff --git a/application/sandbox/daytona.py b/application/sandbox/daytona.py index a4ec1f9d..d9b5f8f3 100644 --- a/application/sandbox/daytona.py +++ b/application/sandbox/daytona.py @@ -561,7 +561,7 @@ class DaytonaSandbox(CodeSandbox): handle.sandbox.fs.create_folder(parent, "755", request_timeout=self._default_timeout) except Exception as folder_exc: # noqa: BLE001 - folder may already exist logger.debug("put_file parent folder create returned: %s", folder_exc) - handle.sandbox.fs.upload_file(data, remote) + handle.sandbox.fs.upload_file(data, remote, timeout=int(self._default_timeout)) def get_file(self, session_id: str, path: str) -> bytes: """Download ``path`` from the session workspace as bytes, capped at ``max_file_bytes``.""" @@ -616,7 +616,9 @@ class DaytonaSandbox(CodeSandbox): size = getattr(info, "size", None) if size is not None and size > self._max_file_bytes: raise IOError(f"file too large: {size} > {self._max_file_bytes} bytes") - return handle.sandbox.fs.download_file(remote) + # download_file takes its timeout positionally and dispatches on + # isinstance(arg, int): a float would be read as a destination path. + return handle.sandbox.fs.download_file(remote, int(self._default_timeout)) def list_files(self, session_id: str) -> List[str]: """List workspace-relative file paths for ``session_id`` (recursive, never escapes the workspace).""" diff --git a/application/storage/db/bootstrap.py b/application/storage/db/bootstrap.py index 34042ef9..6083a6d4 100644 --- a/application/storage/db/bootstrap.py +++ b/application/storage/db/bootstrap.py @@ -153,7 +153,9 @@ def ensure_vector_schema(*, logger: Optional[logging.Logger] = None) -> None: started = time.monotonic() # A plain connection, never the store's pool: this can run pre-fork under # ``gunicorn --preload``, and an inherited pooled socket is a broken one. - conn = psycopg.connect(dsn) + # Bounded: a suspended/unreachable vector cluster must fail this hook + # rather than hang boot until a liveness probe kills the process. + conn = psycopg.connect(dsn, connect_timeout=10) try: cursor = conn.cursor() try: diff --git a/application/vectorstore/pgvector.py b/application/vectorstore/pgvector.py index c49e53fd..c3fe9f6d 100644 --- a/application/vectorstore/pgvector.py +++ b/application/vectorstore/pgvector.py @@ -97,7 +97,12 @@ class PGVectorStore(BaseVectorStore): def _get_connection(self): """Get or create this store's connection, pooled unless pooling is off.""" - if self._connection is None or self._connection.closed: + if self._connection is not None and self._connection.closed: + # Hand the dead connection back before replacing it: psycopg_pool + # never reclaims a checkout that is not returned, so dropping it + # costs the pool a slot for the life of the process. + self.close() + if self._connection is None: if self._pool_max_size > 0: self._connection = _pool_for( self._connection_string, self._pool_max_size diff --git a/tests/graphrag/test_store.py b/tests/graphrag/test_store.py index f0467f11..a01b382e 100644 --- a/tests/graphrag/test_store.py +++ b/tests/graphrag/test_store.py @@ -743,6 +743,29 @@ class TestGraphStorePooling: pooled_conn.rollback.assert_not_called() pool.putconn.assert_called_once_with(pooled_conn) + def test_a_dead_pooled_connection_is_returned_before_being_replaced( + self, monkeypatch + ): + # Same contract as ``PGVectorStore``: a connection that dies while this + # store holds it must go back to the pool, or the slot is lost for the + # life of the process. Extraction holds one store across the whole + # per-chunk LLM loop, which is exactly when a backend gets reaped. + pool, pooled_conn = self._fake_pool() + monkeypatch.setattr(store_module.pgconn, "pool_for", lambda dsn, n: pool) + monkeypatch.setitem(pgconn._POOLS, POOL_DSN, pool) + store = self._store() + store._get_connection() + replacement = MagicMock() + replacement.closed = False + pool.getconn.return_value = replacement + + pooled_conn.closed = True + conn = store._get_connection() + + assert conn is replacement + pool.putconn.assert_called_once_with(pooled_conn) + assert pool.getconn.call_count == 2 + def test_legacy_path_connects_directly_and_closes(self, monkeypatch): def _never(dsn, n): raise AssertionError("pooling is off; no pool must be built") diff --git a/tests/sandbox/test_daytona.py b/tests/sandbox/test_daytona.py index aeef2e4a..ec8e9ef4 100644 --- a/tests/sandbox/test_daytona.py +++ b/tests/sandbox/test_daytona.py @@ -437,6 +437,28 @@ def test_put_file_uploads_under_workspace(sandbox): assert remote.endswith("/docsgpt-sandbox/out/data.csv") +def test_put_file_bounds_the_upload(sandbox): + # Every other fs.* call is bounded at default_timeout; without this the + # upload falls back to the SDK's 30-minute default and a toolbox proxy that + # accepts but never answers pins the chat stream for that long. + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + sandbox.put_file("conv-1", "out/data.csv", b"a,b\n") + assert created.fs.upload_file.call_args.kwargs["timeout"] == 60 + + +def test_get_file_bounds_the_download(sandbox): + # download_file takes its timeout positionally and dispatches on + # isinstance(arg, int) -- a float would be read as a local destination path. + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + created.fs.download_file.return_value = b"hello" + sandbox.get_file("conv-1", "out/data.csv") + args = created.fs.download_file.call_args.args + assert args[1] == 60 + assert isinstance(args[1], int) + + def test_put_file_creates_parent_dirs(sandbox): sandbox.open("conv-1") _, created = sandbox._client.created[0] diff --git a/tests/storage/db/test_bootstrap_vector_schema.py b/tests/storage/db/test_bootstrap_vector_schema.py index a4bf9fc9..ce925786 100644 --- a/tests/storage/db/test_bootstrap_vector_schema.py +++ b/tests/storage/db/test_bootstrap_vector_schema.py @@ -8,6 +8,7 @@ a ``documents`` table built for a different model silently retrieves garbage. from __future__ import annotations +import ast from pathlib import Path from unittest.mock import MagicMock, patch @@ -91,7 +92,11 @@ class TestEnsureVectorSchemaCreates: def test_creates_the_vector_schema_under_the_advisory_lock(self, vector_settings): connect, conn, cursor, vector_schema, graph_schema, _ = self._run() - connect.assert_called_once_with("postgresql://user:pass@localhost/db") + # A bounded connect: an unreachable or suspended vector DB must fail the + # hook, not hang boot until a liveness probe kills the process. + connect.assert_called_once_with( + "postgresql://user:pass@localhost/db", connect_timeout=10 + ) statements = " ".join(str(c) for c in cursor.execute.call_args_list) assert "pg_advisory_xact_lock" in statements vector_schema.assert_called_once() @@ -187,6 +192,29 @@ class TestBootGating: assert "if settings.AUTO_VECTOR_SCHEMA:" in source assert "ensure_vector_schema(" in source + def test_the_boot_hook_cannot_take_the_process_down(self): + """A vector-DB fault must degrade retrieval, not crash-loop the app. + + The hook runs at import time, so an exception escaping it stops + gunicorn and every Celery worker from booting -- taking auth, chat + history and webhooks down with retrieval. ``PGVectorStore`` re-checks + the schema on its write path, so failing soft here loses nothing. + """ + tree = ast.parse(_APP_PY.read_text()) + + def _calls_hook(node) -> bool: + return any( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "ensure_vector_schema" + for n in ast.walk(node) + ) + + assert any( + isinstance(node, ast.Try) and node.handlers and _calls_hook(node) + for node in ast.walk(tree) + ), "ensure_vector_schema() at import time must be wrapped in try/except" + def test_setting_defaults_on(self): from application.core.settings import Settings diff --git a/tests/vectorstore/test_pgvector_schema_and_pool.py b/tests/vectorstore/test_pgvector_schema_and_pool.py index 3b0cae95..9f3c6816 100644 --- a/tests/vectorstore/test_pgvector_schema_and_pool.py +++ b/tests/vectorstore/test_pgvector_schema_and_pool.py @@ -290,6 +290,29 @@ class TestConnectionPooling: pooled_conn.rollback.assert_not_called() fake_pool.putconn.assert_called_once_with(pooled_conn) + def test_a_dead_pooled_connection_is_returned_before_being_replaced( + self, monkeypatch + ): + # psycopg_pool never reclaims a checkout that is not handed back, so + # replacing a connection that died in our hands would burn a pool slot + # permanently: after PGVECTOR_POOL_MAX_SIZE such events every checkout + # blocks for the pool timeout and then fails. + store, fake_pool, pooled_conn = self._pooled_store(monkeypatch) + monkeypatch.setitem( + pgvector_module._POOLS, store._connection_string, fake_pool + ) + store._get_connection() + replacement = MagicMock() + replacement.closed = False + fake_pool.getconn.return_value = replacement + + pooled_conn.closed = True # server terminated the backend mid-idle + conn = store._get_connection() + + assert conn is replacement + fake_pool.putconn.assert_called_once_with(pooled_conn) + assert fake_pool.getconn.call_count == 2 + def test_legacy_path_connects_directly_and_closes(self): store, mock_conn, _, _ = _make_store() store._pool_max_size = 0