Add support for the sslnegotiation parameter (#1187)

Direct TLS connections are already supported via the `direct_tls` argument, however PostgreSQL 17 added native support for this via `sslnegotiation`, so recognize it in DSNs and the environment. I decided not to introduce the `sslnegotiation` connection constructor argument for now, `direct_tls` should continue to be used instead.

Author: elprans

asyncpg/compat.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import enum
 import pathlib
 import platform
 import typing
@@ -78,3 +79,10 @@ def markcoroutinefunction(c): # type: ignore
  from collections.abc import ( # noqa: F401
  Awaitable as Awaitable,
  )
+
+if sys.version_info < (3, 11):
+ class StrEnum(str, enum.Enum):
+ __str__ = str.__str__
+ __repr__ = enum.Enum.__repr__
+else:
+ from enum import StrEnum as StrEnum # noqa: F401

asyncpg/connect_utils.py

@@ -45,6 +45,11 @@ def parse(cls, sslmode):
  return getattr(cls, sslmode.replace('-', '_'))
 
 
+class SSLNegotiation(compat.StrEnum):
+ postgres = "postgres"
+ direct = "direct"
+
+
 _ConnectionParameters = collections.namedtuple(
  'ConnectionParameters',
  [
@@ -53,7 +58,7 @@ def parse(cls, sslmode):
  'database',
  'ssl',
  'sslmode',
- 'direct_tls',
+ 'ssl_negotiation',
  'server_settings',
  'target_session_attrs',
  'krbsrvname',
@@ -269,6 +274,7 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
  auth_hosts = None
  sslcert = sslkey = sslrootcert = sslcrl = sslpassword = None
  ssl_min_protocol_version = ssl_max_protocol_version = None
+ sslnegotiation = None
 
  if dsn:
  parsed = urllib.parse.urlparse(dsn)
@@ -362,6 +368,9 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
  if 'sslrootcert' in query:
  sslrootcert = query.pop('sslrootcert')
 
+ if 'sslnegotiation' in query:
+ sslnegotiation = query.pop('sslnegotiation')
+
  if 'sslcrl' in query:
  sslcrl = query.pop('sslcrl')
 
@@ -503,13 +512,36 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
  if ssl is None and have_tcp_addrs:
  ssl = 'prefer'
 
+ if direct_tls is not None:
+ sslneg = (
+ SSLNegotiation.direct if direct_tls else SSLNegotiation.postgres
+ )
+ else:
+ if sslnegotiation is None:
+ sslnegotiation = os.environ.get("PGSSLNEGOTIATION")
+
+ if sslnegotiation is not None:
+ try:
+ sslneg = SSLNegotiation(sslnegotiation)
+ except ValueError:
+ modes = ', '.join(
+ m.name.replace('_', '-')
+ for m in SSLNegotiation
+ )
+ raise exceptions.ClientConfigurationError(
+ f'`sslnegotiation` parameter must be one of: {modes}'
+ ) from None
+ else:
+ sslneg = SSLNegotiation.postgres
+
  if isinstance(ssl, (str, SSLMode)):
  try:
  sslmode = SSLMode.parse(ssl)
  except AttributeError:
  modes = ', '.join(m.name.replace('_', '-') for m in SSLMode)
  raise exceptions.ClientConfigurationError(
- '`sslmode` parameter must be one of: {}'.format(modes))
+ '`sslmode` parameter must be one of: {}'.format(modes)
+ ) from None
 
  # docs at https://www.postgresql.org/docs/10/static/libpq-connect.html
  if sslmode < SSLMode.allow:
@@ -676,7 +708,7 @@ def _parse_connect_dsn_and_args(*, dsn, host, port, user,
 
  params = _ConnectionParameters(
  user=user, password=password, database=database, ssl=ssl,
- sslmode=sslmode, direct_tls=direct_tls,
+ sslmode=sslmode, ssl_negotiation=sslneg,
  server_settings=server_settings,
  target_session_attrs=target_session_attrs,
  krbsrvname=krbsrvname, gsslib=gsslib)
@@ -882,9 +914,9 @@ async def __connect_addr(
  # UNIX socket
  connector = loop.create_unix_connection(proto_factory, addr)
 
- elif params.ssl and params.direct_tls:
- # if ssl and direct_tls are given, skip STARTTLS and perform direct
- # SSL connection
+ elif params.ssl and params.ssl_negotiation is SSLNegotiation.direct:
+ # if ssl and ssl_negotiation is `direct`, skip STARTTLS and perform
+ # direct SSL connection
  connector = loop.create_connection(
  proto_factory, *addr, ssl=params.ssl
  )

asyncpg/connection.py

@@ -2001,7 +2001,7 @@ async def connect(dsn=None, *,
  max_cacheable_statement_size=1024 * 15,
  command_timeout=None,
  ssl=None,
- direct_tls=False,
+ direct_tls=None,
  connection_class=Connection,
  record_class=protocol.Record,
  server_settings=None,

pyproject.toml

@@ -112,6 +112,15 @@ exclude_lines = [
 show_missing = true
 
 [tool.mypy]
+exclude = [
+ "^.eggs",
+ "^.github",
+ "^.vscode",
+ "^build",
+ "^dist",
+ "^docs",
+ "^tests",
+]
 incremental = true
 strict = true
 implicit_reexport = true

tests/test_connect.py

@@ -592,6 +592,58 @@ class TestConnectParams(tb.TestCase):
  'target_session_attrs': 'any'})
  },
 
+ {
+ 'name': 'params_ssl_negotiation_dsn',
+ 'env': {
+ 'PGSSLNEGOTIATION': 'postgres'
+ },
+
+ 'dsn': 'postgres://u:p@localhost/d?sslnegotiation=direct',
+
+ 'result': ([('localhost', 5432)], {
+ 'user': 'u',
+ 'password': 'p',
+ 'database': 'd',
+ 'ssl_negotiation': 'direct',
+ 'target_session_attrs': 'any',
+ })
+ },
+
+ {
+ 'name': 'params_ssl_negotiation_env',
+ 'env': {
+ 'PGSSLNEGOTIATION': 'direct'
+ },
+
+ 'dsn': 'postgres://u:p@localhost/d',
+
+ 'result': ([('localhost', 5432)], {
+ 'user': 'u',
+ 'password': 'p',
+ 'database': 'd',
+ 'ssl_negotiation': 'direct',
+ 'target_session_attrs': 'any',
+ })
+ },
+
+ {
+ 'name': 'params_ssl_negotiation_params',
+ 'env': {
+ 'PGSSLNEGOTIATION': 'direct'
+ },
+
+ 'dsn': 'postgres://u:p@localhost/d',
+ 'direct_tls': False,
+
+ 'result': ([('localhost', 5432)], {
+ 'user': 'u',
+ 'password': 'p',
+ 'database': 'd',
+ 'ssl_negotiation': 'postgres',
+ 'target_session_attrs': 'any',
+ })
+ },
+
  {
  'name': 'dsn_overrides_env_partially_ssl_prefer',
  'env': {
@@ -1067,6 +1119,7 @@ def run_testcase(self, testcase):
  passfile = testcase.get('passfile')
  database = testcase.get('database')
  sslmode = testcase.get('ssl')
+ direct_tls = testcase.get('direct_tls')
  server_settings = testcase.get('server_settings')
  target_session_attrs = testcase.get('target_session_attrs')
  krbsrvname = testcase.get('krbsrvname')
@@ -1093,7 +1146,7 @@ def run_testcase(self, testcase):
  addrs, params = connect_utils._parse_connect_dsn_and_args(
  dsn=dsn, host=host, port=port, user=user, password=password,
  passfile=passfile, database=database, ssl=sslmode,
- direct_tls=False,
+ direct_tls=direct_tls,
  server_settings=server_settings,
  target_session_attrs=target_session_attrs,
  krbsrvname=krbsrvname, gsslib=gsslib)
@@ -1118,6 +1171,10 @@ def run_testcase(self, testcase):
  # Avoid the hassle of specifying direct_tls
  # unless explicitly tested for
  params.pop('direct_tls', False)
+ if 'ssl_negotiation' not in expected[1]:
+ # Avoid the hassle of specifying sslnegotiation
+ # unless explicitly tested for
+ params.pop('ssl_negotiation', False)
  if 'gsslib' not in expected[1]:
  # Avoid the hassle of specifying gsslib
  # unless explicitly tested for