Beginnings of a table specific column completion.

3 files changed, 45 insertions, 6 deletions

diff --git a/pgcli/packages/parseutils.py b/pgcli/packages/parseutils.py
index e2eb2eee..e3c7be36 100644
--- a/pgcli/packages/parseutils.py
+++ b/pgcli/packages/parseutils.py
@@ -1,4 +1,7 @@
 import re
+import sqlparse
+from sqlparse.sql import IdentifierList, Identifier
+from sqlparse.tokens import Keyword, DML
 
 # This matches only alphanumerics and underscores.
 _LAST_WORD_RE = re.compile(r'(\w+)$')
@@ -48,6 +51,40 @@ def last_word(text, include_special_chars=False):
         else:
             return ''
 
-def extract_tables(sql):
-    return []
+def is_subselect(parsed):
+    if not parsed.is_group():
+        return False
+    for item in parsed.tokens:
+        if item.ttype is DML and item.value.upper() == 'SELECT':
+            return True
+    return False
+
+def extract_from_part(parsed):
+    from_seen = False
+    for item in parsed.tokens:
+        if from_seen:
+            if is_subselect(item):
+                for x in extract_from_part(item):
+                    yield x
+            elif item.ttype is Keyword:
+                raise StopIteration
+            else:
+                yield item
+        elif item.ttype is Keyword and item.value.upper() == 'FROM':
+            from_seen = True
 
+def extract_table_identifiers(token_stream):
+    for item in token_stream:
+        if isinstance(item, IdentifierList):
+            for identifier in item.get_identifiers():
+                yield identifier.get_name()
+        elif isinstance(item, Identifier):
+            yield item.get_name()
+        # It's a bug to check for Keyword here, but in the example
+        # above some tables names are identified as keywords...
+        elif item.ttype is Keyword:
+            yield item.value
+
+def extract_tables(sql):
+    stream = extract_from_part(sqlparse.parse(sql)[0])
+    return list(extract_table_identifiers(stream))
diff --git a/pgcli/pgcompleter.py b/pgcli/pgcompleter.py
index 533b7f33..b4772625 100644
--- a/pgcli/pgcompleter.py
+++ b/pgcli/pgcompleter.py
@@ -32,9 +32,8 @@ class PGCompleter(Completer):
     databases = []
     tables = []
     # This will create a defaultdict which is initialized with a list that has
-    # a '*' by default. Inspired by the constant_factory example in
-    # https://docs.python.org/2/library/collections.html#defaultdict-examples
-    columns = defaultdict(itertools.repeat(['*']).next)
+    # a '*' by default.
+    columns = defaultdict(lambda: ['*'])
     all_completions = set(keywords)
 
     def __init__(self, smart_completion=True):
@@ -88,6 +87,8 @@ class PGCompleter(Completer):
             scoped_cols = []
             for table in scope:
                 scoped_cols.extend(self.columns[table])
+            print ('scope:', scope)
+            print ('scoped_cols:', scoped_cols)
             return self.find_matches(word_before_cursor, scoped_cols)
         elif category == 'columns-and-functions':
             scoped_cols = []
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 6bfa5103..61c363fe 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -120,7 +120,8 @@ class PGExecute(object):
     def columns(self, table):
         with self.conn.cursor() as cur:
             cur.execute(self.columns_query, (table,))
-            return [x[0] for x in cur.fetchall()]
+            cols = [x[0] for x in cur.fetchall()]
+            return cols
 
     def all_columns(self):
         columns = set()