Fix another corner case for table name extraction in nested sql.

2 files changed, 18 insertions, 15 deletions

diff --git a/pgcli/packages/parseutils.py b/pgcli/packages/parseutils.py
index b8094076..b3086a63 100644
--- a/pgcli/packages/parseutils.py
+++ b/pgcli/packages/parseutils.py
@@ -92,7 +92,8 @@ def extract_from_part(parsed, stop_at_punctuation=True):
         # So this check here is necessary.
         elif isinstance(item, IdentifierList):
             for identifier in item.get_identifiers():
-                if identifier.ttype is Keyword and identifier.value.upper() == 'FROM':
+                if (identifier.ttype is Keyword and
+                        identifier.value.upper() == 'FROM'):
                     tbl_prefix_seen = True
                     break
 
@@ -100,7 +101,12 @@ def extract_table_identifiers(token_stream):
     for item in token_stream:
         if isinstance(item, IdentifierList):
             for identifier in item.get_identifiers():
-                real_name = identifier.get_real_name()
+                # Sometimes Keywords (such as FROM ) are classified as
+                # identifiers which don't have the get_real_name() method.
+                try:
+                    real_name = identifier.get_real_name()
+                except AttributeError:
+                    continue
                 if real_name:
                     yield (real_name, identifier.get_alias() or real_name)
         elif isinstance(item, Identifier):
@@ -109,11 +115,8 @@ def extract_table_identifiers(token_stream):
                 yield (real_name, item.get_alias() or real_name)
         elif isinstance(item, Function):
             yield (item.get_name(), 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, item.value)
 
+# extract_tables is inspired from examples in the sqlparse lib.
 def extract_tables(sql, include_alias=False):
     parsed = sqlparse.parse(sql)
     if not parsed:
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index 7e876e70..860fabb6 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -81,15 +81,15 @@ def test_sub_select_table_name_completion():
             'SELECT * FROM (SELECT * FROM ')
     assert suggestion == ('tables', [])
 
-#def test_sub_select_col_name_completion():
-    #suggestion = suggest_type('SELECT * FROM (SELECT  FROM abc',
-            #'SELECT * FROM (SELECT ')
-    #assert suggestion == ('columns', ['abc'])
-
-#def test_sub_select_multiple_col_name_completion():
-    #suggestion = suggest_type('SELECT * FROM (SELECT a, FROM abc',
-            #'SELECT * FROM (SELECT a, ')
-    #assert suggestion == ('columns', ['abc'])
+def test_sub_select_col_name_completion():
+    suggestion = suggest_type('SELECT * FROM (SELECT  FROM abc',
+            'SELECT * FROM (SELECT ')
+    assert suggestion == ('columns-and-functions', ['abc'])
+
+def test_sub_select_multiple_col_name_completion():
+    suggestion = suggest_type('SELECT * FROM (SELECT a, FROM abc',
+            'SELECT * FROM (SELECT a, ')
+    assert suggestion == ('columns-and-functions', ['abc'])
 
 def test_sub_select_dot_col_name_completion():
     suggestion = suggest_type('SELECT * FROM (SELECT t. FROM tabl t',