From f900b73714f20470a566efd4d0b93ec0a565a93f Mon Sep 17 00:00:00 2001
From: Elvis Pranskevichus <elvis@edgedb.com>
Date: Mon, 15 Nov 2021 22:13:35 -0800
Subject: [PATCH] Fix parsing of IPv6 addresses in the connection URI (#845)

Plain IPv6 addresses specified in square brackets in the connection URI
are now parsed correctly.

Fixes: #838.
---
 asyncpg/connect_utils.py | 20 +++++++++++++++++---
 asyncpg/connection.py    |  8 +++++++-
 tests/test_connect.py    | 28 ++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py
index f6f9d651..f98935f5 100644
--- a/asyncpg/connect_utils.py
+++ b/asyncpg/connect_utils.py
@@ -197,11 +197,25 @@ def _parse_hostlist(hostlist, port, *, unquote=False):
         port = _validate_port_spec(hostspecs, port)
 
     for i, hostspec in enumerate(hostspecs):
-        if not hostspec.startswith('/'):
-            addr, _, hostspec_port = hostspec.partition(':')
-        else:
+        if hostspec[0] == '/':
+            # Unix socket
             addr = hostspec
             hostspec_port = ''
+        elif hostspec[0] == '[':
+            # IPv6 address
+            m = re.match(r'(?:\[([^\]]+)\])(?::([0-9]+))?', hostspec)
+            if m:
+                addr = m.group(1)
+                hostspec_port = m.group(2)
+            else:
+                raise ValueError(
+                    'invalid IPv6 address in the connection URI: {!r}'.format(
+                        hostspec
+                    )
+                )
+        else:
+            # IPv4 address
+            addr, _, hostspec_port = hostspec.partition(':')
 
         if unquote:
             addr = urllib.parse.unquote(addr)
diff --git a/asyncpg/connection.py b/asyncpg/connection.py
index ac2b5e31..09aa3dac 100644
--- a/asyncpg/connection.py
+++ b/asyncpg/connection.py
@@ -1813,7 +1813,13 @@ async def connect(dsn=None, *,
         .. note::
 
            The URI must be *valid*, which means that all components must
-           be properly quoted with :py:func:`urllib.parse.quote`.
+           be properly quoted with :py:func:`urllib.parse.quote`, and
+           any literal IPv6 addresses must be enclosed in square brackets.
+           For example:
+
+           .. code-block:: text
+
+              postgres://dbuser@[fe80::1ff:fe23:4567:890a%25eth0]/dbname
 
     :param host:
         Database host address as one of the following:
diff --git a/tests/test_connect.py b/tests/test_connect.py
index f0b1ca07..d66a087b 100644
--- a/tests/test_connect.py
+++ b/tests/test_connect.py
@@ -511,6 +511,34 @@ class TestConnectParams(tb.TestCase):
             })
         },
 
+        {
+            'name': 'dsn_ipv6_multi_host',
+            'dsn': 'postgresql://user@[2001:db8::1234%25eth0],[::1]/db',
+            'result': ([('2001:db8::1234%eth0', 5432), ('::1', 5432)], {
+                'database': 'db',
+                'user': 'user',
+            })
+        },
+
+        {
+            'name': 'dsn_ipv6_multi_host_port',
+            'dsn': 'postgresql://user@[2001:db8::1234]:1111,[::1]:2222/db',
+            'result': ([('2001:db8::1234', 1111), ('::1', 2222)], {
+                'database': 'db',
+                'user': 'user',
+            })
+        },
+
+        {
+            'name': 'dsn_ipv6_multi_host_query_part',
+            'dsn': 'postgresql:///db?user=user&host=[2001:db8::1234],[::1]',
+            'result': ([('2001:db8::1234', 5432), ('::1', 5432)], {
+                'database': 'db',
+                'user': 'user',
+            })
+        },
+
+
         {
             'name': 'dsn_combines_env_multi_host',
             'env': {