From b37dc1a35b8c24707bc844f8cc8ee965121ba25b Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Mon, 10 Mar 2025 23:53:31 +0100 Subject: [PATCH] fix: implement libpq semantic for target_session_attrs=prefer-standby First attempt all the servers in standby mode, then fall back to any mode. Fix #1021. --- docs/news.rst | 7 +++++++ psycopg/psycopg/_conninfo_attempts.py | 12 ++++++++++++ psycopg/psycopg/_conninfo_attempts_async.py | 10 ++++++++++ tests/test_connection.py | 17 +++++++++++++++++ tests/test_connection_async.py | 17 +++++++++++++++++ tests/test_conninfo_attempts.py | 20 ++++++++++++++++++++ tests/test_conninfo_attempts_async.py | 20 ++++++++++++++++++++ 7 files changed, 103 insertions(+) diff --git a/docs/news.rst b/docs/news.rst index eb4cdc2c7..dbbd92d97 100644 --- a/docs/news.rst +++ b/docs/news.rst @@ -7,6 +7,13 @@ ``psycopg`` release notes ========================= +Python 3.2.6 (unreleased) +^^^^^^^^^^^^^^^^^^^^^^^^^ + +- Fix connection semantic when using ``target_session_attrs=prefer-standby`` + (:ticket:`#1021`). + + Current release --------------- diff --git a/psycopg/psycopg/_conninfo_attempts.py b/psycopg/psycopg/_conninfo_attempts.py index 37809a8d4..7bc96dd99 100644 --- a/psycopg/psycopg/_conninfo_attempts.py +++ b/psycopg/psycopg/_conninfo_attempts.py @@ -35,6 +35,11 @@ def conninfo_attempts(params: ConnMapping) -> list[ConnDict]: """ last_exc = None attempts = [] + if prefer_standby := ( + get_param(params, "target_session_attrs") == "prefer-standby" + ): + params = {k: v for k, v in params.items() if k != "target_session_attrs"} + for attempt in split_attempts(params): try: attempts.extend(_resolve_hostnames(attempt)) @@ -50,6 +55,13 @@ def conninfo_attempts(params: ConnMapping) -> list[ConnDict]: if get_param(params, "load_balance_hosts") == "random": shuffle(attempts) + # Order matters: first try all the load-balanced host in standby mode, + # then allow primary + if prefer_standby: + attempts = [ + {**a, "target_session_attrs": "standby"} for a in attempts + ] + attempts + return attempts diff --git a/psycopg/psycopg/_conninfo_attempts_async.py b/psycopg/psycopg/_conninfo_attempts_async.py index 5605f23ad..e50e4f95e 100644 --- a/psycopg/psycopg/_conninfo_attempts_async.py +++ b/psycopg/psycopg/_conninfo_attempts_async.py @@ -35,6 +35,9 @@ async def conninfo_attempts_async(params: ConnMapping) -> list[ConnDict]: """ last_exc = None attempts = [] + if prefer_standby := get_param(params, "target_session_attrs") == "prefer-standby": + params = {k: v for k, v in params.items() if k != "target_session_attrs"} + for attempt in split_attempts(params): try: attempts.extend(await _resolve_hostnames(attempt)) @@ -50,6 +53,13 @@ async def conninfo_attempts_async(params: ConnMapping) -> list[ConnDict]: if get_param(params, "load_balance_hosts") == "random": shuffle(attempts) + # Order matters: first try all the load-balanced host in standby mode, + # then allow primary + if prefer_standby: + attempts = [ + {**a, "target_session_attrs": "standby"} for a in attempts + ] + attempts + return attempts diff --git a/tests/test_connection.py b/tests/test_connection.py index 39f0c2b98..4e477f6f0 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -903,3 +903,20 @@ def test_right_exception_on_session_timeout(conn): # This check is here to monitor if the behaviour on Window chamge. # Rreceiving the same exception of other platform will be acceptable. assert type(ex.value) is want_ex + + +# NOTE: these "tsa" tests assume that the database we use for tests is a primary. + + +@pytest.mark.parametrize("mode", ["any", "read-write", "primary", "prefer-standby"]) +def test_connect_tsa(conn_cls, dsn, mode): + params = conninfo_to_dict(dsn, target_session_attrs=mode) + with conn_cls.connect(**params) as conn: + assert conn.pgconn.status == pq.ConnStatus.OK + + +@pytest.mark.parametrize("mode", ["read-only", "standby", "nosuchmode"]) +def test_connect_tsa_bad(conn_cls, dsn, mode): + params = conninfo_to_dict(dsn, target_session_attrs=mode) + with pytest.raises(psycopg.OperationalError, match=mode): + conn_cls.connect(**params) diff --git a/tests/test_connection_async.py b/tests/test_connection_async.py index 3cffd8992..7662ea9e5 100644 --- a/tests/test_connection_async.py +++ b/tests/test_connection_async.py @@ -909,3 +909,20 @@ async def test_right_exception_on_session_timeout(aconn): # This check is here to monitor if the behaviour on Window chamge. # Rreceiving the same exception of other platform will be acceptable. assert type(ex.value) is want_ex + + +# NOTE: these "tsa" tests assume that the database we use for tests is a primary. + + +@pytest.mark.parametrize("mode", ["any", "read-write", "primary", "prefer-standby"]) +async def test_connect_tsa(aconn_cls, dsn, mode): + params = conninfo_to_dict(dsn, target_session_attrs=mode) + async with await aconn_cls.connect(**params) as aconn: + assert aconn.pgconn.status == pq.ConnStatus.OK + + +@pytest.mark.parametrize("mode", ["read-only", "standby", "nosuchmode"]) +async def test_connect_tsa_bad(aconn_cls, dsn, mode): + params = conninfo_to_dict(dsn, target_session_attrs=mode) + with pytest.raises(psycopg.OperationalError, match=mode): + await aconn_cls.connect(**params) diff --git a/tests/test_conninfo_attempts.py b/tests/test_conninfo_attempts.py index d081f22af..25bc0dfd7 100644 --- a/tests/test_conninfo_attempts.py +++ b/tests/test_conninfo_attempts.py @@ -49,6 +49,16 @@ pytestmark = pytest.mark.anyio ], None, ), + ( + "host=1.1.1.1,2.2.2.2 target_session_attrs=prefer-standby", + [ + "host=1.1.1.1 hostaddr=1.1.1.1 target_session_attrs=standby", + "host=2.2.2.2 hostaddr=2.2.2.2 target_session_attrs=standby", + "host=1.1.1.1 hostaddr=1.1.1.1", + "host=2.2.2.2 hostaddr=2.2.2.2", + ], + None, + ), ( "port=5432", [ @@ -121,6 +131,16 @@ def test_conninfo_attempts_no_resolve(setpgenv, conninfo, want, env, fail_resolv ["host=dup.com hostaddr=3.3.3.3", "host=dup.com hostaddr=3.3.3.4"], None, ), + ( + "host=dup.com target_session_attrs=prefer-standby", + [ + "host=dup.com hostaddr=3.3.3.3 target_session_attrs=standby", + "host=dup.com hostaddr=3.3.3.4 target_session_attrs=standby", + "host=dup.com hostaddr=3.3.3.3", + "host=dup.com hostaddr=3.3.3.4", + ], + None, + ), ], ) def test_conninfo_attempts(conninfo, want, env, fake_resolve): diff --git a/tests/test_conninfo_attempts_async.py b/tests/test_conninfo_attempts_async.py index 403795481..dcf1a51f7 100644 --- a/tests/test_conninfo_attempts_async.py +++ b/tests/test_conninfo_attempts_async.py @@ -46,6 +46,16 @@ pytestmark = pytest.mark.anyio ], None, ), + ( + "host=1.1.1.1,2.2.2.2 target_session_attrs=prefer-standby", + [ + "host=1.1.1.1 hostaddr=1.1.1.1 target_session_attrs=standby", + "host=2.2.2.2 hostaddr=2.2.2.2 target_session_attrs=standby", + "host=1.1.1.1 hostaddr=1.1.1.1", + "host=2.2.2.2 hostaddr=2.2.2.2", + ], + None, + ), ( "port=5432", [ @@ -128,6 +138,16 @@ async def test_conninfo_attempts_no_resolve( ["host=dup.com hostaddr=3.3.3.3", "host=dup.com hostaddr=3.3.3.4"], None, ), + ( + "host=dup.com target_session_attrs=prefer-standby", + [ + "host=dup.com hostaddr=3.3.3.3 target_session_attrs=standby", + "host=dup.com hostaddr=3.3.3.4 target_session_attrs=standby", + "host=dup.com hostaddr=3.3.3.3", + "host=dup.com hostaddr=3.3.3.4", + ], + None, + ), ], ) async def test_conninfo_attempts(conninfo, want, env, fake_resolve): -- 2.47.2