attempt to suggest column name when function is typed

3 files changed, 24 insertions, 2 deletions

diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index 6e45f435..3d46e85c 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -29,7 +29,12 @@ def suggest_type(full_text, text_before_cursor):
         last_token = parsed[0].token_prev(len(parsed[0].tokens))
         last_token = last_token.value if last_token else ''
 
-    if last_token.lower() in ('set', 'by'):
+    def is_function_word(word):
+        return word and len(word) > 1 and word[-1] == '('
+
+    if is_function_word(word_before_cursor):
+        return ('columns', extract_tables(full_text))
+    elif last_token.lower() in ('set', 'by', 'distinct'):
         return ('columns', extract_tables(full_text))
     elif last_token.lower() in ('select', 'where', 'having'):
         return ('columns-and-functions', extract_tables(full_text))
diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py
index 8cc2e9d9..12727d50 100644
--- a/pgcli/pgcompleter.py
+++ b/pgcli/pgcompleter.py
@@ -1,6 +1,5 @@
 from __future__ import print_function
 from collections import defaultdict
-import itertools
 from prompt_toolkit.completion import Completer, Completion
 from .packages.sqlcompletion import suggest_type
 
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
new file mode 100644
index 00000000..d31042f0
--- /dev/null
+++ b/tests/test_sqlcompletion.py
@@ -0,0 +1,18 @@
+from pgcli.packages.sqlcompletion import suggest_type
+
+class TestSqlCompletion:
+    def test_lparen_suggest_cols(self):
+        suggestion = suggest_type('SELECT MAX( FROM table_name', 'SELECT MAX(')
+        assert (suggestion and suggestion[0] == 'columns')
+
+    def test_select_suggest_cols_and_funcs(self):
+        suggestion = suggest_type('SELECT ', 'SELECT ')
+        assert (suggestion and suggestion[0] == 'columns-and-functions')
+
+    def test_from_suggest_tables(self):
+        suggestion = suggest_type('SELECT * FROM ', 'SELECT * FROM ')
+        assert (suggestion and suggestion[0] == 'tables')
+
+    def test_distinct_suggest_cols(self):
+        suggestion = suggest_type('SELECT DISTINCT ', 'SELECT DISTINCT ')
+        assert (suggestion and suggestion[0] == 'columns')