mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
6b60903d73
* 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.
273 lines
11 KiB
Python
273 lines
11 KiB
Python
"""Tests for the PostgreSQL schema helpers (Issue #3380)."""
|
|
|
|
import asyncio
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from urllib.parse import parse_qs, urlsplit
|
|
|
|
import pytest
|
|
|
|
from deerflow.persistence.postgres_schema import (
|
|
build_asyncpg_connect_args,
|
|
build_psycopg_options,
|
|
create_schema_sql,
|
|
dsn_with_search_path,
|
|
ensure_postgres_schema_async,
|
|
normalize_libpq_dsn,
|
|
)
|
|
|
|
|
|
class TestBuildAsyncpgConnectArgs:
|
|
def test_sets_search_path_for_schema(self):
|
|
assert build_asyncpg_connect_args("deerflow") == {"server_settings": {"search_path": "deerflow"}}
|
|
|
|
def test_empty_schema_returns_empty_dict(self):
|
|
assert build_asyncpg_connect_args("") == {}
|
|
|
|
|
|
class TestBuildPsycopgOptions:
|
|
def test_builds_libpq_options(self):
|
|
assert build_psycopg_options("deerflow") == "-c search_path=deerflow"
|
|
|
|
def test_empty_schema_returns_none(self):
|
|
assert build_psycopg_options("") is None
|
|
|
|
|
|
class TestCreateSchemaSql:
|
|
def test_builds_create_schema_statement(self):
|
|
assert create_schema_sql("deerflow") == 'CREATE SCHEMA IF NOT EXISTS "deerflow"'
|
|
|
|
def test_empty_schema_returns_none(self):
|
|
assert create_schema_sql("") is None
|
|
|
|
@pytest.mark.parametrize("schema", ['a"; DROP SCHEMA public; --', "MySchema", "a b", "deerflow\n"])
|
|
def test_rejects_non_plain_identifier(self, schema):
|
|
# Defense-in-depth: the SQL-emitting boundary re-validates so a caller
|
|
# that bypasses the pydantic config validator cannot inject.
|
|
with pytest.raises(ValueError):
|
|
create_schema_sql(schema)
|
|
|
|
|
|
class TestDsnWithSearchPath:
|
|
def test_empty_schema_returns_dsn_unchanged(self):
|
|
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")
|
|
# libpq only decodes %XX in URI query values; '+' is NOT treated as a
|
|
# space. The space MUST therefore be encoded as %20, never as '+'.
|
|
assert "+" not in out
|
|
assert "options=-c%20search_path%3Ddeerflow" in out
|
|
parts = urlsplit(out)
|
|
query = parse_qs(parts.query)
|
|
assert query["options"] == ["-c search_path=deerflow"]
|
|
|
|
def test_merges_with_existing_query(self):
|
|
dsn = "postgresql://u:p@h:5432/db?sslmode=require"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
query = parse_qs(urlsplit(out).query)
|
|
assert query["sslmode"] == ["require"]
|
|
assert query["options"] == ["-c search_path=deerflow"]
|
|
|
|
@staticmethod
|
|
def _raw_query(out: str) -> dict[str, str]:
|
|
# Read the pairs back undecoded: parse_qs applies the HTML-form rule that
|
|
# turns '+' into a space, which is exactly the rule libpq does NOT apply.
|
|
return dict(pair.split("=", 1) for pair in urlsplit(out).query.split("&") if pair)
|
|
|
|
def test_preserves_plus_in_existing_query_value(self):
|
|
# A '+' in a parameter value is a literal plus for libpq, so rewriting it
|
|
# into '%20' silently changes the value the server receives: 'web+app'
|
|
# would arrive as 'web app'.
|
|
dsn = "postgresql://u:p@h:5432/db?application_name=web+app"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
assert self._raw_query(out)["application_name"] == "web+app"
|
|
assert parse_qs(urlsplit(out).query)["options"] == ["-c search_path=deerflow"]
|
|
|
|
def test_preserves_plus_in_existing_options_value(self):
|
|
dsn = "postgresql://u:p@h:5432/db?options=-c%20timezone%3DUTC+8"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
assert self._raw_query(out)["options"] == "-c%20timezone%3DUTC%2B8%20-c%20search_path%3Ddeerflow"
|
|
assert parse_qs(urlsplit(out).query, keep_blank_values=True)["options"] == ["-c timezone=UTC+8 -c search_path=deerflow"]
|
|
|
|
def test_does_not_double_encode_an_existing_percent_escape(self):
|
|
dsn = "postgresql://u:p@h:5432/db?application_name=web%2Bapp"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
assert self._raw_query(out)["application_name"] == "web%2Bapp"
|
|
|
|
def test_replaces_existing_options_query(self):
|
|
dsn = "postgresql://u:p@h:5432/db?options=-c%20search_path%3Dpublic"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
query = parse_qs(urlsplit(out).query)
|
|
assert query["options"] == ["-c search_path=deerflow"]
|
|
|
|
def test_preserves_existing_options_query(self):
|
|
dsn = "postgresql://u:p@h:5432/db?options=-c%20statement_timeout%3D5000"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
query = parse_qs(urlsplit(out).query)
|
|
assert query["options"] == ["-c statement_timeout=5000 -c search_path=deerflow"]
|
|
|
|
def test_replaces_only_existing_search_path_option(self):
|
|
dsn = "postgresql://u:p@h:5432/db?options=-c%20statement_timeout%3D5000%20-c%20search_path%3Dpublic"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
query = parse_qs(urlsplit(out).query)
|
|
assert query["options"] == ["-c statement_timeout=5000 -c search_path=deerflow"]
|
|
|
|
def test_supports_keyword_dsn(self):
|
|
pytest.importorskip("psycopg")
|
|
from psycopg.conninfo import conninfo_to_dict
|
|
|
|
dsn = "host=localhost dbname=deerflow user=postgres"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
assert conninfo_to_dict(out) == {
|
|
"host": "localhost",
|
|
"dbname": "deerflow",
|
|
"user": "postgres",
|
|
"options": "-c search_path=deerflow",
|
|
}
|
|
|
|
def test_preserves_keyword_dsn_options(self):
|
|
pytest.importorskip("psycopg")
|
|
from psycopg.conninfo import conninfo_to_dict
|
|
|
|
dsn = "host=localhost dbname=deerflow options='-c statement_timeout=5000'"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
assert conninfo_to_dict(out)["options"] == "-c statement_timeout=5000 -c search_path=deerflow"
|
|
|
|
def test_normalizes_sqlalchemy_driver_scheme(self):
|
|
# DatabaseConfig.postgres_url may carry a +asyncpg suffix; the libpq DSN
|
|
# produced for psycopg must drop the driver and still inject search_path.
|
|
dsn = "postgresql+asyncpg://u:p@h:5432/db"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
parts = urlsplit(out)
|
|
assert parts.scheme == "postgresql"
|
|
query = parse_qs(parts.query)
|
|
assert query["options"] == ["-c search_path=deerflow"]
|
|
|
|
def test_rejects_non_postgres_url_scheme(self):
|
|
try:
|
|
dsn_with_search_path("mysql://localhost/db", "deerflow")
|
|
except ValueError as exc:
|
|
assert "Unsupported PostgreSQL DSN scheme" in str(exc)
|
|
else:
|
|
raise AssertionError("Expected ValueError")
|
|
|
|
def test_roundtrip_preserves_host_and_db(self):
|
|
dsn = "postgresql://u:p@h:5432/db"
|
|
out = dsn_with_search_path(dsn, "deerflow")
|
|
parts = urlsplit(out)
|
|
assert parts.hostname == "h"
|
|
assert parts.port == 5432
|
|
assert parts.path == "/db"
|
|
|
|
def test_preserves_option_value_containing_space(self):
|
|
# libpq's options parameter separates args on spaces unless they are
|
|
# backslash-escaped. shlex.join would emit single-quotes, which libpq
|
|
# treats as literal characters and would corrupt the option. A token
|
|
# carrying a space must round-trip as a single backslash-escaped token.
|
|
from deerflow.persistence.postgres_schema import _merge_search_path_option
|
|
|
|
merged = _merge_search_path_option(r"-c application_name=My\ App", "deerflow")
|
|
assert "'" not in merged
|
|
assert r"application_name=My\ App" in merged
|
|
assert merged.endswith("-c search_path=deerflow")
|
|
|
|
def test_preserves_option_value_containing_tab(self):
|
|
# Non-space whitespace (TAB/CR/LF) inside an existing escaped token must
|
|
# also be re-escaped on re-join, otherwise libpq re-tokenizes on the bare
|
|
# whitespace byte and the round-trip is lossy.
|
|
from deerflow.persistence.postgres_schema import (
|
|
_merge_search_path_option,
|
|
_split_libpq_options,
|
|
)
|
|
|
|
merged = _merge_search_path_option("-c application_name=My\\\tApp", "deerflow")
|
|
assert "'" not in merged
|
|
# The tab-bearing value must round-trip back to a single token.
|
|
tokens = _split_libpq_options(merged)
|
|
assert "application_name=My\tApp" in tokens
|
|
assert merged.endswith("-c search_path=deerflow")
|
|
|
|
|
|
class TestNormalizeLibpqDsn:
|
|
def test_strips_asyncpg_driver_suffix(self):
|
|
assert normalize_libpq_dsn("postgresql+asyncpg://u:p@h:5432/db") == "postgresql://u:p@h:5432/db"
|
|
|
|
def test_leaves_bare_postgres_scheme_unchanged(self):
|
|
dsn = "postgresql://u:p@h:5432/db"
|
|
assert normalize_libpq_dsn(dsn) == dsn
|
|
|
|
def test_leaves_keyword_dsn_unchanged(self):
|
|
dsn = "host=localhost dbname=deerflow"
|
|
assert normalize_libpq_dsn(dsn) == dsn
|
|
|
|
def test_rejects_non_postgres_scheme(self):
|
|
with pytest.raises(ValueError, match="Unsupported PostgreSQL DSN scheme"):
|
|
normalize_libpq_dsn("mysql://localhost/db")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_schema_close_drains_across_repeated_cancellation(monkeypatch) -> None:
|
|
class _BlockingConnection:
|
|
def __init__(self) -> None:
|
|
self.execute_started = asyncio.Event()
|
|
self.allow_execute = asyncio.Event()
|
|
self.close_started = asyncio.Event()
|
|
self.allow_close = asyncio.Event()
|
|
self.close_finished = asyncio.Event()
|
|
|
|
async def execute(self, _statement: str) -> None:
|
|
self.execute_started.set()
|
|
await self.allow_execute.wait()
|
|
|
|
async def close(self) -> None:
|
|
self.close_started.set()
|
|
await self.allow_close.wait()
|
|
self.close_finished.set()
|
|
|
|
conn = _BlockingConnection()
|
|
|
|
class _AsyncConnection:
|
|
@staticmethod
|
|
async def connect(_dsn: str, *, autocommit: bool):
|
|
assert autocommit is True
|
|
return conn
|
|
|
|
monkeypatch.setitem(sys.modules, "psycopg", SimpleNamespace(AsyncConnection=_AsyncConnection))
|
|
task: asyncio.Task[None] | None = None
|
|
|
|
try:
|
|
task = asyncio.create_task(
|
|
ensure_postgres_schema_async(
|
|
"postgresql://user:pass@localhost/deerflow",
|
|
"deerflow",
|
|
install_hint="install postgres extras",
|
|
)
|
|
)
|
|
await asyncio.wait_for(conn.execute_started.wait(), timeout=1)
|
|
|
|
task.cancel()
|
|
await asyncio.wait_for(conn.close_started.wait(), timeout=1)
|
|
|
|
task.cancel()
|
|
for _ in range(5):
|
|
await asyncio.sleep(0)
|
|
assert not task.done(), "schema setup returned before psycopg connection close finished"
|
|
assert not conn.close_finished.is_set()
|
|
|
|
conn.allow_close.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert conn.close_finished.is_set()
|
|
finally:
|
|
conn.allow_execute.set()
|
|
conn.allow_close.set()
|
|
if task is not None and not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|