fix: make is_dsn_multihost agree with the connection plan

db_retry/connections.py

@@ -27,8 +27,17 @@ class ConnectionPlan:
     failover: tuple[tuple[str, int], ...]
 
 
+def parse_connect_args(url: sqlalchemy.URL) -> tuple[dict[str, typing.Any], list[tuple[str, int]]]:
+    connect_args: typing.Final[dict[str, typing.Any]] = PGDialect_asyncpg().create_connect_args(url)[1]
+    hosts: typing.Final = connect_args.get("host")
+    ports: typing.Final = connect_args.get("port")
+    if isinstance(hosts, list) and isinstance(ports, list):
+        return connect_args, list(zip(hosts, ports, strict=True))
+    return connect_args, []
+
+
 def build_connection_plan(url: sqlalchemy.URL) -> ConnectionPlan:
-    connect_args: dict[str, typing.Any] = PGDialect_asyncpg().create_connect_args(url)[1]
+    connect_args, hosts_and_ports = parse_connect_args(url)
     raw_target_session_attrs: str | None = connect_args.pop("target_session_attrs", None)
     target_session_attrs: SessionAttribute | None = (
         SessionAttribute(raw_target_session_attrs) if raw_target_session_attrs else None
@@ -38,8 +47,7 @@ def build_connection_plan(url: sqlalchemy.URL) -> ConnectionPlan:
     primary_host: str | list[str]
     primary_port: int | list[int] | None
     failover: tuple[tuple[str, int], ...]
-    if isinstance(raw_hosts, list) and isinstance(raw_ports, list):
-        hosts_and_ports: list[tuple[str, int]] = list(zip(raw_hosts, raw_ports, strict=True))
+    if hosts_and_ports:
         random.shuffle(hosts_and_ports)
         primary_host = list(map(itemgetter(0), hosts_and_ports))
         primary_port = list(map(itemgetter(1), hosts_and_ports))

db_retry/dsn.py

@@ -2,6 +2,8 @@
 
 import sqlalchemy as sa
 
+from db_retry.connections import parse_connect_args
+
 
 def build_db_dsn(
     db_dsn: str,
@@ -11,7 +13,7 @@ def build_db_dsn(
 ) -> sa.URL:
     """Parse DSN variable and replace some parts.
 
-    - DSN stored in format postgresql://login:password@/db_placeholder?host=host1&host=host2
+    - DSN stored in format postgresql://login:password@/db_placeholder?host=host1:5432&host=host2:5432
     https://docs.sqlalchemy.org/en/20/dialects/postgresql.html#specifying-multiple-fallback-hosts
     - 'db_placeholder' is replaced here with service database name
     - `target_session_attrs` is chosen based on `use_replica` arg
@@ -29,6 +31,4 @@ def build_db_dsn(
 
 
 def is_dsn_multihost(db_dsn: str) -> bool:
-    parsed_db_dsn: typing.Final = sa.make_url(db_dsn)
-    db_dsn_query: typing.Final[dict[str, typing.Any]] = dict(parsed_db_dsn.query or {})
-    return bool((hosts := db_dsn_query.get("host")) and isinstance(hosts, tuple) and len(hosts) > 1)
+    return bool(parse_connect_args(sa.make_url(db_dsn))[1])

tests/test_dsn.py

@@ -1,13 +1,18 @@
 import typing
 
+import pytest
+import sqlalchemy
+from sqlalchemy import exc as sa_exc
+
 from db_retry import build_db_dsn, is_dsn_multihost
+from db_retry.connections import build_connection_plan
 
 
 def test_build_db_dsn() -> None:
     database_name: typing.Final = "new_db_name"
     drivername: typing.Final = "postgresql+asyncpg"
     result_dsn: typing.Final = build_db_dsn(
-        db_dsn="postgresql://login:password@/db_placeholder?host=host1&host=host2",
+        db_dsn="postgresql://login:password@/db_placeholder?host=host1:5432&host=host2:5432",
         database_name=database_name,
         drivername=drivername,
     )
@@ -30,7 +35,7 @@ def test_use_replica_overrides_a_target_session_attrs_already_in_the_dsn() -> No
     """
     dsn: typing.Final = (
         "postgresql://login:password@/db_placeholder"
-        "?host=host1&host=host2&target_session_attrs=read-write&application_name=svc"
+        "?host=host1:5432&host=host2:5432&target_session_attrs=read-write&application_name=svc"
     )
 
     replica: typing.Final = build_db_dsn(db_dsn=dsn, database_name="db", use_replica=True)
@@ -43,7 +48,29 @@ def test_use_replica_overrides_a_target_session_attrs_already_in_the_dsn() -> No
     assert primary.query["target_session_attrs"] == "read-write"
 
 
+@pytest.mark.parametrize(
+    "db_dsn",
+    [
+        "postgresql://login:password@/db_placeholder?host=host1:5432&host=host2:5433",
+        "postgresql://login:password@/db_placeholder?host=host1,host2&port=5432,5433",
+        "postgresql://login:password@/db_placeholder?host=host1",
+        "postgresql://login:password@/db_placeholder?host=host1:5432",
+        "postgresql://login:password@host1:5432/db_placeholder",
+        "postgresql+asyncpg://login:password@/db_placeholder?host=host1,host2&port=5432,5433",
+    ],
+)
+def test_is_dsn_multihost_agrees_with_connection_plan(db_dsn: str) -> None:
+    assert is_dsn_multihost(db_dsn) == bool(build_connection_plan(sqlalchemy.make_url(db_dsn)).failover)
+
+
 def test_is_dsn_multihost() -> None:
-    assert is_dsn_multihost("postgresql://login:password@/db_placeholder?host=host1&host=host2")
+    assert is_dsn_multihost("postgresql://login:password@/db_placeholder?host=host1:5432&host=host2:5432")
+    assert is_dsn_multihost("postgresql://login:password@/db_placeholder?host=host1,host2&port=5432,5432")
     assert not is_dsn_multihost("postgresql://login:password@/db_placeholder?host=host1")
     assert not is_dsn_multihost("postgresql://login:password@host/db_placeholder")
+    assert not is_dsn_multihost("postgresql://login:password@/db_placeholder")
+
+
+def test_is_dsn_multihost_rejects_what_the_connection_factory_rejects() -> None:
+    with pytest.raises(sa_exc.ArgumentError):
+        is_dsn_multihost("postgresql://login:password@/db_placeholder?host=host1&host=host2")