Update based on review.

Author: Jesse De Loore

asyncpg/connect_utils.py

@@ -57,7 +57,7 @@ def parse(cls, sslmode):
         'direct_tls',
         'connect_timeout',
         'server_settings',
-        'target_session_attribute',
+        'target_session_attrs',
     ])
 
 
@@ -258,7 +258,7 @@ def _dot_postgresql_path(filename) -> pathlib.Path:
 def _parse_connect_dsn_and_args(*, dsn, host, port, user,
                                 password, passfile, database, ssl,
                                 direct_tls, connect_timeout, server_settings,
-                                target_session_attribute):
+                                target_session_attrs):
     # `auth_hosts` is the version of host information for the purposes
     # of reading the pgpass file.
     auth_hosts = None
@@ -595,11 +595,24 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
             'server_settings is expected to be None or '
             'a Dict[str, str]')
 
+    if target_session_attrs is None:
+
+        target_session_attrs = os.getenv("PGTARGETSESSIONATTRS", SessionAttribute.any)
+    try:
+
+        target_session_attrs = SessionAttribute(target_session_attrs)
+    except ValueError as exc:
+        raise exceptions.InterfaceError(
+            "target_session_attrs is expected to be one of "
+            "{!r}"
+            ", got {!r}".format(SessionAttribute.__members__.values, target_session_attrs)
+        ) from exc
+
     params = _ConnectionParameters(
         user=user, password=password, database=database, ssl=ssl,
         sslmode=sslmode, direct_tls=direct_tls,
         connect_timeout=connect_timeout, server_settings=server_settings,
-        target_session_attribute=target_session_attribute)
+        target_session_attrs=target_session_attrs)
 
     return addrs, params
 
@@ -610,7 +623,7 @@ def _parse_connect_arguments(*, dsn, host, port, user, password, passfile,
                              max_cached_statement_lifetime,
                              max_cacheable_statement_size,
                              ssl, direct_tls, server_settings,
-                             target_session_attribute):
+                             target_session_attrs):
     local_vars = locals()
     for var_name in {'max_cacheable_statement_size',
                      'max_cached_statement_lifetime',
@@ -639,7 +652,7 @@ def _parse_connect_arguments(*, dsn, host, port, user, password, passfile,
         password=password, passfile=passfile, ssl=ssl,
         direct_tls=direct_tls, database=database,
         connect_timeout=timeout, server_settings=server_settings,
-        target_session_attribute=target_session_attribute)
+        target_session_attrs=target_session_attrs)
 
     config = _ClientConfiguration(
         command_timeout=command_timeout,
@@ -941,7 +954,7 @@ async def _connect(*, loop, timeout, connection_class, record_class, **kwargs):
         loop = asyncio.get_event_loop()
 
     addrs, params, config = _parse_connect_arguments(timeout=timeout, **kwargs)
-    target_attr = params.target_session_attribute
+    target_attr = params.target_session_attrs
 
     candidates = []
     chosen_connection = None

asyncpg/connection.py

@@ -1794,7 +1794,7 @@ async def connect(dsn=None, *,
                   connection_class=Connection,
                   record_class=protocol.Record,
                   server_settings=None,
-                  target_session_attribute=SessionAttribute.any):
+                  target_session_attrs=None):
     r"""A coroutine to establish a connection to a PostgreSQL server.
 
     The connection parameters may be specified either as a connection
@@ -2005,16 +2005,21 @@ async def connect(dsn=None, *,
         this connection object.  Must be a subclass of
         :class:`~asyncpg.Record`.
 
-    :param SessionAttribute target_session_attribute:
+    :param SessionAttribute target_session_attrs:
         If specified, check that the host has the correct attribute.
         Can be one of:
             "any": the first successfully connected host
             "primary": the host must NOT be in hot standby mode
             "standby": the host must be in hot standby mode
+            "read-write": the host must allow writes
+            "read-only": the host most NOT allow writes
             "prefer-standby": first try to find a standby host, but if
                             none of the listed hosts is a standby server,
                             return any of them.
 
+        If not specified will try to use PGTARGETSESSIONATTRS from the environment.
+        Defaults to "any" if no value is set.
+
     :return: A :class:`~asyncpg.connection.Connection` instance.
 
     Example:
@@ -2099,15 +2104,6 @@ async def connect(dsn=None, *,
     if record_class is not protocol.Record:
         _check_record_class(record_class)
 
-    try:
-        target_session_attribute = SessionAttribute(target_session_attribute)
-    except ValueError as exc:
-        raise exceptions.InterfaceError(
-            "target_session_attribute is expected to be one of "
-            "'any', 'primary', 'standby' or 'prefer-standby'"
-            ", got {!r}".format(target_session_attribute)
-        ) from exc
-
     if loop is None:
         loop = asyncio.get_event_loop()
 
@@ -2130,7 +2126,7 @@ async def connect(dsn=None, *,
         statement_cache_size=statement_cache_size,
         max_cached_statement_lifetime=max_cached_statement_lifetime,
         max_cacheable_statement_size=max_cacheable_statement_size,
-        target_session_attribute=target_session_attribute
+        target_session_attrs=target_session_attrs
     )
 
 

tests/test_connect.py

@@ -788,7 +788,7 @@ def run_testcase(self, testcase):
         database = testcase.get('database')
         sslmode = testcase.get('ssl')
         server_settings = testcase.get('server_settings')
-        target_session_attribute = testcase.get('target_session_attribute')
+        target_session_attrs = testcase.get('target_session_attrs')
 
         expected = testcase.get('result')
         expected_error = testcase.get('error')
@@ -813,7 +813,7 @@ def run_testcase(self, testcase):
                 passfile=passfile, database=database, ssl=sslmode,
                 direct_tls=False, connect_timeout=None,
                 server_settings=server_settings,
-                target_session_attribute=target_session_attribute)
+                target_session_attrs=target_session_attrs)
 
             params = {
                 k: v for k, v in params._asdict().items()
@@ -1750,7 +1750,7 @@ class TestConnectionAttributes(tb.HotStandbyTestCase):
     async def _run_connection_test(
         self, connect, target_attribute, expected_port
     ):
-        conn = await connect(target_session_attribute=target_attribute)
+        conn = await connect(target_session_attrs=target_attribute)
         self.assertTrue(_get_connected_host(conn).endswith(expected_port))
         await conn.close()
 
@@ -1790,7 +1790,7 @@ async def test_target_attribute_not_matched(self):
 
         for connect, target_attr in tests:
             with self.assertRaises(exceptions.TargetServerAttributeNotMatched):
-                await connect(target_session_attribute=target_attr)
+                await connect(target_session_attrs=target_attr)
 
         if self.master_cluster.get_pg_version()[0] < 14:
             self.skipTest("PostgreSQL<14 does not support these features")
@@ -1801,12 +1801,12 @@ async def test_target_attribute_not_matched(self):
 
         for connect, target_attr in tests:
             with self.assertRaises(exceptions.TargetServerAttributeNotMatched):
-                await connect(target_session_attribute=target_attr)
+                await connect(target_session_attrs=target_attr)
 
     async def test_prefer_standby_when_standby_is_up(self):
         if self.master_cluster.get_pg_version()[0] == 11:
             self.skipTest("PostgreSQL 11 seems to have issues with this test")
-        con = await self.connect(target_session_attribute='prefer-standby')
+        con = await self.connect(target_session_attrs='prefer-standby')
         standby_port = self.standby_cluster.get_connection_spec()['port']
         connected_host = _get_connected_host(con)
         self.assertTrue(connected_host.endswith(standby_port))
@@ -1824,7 +1824,7 @@ async def test_prefer_standby_picks_master_when_standby_is_down(self):
             'port': [primary_spec['port'], 15345],
             'database': primary_spec['database'],
             'user': primary_spec['user'],
-            'target_session_attribute': 'prefer-standby'
+            'target_session_attrs': 'prefer-standby'
         }
 
         con = await self.connect(**connection_spec)