Add command-line option --single-connection

This option makes it so that the executor and the completer use the same database connection. This is needed for the completer to function when using one-time passwords for connecting to the database.

2 files changed, 18 insertions, 8 deletions

diff --git a/pgcli/completion_refresher.py b/pgcli/completion_refresher.py
index b4a89d91..e1e05e5a 100644
--- a/pgcli/completion_refresher.py
+++ b/pgcli/completion_refresher.py
@@ -49,12 +49,17 @@ class CompletionRefresher(object):
 
     def _bg_refresh(self, pgexecute, special, callbacks, history=None,
       settings=None):
+        settings = settings or {}
         completer = PGCompleter(smart_completion=True, pgspecial=special,
             settings=settings)
 
-        # 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)
+        if settings.get('single_connection'):
+            executor = pgexecute
+        else:
+            # 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)
 
         # 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 81418194..0f81c9d3 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -101,7 +101,8 @@ class PGCli(object):
         os.environ['LESS'] = '-SRXF'
 
     def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False,
-                 pgexecute=None, pgclirc_file=None, row_limit=None):
+                 pgexecute=None, pgclirc_file=None, row_limit=None,
+                 single_connection=False):
 
         self.force_passwd_prompt = force_passwd_prompt
         self.never_passwd_prompt = never_passwd_prompt
@@ -142,7 +143,8 @@ class PGCli(object):
         self.settings = {'casing_file': get_casing_file(c),
           'generate_casing_file': c['main'].as_bool('generate_casing_file'),
           'generate_aliases': c['main'].as_bool('generate_aliases'),
-          'asterisk_column_order': c['main']['asterisk_column_order']}
+          'asterisk_column_order': c['main']['asterisk_column_order'],
+          'single_connection': single_connection}
         completer = PGCompleter(smart_completion, pgspecial=self.pgspecial,
             settings=self.settings)
         self.completer = completer
@@ -670,6 +672,9 @@ class PGCli(object):
         help='Force password prompt.')
 @click.option('-w', '--no-password', 'never_prompt', is_flag=True,
         default=False, help='Never prompt for password.')
+@click.option('--single-connection', 'single_connection', is_flag=True,
+        default=False,
+        help='Don''t use a separate connection for completions.')
 @click.option('-v', '--version', is_flag=True, help='Version of pgcli.')
 @click.option('-d', '--dbname', default='', envvar='PGDATABASE',
         help='database name to connect to.')
@@ -681,8 +686,8 @@ class PGCli(object):
         help='Set threshold for row limit prompt. Use 0 to disable prompt.')
 @click.argument('database', default=lambda: None, envvar='PGDATABASE', nargs=1)
 @click.argument('username', default=lambda: None, envvar='PGUSER', nargs=1)
-def cli(database, user, host, port, prompt_passwd, never_prompt, dbname,
-        username, version, pgclirc, dsn, row_limit):
+def cli(database, user, host, port, prompt_passwd, never_prompt,
+    single_connection, dbname, username, version, pgclirc, dsn, row_limit):
 
     if version:
         print('Version:', __version__)
@@ -705,7 +710,7 @@ def cli(database, user, host, port, prompt_passwd, never_prompt, dbname,
                    config_full_path)
 
     pgcli = PGCli(prompt_passwd, never_prompt, pgclirc_file=pgclirc,
-                  row_limit=row_limit)
+                  row_limit=row_limit, single_connection=single_connection)
 
     # Choose which ever one has a valid value.
     database = database or dbname