Complete keywords depending on previous token

Keywords list is based on
https://www.postgresql.org/docs/9.6/static/sql-commands.html.

7 files changed, 179 insertions, 41 deletions

diff --git a/changelog.rst b/changelog.rst
index 1a64db29..ab41f20c 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -7,6 +7,7 @@ Features:
 * Casing for column headers (Thanks: `Joakim Koljonen`_)
 * Allow configurable character to be used for multi-line query continuations. (Thanks: `Owen Stephens`_)
 * 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`_)
 
 Bug fixes:
 ----------
@@ -675,3 +676,4 @@ Improvements:
 .. _`Owen Stephens`: https://github.com/owst
 .. _`Russell Davies`: https://github.com/russelldavies
 .. _`Dick Marinus`: https://github.com/meeuw
+.. _`Étienne Bersac`: https://github.com/bersace
diff --git a/pgcli/packages/pgliterals/pgliterals.json b/pgcli/packages/pgliterals/pgliterals.json
index c7b9b7f6..9021a9b9 100644
--- a/pgcli/packages/pgliterals/pgliterals.json
+++ b/pgcli/packages/pgliterals/pgliterals.json
@@ -3,12 +3,45 @@
         "ACCESS": [],
         "ADD": [],
         "ALL": [],
-        "ALTER": [],
+        "ALTER": [
+            "AGGREGATE",
+            "COLLATION",
+            "COLUMN",
+            "CONVERSION",
+            "DATABASE",
+            "DEFAULT",
+            "DOMAIN",
+            "EVENT TRIGGER",
+            "EXTENSION",
+            "FOREIGN",
+            "FUNCTION",
+            "GROUP",
+            "INDEX",
+            "LANGUAGE",
+            "LARGE OBJECT",
+            "MATERIALIZED VIEW",
+            "OPERATOR",
+            "POLICY",
+            "ROLE",
+            "RULE",
+            "SCHEMA",
+            "SEQUENCE",
+            "SERVER",
+            "SYSTEM",
+            "TABLE",
+            "TABLESPACE",
+            "TEXT SEARCH",
+            "TRIGGER",
+            "TYPE",
+            "USER",
+            "VIEW"
+        ],
         "AND": [],
         "ANY": [],
         "AS": [],
         "ASC": [],
         "AUDIT": [],
+        "BEGIN": [],
         "BETWEEN": [],
         "BY": [],
         "CASE": [],
@@ -21,7 +54,46 @@
         "CONCURRENTLY": [],
         "CONNECT": [],
         "COPY": [],
-        "CREATE": [],
+        "CREATE": [
+            "ACCESS METHOD",
+            "AGGREGATE",
+            "CAST",
+            "COLLATION",
+            "CONVERSION",
+            "DATABASE",
+            "DOMAIN",
+            "EVENT TRIGGER",
+            "EXTENSION",
+            "FOREIGN DATA WRAPPER",
+            "FOREIGN EXTENSION",
+            "FUNCTION",
+            "GLOBAL",
+            "GROUP",
+            "IF NOT EXISTS",
+            "INDEX",
+            "LANGUAGE",
+            "LOCAL",
+            "MATERIALIZED VIEW",
+            "OPERATOR",
+            "OR REPLACE",
+            "POLICY",
+            "ROLE",
+            "RULE",
+            "SCHEMA",
+            "SEQUENCE",
+            "SERVER",
+            "TABLE",
+            "TABLESPACE",
+            "TEMPORARY",
+            "TEXT SEARCH",
+            "TRIGGER",
+            "TYPE",
+            "UNIQUE",
+            "UNLOGGED",
+            "USER",
+            "USER MAPPING",
+            "VIEW"
+        ],
         "CURRENT": [],
         "DATABASE": [],
         "DATE": [],
@@ -32,7 +104,41 @@
         "DESC": [],
         "DESCRIBE": [],
         "DISTINCT": [],
-        "DROP": [],
+        "DROP": [
+            "ACCESS METHOD",
+            "AGGREGATE",
+            "CAST",
+            "COLLATION",
+            "CONVERSION",
+            "DATABASE",
+            "DOMAIN",
+            "EVENT TRIGGER",
+            "EXTENSION",
+            "FOREIGN DATA WRAPPER",
+            "FOREIGN TABLE",
+            "FUNCTION",
+            "GROUP",
+            "INDEX",
+            "LANGUAGE",
+            "MATERIALIZED VIEW",
+            "OPERATOR",
+            "OWNED",
+            "POLICY",
+            "ROLE",
+            "RULE",
+            "SCHEMA",
+            "SEQUENCE",
+            "SERVER",
+            "TABLE",
+            "TABLESPACE",
+            "TEXT SEARCH",
+            "TRANSFORM",
+            "TRIGGER",
+            "TYPE",
+            "USER",
+            "USER MAPPING",
+            "VIEW"
+        ],
         "EXPLAIN": [],
         "ELSE": [],
         "ENCODING": [],
@@ -105,6 +211,7 @@
         "RAISE": [],
         "RENAME": [],
         "REPLACE": [],
+        "RESET": ["ALL"],
         "RAW": [],
         "REFRESH MATERIALIZED VIEW": [],
         "RESOURCE": [],
diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index bbda0f92..6bad7c38 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -46,7 +46,8 @@ Column = namedtuple(
 )
 Column.__new__.__defaults__ = (None, None, tuple(), False)
 
-Keyword = namedtuple('Keyword', [])
+Keyword = namedtuple('Keyword', ['last_token'])
+Keyword.__new__.__defaults__ = (None,)
 NamedQuery = namedtuple('NamedQuery', [])
 Datatype = namedtuple('Datatype', ['schema'])
 Alias = namedtuple('Alias', ['aliases'])
@@ -394,7 +395,7 @@ def suggest_based_on_last_token(token, stmt):
             return (Column(table_refs=tables, local_tables=stmt.local_tables,
                            qualifiable=True),
                     Function(schema=None),
-                    Keyword(),)
+                    Keyword(token_v.upper()),)
     elif token_v == 'as':
         # Don't suggest anything for aliases
         return ()
@@ -492,7 +493,7 @@ def suggest_based_on_last_token(token, stmt):
             suggestions.append(Schema())
         return tuple(suggestions)
     elif token_v in {'alter', 'create', 'drop'}:
-        return (Keyword(),)
+        return (Keyword(token_v.upper()),)
     elif token.is_keyword:
         # token is a keyword we haven't implemented any special handling for
         # go backwards in the query until we find one we do recognize
@@ -500,7 +501,7 @@ def suggest_based_on_last_token(token, stmt):
         if prev_keyword:
             return suggest_based_on_last_token(prev_keyword, stmt)
         else:
-            return (Keyword(),)
+            return (Keyword(token_v.upper()),)
     else:
         return (Keyword(),)
 
diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py
index 6c0b3754..eca8b86a 100644
--- a/pgcli/pgcompleter.py
+++ b/pgcli/pgcompleter.py
@@ -51,8 +51,10 @@ def generate_alias(tbl):
         [l for l, prev in zip(tbl,  '_' + tbl) if prev == '_' and l != '_'])
 
 class PGCompleter(Completer):
+    # keywords_tree: A dict mapping keywords to well known following keywords.
+    # e.g. 'CREATE': ['TABLE', 'USER', ...],
     keywords_tree = get_literals('keywords', type_=dict)
-    keywords = tuple(chain(keywords_tree.keys(), *keywords_tree.values()))
+    keywords = tuple(set(chain(keywords_tree.keys(), *keywords_tree.values())))
     functions = get_literals('functions')
     datatypes = get_literals('datatypes')
 
@@ -650,7 +652,14 @@ class PGCompleter(Completer):
         return self.find_matches(word_before_cursor, self.databases,
                                  meta='database')
 
-    def get_keyword_matches(self, _, word_before_cursor):
+    def get_keyword_matches(self, suggestion, word_before_cursor):
+        keywords = self.keywords_tree.keys()
+        # Get well known following keywords for the last token. If any, narrow
+        # candidates to this list.
+        next_keywords = self.keywords_tree.get(suggestion.last_token, [])
+        if next_keywords:
+            keywords = next_keywords
+
         casing = self.keyword_casing
         if casing == 'auto':
             if word_before_cursor and word_before_cursor[-1].islower():
@@ -659,9 +668,9 @@ class PGCompleter(Completer):
                 casing = 'upper'
 
         if casing == 'upper':
-            keywords = [k.upper() for k in self.keywords]
+            keywords = [k.upper() for k in keywords]
         else:
-            keywords = [k.lower() for k in self.keywords]
+            keywords = [k.lower() for k in keywords]
 
         return self.find_matches(word_before_cursor, keywords,
                                  mode='strict', meta='keyword')
diff --git a/tests/metadata.py b/tests/metadata.py
index 15b4bc2b..37c8ff96 100644
--- a/tests/metadata.py
+++ b/tests/metadata.py
@@ -66,7 +66,7 @@ class MetaData(object):
         return [datatype(dt, pos) for dt in self.completer.datatypes]
 
     def keywords(self, pos=0):
-        return [keyword(kw, pos) for kw in self.completer.keywords]
+        return [keyword(kw, pos) for kw in self.completer.keywords_tree.keys()]
 
     def columns(self, tbl, parent='public', typ='tables', pos=0):
         if typ == 'functions':
diff --git a/tests/test_naive_completion.py b/tests/test_naive_completion.py
index 986933f5..4494b358 100644
--- a/tests/test_naive_completion.py
+++ b/tests/test_naive_completion.py
@@ -56,3 +56,18 @@ def test_paths_completion(completer, complete_event):
         complete_event,
         smart_completion=True))
     assert result > set([Completion(text="setup.py", start_position=0)])
+
+
+def test_alter_well_known_keywords_completion(completer, complete_event):
+    text = 'ALTER '
+    position = len(text)
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event,
+        smart_completion=True))
+    assert result > set([
+        Completion(text="DATABASE", display_meta='keyword'),
+        Completion(text="TABLE", display_meta='keyword'),
+        Completion(text="SYSTEM", display_meta='keyword'),
+    ])
+    assert Completion(text="CREATE", display_meta="keyword") not in result
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index c4a7da11..23064459 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -4,23 +4,24 @@ from pgcli.packages.sqlcompletion import (
 from pgcli.packages.parseutils.tables import TableReference
 import pytest
 
-def cols_etc(table, schema=None, alias=None, is_function=False, parent=None):
+def cols_etc(table, schema=None, alias=None, is_function=False, parent=None,
+             last_keyword=None):
     """Returns the expected select-clause suggestions for a single-table
     select."""
     return set([
         Column(table_refs=(TableReference(schema, table, alias, is_function),),
                qualifiable=True),
         Function(schema=parent),
-        Keyword()])
+        Keyword(last_keyword)])
 
 def test_select_suggests_cols_with_visible_table_scope():
     suggestions = suggest_type('SELECT  FROM tabl', 'SELECT ')
-    assert set(suggestions) == cols_etc('tabl')
+    assert set(suggestions) == cols_etc('tabl', last_keyword='SELECT')
 
 
 def test_select_suggests_cols_with_qualified_table_scope():
     suggestions = suggest_type('SELECT  FROM sch.tabl', 'SELECT ')
-    assert set(suggestions) == cols_etc('tabl', 'sch')
+    assert set(suggestions) == cols_etc('tabl', 'sch', last_keyword='SELECT')
 
 
 def test_cte_does_not_crash():
@@ -34,7 +35,7 @@ def test_cte_does_not_crash():
 ])
 def test_where_suggests_columns_functions_quoted_table(expression):
     suggestions = suggest_type(expression, expression)
-    assert set(suggestions) == cols_etc('tabl', alias='"tabl"')
+    assert set(suggestions) == cols_etc('tabl', alias='"tabl"', last_keyword='WHERE')
 
 
 @pytest.mark.parametrize('expression', [
@@ -53,7 +54,7 @@ def test_where_suggests_columns_functions_quoted_table(expression):
 ])
 def test_where_suggests_columns_functions(expression):
     suggestions = suggest_type(expression, expression)
-    assert set(suggestions) == cols_etc('tabl')
+    assert set(suggestions) == cols_etc('tabl', last_keyword='WHERE')
 
 
 @pytest.mark.parametrize('expression', [
@@ -62,7 +63,7 @@ def test_where_suggests_columns_functions(expression):
 ])
 def test_where_in_suggests_columns(expression):
     suggestions = suggest_type(expression, expression)
-    assert set(suggestions) == cols_etc('tabl')
+    assert set(suggestions) == cols_etc('tabl', last_keyword='WHERE')
 
 @pytest.mark.parametrize('expression', [
     'SELECT 1 AS ',
@@ -76,7 +77,7 @@ def test_after_as(expression):
 def test_where_equals_any_suggests_columns_or_keywords():
     text = 'SELECT * FROM tabl WHERE foo = ANY('
     suggestions = suggest_type(text, text)
-    assert set(suggestions) == cols_etc('tabl')
+    assert set(suggestions) == cols_etc('tabl', last_keyword='WHERE')
 
 
 def test_lparen_suggests_cols():
@@ -90,7 +91,7 @@ def test_select_suggests_cols_and_funcs():
     assert set(suggestions) == set([
          Column(table_refs=(), qualifiable=True),
          Function(schema=None),
-         Keyword(),
+         Keyword('SELECT'),
     ])
 
 
@@ -217,21 +218,24 @@ def test_distinct_suggests_cols(text):
     assert set(suggestions) == set([
         Column(table_refs=(), local_tables=(), qualifiable=True),
         Function(schema=None),
-        Keyword()
+        Keyword('DISTINCT')
     ])
 
 
-@pytest.mark.parametrize('text, text_before', [
+@pytest.mark.parametrize('text, text_before, last_keyword', [
     (
         'SELECT DISTINCT FROM tbl x JOIN tbl1 y',
-        'SELECT DISTINCT'
+        'SELECT DISTINCT',
+        'SELECT',
     ),
     (
         'SELECT * FROM tbl x JOIN tbl1 y ORDER BY ',
-        'SELECT * FROM tbl x JOIN tbl1 y ORDER BY '
+        'SELECT * FROM tbl x JOIN tbl1 y ORDER BY ',
+        'BY',
     )
 ])
-def test_distinct_and_order_by_suggestions_with_aliases(text, text_before):
+def test_distinct_and_order_by_suggestions_with_aliases(text, text_before,
+                                                        last_keyword):
     suggestions = suggest_type(text, text_before)
     assert set(suggestions) == set([
         Column(
@@ -243,7 +247,7 @@ def test_distinct_and_order_by_suggestions_with_aliases(text, text_before):
             qualifiable=True
         ),
         Function(schema=None),
-        Keyword()
+        Keyword(last_keyword)
     ])
 
 
@@ -275,7 +279,7 @@ def test_col_comma_suggests_cols():
     assert set(suggestions) == set([
         Column(table_refs=((None, 'tbl', None, False),), qualifiable=True),
         Function(schema=None),
-        Keyword(),
+        Keyword('SELECT'),
     ])
 
 
@@ -318,7 +322,7 @@ def test_insert_into_lparen_comma_suggests_cols():
 def test_partially_typed_col_name_suggests_col_names():
     suggestions = suggest_type('SELECT * FROM tabl WHERE col_n',
             'SELECT * FROM tabl WHERE col_n')
-    assert set(suggestions) == cols_etc('tabl')
+    assert set(suggestions) == cols_etc('tabl', last_keyword='WHERE')
 
 
 def test_dot_suggests_cols_of_a_table_or_schema_qualified_table():
@@ -435,7 +439,7 @@ def test_sub_select_col_name_completion():
     assert set(suggestions) == set([
         Column(table_refs=((None, 'abc', None, False),), qualifiable=True),
         Function(schema=None),
-        Keyword(),
+        Keyword('SELECT'),
     ])
 
 
@@ -599,7 +603,7 @@ def test_2_statements_2nd_current():
     assert set(suggestions) == set([
         Column(table_refs=((None, 'b', None, False),), qualifiable=True),
         Function(schema=None),
-        Keyword()
+        Keyword('SELECT')
     ])
 
     # Should work even if first statement is invalid
@@ -621,7 +625,7 @@ def test_2_statements_1st_current():
 
     suggestions = suggest_type('select  from a; select * from b',
                                'select ')
-    assert set(suggestions) == cols_etc('a')
+    assert set(suggestions) == cols_etc('a', last_keyword='SELECT')
 
 
 def test_3_statements_2nd_current():
@@ -634,7 +638,7 @@ def test_3_statements_2nd_current():
 
     suggestions = suggest_type('select * from a; select  from b; select * from c',
                                'select * from a; select ')
-    assert set(suggestions) == cols_etc('b')
+    assert set(suggestions) == cols_etc('b', last_keyword='SELECT')
 
 @pytest.mark.parametrize('text', [
 '''
@@ -685,7 +689,7 @@ def test_statements_in_function_body(text):
     assert set(suggestions) == set([
         Column(table_refs=((None, 'foo', None, False),), qualifiable=True),
         Function(schema=None),
-        Keyword()
+        Keyword('SELECT'),
     ])
 
 functions = [
@@ -709,12 +713,12 @@ SELECT 1 FROM foo;
 @pytest.mark.parametrize('text', functions)
 def test_statements_with_cursor_after_function_body(text):
     suggestions = suggest_type(text, text[:text.find('; ') + 1])
-    assert set(suggestions) == set([Keyword()])
+    assert set(suggestions) == set([Keyword(), Special()])
 
 @pytest.mark.parametrize('text', functions)
 def test_statements_with_cursor_before_function_body(text):
     suggestions = suggest_type(text, '')
-    assert set(suggestions) == set([Keyword()])
+    assert set(suggestions) == set([Keyword(), Special()])
 
 def test_create_db_with_template():
     suggestions = suggest_type('create database foo with template ',
@@ -828,16 +832,16 @@ def test_invalid_sql():
 def test_suggest_where_keyword(text):
     # https://github.com/dbcli/mycli/issues/135
     suggestions = suggest_type(text, text)
-    assert set(suggestions) == cols_etc('foo')
+    assert set(suggestions) == cols_etc('foo', last_keyword='WHERE')
 
 
 @pytest.mark.parametrize('text, before, expected', [
     ('\\ns abc SELECT ', 'SELECT ', [
         Column(table_refs=(), qualifiable=True),
         Function(schema=None),
-        Keyword()
+        Keyword('SELECT')
     ]),
-    ('\\ns abc SELECT foo ', 'SELECT foo ',(Keyword(),)),
+    ('\\ns abc SELECT foo ', 'SELECT foo ', (Keyword(),)),
     ('\\ns abc SELECT t1. FROM tabl1 t1', 'SELECT t1.', [
         Table(schema='t1'),
         View(schema='t1'),
@@ -853,7 +857,7 @@ def test_named_query_completion(text, before, expected):
 def test_select_suggests_fields_from_function():
     suggestions = suggest_type('SELECT  FROM func()', 'SELECT ')
     assert set(suggestions) == cols_etc(
-        'func', is_function=True)
+        'func', is_function=True, last_keyword='SELECT')
 
 
 @pytest.mark.parametrize('sql', [
@@ -900,4 +904,4 @@ def test_handle_unrecognized_kw_generously():
     'ALTER TABLE foo ALTER ',
 ])
 def test_keyword_after_alter(sql):
-    assert Keyword() in set(suggest_type(sql, sql))
+    assert Keyword('ALTER') in set(suggest_type(sql, sql))