Make parseutils to work with hanging comma and insert statements.

2 files changed, 21 insertions, 16 deletions

diff --git a/pgcli/packages/parseutils.py b/pgcli/packages/parseutils.py
index f4a5f333..4b5564da 100644
--- a/pgcli/packages/parseutils.py
+++ b/pgcli/packages/parseutils.py
@@ -1,6 +1,6 @@
 import re
 import sqlparse
-from sqlparse.sql import IdentifierList, Identifier
+from sqlparse.sql import IdentifierList, Identifier, Function
 from sqlparse.tokens import Keyword, DML
 
 cleanup_regex = {
@@ -68,9 +68,9 @@ def is_subselect(parsed):
     return False
 
 def extract_from_part(parsed):
-    from_seen = False
+    tbl_prefix_seen = False
     for item in parsed.tokens:
-        if from_seen:
+        if tbl_prefix_seen:
             if is_subselect(item):
                 for x in extract_from_part(item):
                     yield x
@@ -80,15 +80,20 @@ def extract_from_part(parsed):
                 yield item
         elif item.ttype is Keyword and item.value.upper() in ('FROM', 'INTO',
                 'UPDATE', 'TABLE', ):
-            from_seen = True
+            tbl_prefix_seen = True
+        elif isinstance(item, IdentifierList):
+            for identifier in item.get_identifiers():
+                if identifier.ttype is Keyword and identifier.value.upper() == 'FROM':
+                    tbl_prefix_seen = True
+                    break
 
 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()
+                yield identifier.get_real_name()
+        elif isinstance(item, Identifier) or isinstance(item, Function):
+            yield item.get_real_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:
diff --git a/tests/test_parse_utils.py b/tests/test_parseutils.py
index 6b00676c..63198d57 100644
--- a/tests/test_parse_utils.py
+++ b/tests/test_parseutils.py
@@ -16,14 +16,14 @@ def test_simple_select_with_cols_multiple_tables():
     tables = extract_tables('select a,b from abc, def')
     assert tables == ['abc', 'def']
 
-#def test_select_with_hanging_comma_single_table():
-    #tables = extract_tables('select a, from abc')
-    #assert tables == ['abc']
+def test_select_with_hanging_comma_single_table():
+    tables = extract_tables('select a, from abc')
+    assert tables == ['abc']
 
-#def test_select_with_hanging_comma_multiple_tables():
-    #tables = extract_tables('select a, from abc, def')
-    #assert tables == ['abc']
+def test_select_with_hanging_comma_multiple_tables():
+    tables = extract_tables('select a, from abc, def')
+    assert tables == ['abc', 'def']
 
-#def test_simple_insert_single_table():
-    #tables = extract_tables('insert into abc (id, name) values (1, "def")')
-    #assert tables == ['abc']
+def test_simple_insert_single_table():
+    tables = extract_tables('insert into abc (id, name) values (1, "def")')
+    assert tables == ['abc']