diff --git a/backend/app/gateway/health.py b/backend/app/gateway/health.py index 3e850b5102..729e0340ef 100644 --- a/backend/app/gateway/health.py +++ b/backend/app/gateway/health.py @@ -219,9 +219,9 @@ async def _probe_postgres_backend(conn_string: str, schema: str) -> str: logger.error("Readiness probe: psycopg is not installed for the postgres checkpointer backend") return PROBE_UNREACHABLE try: - from deerflow.persistence.postgres_schema import dsn_with_search_path, normalize_libpq_dsn + from deerflow.persistence.postgres_schema import dsn_with_search_path - dsn = dsn_with_search_path(normalize_libpq_dsn(conn_string), schema) + dsn = dsn_with_search_path(conn_string, schema) async with asyncio.timeout(_PROBE_TIMEOUT_SECONDS): connection = await AsyncConnection.connect(dsn, connect_timeout=int(_PROBE_TIMEOUT_SECONDS)) try: diff --git a/backend/packages/harness/deerflow/persistence/postgres_schema.py b/backend/packages/harness/deerflow/persistence/postgres_schema.py index 35167cc1a8..082d4e9921 100644 --- a/backend/packages/harness/deerflow/persistence/postgres_schema.py +++ b/backend/packages/harness/deerflow/persistence/postgres_schema.py @@ -178,9 +178,11 @@ def dsn_with_search_path(dsn: str, schema: str) -> str: does NOT treat ``+`` as a space (that is an HTML-form convention). So the space MUST be encoded as ``%20`` rather than ``+``, otherwise libpq sees a single broken token ``-c+search_path=...`` and the search_path is - never applied. Existing query parameters are preserved. Empty *schema* - returns *dsn* unchanged. + never applied. Existing query parameters are preserved. The DSN first goes + through :func:`normalize_libpq_dsn`, so a SQLAlchemy ``+driver`` suffix is + dropped even when *schema* is empty. Empty *schema* adds no option. """ + dsn = normalize_libpq_dsn(dsn) if not schema: return dsn parts = urlsplit(dsn) @@ -192,14 +194,6 @@ def dsn_with_search_path(dsn: str, schema: str) -> str: params["options"] = _merge_search_path_option(params.get("options", ""), schema) return make_conninfo(**params) - # DatabaseConfig.postgres_url may carry a SQLAlchemy driver suffix such as - # ``postgresql+asyncpg://``. psycopg's libpq only understands the bare - # ``postgres``/``postgresql`` scheme, so accept the compound form but emit - # a psycopg-consumable DSN by dropping the ``+driver`` part. - scheme_base = parts.scheme.split("+", 1)[0] - if scheme_base not in {"postgres", "postgresql"}: - raise ValueError(f"Unsupported PostgreSQL DSN scheme for schema injection: {parts.scheme!r}") - # Rebuild the query by hand rather than with parse_qsl/urlencode: parse_qsl # decodes '+' as a space (an HTML-form rule libpq does not apply), so # round-tripping an existing parameter would rewrite a literal plus as '%20' @@ -220,7 +214,7 @@ def dsn_with_search_path(dsn: str, schema: str) -> str: # quote() encodes space as %20 (libpq-safe) and a literal '+' as %2B, so the # merged value keeps both distinctions. query = "&".join([*preserved, f"options={quote(options, safe='')}"]) - return urlunsplit((scheme_base, parts.netloc, parts.path, query, parts.fragment)) + return urlunsplit((parts.scheme, parts.netloc, parts.path, query, parts.fragment)) def ensure_postgres_schema(conn_string: str, schema: str, *, install_hint: str) -> None: diff --git a/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py b/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py index 0e73dff1a5..fe13152c15 100644 --- a/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py +++ b/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py @@ -25,7 +25,7 @@ from collections.abc import AsyncIterator from langgraph.types import Checkpointer from deerflow.config.app_config import AppConfig, get_app_config -from deerflow.persistence.postgres_schema import create_schema_sql, dsn_with_search_path, normalize_libpq_dsn +from deerflow.persistence.postgres_schema import create_schema_sql, dsn_with_search_path from deerflow.runtime.cancellation import drained_async_context from deerflow.runtime.checkpointer.provider import ( POSTGRES_CONN_REQUIRED, @@ -68,7 +68,7 @@ def _build_postgres_pool(conn_string: str, schema: str = ""): # *on top of* the conninfo and would silently drop a DSN-supplied option # such as statement_timeout. This also strips a SQLAlchemy ``+driver`` # suffix so libpq can parse the DSN. Matches the sync/DSN paths. - dsn = dsn_with_search_path(normalize_libpq_dsn(conn_string), schema) + dsn = dsn_with_search_path(conn_string, schema) return AsyncConnectionPool( dsn, diff --git a/backend/tests/test_checkpointer.py b/backend/tests/test_checkpointer.py index 5aa48e7a82..2454cbc5af 100644 --- a/backend/tests/test_checkpointer.py +++ b/backend/tests/test_checkpointer.py @@ -939,6 +939,26 @@ class TestCheckpointerDatabaseConfig: assert resolved.type == "sqlite" assert resolved.connection_string == db_config.checkpointer_sqlite_path + def test_sync_checkpointer_strips_driver_suffix_without_schema(self): + """A SQLAlchemy ``+asyncpg`` URL must reach libpq as a bare DSN when no schema is set.""" + from deerflow.runtime.checkpointer.provider import checkpointer_context + + app_config = SimpleNamespace( + checkpointer=None, + database=DatabaseConfig(backend="postgres", postgres_url="postgresql+asyncpg://localhost/db"), + ) + mock_saver_cls = MagicMock() + mock_saver_cls.from_conn_string.return_value = nullcontext(MagicMock()) + + with ( + patch("deerflow.runtime.checkpointer.provider.get_app_config", return_value=app_config), + patch.dict(sys.modules, {"langgraph.checkpoint.postgres": MagicMock(PostgresSaver=mock_saver_cls)}), + checkpointer_context(), + ): + pass + + mock_saver_cls.from_conn_string.assert_called_once_with("postgresql://localhost/db") + def test_sync_checkpointer_singleton_uses_database_config(self): """The cached sync checkpointer factory must resolve database config before locking.""" app_config = SimpleNamespace( @@ -1045,6 +1065,27 @@ class TestStoreDatabaseConfig: mock_store.setup.assert_awaited_once() assert "No 'checkpointer' section" not in caplog.text + @pytest.mark.anyio + async def test_async_store_strips_driver_suffix_without_schema(self): + """A SQLAlchemy ``+asyncpg`` URL must reach libpq as a bare DSN when no schema is set.""" + from deerflow.runtime.store.async_provider import make_store + + app_config = SimpleNamespace( + checkpointer=None, + database=DatabaseConfig(backend="postgres", postgres_url="postgresql+asyncpg://localhost/db"), + ) + mock_cm = AsyncMock() + mock_cm.__aenter__.return_value = AsyncMock() + mock_cm.__aexit__.return_value = False + mock_store_cls = MagicMock() + mock_store_cls.from_conn_string.return_value = mock_cm + + with patch.dict(sys.modules, {"langgraph.store.postgres.aio": MagicMock(AsyncPostgresStore=mock_store_cls)}): + async with make_store(app_config): + pass + + mock_store_cls.from_conn_string.assert_called_once_with("postgresql://localhost/db") + @pytest.mark.anyio async def test_async_sqlite_store_uses_unified_database_path(self, tmp_path): """Unified database SQLite config must use the shared deerflow.db path.""" @@ -1100,6 +1141,26 @@ class TestStoreDatabaseConfig: assert resolved.type == "postgres" assert resolved.connection_string == "postgresql://localhost/db" + def test_sync_store_strips_driver_suffix_without_schema(self): + """A SQLAlchemy ``+asyncpg`` URL must reach libpq as a bare DSN when no schema is set.""" + from deerflow.runtime.store.provider import store_context + + app_config = SimpleNamespace( + checkpointer=None, + database=DatabaseConfig(backend="postgres", postgres_url="postgresql+asyncpg://localhost/db"), + ) + mock_store_cls = MagicMock() + mock_store_cls.from_conn_string.return_value = nullcontext(MagicMock()) + + with ( + patch("deerflow.runtime.store.provider.get_app_config", return_value=app_config), + patch.dict(sys.modules, {"langgraph.store.postgres": MagicMock(PostgresStore=mock_store_cls)}), + store_context(), + ): + pass + + mock_store_cls.from_conn_string.assert_called_once_with("postgresql://localhost/db") + def test_sync_store_singleton_uses_database_config(self): """The cached sync Store factory must resolve database config before locking.""" app_config = SimpleNamespace( diff --git a/backend/tests/test_postgres_schema_helper.py b/backend/tests/test_postgres_schema_helper.py index 0ede45ab91..66046fb7b3 100644 --- a/backend/tests/test_postgres_schema_helper.py +++ b/backend/tests/test_postgres_schema_helper.py @@ -53,6 +53,9 @@ class TestDsnWithSearchPath: dsn = "postgresql://u:p@h:5432/db" assert dsn_with_search_path(dsn, "") == dsn + def test_empty_schema_still_drops_driver_suffix(self): + assert dsn_with_search_path("postgresql+asyncpg://u:p@h:5432/db", "") == "postgresql://u:p@h:5432/db" + def test_appends_options_query_encoded(self): dsn = "postgresql://u:p@h:5432/db" out = dsn_with_search_path(dsn, "deerflow")