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.

diff --git a/docs/news.rst b/docs/news.rst
index eb4cdc2c734d18c3f90a939f0267234f6dc27dcb..dbbd92d977b567dc5e241181fdc8d5ca654749fd 100644 (file)
--- 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 37809a8d4e4f3cb4f01cf4a81076a7feea6fba62..7bc96dd9976b0e2e659d2cde11c9242bf1b3822c 100644 (file)
--- 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 5605f23ad164f71f63ca73362977f28f623f319c..e50e4f95efe534f4a6cb1b42d8ff4c9e47267e21 100644 (file)
--- 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 39f0c2b98cdb638f7561d49ae66161de05ccf0e8..4e477f6f0b05971a19c38f8139d82c2ec7abb089 100644 (file)
--- 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 3cffd899240d6325c6c44509d232a31abd516954..7662ea9e51788e0942dc3c945adf3d1b5dbc7760 100644 (file)
--- 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 d081f22af21d8d788f97dbc51fd26cd2cf083840..25bc0dfd74208bf09be992154a49c432040b2427 100644 (file)
--- 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 403795481af80bf1257daf13d35836a16ee56b16..dcf1a51f74548d5815d7869471cfffdb4a1935b5 100644 (file)
--- 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):