Merge pull request #721 from smarkets/pgcli-ssl-url-support

Make it possible to use SSL with pgcli

7 files changed, 69 insertions, 23 deletions

diff --git a/AUTHORS b/AUTHORS
index 9863c81e..a35ca06b 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -1,6 +1,6 @@
-Many thanks to the following contributors. 
+Many thanks to the following contributors.
 
-Project Lead: 
+Project Lead:
 -------------
     * Irina Truong
 
@@ -13,7 +13,7 @@ Core Devs:
     * Daniel Rocco
     * Karl-Aksel Puulmann
     * Dick Marinus
-    
+
 Contributors:
 -------------
     * Brett
@@ -64,6 +64,7 @@ Contributors:
     * Hraban Luyat
     * Jackson Popkin
     * Gustavo Castro
+    * Alexander Schmolck
 
 Creator:
 --------
diff --git a/changelog.rst b/changelog.rst
index 49db98a4..cff7db42 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -10,6 +10,7 @@ Features:
 * Completions after ORDER BY and DISTINCT now take account of table aliases. (Thanks: `Owen Stephens`_)
 * Narrow keyword candidates based on previous keyword. (Thanks: `Étienne Bersac`_)
 * Opening an external editor will edit the last-run query. (Thanks: `Thomas Roten`_)
+* Support query options in postgres URIs such as ?sslcert=foo.pem (Thanks: `Alexander Schmolck`_)
 
 Bug fixes:
 ----------
@@ -690,3 +691,4 @@ Improvements:
 .. _`Étienne Bersac`: https://github.com/bersace
 .. _`Thomas Roten`: https://github.com/tsroten
 .. _`Gustavo Castro`: https://github.com/gustavo-castro
+.. _`Alexander Schmolck`: https://github.com/aschmolck
diff --git a/pgcli/completion_refresher.py b/pgcli/completion_refresher.py
index e1e05e5a..9944a701 100644
--- a/pgcli/completion_refresher.py
+++ b/pgcli/completion_refresher.py
@@ -59,7 +59,8 @@ class CompletionRefresher(object):
             # Create a new pgexecute method to popoulate the completions.
             e = pgexecute
             executor = PGExecute(
-                e.dbname, e.user, e.password, e.host, e.port, e.dsn)
+                e.dbname, e.user, e.password, e.host, e.port, e.dsn,
+                **e.extra_args)
 
         # If callbacks is a single function then push it into a list.
         if callable(callbacks):
diff --git a/pgcli/main.py b/pgcli/main.py
index 9f57aa52..8970c43f 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -53,9 +53,9 @@ from .__init__ import __version__
 click.disable_unicode_literals_warning = True
 
 try:
-    from urlparse import urlparse, unquote
+    from urlparse import urlparse, unquote, parse_qs
 except ImportError:
-    from urllib.parse import urlparse, unquote
+    from urllib.parse import urlparse, unquote, parse_qs
 
 from getpass import getuser
 from psycopg2 import OperationalError
@@ -292,13 +292,26 @@ class PGCli(object):
     def connect_uri(self, uri):
         uri = urlparse(uri)
         database = uri.path[1:]  # ignore the leading fwd slash
-        arguments = [database, uri.hostname, uri.username,
-                     uri.port, uri.password]
-        # unquote each URI part (they may be percent encoded)
-        self.connect(*list(map(lambda p: unquote(str(p)) if p else p, arguments)))
+
+        def fixup_possible_percent_encoding(s):
+            return unquote(str(s)) if s else s
+
+        arguments = dict(database=fixup_possible_percent_encoding(database),
+                         host=fixup_possible_percent_encoding(uri.hostname),
+                         user=fixup_possible_percent_encoding(uri.username),
+                         port=fixup_possible_percent_encoding(uri.port),
+                         passwd=fixup_possible_percent_encoding(uri.password))
+        # Deal with extra params e.g. ?sslmode=verify-ca&ssl-cert=/mycert
+        if uri.query:
+            arguments = dict(
+                {k: v for k, (v,) in parse_qs(uri.query).items()},
+                **arguments)
+
+        # unquote str(each URI part (they may be percent encoded)
+        self.connect(**arguments)
 
     def connect(self, database='', host='', user='', port='', passwd='',
-                dsn=''):
+                dsn='', **kwargs):
         # Connect to the database.
 
         if not user:
@@ -331,14 +344,15 @@ class PGCli(object):
         # a password (no -w flag), prompt for a passwd and try again.
         try:
             try:
-                pgexecute = PGExecute(database, user, passwd, host, port, dsn)
+                pgexecute = PGExecute(database, user, passwd, host, port, dsn,
+                                      **kwargs)
             except OperationalError as e:
                 if ('no password supplied' in utf8tounicode(e.args[0]) and
                         auto_passwd_prompt):
                     passwd = click.prompt('Password', hide_input=True,
                                           show_default=False, type=str)
                     pgexecute = PGExecute(database, user, passwd, host, port,
-                                          dsn)
+                                          dsn, **kwargs)
                 else:
                     raise e
 
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 73502324..47b17910 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -146,17 +146,18 @@ class PGExecute(object):
         FROM pg_catalog.pg_database d
         ORDER BY 1'''
 
-    def __init__(self, database, user, password, host, port, dsn):
+    def __init__(self, database, user, password, host, port, dsn, **kwargs):
         self.dbname = database
         self.user = user
         self.password = password
         self.host = host
         self.port = port
         self.dsn = dsn
+        self.extra_args = {k: unicode2utf8(v) for k, v in kwargs.items()}
         self.connect()
 
     def connect(self, database=None, user=None, password=None, host=None,
-            port=None, dsn=None):
+                port=None, dsn=None, **kwargs):
 
         db = (database or self.dbname)
         user = (user or self.user)
@@ -164,6 +165,7 @@ class PGExecute(object):
         host = (host or self.host)
         port = (port or self.port)
         dsn = (dsn or self.dsn)
+        kwargs = (kwargs or self.extra_args)
         pid = -1
         if dsn:
             if password:
@@ -172,11 +174,12 @@ class PGExecute(object):
             cursor = conn.cursor()
         else:
             conn = psycopg2.connect(
-                    database=unicode2utf8(db),
-                    user=unicode2utf8(user),
-                    password=unicode2utf8(password),
-                    host=unicode2utf8(host),
-                    port=unicode2utf8(port))
+                database=unicode2utf8(db),
+                user=unicode2utf8(user),
+                password=unicode2utf8(password),
+                host=unicode2utf8(host),
+                port=unicode2utf8(port),
+                **kwargs)
 
             cursor = conn.cursor()
 
@@ -653,4 +656,3 @@ class PGExecute(object):
             cur.execute(query)
             for row in cur:
                 yield row[0]
-
diff --git a/tests/test_completion_refresher.py b/tests/test_completion_refresher.py
index f47a9745..08f7811e 100644
--- a/tests/test_completion_refresher.py
+++ b/tests/test_completion_refresher.py
@@ -79,6 +79,7 @@ def test_refresh_with_callbacks(refresher):
     callbacks = [Mock()]
     pgexecute_class = Mock()
     pgexecute = Mock()
+    pgexecute.extra_args = {}
     special = Mock()
 
     with patch('pgcli.completion_refresher.PGExecute', pgexecute_class):
diff --git a/tests/test_main.py b/tests/test_main.py
index b705b730..b380d4db 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -91,11 +91,36 @@ def test_quoted_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%5E:%5Dfoo@baz.com/testdb%5B')
-    mock_connect.assert_called_with('testdb[', 'baz.com', 'bar^', None, ']foo')
+    mock_connect.assert_called_with(database='testdb[',
+                                    port=None,
+                                    host='baz.com',
+                                    user='bar^',
+                                    passwd=']foo')
+
+
+def test_ssl_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%5E:%5Dfoo@baz.com/testdb%5B?'
+            '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',
+                                    sslcert='my.pem',
+                                    sslkey='my-key.pem',
+                                    sslrootcert='ca.pem')
 
 
 def test_port_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@baz.com:2543/testdb')
-    mock_connect.assert_called_with('testdb', 'baz.com', 'bar', '2543', 'foo')
+    mock_connect.assert_called_with(database='testdb',
+                                    host='baz.com',
+                                    user='bar',
+                                    passwd='foo',
+                                    port='2543')