Merge branch 'master' into master

4 files changed, 71 insertions, 41 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index b0096dc4..1ce3f4b0 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -40,6 +40,7 @@ from pygments.token import Token
 
 from pgspecial.main import (PGSpecial, NO_QUERY, PAGER_OFF)
 import pgspecial as special
+import keyring
 from .pgcompleter import PGCompleter
 from .pgtoolbar import create_toolbar_tokens_func
 from .pgstyle import style_factory, style_factory_output
@@ -91,6 +92,10 @@ OutputSettings.__new__.__defaults__ = (
 )
 
 
+class PgCliQuitError(Exception):
+    pass
+
+
 class PGCli(object):
 
     default_prompt = '\\u@\\h:\\d> '
@@ -202,6 +207,9 @@ class PGCli(object):
         self.eventloop = create_eventloop()
         self.cli = None
 
+    def quit(self):
+        raise PgCliQuitError
+
     def register_special_commands(self):
 
         self.pgspecial.register(
@@ -211,6 +219,12 @@ class PGCli(object):
         refresh_callback = lambda: self.refresh_completions(
             persist_priorities='all')
 
+        self.pgspecial.register(self.quit, '\\q', '\\q',
+                                'Quit pgcli.', arg_type=NO_QUERY, case_sensitive=True,
+                                aliases=(':q',))
+        self.pgspecial.register(self.quit, 'quit', 'quit',
+                                'Quit pgcli.', arg_type=NO_QUERY, case_sensitive=False,
+                                aliases=('exit',))
         self.pgspecial.register(refresh_callback, '\\#', '\\#',
                                 'Refresh auto-completions.', arg_type=NO_QUERY)
         self.pgspecial.register(refresh_callback, '\\refresh', '\\refresh',
@@ -367,7 +381,7 @@ class PGCli(object):
                          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
+        # Deal with extra params e.g. ?sslmode=verify-ca&sslrootcert=/myrootcert
         if uri.query:
             arguments = dict(
                 {k: v for k, (v,) in parse_qs(uri.query).items()},
@@ -391,6 +405,11 @@ class PGCli(object):
         if not self.force_passwd_prompt and not passwd:
             passwd = os.environ.get('PGPASSWORD', '')
 
+        # Find password from store
+        key = '%s@%s' % (user, host)
+        if not passwd:
+            passwd = keyring.get_password('pgcli', key)
+
         # Prompt for a password immediately if requested via the -W flag. This
         # avoids wasting time trying to connect to the database and catching a
         # no-password exception.
@@ -412,6 +431,8 @@ class PGCli(object):
             try:
                 pgexecute = PGExecute(database, user, passwd, host, port, dsn,
                                       application_name='pgcli', **kwargs)
+                if passwd:
+                    keyring.set_password('pgcli', key, passwd)
             except (OperationalError, InterfaceError) as e:
                 if ('no password supplied' in utf8tounicode(e.args[0]) and
                         auto_passwd_prompt):
@@ -421,6 +442,8 @@ class PGCli(object):
                     pgexecute = PGExecute(database, user, passwd, host, port,
                                           dsn, application_name='pgcli',
                                           **kwargs)
+                    if passwd:
+                        keyring.set_password('pgcli', key, passwd)
                 else:
                     raise e
 
@@ -449,19 +472,22 @@ class PGCli(object):
         # It's internal api of prompt_toolkit that may change. This was added to fix #668.
         # We may find a better way to do it in the future.
         saved_callables = cli.application.pre_run_callables
-        while special.editor_command(document.text):
-            filename = special.get_filename(document.text)
-            query = (special.get_editor_query(document.text) or
-                     self.get_last_query())
-            sql, message = special.open_external_editor(filename, sql=query)
-            if message:
-                # Something went wrong. Raise an exception and bail.
-                raise RuntimeError(message)
-            cli.current_buffer.document = Document(sql, cursor_position=len(sql))
-            cli.application.pre_run_callables = []
-            document = cli.run()
-            continue
-        cli.application.pre_run_callables = saved_callables
+        try:
+            while special.editor_command(document.text):
+                filename = special.get_filename(document.text)
+                query = (special.get_editor_query(document.text) or
+                         self.get_last_query())
+                sql, message = special.open_external_editor(
+                    filename, sql=query)
+                if message:
+                    # Something went wrong. Raise an exception and bail.
+                    raise RuntimeError(message)
+                cli.current_buffer.document = Document(sql,
+                                                       cursor_position=len(sql))
+                cli.application.pre_run_callables = []
+                document = cli.run()
+        finally:
+            cli.application.pre_run_callables = saved_callables
         return document
 
     def execute_command(self, text, query):
@@ -489,6 +515,8 @@ class PGCli(object):
             logger.error("sql: %r, error: %r", text, e)
             logger.error("traceback: %r", traceback.format_exc())
             self._handle_server_closed_connection()
+        except PgCliQuitError as e:
+            raise
         except Exception as e:
             logger.error("sql: %r, error: %r", text, e)
             logger.error("traceback: %r", traceback.format_exc())
@@ -555,13 +583,6 @@ class PGCli(object):
             while True:
                 document = self.cli.run()
 
-                # The reason we check here instead of inside the pgexecute is
-                # because we want to raise the Exit exception which will be
-                # caught by the try/except block that wraps the pgexecute.run()
-                # statement.
-                if quit_command(document.text):
-                    raise EOFError
-
                 try:
                     document = self.handle_editor_command(self.cli, document)
                 except RuntimeError as e:
@@ -573,15 +594,19 @@ class PGCli(object):
                 # Initialize default metaquery in case execution fails
                 query = MetaQuery(query=document.text, successful=False)
 
-                watch_command, timing = special.get_watch_command(document.text)
-                if watch_command:
-                    while watch_command:
+                self.watch_command, timing = special.get_watch_command(
+                        document.text)
+                if self.watch_command:
+                    while self.watch_command:
                         try:
-                            query = self.execute_command(watch_command, query)
-                            click.echo('Waiting for {0} seconds before repeating'.format(timing))
+                            query = self.execute_command(
+                                    self.watch_command, query)
+                            click.echo(
+                                    'Waiting for {0} seconds before repeating'
+                                    .format(timing))
                             sleep(timing)
                         except KeyboardInterrupt:
-                            watch_command = None
+                            self.watch_command = None
                 else:
                     query = self.execute_command(document.text, query)
 
@@ -593,7 +618,7 @@ class PGCli(object):
 
                 self.query_history.append(query)
 
-        except EOFError:
+        except PgCliQuitError:
             if not self.less_chatty:
                 print ('Goodbye!')
 
@@ -849,7 +874,7 @@ class PGCli(object):
         return self.query_history[-1][0] if self.query_history else None
 
     def echo_via_pager(self, text, color=None):
-        if self.pgspecial.pager_config == PAGER_OFF:
+        if self.pgspecial.pager_config == PAGER_OFF or self.watch_command:
             click.echo(text, color=color)
         else:
             click.echo_via_pager(text, color)
@@ -1044,13 +1069,6 @@ def is_select(status):
     return status.split(None, 1)[0].lower() == 'select'
 
 
-def quit_command(sql):
-    return (sql.strip().lower() == 'exit'
-            or sql.strip().lower() == 'quit'
-            or sql.strip() == r'\q'
-            or sql.strip() == ':q')
-
-
 def exception_formatter(e):
     return click.style(utf8tounicode(str(e)), fg='red')
 
diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index f92a8a71..aa829853 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -28,6 +28,7 @@ Schema.__new__.__defaults__ = (False,)
 # used to ensure that the alias we suggest is unique
 FromClauseItem = namedtuple('FromClauseItem', 'schema table_refs local_tables')
 Table = namedtuple('Table', ['schema', 'table_refs', 'local_tables'])
+TableFormat = namedtuple('TableFormat', [])
 View = namedtuple('View', ['schema', 'table_refs'])
 # JoinConditions are suggested after ON, e.g. 'foo.barid = bar.barid'
 JoinCondition = namedtuple('JoinCondition', ['table_refs', 'parent'])
@@ -252,6 +253,9 @@ def suggest_special(text):
     if cmd in ('\\c', '\\connect'):
         return (Database(),)
 
+    if cmd == '\\T':
+        return (TableFormat(),)
+
     if cmd == '\\dn':
         return (Schema(),)
 
diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py
index cba9fb7c..72d54d82 100644
--- a/pgcli/pgcompleter.py
+++ b/pgcli/pgcompleter.py
@@ -4,13 +4,15 @@ import re
 from itertools import count, repeat, chain
 import operator
 from collections import namedtuple, defaultdict, OrderedDict
+from cli_helpers.tabular_output import TabularOutputFormatter
 from pgspecial.namedqueries import NamedQueries
 from prompt_toolkit.completion import Completer, Completion
 from prompt_toolkit.contrib.completers import PathCompleter
 from prompt_toolkit.document import Document
-from .packages.sqlcompletion import (FromClauseItem,
-    suggest_type, Special, Database, Schema, Table, Function, Column, View,
-    Keyword, NamedQuery, Datatype, Alias, Path, JoinCondition, Join)
+from .packages.sqlcompletion import (
+    FromClauseItem, suggest_type, Special, Database, Schema, Table,
+    TableFormat, Function, Column, View, Keyword, NamedQuery,
+    Datatype, Alias, Path, JoinCondition, Join)
 from .packages.parseutils.meta import ColumnMetadata, ForeignKey
 from .packages.parseutils.utils import last_word
 from .packages.parseutils.tables import TableReference
@@ -321,7 +323,8 @@ class PGCompleter(Completer):
             return []
         prio_order = [
             'keyword', 'function', 'view', 'table', 'datatype', 'database',
-            'schema', 'column', 'table alias', 'join', 'name join', 'fk join'
+            'schema', 'column', 'table alias', 'join', 'name join', 'fk join',
+            'table format'
         ]
         type_priority = prio_order.index(meta) if meta in prio_order else -1
         text = last_word(text, include='most_punctuations').lower()
@@ -778,6 +781,9 @@ class PGCompleter(Completer):
         tables = [self._make_cand(t, alias, suggestion) for t in tables]
         return self.find_matches(word_before_cursor, tables, meta='table')
 
+    def get_table_formats(self, _, word_before_cursor):
+        formats = TabularOutputFormatter().supported_formats
+        return self.find_matches(word_before_cursor, formats, meta='table format')
 
     def get_view_matches(self, suggestion, word_before_cursor, alias=False):
         views = self.populate_schema_objects(suggestion.schema, 'views')
@@ -861,6 +867,7 @@ class PGCompleter(Completer):
         Function: get_function_matches,
         Schema: get_schema_matches,
         Table: get_table_matches,
+        TableFormat: get_table_formats,
         View: get_view_matches,
         Alias: get_alias_matches,
         Database: get_database_matches,
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 74084fe6..2158f909 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -346,7 +346,8 @@ class PGExecute(object):
         """
         return (isinstance(e, psycopg2.OperationalError) and
                 (not e.pgcode or
-                 psycopg2.errorcodes.lookup(e.pgcode) != 'LOCK_NOT_AVAILABLE'))
+                 psycopg2.errorcodes.lookup(e.pgcode) not in
+                 ('LOCK_NOT_AVAILABLE', 'CANT_CHANGE_RUNTIME_PARAM')))
 
     def execute_normal_sql(self, split_sql):
         """Returns tuple (title, rows, headers, status)"""