diff --git a/db_retry/connections.py b/db_retry/connections.py index 740240a..de1c9ea 100644 --- a/db_retry/connections.py +++ b/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)) diff --git a/db_retry/dsn.py b/db_retry/dsn.py index 273b612..17c69fe 100644 --- a/db_retry/dsn.py +++ b/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]) diff --git a/tests/test_dsn.py b/tests/test_dsn.py index 6283baa..60bfdd0 100644 --- a/tests/test_dsn.py +++ b/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")