repair psycopg2 (and psycopg) multiple hosts format

Fixed issue in psycopg2 dialect where the "multiple hosts" feature
implemented for :ticket:`4392`, where multiple ``host:port`` pairs could be
passed in the query string as
``?host=host1:port1&host=host2:port2&host=host3:port3`` was not implemented
correctly, as it did not propagate the "port" parameter appropriately.
Connections that didn't use a different "port" likely worked without issue,
and connections that had "port" for some of the entries may have
incorrectly passed on that hostname. The format is now corrected to pass
hosts/ports appropriately.

As part of this change, maintained support for another multihost style that
worked unintentionally, which is comma-separated
``?host=h1,h2,h3&port=p1,p2,p3``. This format is more consistent with
libpq's query-string format, whereas the previous format is inspired by a
different aspect of libpq's URI format but is not quite the same thing.

If the two styles are mixed together, an error is raised as this is
ambiguous.

Fixes: #4392
Change-Id: Ic9cc0b0e6e90725e158d9efe73e042853dd1263f

1 files changed, 116 insertions, 14 deletions

diff --git a/test/dialect/postgresql/test_dialect.py b/test/dialect/postgresql/test_dialect.py
index d55aa8203..1ffd82ae4 100644
--- a/test/dialect/postgresql/test_dialect.py
+++ b/test/dialect/postgresql/test_dialect.py
@@ -43,6 +43,7 @@ from sqlalchemy.engine import url
 from sqlalchemy.sql.selectable import LABEL_STYLE_TABLENAME_PLUS_COL
 from sqlalchemy.testing import config
 from sqlalchemy.testing import engines
+from sqlalchemy.testing import expect_raises_message
 from sqlalchemy.testing import fixtures
 from sqlalchemy.testing import is_
 from sqlalchemy.testing import is_false
@@ -200,24 +201,93 @@ $$ LANGUAGE plpgsql;"""
         eq_(cargs, [])
         eq_(cparams, {"host": "somehost", "any_random_thing": "yes"})
 
-    def test_psycopg2_nonempty_connection_string_w_query_two(self):
-        dialect = psycopg2_dialect.dialect()
-        url_string = "postgresql+psycopg2://USER:PASS@/DB?host=hostA"
-        u = url.make_url(url_string)
-        cargs, cparams = dialect.create_connect_args(u)
-        eq_(cargs, [])
-        eq_(cparams["host"], "hostA")
-
-    def test_psycopg2_nonempty_connection_string_w_query_three(self):
-        dialect = psycopg2_dialect.dialect()
-        url_string = (
+    @testing.combinations(
+        (
+            "postgresql+psycopg2://USER:PASS@/DB?host=hostA",
+            {
+                "dbname": "DB",
+                "user": "USER",
+                "password": "PASS",
+                "host": "hostA",
+            },
+        ),
+        (
             "postgresql+psycopg2://USER:PASS@/DB"
-            "?host=hostA:portA&host=hostB&host=hostC"
-        )
+            "?host=hostA&host=hostB&host=hostC",
+            {
+                "dbname": "DB",
+                "user": "USER",
+                "password": "PASS",
+                "host": "hostA,hostB,hostC",
+                "port": ",,",
+            },
+        ),
+        (
+            "postgresql+psycopg2://USER:PASS@/DB"
+            "?host=hostA&host=hostB:portB&host=hostC:portC",
+            {
+                "dbname": "DB",
+                "user": "USER",
+                "password": "PASS",
+                "host": "hostA,hostB,hostC",
+                "port": ",portB,portC",
+            },
+        ),
+        (
+            "postgresql+psycopg2://USER:PASS@/DB?"
+            "host=hostA:portA&host=hostB:portB&host=hostC:portC",
+            {
+                "dbname": "DB",
+                "user": "USER",
+                "password": "PASS",
+                "host": "hostA,hostB,hostC",
+                "port": "portA,portB,portC",
+            },
+        ),
+        (
+            "postgresql+psycopg2:///"
+            "?host=hostA:portA&host=hostB:portB&host=hostC:portC",
+            {"host": "hostA,hostB,hostC", "port": "portA,portB,portC"},
+        ),
+        (
+            "postgresql+psycopg2:///"
+            "?host=hostA:portA&host=hostB:portB&host=hostC:portC",
+            {"host": "hostA,hostB,hostC", "port": "portA,portB,portC"},
+        ),
+        (
+            "postgresql+psycopg2:///"
+            "?host=hostA,hostB,hostC&port=portA,portB,portC",
+            {"host": "hostA,hostB,hostC", "port": "portA,portB,portC"},
+        ),
+        argnames="url_string,expected",
+    )
+    @testing.combinations(
+        psycopg2_dialect.dialect(),
+        psycopg_dialect.dialect(),
+        argnames="dialect",
+    )
+    def test_psycopg_multi_hosts(self, dialect, url_string, expected):
         u = url.make_url(url_string)
         cargs, cparams = dialect.create_connect_args(u)
         eq_(cargs, [])
-        eq_(cparams["host"], "hostA:portA,hostB,hostC")
+        eq_(cparams, expected)
+
+    @testing.combinations(
+        "postgresql+psycopg2:///?host=H&host=H&port=5432,5432",
+        "postgresql+psycopg2://user:pass@/dbname?host=H&host=H&port=5432,5432",
+        argnames="url_string",
+    )
+    @testing.combinations(
+        psycopg2_dialect.dialect(),
+        psycopg_dialect.dialect(),
+        argnames="dialect",
+    )
+    def test_psycopg_no_mix_hosts(self, dialect, url_string):
+        with expect_raises_message(
+            exc.ArgumentError, "Can't mix 'multihost' formats together"
+        ):
+            u = url.make_url(url_string)
+            dialect.create_connect_args(u)
 
     def test_psycopg2_disconnect(self):
         class Error(Exception):
@@ -256,6 +326,38 @@ $$ LANGUAGE plpgsql;"""
         eq_(dialect.is_disconnect("not an error", None, None), False)
 
 
+class BackendDialectTest(fixtures.TestBase):
+    __backend__ = True
+
+    @testing.only_on(["+psycopg", "+psycopg2"])
+    @testing.combinations(
+        "host=H:P&host=H:P&host=H:P",
+        "host=H:P&host=H&host=H",
+        "host=H:P&host=H&host=H:P",
+        "host=H&host=H:P&host=H",
+        "host=H,H,H&port=P,P,P",
+    )
+    def test_connect_psycopg_multiple_hosts(self, pattern):
+        """test the fix for #4392"""
+
+        tdb_url = testing.db.url
+
+        host = tdb_url.host
+        if host == "127.0.0.1":
+            host = "localhost"
+        port = str(tdb_url.port) if tdb_url.port else "5432"
+
+        query_str = pattern.replace("H", host).replace("P", port)
+        url_string = (
+            f"{tdb_url.drivername}://{tdb_url.username}:"
+            f"{tdb_url.password}@/{tdb_url.database}?{query_str}"
+        )
+
+        e = create_engine(url_string)
+        with e.connect() as conn:
+            eq_(conn.exec_driver_sql("select 1").scalar(), 1)
+
+
 class PGCodeTest(fixtures.TestBase):
     __only_on__ = "postgresql"