First attempt at multiple column suggestion.

1 files changed, 34 insertions, 15 deletions

diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index c6029ae0..b56a5008 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -11,6 +11,8 @@ def suggest_type(full_text, text_before_cursor):
     A scope for a column category will be a list of tables.
     """
 
+    tables = extract_tables(full_text)
+
     word_before_cursor = last_word(text_before_cursor,
             include='all_punctuations')
 
@@ -25,25 +27,42 @@ def suggest_type(full_text, text_before_cursor):
     else:
         parsed = sqlparse.parse(text_before_cursor)
 
-    last_token = ''
-    if parsed:
-        last_token = parsed[0].token_prev(len(parsed[0].tokens))
-        last_token = last_token.value if last_token else ''
+    # Need to check if `p` is not empty, since an empty string will result in
+    # an empty tuple.
+    p = parsed[0] if parsed else None
+    n = p and len(p.tokens) or 0
+    last_token = p and p.token_prev(n) or ''
+    last_token_v = last_token.value if last_token else ''
 
     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))
-    elif last_token.lower() in ('from', 'update', 'into', 'describe'):
-        return ('tables', [])
-    elif last_token in ('d',):  # \d
-        return ('tables', [])
-    elif last_token.lower() in ('c', 'use'):  # \c
-        return ('databases', [])
+    elif last_token_v.lower().endswith('('):
+        return ('columns', tables)
+    if last_token_v.lower() in ('set', 'by', 'distinct'):
+        return ('columns', tables)
+    elif last_token_v.lower() in ('select', 'where', 'having'):
+        return ('columns-and-functions', tables)
+    elif last_token_v.lower() in ('from', 'update', 'into', 'describe'):
+        return ('tables', tables)
+    elif last_token_v in ('d',):  # \d
+        return ('tables', tables)
+    elif last_token_v.lower() in ('c', 'use'):  # \c
+        return ('databases', tables)
+    elif last_token_v == ',':
+        prev_keyword = find_prev_keyword(text_before_cursor)
+        if prev_keyword in('select', 'where', 'having'):
+            return ('columns-and-functions', tables)
+        return ('keywords', tables)
     else:
-        return ('keywords', [])
+        return ('keywords', tables)
+
+def find_prev_keyword(sql):
+    if not sql.strip():
+        return None
+
+    for t in sqlparse.parse(sql)[0].flatten():
+        if t.is_keyword:
+            return t.value.lower()