Support multihost connection string (#978)

* Switch to psycopg2 parse_dsn instead of urlparse

* Added wronglink to contributors and updated changelog

* Fix test codestyle

* Support for PGPORT customization in tests

* Support for PGPORT customization in tests

* Refactored PGExecute init and moved short_host generation to object property

* Fix test util codestyle

* Fix local tests run

* Store PGExecute initial params in  _conn_params and added PGExecute.copy method

* Fix codestyle

* Added docstring to PGExecute.copy() method

4 files changed, 46 insertions, 6 deletions

diff --git a/tests/conftest.py b/tests/conftest.py
index 9554623a..26a40816 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -2,7 +2,7 @@ from __future__ import print_function
 
 import os
 import pytest
-from utils import (POSTGRES_HOST, POSTGRES_USER, POSTGRES_PASSWORD, create_db, db_connection,
+from utils import (POSTGRES_HOST, POSTGRES_PORT, POSTGRES_USER, POSTGRES_PASSWORD, create_db, db_connection,
                    drop_tables)
 import pgcli.pgexecute
 
@@ -25,8 +25,8 @@ def cursor(connection):
 
 @pytest.fixture
 def executor(connection):
-    return pgcli.pgexecute.PGExecute(database='_test_db', user=POSTGRES_USER,
-                                     host=POSTGRES_HOST, password=POSTGRES_PASSWORD, port=None, dsn=None)
+    return pgcli.pgexecute.PGExecute(database='_test_db', user=POSTGRES_USER, host=POSTGRES_HOST,
+                                     password=POSTGRES_PASSWORD, port=POSTGRES_PORT, dsn=None)
 
 
 @pytest.fixture
diff --git a/tests/test_main.py b/tests/test_main.py
index fc4de7d9..b1c40356 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -267,7 +267,6 @@ def test_quoted_db_uri(tmpdir):
         cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile")))
         cli.connect_uri('postgres://bar%5E:%5Dfoo@baz.com/testdb%5B')
     mock_connect.assert_called_with(database='testdb[',
-                                    port=None,
                                     host='baz.com',
                                     user='bar^',
                                     passwd=']foo')
@@ -281,7 +280,6 @@ def test_ssl_db_uri(tmpdir):
             'sslmode=verify-full&sslcert=m%79.pem&sslkey=my-key.pem&sslrootcert=c%61.pem')
     mock_connect.assert_called_with(database='testdb[',
                                     host='baz.com',
-                                    port=None,
                                     user='bar^',
                                     passwd=']foo',
                                     sslmode='verify-full',
@@ -299,3 +297,15 @@ def test_port_db_uri(tmpdir):
                                     user='bar',
                                     passwd='foo',
                                     port='2543')
+
+
+def test_multihost_db_uri(tmpdir):
+    with mock.patch.object(PGCli, 'connect') as mock_connect:
+        cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile")))
+        cli.connect_uri(
+            'postgres://bar:foo@baz1.com:2543,baz2.com:2543,baz3.com:2543/testdb')
+    mock_connect.assert_called_with(database='testdb',
+                                    host='baz1.com,baz2.com,baz3.com',
+                                    user='bar',
+                                    passwd='foo',
+                                    port='2543,2543,2543')
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index a2340547..39788950 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -34,6 +34,21 @@ def test_conn(executor):
 
 
 @dbtest
+def test_copy(executor):
+    executor_copy = executor.copy()
+    run(executor_copy, '''create table test(a text)''')
+    run(executor_copy, '''insert into test values('abc')''')
+    assert run(executor_copy, '''select * from test''', join=True) == dedent("""\
+        +-----+
+        | a   |
+        |-----|
+        | abc |
+        +-----+
+        SELECT 1""")
+
+
+
+@dbtest
 def test_bools_are_treated_as_strings(executor):
     run(executor, '''create table test(a boolean)''')
     run(executor, '''insert into test values(True)''')
@@ -400,3 +415,13 @@ def test_nonexistent_view_definition(executor):
         result = executor.view_definition('there_is_no_such_view')
     with pytest.raises(RuntimeError):
         result = executor.view_definition('mvw1')
+
+
+@dbtest
+def test_short_host(executor):
+    with patch.object(executor, 'host', 'localhost'):
+        assert executor.short_host == 'localhost'
+    with patch.object(executor, 'host', 'localhost.example.org'):
+        assert executor.short_host == 'localhost'
+    with patch.object(executor, 'host', 'localhost1.example.org,localhost2.example.org'):
+        assert executor.short_host == 'localhost1'
diff --git a/tests/utils.py b/tests/utils.py
index 31fd12a3..62b06602 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -7,11 +7,16 @@ from os import getenv
 
 POSTGRES_USER = getenv('PGUSER', 'postgres')
 POSTGRES_HOST = getenv('PGHOST', 'localhost')
+POSTGRES_PORT = getenv('PGPORT', 5432)
 POSTGRES_PASSWORD = getenv('PGPASSWORD', '')
 
 
 def db_connection(dbname=None):
-    conn = psycopg2.connect(user=POSTGRES_USER, host=POSTGRES_HOST, database=dbname)
+    conn = psycopg2.connect(user=POSTGRES_USER,
+                            host=POSTGRES_HOST,
+                            password=POSTGRES_PASSWORD,
+                            port=POSTGRES_PORT,
+                            database=dbname)
     conn.autocommit = True
     return conn