diff --git a/backend/packages/harness/deerflow/persistence/postgres_schema.py b/backend/packages/harness/deerflow/persistence/postgres_schema.py index 9d07f81b56..35167cc1a8 100644 --- a/backend/packages/harness/deerflow/persistence/postgres_schema.py +++ b/backend/packages/harness/deerflow/persistence/postgres_schema.py @@ -20,7 +20,7 @@ defense-in-depth; connection-argument helpers only assemble driver payloads. from __future__ import annotations import re -from urllib.parse import parse_qsl, quote, urlencode, urlsplit, urlunsplit +from urllib.parse import quote, unquote, urlsplit, urlunsplit from deerflow.utils.file_io import await_drained @@ -200,18 +200,26 @@ def dsn_with_search_path(dsn: str, schema: str) -> str: 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' + # and change the value the server sees. Everything but ``options`` is carried + # over exactly as written. options_values: list[str] = [] - query_pairs = [] - for key, value in parse_qsl(parts.query, keep_blank_values=True): - if key == "options": - options_values.append(value) + preserved: list[str] = [] + for pair in parts.query.split("&"): + if not pair: + continue + raw_key, _, raw_value = pair.partition("=") + if unquote(raw_key) == "options": + options_values.append(unquote(raw_value)) else: - query_pairs.append((key, value)) + preserved.append(pair) options = _merge_search_path_option(" ".join(options_values), schema) - query_pairs.append(("options", options)) - # quote_via=quote encodes space as %20 (libpq-safe), not + (form-style). - query = urlencode(query_pairs, quote_via=quote) + # 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)) diff --git a/backend/tests/test_postgres_schema_helper.py b/backend/tests/test_postgres_schema_helper.py index bf25ec2290..0ede45ab91 100644 --- a/backend/tests/test_postgres_schema_helper.py +++ b/backend/tests/test_postgres_schema_helper.py @@ -71,6 +71,32 @@ class TestDsnWithSearchPath: 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")