From 6b60903d73d31eee23ed3b52a8c971f0392a1705 Mon Sep 17 00:00:00 2001 From: Chaoqi <136579807+Chaoqi31@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:45:09 -0400 Subject: [PATCH] fix(runtime): strip the SQLAlchemy driver suffix from sync and store DSNs (#6917) * test(runtime): cover the driver suffix in sync and store DSNs * fix(runtime): strip the SQLAlchemy driver suffix from sync and store DSNs With `postgres_url: postgresql+asyncpg://...` and the default empty `postgres_schema`, the sync checkpointer, the sync store and the async store passed the DSN to psycopg's `from_conn_string` unchanged, and libpq cannot parse the `+asyncpg` scheme. `dsn_with_search_path` only drops the suffix when it adds a schema, so a non-empty schema hid the bug. The async checkpointer already calls `normalize_libpq_dsn` first. Call `normalize_libpq_dsn` at these three call sites as well. * refactor(persistence): drop the driver suffix inside dsn_with_search_path dsn_with_search_path dropped a SQLAlchemy `+driver` suffix only when it added a search_path, so every libpq caller had to wrap the DSN in normalize_libpq_dsn to cover an empty schema, and flattening that nesting later would bring the bug back. Normalize at the top of dsn_with_search_path so both paths drop the suffix and check the scheme in one place, and remove the outer call from its five callers. The sync checkpointer and both stores are back to their code on main. --- backend/app/gateway/health.py | 4 +- .../deerflow/persistence/postgres_schema.py | 16 ++--- .../runtime/checkpointer/async_provider.py | 4 +- backend/tests/test_checkpointer.py | 61 +++++++++++++++++++ backend/tests/test_postgres_schema_helper.py | 3 + 5 files changed, 73 insertions(+), 15 deletions(-) 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")